From e4343424263b76b4176fd0846864e666a2e6eb08 Mon Sep 17 00:00:00 2001 From: Kilian Singer Date: Wed, 13 May 2026 09:49:02 +0200 Subject: [PATCH 1/6] Zeeman Matrix elements are off-diagonal with with small magnetic field present as it typically dominates Spin-Orbit coupling due to the latter scaling with 1/n^3. So (ml1 + gs * ms1) is off-diagonal due to |n l j mj> are containing different admixtures of ml and ms states given by Clebsch-Gordan coefficients. Modified class StarkMap to take this into consideration. The error is still present in class LevelPlot. --- arc/alkali_atom_functions.py | 43 +++++++++++++++++ arc/calculations_atom_single.py | 83 ++++++++++++++++++++++++++------- 2 files changed, 109 insertions(+), 17 deletions(-) diff --git a/arc/alkali_atom_functions.py b/arc/alkali_atom_functions.py index f4dacd0..c9c1053 100644 --- a/arc/alkali_atom_functions.py +++ b/arc/alkali_atom_functions.py @@ -2679,7 +2679,50 @@ def getZeemanEnergyShift( ms = mj - ml sumOverMl += (ml + gs * ms) * abs(CG(l, ml, s, ms, j, mj)) ** 2 return prefactor * sumOverMl + + def getZeemanEnergyShiftOffDiagonal( + self, + l: int, + j1: float, + mj1: float, + j2: float, + mj2: float, + magneticFieldBz: float, + s: float = 0.5, + ) -> float: + r""" + Retuns off diagonal linear (paramagnetic) Zeeman shift. + :math:`\mathcal{H}_P=\frac{\mu_B B_z}{\hbar}(\hat{L}_{\rm z}+\ + g_{\rm S}S_{\rm z})` + + Args: + l (int): orbital angular momentum + j1 (float): total angular momentum of first state + mj1 (float): projection of total angular momentum of first state along z-axis + j2 (float): total angular momentum of second state + mj2 (float): projection of total angular momentum of second state along z-axis + magneticFieldBz (float): applied magnetic field (along z-axis + only) in units of T (Tesla) + s (float): optional, total spin angular momentum of state. + By default 0.5 for Alkali atoms. + + Returns: + float: energy offset of the state (in J) + """ + prefactor = physical_constants["Bohr magneton"][0] * magneticFieldBz + gs = -physical_constants["electron g factor"][0] + sumOverMl = 0 + + for ml1 in np.linspace(mj1 - s, mj1 + s, round(2 * s + 1)): + for ml2 in np.linspace(mj2 - s, mj2 + s, round(2 * s + 1)): + if abs(ml1) <= l + 0.1 and abs(ml2) <= l + 0.1: + ms1 = mj1 - ml1 + ms2 = mj2 - ml2 + if(ms1==ms2 and ml1==ml2): + sumOverMl += (ml1 + gs * ms1) * CG(l, ml1, s, ms1, j1, mj1)*CG(l, ml1, s, ms1, j2, mj2) + return prefactor * sumOverMl + def _getRadialDipoleSemiClassical( self, n1: int, diff --git a/arc/calculations_atom_single.py b/arc/calculations_atom_single.py index 467d4e7..8cce85b 100644 --- a/arc/calculations_atom_single.py +++ b/arc/calculations_atom_single.py @@ -584,7 +584,15 @@ def __init__(self, atom): off-diagonal elements of Stark-matrix divided by electric field value. To get off diagonal elemements multiply this matrix with electric field value. Full Stark matrix is obtained as - `fullStarkMatrix` = :obj:`mat1` + :obj:`mat2` *`eField`. Calculated by + `fullStarkMatrix` = :obj:`mat1` + :obj:`mat2` *`eField` + :obj:`mat3`. Calculated by + :obj:`defineBasis` in the basis :obj:`basisStates`. + """ + self.mat3 = [] + """ + off-diagonal elements of Zeeman-Matrix divided multiplied by magnetic + field value. Only relevant if Bz is not zero. + Full Stark matrix is obtained as + `fullStarkMatrix` = :obj:`mat1` + :obj:`mat2` *`eField` + :obj:`mat3`. Calculated by :obj:`defineBasis` in the basis :obj:`basisStates`. """ self.indexOfCoupledState = [] @@ -697,7 +705,10 @@ def defineBasis( second part :obj:`mat2` corresponds to off-diagonal elements that are propotional to electric field. Overall interaction matrix for electric field `eField` can be then obtained as + with Bz equal to zero: `fullStarkMatrix` = :obj:`mat1` + :obj:`mat2` *`eField` + with Bz not equal to zero: + `fullStarkMatrix` = :obj:`mat1` + :obj:`mat2` *`eField` + :obj:`mat3` Args: n (int): principal quantum number of the state @@ -742,15 +753,25 @@ def defineBasis( self.Bz = Bz self.s = s # save calculation details END - - for tn in xrange(nMin, nMax + 1): - for tl in xrange(min(maxL + 1, tn)): - for tj in np.linspace(tl - s, tl + s, round(2 * s + 1)): - if (abs(mj) - 0.1 <= tj) and ( - tn >= self.atom.groundStateN - or [tn, tl, tj] in self.atom.extraLevels - ): - states.append([tn, tl, tj, mj]) + if Bz==0: + for tn in xrange(nMin, nMax + 1): + for tl in xrange(min(maxL + 1, tn)): + for tj in np.linspace(tl - s, tl + s, round(2 * s + 1)): + if (abs(mj) - 0.1 <= tj) and ( + tn >= self.atom.groundStateN + or [tn, tl, tj] in self.atom.extraLevels + ): + states.append([tn, tl, tj, mj]) + else: + for tn in xrange(nMin, nMax + 1): + for tl in xrange(min(maxL + 1, tn)): + for tj in np.linspace(tl - s, tl + s, round(2 * s + 1)): + for tmj in np.linspace(mj-1,mj+1,round(2 * s + 2)): # we need a bit more mj states as we do not know ml + if (abs(tmj) - 0.1 <= tj) and ( + tn >= self.atom.groundStateN + or [tn, tl, tj] in self.atom.extraLevels + ): + states.append([tn, tl, tj, tmj]) dimension = len(states) if progressOutput: @@ -777,6 +798,8 @@ def defineBasis( self.mat1 = np.zeros((dimension, dimension), dtype=np.double) self.mat2 = np.zeros((dimension, dimension), dtype=np.double) + if Bz!=0: + self.mat3 = np.zeros((dimension, dimension), dtype=np.double) self.basisStates = states self.indexOfCoupledState = indexOfCoupledState @@ -819,11 +842,11 @@ def defineBasis( states[ii][0], states[ii][1], states[ii][2], - mj, + states[ii][3], states[jj][0], states[jj][1], states[jj][2], - mj, + states[jj][3], s=self.s, ) * 1.0e-9 @@ -832,10 +855,29 @@ def defineBasis( self.mat2[jj][ii] = coupling self.mat2[ii][jj] = coupling + if Bz!=0: + coupling2=0 + if states[ii][0]==states[jj][0] and states[ii][1]==states[jj][1]: + coupling2 = self.atom.getZeemanEnergyShiftOffDiagonal( + states[ii][1], + states[ii][2], + states[ii][3], + states[jj][2], + states[jj][3], + self.Bz, + s=self.s + )/ C_h* 1.0e-9 + + self.mat3[jj][ii] = coupling + self.mat3[ii][jj] = coupling + if progressOutput: print("\n") if debugOutput: - print(self.mat1 + self.mat2) + if Bz==0: + print(self.mat1 + self.mat2) + else: + print(self.mat1 + self.mat2 + self.mat3) print(self.mat2[0]) self.atom.updateDipoleMatrixElementsFile() @@ -980,8 +1022,10 @@ def diagonalise( "\r%d%%" % (float(progress) / float(len(eFieldList)) * 100) ) sys.stdout.flush() - - m = self.mat1 + self.mat2 * eField + if self.Bz==0: + m = self.mat1 + self.mat2 * eField + else: + m = self.mat1 + self.mat2 * eField + self.mat3 ev, egvector = eigh(m) @@ -1583,6 +1627,7 @@ def getState( maxL, accountForAmplitude=0.95, debugOutput=False, + Bz=0 ): r""" Returns basis states and coefficients that make up for a given electric @@ -1610,6 +1655,7 @@ def getState( for 95\% of the state amplitude. debugOutput (bool): optional, prints additional debug information if True. Default False. + Bz (float): Magnetic field in z direction. Returns: **array of states** in format [[n1, l1, j1, mj1], ...] and @@ -1621,10 +1667,13 @@ def getState( """ self.defineBasis( - state[0], state[1], state[2], state[3], minN, maxN, maxL + state[0], state[1], state[2], state[3], minN, maxN, maxL, Bz ) - m = self.mat1 + self.mat2 * electricField + if Bz==0: + m = self.mat1 + self.mat2 * electricField + else: + m = self.mat1 + self.mat2 * electricField + self.mat3 ev, egvector = eigh(m) # find which state in the electric field has strongest contribution From 3362fd015307d2a447fed5988f6229f3400bc5fb Mon Sep 17 00:00:00 2001 From: Stefan Aull Date: Sat, 16 May 2026 16:02:31 +0200 Subject: [PATCH 2/6] Fix mat3: use coupling2 instead of coupling --- arc/calculations_atom_single.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/arc/calculations_atom_single.py b/arc/calculations_atom_single.py index 8cce85b..3ad9296 100644 --- a/arc/calculations_atom_single.py +++ b/arc/calculations_atom_single.py @@ -868,8 +868,8 @@ def defineBasis( s=self.s )/ C_h* 1.0e-9 - self.mat3[jj][ii] = coupling - self.mat3[ii][jj] = coupling + self.mat3[jj][ii] = coupling2 + self.mat3[ii][jj] = coupling2 if progressOutput: print("\n") From d133deb00399fecb113b3821b29c808295e4d2ca Mon Sep 17 00:00:00 2001 From: Stefan Aull Date: Tue, 19 May 2026 18:19:40 +0200 Subject: [PATCH 3/6] Cleaned up Zeeman implementation: removed basis extention, removed mat3, fixed selection rule mj1=mj2 in getZeemanEnergyShiftOffDiagonal --- arc/alkali_atom_functions.py | 25 +++++---- arc/calculations_atom_single.py | 95 ++++++++++++--------------------- 2 files changed, 49 insertions(+), 71 deletions(-) diff --git a/arc/alkali_atom_functions.py b/arc/alkali_atom_functions.py index c9c1053..6a3cc3a 100644 --- a/arc/alkali_atom_functions.py +++ b/arc/alkali_atom_functions.py @@ -2679,7 +2679,7 @@ def getZeemanEnergyShift( ms = mj - ml sumOverMl += (ml + gs * ms) * abs(CG(l, ml, s, ms, j, mj)) ** 2 return prefactor * sumOverMl - + def getZeemanEnergyShiftOffDiagonal( self, l: int, @@ -2710,19 +2710,24 @@ def getZeemanEnergyShiftOffDiagonal( Returns: float: energy offset of the state (in J) """ + if abs(mj1 - mj2) > 0.1: + return 0.0 + prefactor = physical_constants["Bohr magneton"][0] * magneticFieldBz gs = -physical_constants["electron g factor"][0] - sumOverMl = 0 + sumOverMl = 0.0 - for ml1 in np.linspace(mj1 - s, mj1 + s, round(2 * s + 1)): - for ml2 in np.linspace(mj2 - s, mj2 + s, round(2 * s + 1)): - if abs(ml1) <= l + 0.1 and abs(ml2) <= l + 0.1: - ms1 = mj1 - ml1 - ms2 = mj2 - ml2 - if(ms1==ms2 and ml1==ml2): - sumOverMl += (ml1 + gs * ms1) * CG(l, ml1, s, ms1, j1, mj1)*CG(l, ml1, s, ms1, j2, mj2) + for ml in np.linspace(mj1 - s, mj1 + s, round(2 * s + 1)): + if abs(ml) <= l + 0.1: + ms = mj1 - ml + if abs(ms) <= s + 0.1: + sumOverMl += ( + (self.gL * ml + gs * ms) + * CG(l, ml, s, ms, j1, mj1) + * CG(l, ml, s, ms, j2, mj2) + ) return prefactor * sumOverMl - + def _getRadialDipoleSemiClassical( self, n1: int, diff --git a/arc/calculations_atom_single.py b/arc/calculations_atom_single.py index 3ad9296..916a0ea 100644 --- a/arc/calculations_atom_single.py +++ b/arc/calculations_atom_single.py @@ -584,15 +584,7 @@ def __init__(self, atom): off-diagonal elements of Stark-matrix divided by electric field value. To get off diagonal elemements multiply this matrix with electric field value. Full Stark matrix is obtained as - `fullStarkMatrix` = :obj:`mat1` + :obj:`mat2` *`eField` + :obj:`mat3`. Calculated by - :obj:`defineBasis` in the basis :obj:`basisStates`. - """ - self.mat3 = [] - """ - off-diagonal elements of Zeeman-Matrix divided multiplied by magnetic - field value. Only relevant if Bz is not zero. - Full Stark matrix is obtained as - `fullStarkMatrix` = :obj:`mat1` + :obj:`mat2` *`eField` + :obj:`mat3`. Calculated by + `fullStarkMatrix` = :obj:`mat1` + :obj:`mat2` *`eField`. Calculated by :obj:`defineBasis` in the basis :obj:`basisStates`. """ self.indexOfCoupledState = [] @@ -705,10 +697,7 @@ def defineBasis( second part :obj:`mat2` corresponds to off-diagonal elements that are propotional to electric field. Overall interaction matrix for electric field `eField` can be then obtained as - with Bz equal to zero: `fullStarkMatrix` = :obj:`mat1` + :obj:`mat2` *`eField` - with Bz not equal to zero: - `fullStarkMatrix` = :obj:`mat1` + :obj:`mat2` *`eField` + :obj:`mat3` Args: n (int): principal quantum number of the state @@ -753,25 +742,14 @@ def defineBasis( self.Bz = Bz self.s = s # save calculation details END - if Bz==0: - for tn in xrange(nMin, nMax + 1): - for tl in xrange(min(maxL + 1, tn)): - for tj in np.linspace(tl - s, tl + s, round(2 * s + 1)): - if (abs(mj) - 0.1 <= tj) and ( - tn >= self.atom.groundStateN - or [tn, tl, tj] in self.atom.extraLevels - ): - states.append([tn, tl, tj, mj]) - else: - for tn in xrange(nMin, nMax + 1): - for tl in xrange(min(maxL + 1, tn)): - for tj in np.linspace(tl - s, tl + s, round(2 * s + 1)): - for tmj in np.linspace(mj-1,mj+1,round(2 * s + 2)): # we need a bit more mj states as we do not know ml - if (abs(tmj) - 0.1 <= tj) and ( - tn >= self.atom.groundStateN - or [tn, tl, tj] in self.atom.extraLevels - ): - states.append([tn, tl, tj, tmj]) + for tn in xrange(nMin, nMax + 1): + for tl in xrange(min(maxL + 1, tn)): + for tj in np.linspace(tl - s, tl + s, round(2 * s + 1)): + if (abs(mj) - 0.1 <= tj) and ( + tn >= self.atom.groundStateN + or [tn, tl, tj] in self.atom.extraLevels + ): + states.append([tn, tl, tj, mj]) dimension = len(states) if progressOutput: @@ -798,8 +776,6 @@ def defineBasis( self.mat1 = np.zeros((dimension, dimension), dtype=np.double) self.mat2 = np.zeros((dimension, dimension), dtype=np.double) - if Bz!=0: - self.mat3 = np.zeros((dimension, dimension), dtype=np.double) self.basisStates = states self.indexOfCoupledState = indexOfCoupledState @@ -837,6 +813,28 @@ def defineBasis( # add off-diagonal element for jj in xrange(ii + 1, dimension): + if ( + Bz != 0 + and states[ii][0] == states[jj][0] + and states[ii][1] == states[jj][1] + and abs(states[ii][3] - states[jj][3]) < 0.1 + ): + zeemanCoupling = ( + self.atom.getZeemanEnergyShiftOffDiagonal( + states[ii][1], + states[ii][2], + states[ii][3], + states[jj][2], + states[jj][3], + self.Bz, + s=self.s, + ) + / C_h + * 1.0e-9 + ) + self.mat1[jj][ii] = zeemanCoupling + self.mat1[ii][jj] = zeemanCoupling + coupling = ( self._eFieldCouplingDivE( states[ii][0], @@ -855,29 +853,10 @@ def defineBasis( self.mat2[jj][ii] = coupling self.mat2[ii][jj] = coupling - if Bz!=0: - coupling2=0 - if states[ii][0]==states[jj][0] and states[ii][1]==states[jj][1]: - coupling2 = self.atom.getZeemanEnergyShiftOffDiagonal( - states[ii][1], - states[ii][2], - states[ii][3], - states[jj][2], - states[jj][3], - self.Bz, - s=self.s - )/ C_h* 1.0e-9 - - self.mat3[jj][ii] = coupling2 - self.mat3[ii][jj] = coupling2 - if progressOutput: print("\n") if debugOutput: - if Bz==0: - print(self.mat1 + self.mat2) - else: - print(self.mat1 + self.mat2 + self.mat3) + print(self.mat1 + self.mat2) print(self.mat2[0]) self.atom.updateDipoleMatrixElementsFile() @@ -1022,10 +1001,7 @@ def diagonalise( "\r%d%%" % (float(progress) / float(len(eFieldList)) * 100) ) sys.stdout.flush() - if self.Bz==0: - m = self.mat1 + self.mat2 * eField - else: - m = self.mat1 + self.mat2 * eField + self.mat3 + m = self.mat1 + self.mat2 * eField ev, egvector = eigh(m) @@ -1670,10 +1646,7 @@ def getState( state[0], state[1], state[2], state[3], minN, maxN, maxL, Bz ) - if Bz==0: - m = self.mat1 + self.mat2 * electricField - else: - m = self.mat1 + self.mat2 * electricField + self.mat3 + m = self.mat1 + self.mat2 * electricField ev, egvector = eigh(m) # find which state in the electric field has strongest contribution From 55c8afef6d8bce035e80c3d27a640557e7a54101 Mon Sep 17 00:00:00 2001 From: Kilian Singer Date: Wed, 20 May 2026 10:40:44 +0200 Subject: [PATCH 4/6] made comparison more robust --- arc/alkali_atom_functions.py | 2 +- arc/calculations_atom_single.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/arc/alkali_atom_functions.py b/arc/alkali_atom_functions.py index c9c1053..7eaa3f7 100644 --- a/arc/alkali_atom_functions.py +++ b/arc/alkali_atom_functions.py @@ -2719,7 +2719,7 @@ def getZeemanEnergyShiftOffDiagonal( if abs(ml1) <= l + 0.1 and abs(ml2) <= l + 0.1: ms1 = mj1 - ml1 ms2 = mj2 - ml2 - if(ms1==ms2 and ml1==ml2): + if(abs(ms1-ms2)<0.1 and abs(ml1-ml2)<0.1): sumOverMl += (ml1 + gs * ms1) * CG(l, ml1, s, ms1, j1, mj1)*CG(l, ml1, s, ms1, j2, mj2) return prefactor * sumOverMl diff --git a/arc/calculations_atom_single.py b/arc/calculations_atom_single.py index 3ad9296..e986bc6 100644 --- a/arc/calculations_atom_single.py +++ b/arc/calculations_atom_single.py @@ -857,7 +857,7 @@ def defineBasis( if Bz!=0: coupling2=0 - if states[ii][0]==states[jj][0] and states[ii][1]==states[jj][1]: + if abs(states[ii][0]-states[jj][0])<0.1 and abs(states[ii][1]-states[jj][1])<0.1: coupling2 = self.atom.getZeemanEnergyShiftOffDiagonal( states[ii][1], states[ii][2], From 554622c46f0905653ba97371db990904d7a6497c Mon Sep 17 00:00:00 2001 From: Kilian Singer Date: Wed, 20 May 2026 17:07:13 +0200 Subject: [PATCH 5/6] removed merge conflict --- arc/alkali_atom_functions.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/arc/alkali_atom_functions.py b/arc/alkali_atom_functions.py index 53422b0..6a3cc3a 100644 --- a/arc/alkali_atom_functions.py +++ b/arc/alkali_atom_functions.py @@ -2717,15 +2717,6 @@ def getZeemanEnergyShiftOffDiagonal( gs = -physical_constants["electron g factor"][0] sumOverMl = 0.0 -<<<<<<< HEAD - for ml1 in np.linspace(mj1 - s, mj1 + s, round(2 * s + 1)): - for ml2 in np.linspace(mj2 - s, mj2 + s, round(2 * s + 1)): - if abs(ml1) <= l + 0.1 and abs(ml2) <= l + 0.1: - ms1 = mj1 - ml1 - ms2 = mj2 - ml2 - if(abs(ms1-ms2)<0.1 and abs(ml1-ml2)<0.1): - sumOverMl += (ml1 + gs * ms1) * CG(l, ml1, s, ms1, j1, mj1)*CG(l, ml1, s, ms1, j2, mj2) -======= for ml in np.linspace(mj1 - s, mj1 + s, round(2 * s + 1)): if abs(ml) <= l + 0.1: ms = mj1 - ml @@ -2735,7 +2726,6 @@ def getZeemanEnergyShiftOffDiagonal( * CG(l, ml, s, ms, j1, mj1) * CG(l, ml, s, ms, j2, mj2) ) ->>>>>>> 0162ccc3e641384c62f84026ae63de8733bfe494 return prefactor * sumOverMl def _getRadialDipoleSemiClassical( From 973469941f640109607bca24273422c156f7f013 Mon Sep 17 00:00:00 2001 From: Stefan Aull Date: Wed, 20 May 2026 18:02:15 +0200 Subject: [PATCH 6/6] Implemented highlighting a state from the uncoupled basis, e.g. $|n, l, m_l, m_s\rangle$ instead of coupled basis. --- arc/calculations_atom_single.py | 90 ++++++++++++++++++++++++++++----- 1 file changed, 76 insertions(+), 14 deletions(-) diff --git a/arc/calculations_atom_single.py b/arc/calculations_atom_single.py index 3f62d5f..83d04ae 100644 --- a/arc/calculations_atom_single.py +++ b/arc/calculations_atom_single.py @@ -636,6 +636,7 @@ def __init__(self, atom): self.fittedCurveY = [] self.drivingFromState = [0, 0, 0, 0, 0] + self.highlightUncoupledState = None self.maxCoupling = 0.0 # STARK memoization @@ -867,6 +868,7 @@ def diagonalise( self, eFieldList, drivingFromState=[0, 0, 0, 0, 0], + highlightUncoupledState=None, progressOutput=False, debugOutput=False, upTo=4, @@ -884,6 +886,11 @@ def diagonalise( eFieldList (array): array of electric field strength (in V/m) for which we want to know energy eigenstates + highlightUncoupledState (array): optional target state in the + uncoupled basis :math:`[n,\\ell,m_\\ell,m_s]`. If provided, + highlighting shows the projection of each eigenstate onto this + uncoupled state. The calculation itself remains in the coupled + :math:`[n,\\ell,j,m_j]` basis. progressOutput (:obj:`bool`, optional): if True prints the progress of calculation; Set to false by default. debugOutput (:obj:`bool`, optional): if True prints additional @@ -901,13 +908,45 @@ def diagonalise( upTo = -1. """ - # if we are driving from some state - # ========= FIND LASER COUPLINGS (START) ======= - coupling = [] dimension = len(self.basisStates) self.maxCoupling = 0.0 self.drivingFromState = drivingFromState + self.highlightUncoupledState = highlightUncoupledState + + if highlightUncoupledState is not None and drivingFromState[0] != 0: + raise ValueError( + "highlightUncoupledState cannot be used together with " + "drivingFromState." + ) + + uncoupledStateVector = None + if highlightUncoupledState is not None: + uncoupledStateVector = np.zeros(dimension, dtype=np.double) + hn = round(highlightUncoupledState[0]) + hl = round(highlightUncoupledState[1]) + hml = highlightUncoupledState[2] + hms = highlightUncoupledState[3] + hmj = hml + hms + + for i, state in enumerate(self.basisStates): + if ( + state[0] == hn + and state[1] == hl + and abs(state[3] - hmj) < 0.1 + ): + uncoupledStateVector[i] = CG( + state[1], hml, self.s, hms, state[2], state[3] + ) + + if np.linalg.norm(uncoupledStateVector) < 0.1: + raise ValueError( + "highlightUncoupledState is not represented in the " + "current coupled basis." + ) + + # if we are driving from some state + # ========= FIND LASER COUPLINGS (START) ======= if self.drivingFromState[0] != 0: if progressOutput: print("Finding driving field coupling...") @@ -1009,7 +1048,12 @@ def diagonalise( sh = [] comp = [] for i in xrange(len(ev)): - sh.append(abs(egvector[indexOfCoupledState, i]) ** 2) + if uncoupledStateVector is None: + sh.append(abs(egvector[indexOfCoupledState, i]) ** 2) + else: + sh.append( + abs(np.vdot(uncoupledStateVector, egvector[:, i])) ** 2 + ) comp.append( self._stateComposition2( egvector[:, i], @@ -1085,10 +1129,19 @@ def exportData(self, fileBase, exportFormat="csv"): % (self.s) ) if self.drivingFromState[0] < 0.1: - commonHeader += ( - " - State highlighting based on the relative contribution \n" - + " of the original state in the eigenstates obtained by diagonalization." - ) + if self.highlightUncoupledState is not None: + state = self.highlightUncoupledState + commonHeader += ( + " - State highlighting based on the relative contribution \n" + + " of the uncoupled state " + + "|n=%d, l=%d, m_l=%.1f, m_s=%.1f> in the eigenstates." + % (state[0], state[1], state[2], state[3]) + ) + else: + commonHeader += ( + " - State highlighting based on the relative contribution \n" + + " of the original state in the eigenstates obtained by diagonalization." + ) else: commonHeader += ( " - State highlighting based on the relative driving strength \n" @@ -1283,10 +1336,18 @@ def plotLevelDiagram( cax = self.fig.add_axes([0.91, 0.1, 0.02, 0.8]) cb = matplotlib.colorbar.ColorbarBase(cax, cmap=cm, norm=cNorm) if self.drivingFromState[0] < 0.1: - cb.set_label( - r"$|\langle %s | \mu \rangle |^2$" - % printStateStringLatex(n, l, j, s=self.s) - ) + if self.highlightUncoupledState is not None: + state = self.highlightUncoupledState + cb.set_label( + r"$|\langle n=%d,\ell=%d,m_\ell=%.1f,m_s=%.1f" + r"|\mu\rangle|^2$" + % (state[0], state[1], state[2], state[3]) + ) + else: + cb.set_label( + r"$|\langle %s | \mu \rangle |^2$" + % printStateStringLatex(n, l, j, s=self.s) + ) else: cb.set_label(r"$( \Omega_\mu | \Omega )^2$") @@ -1483,11 +1544,12 @@ def getPolarizability( float: scalar polarizability in units of MHz cm :math:`^2` / V \ :math:`^2` """ - if self.drivingFromState[0] != 0: + if self.drivingFromState[0] != 0 or self.highlightUncoupledState is not None: raise Exception( "Program can only find Polarizability of the original " + "state if you highlight original state. You can do so by NOT " - + "specifying drivingFromState in diagonalise function." + + "specifying drivingFromState or highlightUncoupledState in " + + "diagonalise function." ) eFieldList = self.eFieldList