[docs]classStarkMap:""" Calculates Stark maps for single atom in a field This initializes calculation for the atom of a given type. For details of calculation see Zimmerman [1]_. For a quick working example see `Stark map example snippet`_. Args: atom (:obj:`AlkaliAtom`): ={ :obj:`alkali_atom_data.Lithium6`, :obj:`alkali_atom_data.Lithium7`, :obj:`alkali_atom_data.Sodium`, :obj:`alkali_atom_data.Potassium39`, :obj:`alkali_atom_data.Potassium40`, :obj:`alkali_atom_data.Potassium41`, :obj:`alkali_atom_data.Rubidium85`, :obj:`alkali_atom_data.Rubidium87`, :obj:`alkali_atom_data.Caesium` } Select the alkali metal for energy level diagram calculation Examples: State :math:`28~S_{1/2}~|m_j|=0.5` polarizability calculation >>> from arc import * >>> calc = StarkMap(Caesium()) >>> calc.defineBasis(28, 0, 0.5, 0.5, 23, 32, 20) >>> calc.diagonalise(np.linspace(00.,6000,600)) >>> print("%.5f MHz cm^2 / V^2 " % calc.getPolarizability()) 0.76705 MHz cm^2 / V^2 Stark map calculation >>> from arc import * >>> calc = StarkMap(Caesium()) >>> calc.defineBasis(28, 0, 0.5, 0.5, 23, 32, 20) >>> calc.diagonalise(np.linspace(00.,60000,600)) >>> calc.plotLevelDiagram() >>> calc.showPlot() << matplotlib plot will open containing a Stark map >> Examples: **Advanced interfacing of Stark map calculations (StarkMap class)** Here we show one easy way to obtain the Stark matrix (from diagonal :obj:`mat1` and off-diagonal part :obj:`mat2` ) and basis states (stored in :obj:`basisStates` ), if this middle-product of the calculation is needed for some code build on top of the existing ARC package. >>> from arc import * >>> calc = StarkMap(Caesium()) >>> calc.defineBasis(28, 0, 0.5, 0.5, 23, 32, 20) >>> # Now we have matrix and basis states, that we can used in our own code >>> # Let's say we want Stark map at electric field of 0.2 V/m >>> eField = 0.2 # V/m >>> # We can easily extract Stark matrix >>> # as diagonal matrix (state detunings) >>> # + off-diagonal matrix (propotional to electric field) >>> matrix = calc.mat1+calc.mat2*eField >>> # and the basis states as array [ [n,l,j,mj] , ...] >>> basisStates = calc.basisStates >>> # you can do your own calculation now... References: .. [1] M. L. Zimmerman et.al, PRA **20**:2251 (1979) https://doi.org/10.1103/PhysRevA.20.2251 .. _`Stark map example snippet`: ./Rydberg_atoms_a_primer.html#Rydberg-Atom-Stark-Shifts """def__init__(self,atom):self.atom=atomself.basisStates=[]""" List of basis states for calculation in the form [ [n,l,j,mj], ...]. Calculated by :obj:`defineBasis` . """self.mat1=[]""" diagonal elements of Stark-matrix (detuning of states) calculated by :obj:`defineBasis` in the basis :obj:`basisStates`. """self.mat2=[]""" 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 :obj:`defineBasis` in the basis :obj:`basisStates`. """self.indexOfCoupledState=[]""" Index of coupled state (initial state passed to :obj:`defineBasis`) in :obj:`basisStates` list of basis states """# finding energy levelsself.eFieldList=[]""" Saves electric field (in units of V/m) for which energy levels are calculated See also: :obj:`y`, :obj:`highlight`, :obj:`diagonalise` """self.y=[]# eigenValues""" `y[i]` is an array of eigenValues corresponding to the energies of the atom states at the electric field `eFieldList[i]`. For example `y[i][j]` is energy of the `j` eigenvalue (energy of the state) measured in cm :math:`{}^{-1}` relative to the ionization threshold. See also: :obj:`eFieldList`, :obj:`highlight`, :obj:`diagonalise` """self.highlight=[]#contribution of initial state there (overlap |<original state | given state>|^2)""" `highlight[i]` is an array of values measuring highlighted feature in the eigenstates at electric field intensity `eFieldList[i]`. E.g. `highlight[i][j]` measures highlighted feature of the state with energy `y[i][j]` at electric field `eFieldList[i]`. What will be highlighted feature is defined in the call of :obj:`diagonalise` (see that part of documentation for details). See also: :obj:`eFieldList`, :obj:`y`, :obj:`diagonalise` """# pointers towards figureself.fig=0self.ax=0# values used for fitting polarizability, and fitself.fitX=[]self.fitY=[]self.fittedCurveY=[]self.drivingFromState=[0,0,0,0,0]self.maxCoupling=0.# STARK memoizationself.eFieldCouplingSaved=Falsedef_eFieldCouplingDivE(self,n1,l1,j1,mj1,n2,l2,j2,mj2):# eFied coupling devided with E (witout actuall multiplication to getE)# delta(mj1,mj2') delta(l1,l2+-1)if((abs(mj1-mj2)>0.1)or(abs(l1-l2)!=1)):return0# matrix elementresult=self.atom.getRadialMatrixElement(n1,l1,j1,n2,l2,j2)*\
physical_constants["Bohr radius"][0]*C_esumPart=self.eFieldCouplingSaved.getAngular(l1,j1,mj1,l2,j2,mj2)returnresult*sumPartdef_eFieldCoupling(self,n1,l1,j1,mj1,n2,l2,j2,mj2,eField):returnself._eFieldCouplingDivE(n1,l1,j1,mj1,n2,l2,j2,mj2)*eField

[docs]defdefineBasis(self,n,l,j,mj,nMin,nMax,maxL,Bz=0,progressOutput=False,debugOutput=False):""" Initializes basis of states around state of interest Defines basis of states for further calculation. :math:`n,l,j,m_j` specify state whose neighbourhood and polarizability we want to explore. Other parameters specify basis of calculations. This method stores basis in :obj:`basisStates`, while corresponding interaction matrix is stored in two parts. First part is diagonal electric-field independent part stored in :obj:`mat1`, while the 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 `fullStarkMatrix` = :obj:`mat1` + :obj:`mat2` *`eField` Args: n (int): principal quantum number of the state l (int): angular orbital momentum of the state j (flaot): total angular momentum of the state mj (float): projection of total angular momentum of the state nMin (int): *minimal* principal quantum number of the states to be included in the basis for calculation nMax (int): *maximal* principal quantum number of the states to be included in the basis for calculation maxL (int): *maximal* value of orbital angular momentum for the states to be included in the basis for calculation Bz (float): optional, magnetic field directed along z-axis in units of Tesla. Calculation will be correct only for weak magnetic fields, where paramagnetic term is much stronger then diamagnetic term. Diamagnetic term is neglected. progressOutput (:obj:`bool`, optional): if True prints the progress of calculation; Set to false by default. debugOutput (:obj:`bool`, optional): if True prints additional information usefull for debuging. Set to false by default. """globalwignerPrecalwignerPrecal=Trueself.eFieldCouplingSaved=_EFieldCoupling()states=[]# save calculation details STARTself.n=n;self.l=l;self.j=jself.mj=mj;self.nMin=nMin;self.nMax=nMax;self.maxL=maxLself.Bz=Bz# save calculation details ENDfortninxrange(nMin,nMax):fortlinxrange(min(maxL+1,tn)):if(abs(mj)-0.1<=float(tl)+0.5):states.append([tn,tl,float(tl)+0.5,mj])if(tl>0)and(abs(mj)-0.1<=float(tl)-0.5):states.append([tn,tl,float(tl)-0.5,mj])dimension=len(states)ifprogressOutput:print("Found ",dimension," states.")ifdebugOutput:print(states)indexOfCoupledState=0index=0forsinstates:if(s[0]==n)and(abs(s[1]-l)<0.1)and(abs(s[2]-j)<0.1)and\
(abs(s[3]-mj)<0.1):indexOfCoupledState=indexindex+=1ifdebugOutput:print("Index of initial state")print(indexOfCoupledState)print("Initial state = ")print(states[indexOfCoupledState])self.mat1=np.zeros((dimension,dimension),dtype=np.double)self.mat2=np.zeros((dimension,dimension),dtype=np.double)self.basisStates=statesself.indexOfCoupledState=indexOfCoupledStateifprogressOutput:print("Generating matrix...")progress=0.foriiinxrange(dimension):ifprogressOutput:progress+=((dimension-ii)*2-1)sys.stdout.write("\r%d%%"%(float(progress)/float(dimension**2)*100))sys.stdout.flush()# add diagonal elementself.mat1[ii][ii]=self.atom.getEnergy(states[ii][0],\
states[ii][1],states[ii][2])\
*C_e/C_h*1e-9 \
+self.atom.getZeemanEnergyShift(states[ii][1],states[ii][2],states[ii][3],self.Bz)/C_h*1.0e-9# add off-diagonal elementforjjinxrange(ii+1,dimension):coupling=self._eFieldCouplingDivE(states[ii][0]\
,states[ii][1],\
states[ii][2],mj,\
states[jj][0],\
states[jj][1],\
states[jj][2],mj)*\
1.e-9/C_hself.mat2[jj][ii]=couplingself.mat2[ii][jj]=couplingifprogressOutput:print("\n")ifdebugOutput:print(self.mat1+self.mat2)print(self.mat2[0])self.atom.updateDipoleMatrixElementsFile()self.eFieldCouplingSaved._closeDatabase()self.eFieldCouplingSaved=Falsereturn0

[docs]defdiagonalise(self,eFieldList,drivingFromState=[0,0,0,0,0],progressOutput=False,debugOutput=False):""" Finds atom eigenstates in a given electric field Eigenstates are calculated for a list of given electric fields. To extract polarizability of the originaly stated state see :obj:`getPolarizability` method. Results are saved in :obj:`eFieldList`, :obj:`y` and :obj:`highlight`. Args: eFieldList (array): array of electric field strength (in V/m) for which we want to know energy eigenstates progressOutput (:obj:`bool`, optional): if True prints the progress of calculation; Set to false by default. debugOutput (:obj:`bool`, optional): if True prints additional information usefull for debuging. Set to false by default. """# if we are driving from some state# ========= FIND LASER COUPLINGS (START) =======coupling=[]dimension=len(self.basisStates)self.maxCoupling=0.self.drivingFromState=drivingFromStateif(self.drivingFromState[0]!=0):ifprogressOutput:print("Finding driving field coupling...")# get first what was the state we are calculating coupling withstate1=drivingFromStaten1=int(round(state1[0]))l1=int(round(state1[1]))j1=state1[2]m1=state1[3]q=state1[4]foriinxrange(dimension):thisCoupling=0.ifprogressOutput:sys.stdout.write("\r%d%%"%(i/float(dimension-1)*100.))sys.stdout.flush()if(int(abs(self.basisStates[i][1]-l1))==1)and\
(int(abs(self.basisStates[i][2]-j1))<=1)and\
(int(abs(self.basisStates[i][3]-m1-q))==0):state2=self.basisStates[i]n2=int(state2[0])l2=int(state2[1])j2=state2[2]m2=state2[3]ifdebugOutput:print(n1," ",l1," ",j1," ",m1," < - ",q," - >",n2," ",\
l2," ",j2," ",m2,"\n")dme=self.atom.getDipoleMatrixElement(n1,l1,j1,m1,\
n2,l2,j2,m2,\
q)thisCoupling+=dmethisCoupling=abs(thisCoupling)**2ifthisCoupling>self.maxCoupling:self.maxCoupling=thisCouplingif(thisCoupling>0.00000001)anddebugOutput:print("coupling = ",thisCoupling)coupling.append(thisCoupling)ifprogressOutput:print("\n")ifself.maxCoupling<0.00000001:raiseException("State that you specified in drivingFromState, for a "+\
"given laser polarization, is uncoupled from the specified Stark "+\
"manifold. If you just want to see the specified Stark manifold "+\
"remove driveFromState optional argument from call of function "+\
"diagonalise. Or specify state and driving that is coupled "+\
"to a given manifold to see coupling strengths.")# ========= FIND LASER COUPLINGS (END) =======indexOfCoupledState=self.indexOfCoupledStateself.eFieldList=eFieldListself.y=[]self.highlight=[]self.composition=[]ifprogressOutput:print("Finding eigenvectors...")progress=0.foreFieldineFieldList:ifprogressOutput:progress+=1.sys.stdout.write("\r%d%%"% \
(float(progress)/float(len(eFieldList))*100))sys.stdout.flush()m=self.mat1+self.mat2*eFieldev,egvector=eigh(m)self.y.append(ev)if(drivingFromState[0]<0.1):sh=[]comp=[]foriinxrange(len(ev)):sh.append(abs(egvector[indexOfCoupledState,i])**2)comp.append(self._stateComposition2(egvector[:,i]))self.highlight.append(sh)self.composition.append(comp)else:sh=[]comp=[]foriinxrange(len(ev)):sumCoupledStates=0.forjinxrange(dimension):sumCoupledStates+=abs(coupling[j]/self.maxCoupling)*\
abs(egvector[j,i]**2)comp.append(self._stateComposition2(egvector[:,i]))sh.append(sumCoupledStates)self.highlight.append(sh)self.composition.append(comp)ifprogressOutput:print("\n")return

[docs]defplotLevelDiagram(self,units=1,highlighState=True,progressOutput=False,\
debugOutput=False,highlightColour='red',addToExistingPlot=False):""" Makes a plot of a stark map of energy levels To save this plot, see :obj:`savePlot`. To print this plot see :obj:`showPlot`. Args: units (:obj:`int`,optional): possible values {1,2} ; if the value is 1 (default) Stark diagram will be plotted in energy units cm :math:`{}^{-1}`; if value is 2, Stark diagram will be plotted as energy :math:`/h` in units of GHz highlightState (:obj:`bool`, optional): False by default. If True, scatter plot colour map will map in red amount of original state for the given eigenState progressOutput (:obj:`bool`, optional): if True prints the progress of calculation; Set to False by default. debugOutput (:obj:`bool`, optional): if True prints additional information usefull for debuging. Set to False by default. addToExistingPlot (:obj:`bool`, optional): if True adds points to existing old plot. Note that then interactive plotting doesn't work. False by default. """rvb=LinearSegmentedColormap.from_list('mymap',\
['0.9',highlightColour,'black'])self.units=unitsself.addToExistingPlot=addToExistingPlotifprogressOutput:print("plotting...")originalState=self.basisStates[self.indexOfCoupledState]n=originalState[0]l=originalState[1]j=originalState[2]existingPlot=Falseif(self.fig==0ornotaddToExistingPlot):if(self.fig!=0):plt.close()self.fig,self.ax=plt.subplots(1,1,figsize=(11.,5))else:existingPlot=TrueeFieldList=[]y=[]yState=[]forbrinxrange(len(self.y)):foriinxrange(len(self.y[br])):eFieldList.append(self.eFieldList[br])y.append(self.y[br][i])yState.append(self.highlight[br][i])yState=np.array(yState)sortOrder=yState.argsort(kind='heapsort')eFieldList=np.array(eFieldList)y=np.array(y)eFieldList=eFieldList[sortOrder]y=y[sortOrder]yState=yState[sortOrder]if(units==1):## in cm^-1ifnothighlighState:self.ax.scatter(eFieldList/100.,y*0.03336,s=1,color="k",picker=5)else:cm=rvbcNorm=matplotlib.colors.Normalize(vmin=0.,vmax=1.)self.ax.scatter(eFieldList/100,y*0.03336,\
c=yState,s=5,norm=cNorm,cmap=cm,lw=0,picker=5)ifnotexistingPlot: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))else:cb.set_label(r"$( \Omega_\mu | \Omega )^2$")else:## in GHzifnothighlighState:self.ax.scatter(eFieldList/100.,y,\
s=1,color="k",picker=5)# in GHzelse:cm=rvbcNorm=matplotlib.colors.Normalize(vmin=0.,vmax=1.)self.ax.scatter(eFieldList/100.,y,c=yState,\
s=5,norm=cNorm,cmap=cm,lw=0,picker=5)ifnotexistingPlot: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))else:cb.set_label(r"$(\Omega_\mu / \Omega )^2$")self.ax.set_xlabel("Electric field (V/cm)")if(units==1):## in cm^{-1}uppery=self.atom.getEnergy(n,l,j)*C_e/C_h*1e-9*0.03336+10lowery=self.atom.getEnergy(n,l,j)*C_e/C_h*1e-9*0.03336-10self.ax.set_ylabel("State energy, $E/(h c)$ (cm$^{-1}$)")else:## in GHzuppery=self.atom.getEnergy(n,l,j)*C_e/C_h*1e-9+5lowery=self.atom.getEnergy(n,l,j)*C_e/C_h*1e-9-5self.ax.set_ylabel(r"State energy, $E/h$ (GHz)")self.ax.set_ylim(lowery,uppery)##self.ax.set_xlim(min(eFieldList)/100.,max(eFieldList)/100.)return0

[docs]defsavePlot(self,filename="StarkMap.pdf"):""" Saves plot made by :obj:`plotLevelDiagram` Args: filename (:obj:`str`, optional): file location where the plot should be saved """if(self.fig!=0):self.fig.savefig(filename,bbox_inches='tight')else:print("Error while saving a plot: nothing is plotted yet")return0

[docs]defshowPlot(self,interactive=True):""" Shows plot made by :obj:`plotLevelDiagram` """if(self.fig!=0):ifinteractive:ifself.addToExistingPlot:print("NOTE: Interactive plotting doesn't work with"" addToExistingPlot option set to True""\nPlease turn off this option in plotLevelDiagram.\n")else:self.ax.set_title("Click on state to see state composition")self.clickedPoint=0self.fig.canvas.draw()self.fig.canvas.mpl_connect('pick_event',self._onPick)plt.show()else:print("Error while showing a plot: nothing is plotted yet")return0

[docs]defgetPolarizability(self,maxField=1.e10,showPlot=False,\
debugOutput=False,minStateContribution=0.0):""" Returns the polarizability of the state (set during the initalization process) Args: maxField (:obj:`float`, optional): maximum field (in V/m) to be used for fitting the polarizability. By default, max field is very large, so it will use eigenvalues calculated in the whole range. showPlot (:obj:`bool`, optional): shows plot of calculated eigenValues of the given state (dots), and the fit (solid line) for extracting polarizability debugOutput (:obj:`bool`, optional): if True prints additional information usefull for debuging. Set to false by default. Returns: float: scalar polarizability in units of MHz cm :math:`^2` / V \ :math:`^2` """if(self.drivingFromState[0]!=0):raiseException("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.")eFieldList=self.eFieldListyState=self.highlighty=self.yoriginalState=self.basisStates[self.indexOfCoupledState]n=originalState[0]l=originalState[1]j=originalState[2]energyOfOriginalState=self.atom.getEnergy(n,l,j)*C_e/C_h*1e-9# in GHzifdebugOutput:print("finding original state for each electric field value")stopFitIndex=0whilestopFitIndex<len(eFieldList)-1and \
eFieldList[stopFitIndex]<maxField:stopFitIndex+=1xOriginalState=[]yOriginalState=[]foriiinxrange(stopFitIndex):maxPortion=0.yval=0.jj=0forjjinxrange(len(y[ii])):ifyState[ii][jj]>maxPortion:maxPortion=yState[ii][jj]yval=y[ii][jj]# measure state energy relative to the original stateif(minStateContribution<maxPortion):xOriginalState.append(eFieldList[ii])yOriginalState.append(yval-energyOfOriginalState)xOriginalState=np.array(xOriginalState)/100.# converts to V/cmyOriginalState=np.array(yOriginalState)# in GHz## in GHzuppery=5.0lowery=-5.0ifdebugOutput:print("found ",len(xOriginalState))ifshowPlot:self.fig,self.ax=plt.subplots(1,1,figsize=(6.5,3))self.ax.scatter(xOriginalState,yOriginalState,s=2,color="k")self.ax.set_xlabel("E field (V/cm)")self.ax.set_ylim(lowery,uppery)self.ax.set_ylabel(r"Energy/$h$ (GHz)")self.ax.set_xlim(xOriginalState[0],\
xOriginalState[-1])defpolarizabilityFit(eField,offset,alpha):returnoffset-0.5*alpha*eField**2try:popt,pcov=curve_fit(polarizabilityFit,\
xOriginalState,\
yOriginalState,\
[0,0])except:print("\nERROR: fitting energy levels for extracting polarizability\ of the state failed. Please check the range of electric \ fields where you are trying to fit polarizability and ensure\ that there is only one state with continuous energy change\ that has dominant contribution of the initial state.\n\n")return0ifdebugOutput:print("Scalar polarizability = ",popt[1]*1.e3," MHz cm^2 / V^2 ")y_fit=[]forvalinxOriginalState:y_fit.append(polarizabilityFit(val,popt[0],popt[1]))y_fit=np.array(y_fit)ifshowPlot:self.ax.plot(xOriginalState,y_fit,"r--")self.ax.legend(("fitted model function","calculated energy level"),\
loc=1,fontsize=10)self.ax.set_ylim(min(yOriginalState),max(yOriginalState))plt.show()self.fitX=xOriginalStateself.fitY=yOriginalStateself.fittedCurveY=y_fitreturnpopt[1]*1.e3# returned value is in MHz cm^2 / V^2

[docs]defmakeLevels(self,nFrom,nTo,lFrom,lTo):""" Constructs energy level diagram in a given range Args: nFrom (int): minimal principal quantum number of the states we are interested in nTo (int): maximal principal quantum number of the states we are interested in lFrom (int): minimal orbital angular momentum of the states we are interested in lTo (int): maximal orbital angular momentum of the states we are interested in """#save local copy of the space restrictionsself.nFrom=nFromself.nTo=nToself.lFrom=lFromself.lTo=lTo# find all the levels within this space restrictionsnFrom=max(nFrom,self.atom.groundStateN)whilenFrom<=nTo:l=lFromwhilel<=min(lTo,4,nFrom-1):if(l>0.5):self.listX.append(l)self.listY.append(self.atom.getEnergy(nFrom,l,l-0.5))self.levelLabel.append([nFrom,l,l-0.5])self.listX.append(l)self.listY.append(self.atom.getEnergy(nFrom,l,l+0.5))self.levelLabel.append([nFrom,l,l+0.5])l=l+1nFrom+=1# if user requested principal quantum nuber below the# ground state principal quantum number# add those L states that are higher in energy then the ground stateforstateinself.atom.extraLevels:ifstate[1]<=lToandstate[0]>=self.nFrom:self.listX.append(state[1])self.listY.append(self.atom.getEnergy(state[0],state[1],state[2]))self.levelLabel.append(state)