Source code for arc.calculations_atom_single

# -*- coding: utf-8 -*-

    This module provides calculations of single-atom properties.

    Included calculations are Stark maps, level plot visualisations,
    lifetimes and radiative decays.


from __future__ import print_function

from math import exp,log,sqrt
import matplotlib.pyplot as plt
from matplotlib.ticker import MultipleLocator, FormatStrFormatter
import numpy as np
import re
from .wigner import Wigner6j,Wigner3j,CG
from scipy.constants import physical_constants, pi , epsilon_0, hbar
from scipy.constants import k as C_k
from scipy.constants import c as C_c
from scipy.constants import h as C_h
from scipy.constants import e as C_e
from scipy.optimize import curve_fit

# for matrices
from numpy import zeros,savetxt, complex64,complex128
from numpy.linalg import eigvalsh,eig,eigh
from import conjugate
from numpy.lib.polynomial import real

from scipy.sparse import lil_matrix,csr_matrix
from scipy.sparse.linalg import eigsh
from scipy.special.specfun import fcoef

import sys
if sys.version_info > (2,):
    xrange = range
from .alkali_atom_functions import printStateString, _EFieldCoupling, printStateLetter,printStateStringLatex

from matplotlib.colors import LinearSegmentedColormap
import matplotlib

import sqlite3
sqlite3.register_adapter(np.float64, float)
sqlite3.register_adapter(np.float32, float)
sqlite3.register_adapter(np.int64, int)
sqlite3.register_adapter(np.int32, int)
import datetime

[docs]class StarkMap: """ 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, PRA **20**:2251 (1979) .. _`Stark map example snippet`: ./Rydberg_atoms_a_primer.html#Rydberg-Atom-Stark-Shifts """ def __init__(self,atom): self.atom = atom self.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 levels self.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 figure self.fig = 0 = 0 # values used for fitting polarizability, and fit self.fitX = [] self.fitY = [] self.fittedCurveY = [] self.drivingFromState = [0,0,0,0,0] self.maxCoupling = 0. # STARK memoization self.eFieldCouplingSaved = False def _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) ): return 0 # matrix element result = self.atom.getRadialMatrixElement(n1,l1,j1,n2,l2,j2)*\ physical_constants["Bohr radius"][0]*C_e sumPart = self.eFieldCouplingSaved.getAngular(l1,j1,mj1,l2,j2,mj2) return result*sumPart def _eFieldCoupling(self,n1,l1,j1,mj1,n2,l2,j2,mj2,eField): return self._eFieldCouplingDivE(n1,l1,j1,mj1,n2,l2,j2,mj2)*eField
[docs] def defineBasis(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. """ global wignerPrecal wignerPrecal = True self.eFieldCouplingSaved = _EFieldCoupling() states = [] # save calculation details START self.n = n; self.l =l; self.j=j self.mj = mj; self.nMin = nMin; self.nMax = nMax; self.maxL = maxL self.Bz = Bz # save calculation details END for tn in xrange(nMin,nMax): for tl in xrange(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) if progressOutput: print("Found ",dimension," states.") if debugOutput: print(states) indexOfCoupledState = 0 index = 0 for s in states: 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 = index index +=1 if debugOutput: 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 = states self.indexOfCoupledState = indexOfCoupledState if progressOutput: print("Generating matrix...") progress = 0. for ii in xrange(dimension): if progressOutput: progress += ((dimension-ii)*2-1) sys.stdout.write("\r%d%%" % (float(progress)/float(dimension**2)*100)) sys.stdout.flush() # add diagonal element self.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 element for jj in xrange(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_h self.mat2[jj][ii] = coupling self.mat2[ii][jj] = coupling if progressOutput: print("\n") if debugOutput: print(self.mat1+self.mat2) print(self.mat2[0]) self.atom.updateDipoleMatrixElementsFile() self.eFieldCouplingSaved._closeDatabase() self.eFieldCouplingSaved = False return 0
[docs] def diagonalise(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 = drivingFromState if (self.drivingFromState[0] != 0): if progressOutput: print("Finding driving field coupling...") # get first what was the state we are calculating coupling with state1 = drivingFromState n1 = int(round(state1[0])) l1 = int(round(state1[1])) j1 = state1[2] m1 = state1[3] q = state1[4] for i in xrange(dimension): thisCoupling = 0. if progressOutput: 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] if debugOutput: 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 += dme thisCoupling = abs(thisCoupling)**2 if thisCoupling > self.maxCoupling: self.maxCoupling = thisCoupling if (thisCoupling >0.00000001) and debugOutput: print("coupling = ",thisCoupling) coupling.append(thisCoupling) if progressOutput: print("\n") if self.maxCoupling<0.00000001: raise Exception("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.indexOfCoupledState self.eFieldList = eFieldList self.y = [] self.highlight = [] self.composition = [] if progressOutput: print("Finding eigenvectors...") progress = 0. for eField in eFieldList: if progressOutput: progress += 1. sys.stdout.write("\r%d%%" % \ (float(progress)/float(len(eFieldList))*100)) sys.stdout.flush() m = self.mat1+self.mat2*eField ev,egvector = eigh(m) self.y.append(ev) if (drivingFromState[0]<0.1): sh = [] comp = [] for i in xrange(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 = [] for i in xrange(len(ev)): sumCoupledStates = 0. for j in xrange(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) if progressOutput: print("\n") return
[docs] def exportData(self,fileBase,exportFormat = "csv"): """ Exports StarkMap calculation data. Only supported format (selected by default) is .csv in a human-readable form with a header that saves details of calculation. Function saves three files: 1) `filebase` _eField.csv; 2) `filebase` _energyLevels 3) `filebase` _highlight For more details on the format, see header of the saved files. Args: filebase (string): filebase for the names of the saved files without format extension. Add as a prefix a directory path if necessary (e.g. saving outside the current working directory) exportFormat (string): optional. Format of the exported file. Currently only .csv is supported but this can be extended in the future. """ fmt='on %Y-%m-%d @ %H:%M:%S' ts = commonHeader = "Export from Alkali Rydberg Calculator (ARC) %s.\n" % ts commonHeader += ("\n *** Stark Map for %s %s m_j = %d/2. ***\n\n" % (self.atom.elementName, printStateString(self.n, self.l, self.j), int(round(2.*self.mj)) ) ) commonHeader += (" - Included states - principal quantum number (n) range [%d-%d].\n" %\ (self.nMin, self.nMax)) commonHeader += (" - Included states with orbital momentum (l) in range [%d,%d] (i.e. %s-%s).\n"%\ (0, self.maxL, printStateLetter(0), printStateLetter(self.maxL))) 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." else: commonHeader += (" - State highlighting based on the relative driving strength \n"+\ " to a given energy eigenstate (energy level) from state\n"+\ " %s m_j =%d/2 with polarization q=%d.\n"%\ ( printStateString(*self.drivingFromState[0:3]),\ int(round(2.*self.drivingFromState[3])), self.drivingFromState[4])) if exportFormat=="csv": print("Exporting StarkMap calculation results as .csv ...") commonHeader += " - Export consists of three (3) files:\n" commonHeader += (" 1) %s,\n" % (fileBase+"_eField."+exportFormat)) commonHeader += (" 2) %s,\n" % (fileBase+"_energyLevels."+exportFormat)) commonHeader += (" 3) %s.\n\n" % (fileBase+"_highlight."+exportFormat)) filename = fileBase+"_eField."+exportFormat np.savetxt(filename, \ self.eFieldList, fmt='%.18e', delimiter=', ',\ newline='\n', \ header=(commonHeader + " - - - eField (V/m) - - -"),\ comments='# ') print(" Electric field values (V/m) saved in %s" % filename) filename = fileBase+"_energyLevels."+exportFormat headerDetails = " NOTE : Each row corresponds to eigenstates for a single specified electric field" np.savetxt(filename, \ self.y, fmt='%.18e', delimiter=', ',\ newline='\n', \ header=(commonHeader + ' - - - Energy (GHz) - - -\n' + headerDetails),\ comments='# ') print(" Lists of energies (in GHz relative to ionisation) saved in %s" % filename) filename = fileBase+"_highlight."+exportFormat np.savetxt(filename, \ self.highlight, fmt='%.18e', delimiter=', ',\ newline='\n', \ header=(commonHeader + ' - - - Highlight value (rel.units) - - -\n'+ headerDetails),\ comments='# ') print(" Highlight values saved in %s" % filename) print("... data export finished!") else: raise ValueError("Unsupported export format (.%s)." % format)
[docs] def plotLevelDiagram(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 = units self.addToExistingPlot = addToExistingPlot if progressOutput: print("plotting...") originalState = self.basisStates[self.indexOfCoupledState] n = originalState[0] l = originalState[1] j = originalState[2] existingPlot = False if (self.fig == 0 or not addToExistingPlot): if (self.fig != 0): plt.close() self.fig, = plt.subplots(1,1,figsize=(11.,5)) else: existingPlot = True eFieldList = [] y =[] yState = [] for br in xrange(len(self.y)): for i in xrange(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^-1 if not highlighState:,y*0.03336,s=1,color="k",picker=5) else: cm = rvb cNorm = matplotlib.colors.Normalize(vmin=0., vmax=1.),y*0.03336,\ c=yState,s=5,norm=cNorm, cmap=cm,lw=0,picker=5) if not existingPlot: 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 GHz if not highlighState:,y,\ s=1,color="k",picker=5) # in GHz else: cm = rvb cNorm = matplotlib.colors.Normalize(vmin=0., vmax=1.),y,c=yState,\ s=5,norm=cNorm, cmap=cm,lw=0,picker=5) if not existingPlot: 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$")"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+10 lowery = self.atom.getEnergy(n,l,j)*C_e/C_h*1e-9*0.03336-10"State energy, $E/(h c)$ (cm$^{-1}$)") else: ## in GHz uppery = self.atom.getEnergy(n,l,j)*C_e/C_h*1e-9+5 lowery = self.atom.getEnergy(n,l,j)*C_e/C_h*1e-9-5"State energy, $E/h$ (GHz)"),uppery) ##,max(eFieldList)/100.) return 0
[docs] def savePlot(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") return 0
[docs] def showPlot(self, interactive = True): """ Shows plot made by :obj:`plotLevelDiagram` """ if (self.fig != 0): if interactive: if self.addToExistingPlot: print("NOTE: Interactive plotting doesn't work with" " addToExistingPlot option set to True" "\nPlease turn off this option in plotLevelDiagram.\n") else:"Click on state to see state composition") self.clickedPoint = 0 self.fig.canvas.draw() self.fig.canvas.mpl_connect('pick_event', self._onPick) else: print("Error while showing a plot: nothing is plotted yet") return 0
def _onPick(self,event): if isinstance(event.artist, matplotlib.collections.PathCollection): if (self.units==1): scaleFactor = 0.03336 else: scaleFactor = 1.0 x = event.mouseevent.xdata*100. y = event.mouseevent.ydata/scaleFactor i = np.searchsorted(self.eFieldList,x) if i == len(self.eFieldList): i -= 1 if ((i>0) and (abs(self.eFieldList[i-1]-x)<abs(self.eFieldList[i]-x))): i -=1 j = 0 for jj in xrange(len(self.y[i])): if (abs(self.y[i][jj]-y) < abs(self.y[i][j]-y)): j = jj # now choose the most higlighted state in this area distance = abs(self.y[i][j]-y)*1.5 for jj in xrange(len(self.y[i])): if (abs(self.y[i][jj]-y) < distance and \ (abs(self.highlight[i][jj])>abs(self.highlight[i][j]))): j = jj if (self.clickedPoint!=0): self.clickedPoint.remove() self.clickedPoint, =[self.eFieldList[i]/100.],\ [self.y[i][j]*scaleFactor],"bs",\ linewidth=0,zorder=3)"[%s] = " % self.atom.elementName)+\ self._stateComposition(self.composition[i][j])+\ (" Colourbar value = %.2f"% self.highlight[i][j]), fontsize=11) event.canvas.draw() def _stateComposition(self,stateVector): i = 0 totalContribution = 0 value = "$" while (i<len(stateVector)) and (totalContribution<0.95): if (i!=0 and stateVector[i][0]>0): value+= "+" value = value+ ("%.2f" % stateVector[i][0])+\ self._addState(*self.basisStates[stateVector[i][1]]) totalContribution += abs(stateVector[i][0])**2 i += 1 if totalContribution<0.999: value+="+\\ldots" return value+"$" def _stateComposition2(self,stateVector,upTo=4): contribution = np.absolute(stateVector) order = np.argsort(contribution,kind='heapsort') index = -1 totalContribution = 0 mainStates = [] #[state Value, state index] while (index>-upTo) and (totalContribution<0.95): i = order[index] mainStates.append([stateVector[i],i]) totalContribution += contribution[i]**2 index -= 1 return mainStates def _addState(self,n1,l1,j1,mj1): return "|%s m_j=%d/2\\rangle" %\ (printStateStringLatex(n1, l1, j1),int(2*mj1))
[docs] def getPolarizability(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): 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.") eFieldList = self.eFieldList yState = self.highlight y = self.y originalState = 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 GHz if debugOutput: print("finding original state for each electric field value") stopFitIndex = 0 while stopFitIndex<len(eFieldList)-1 and \ eFieldList[stopFitIndex]<maxField: stopFitIndex += 1 xOriginalState = [] yOriginalState = [] for ii in xrange(stopFitIndex): maxPortion = 0. yval = 0. jj=0 for jj in xrange(len(y[ii])): if yState[ii][jj]>maxPortion: maxPortion = yState[ii][jj] yval = y[ii][jj] # measure state energy relative to the original state if (minStateContribution<maxPortion): xOriginalState.append(eFieldList[ii]) yOriginalState.append(yval-energyOfOriginalState) xOriginalState = np.array(xOriginalState)/100. # converts to V/cm yOriginalState = np.array(yOriginalState) # in GHz ## in GHz uppery = 5.0 lowery = -5.0 if debugOutput: print("found ",len(xOriginalState)) if showPlot: self.fig, = plt.subplots(1, 1,figsize=(6.5, 3)),yOriginalState,s=2,color="k")"E field (V/cm)"),uppery)"Energy/$h$ (GHz)")[0],\ xOriginalState[-1]) def polarizabilityFit(eField,offset,alpha): return offset-0.5*alpha*eField**2 try: 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") return 0 if debugOutput: print("Scalar polarizability = ",popt[1]*1.e3," MHz cm^2 / V^2 ") y_fit = [] for val in xOriginalState: y_fit.append(polarizabilityFit(val,popt[0],popt[1])) y_fit = np.array(y_fit) if showPlot:,y_fit,"r--")"fitted model function","calculated energy level"),\ loc=1,fontsize=10),max(yOriginalState)) self.fitX = xOriginalState self.fitY = yOriginalState self.fittedCurveY = y_fit return popt[1]*1.e3 # returned value is in MHz cm^2 / V^2
# ================= Level plots, decays, cascades etc =======================
[docs]class LevelPlot: """ Single atom level plots and decays For an example see `Rydberg energy levels example snippet`_. .. _`Rydberg energy levels example snippet`: ./Rydberg_atoms_a_primer.html#Rydberg-Atom-Energy-Levels 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` } Alkali atom type whose levels we want to examine """ def __init__(self,atomType ): self.atom = atomType self.nFrom = 0 self.nTo = 0 self.lFrom = 0 self.lTo = 0 self.listX = [] # list of l self.listY = [] # list of energies self.levelLabel = [] self.fig = 0 = 0 self.width =0.2 self.state1=[0,0,0] self.state2 =[0,-1,0] self.transitionMatrix = [] self.populations = [] self.transitionMatrixWavelength3 = [] # characterization of the graph self.spectraX = [] self.spectraY = [] self.spectraLine = []
[docs] def makeLevels(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 restrictions self.nFrom = nFrom self.nTo = nTo self.lFrom = lFrom self.lTo = lTo # find all the levels within this space restrictions nFrom = max(nFrom,self.atom.groundStateN) while nFrom<=nTo: l = lFrom while l<=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+1 nFrom += 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 state for state in self.atom.extraLevels: if state[1]<=lTo and state[0]>=self.nFrom: self.listX.append(state[1]) self.listY.append(self.atom.getEnergy(state[0],state[1],state[2])) self.levelLabel.append(state)
def makeTransitionMatrix(self,environmentTemperature = 0.0,printDecays=True): self.transitionMatrix =[] for i in xrange(len(self.levelLabel)): state1 = self.levelLabel[i] transitionVector = [] # decay of the stay decay = 0.0 for state2 in self.levelLabel: dipoleAllowed = (abs(state1[1]-state2[1])==1)and\ (abs(state1[2]-state2[2])<=1.01) if (dipoleAllowed): # decay to this state rate = self.atom.getTransitionRate(state2[0],state2[1],state2[2],\ state1[0],state1[1],state1[2],\ temperature=environmentTemperature) transitionVector.append(rate) # decay from this state rate = self.atom.getTransitionRate(state1[0],state1[1],state1[2],\ state2[0],state2[1],state2[2],\ temperature=environmentTemperature) decay = decay-rate else: transitionVector.append(0.0) transitionVector[i] = decay if printDecays: print("Decay time of ") printState(state1[0], state1[1], state1[2]) if decay < -1e-20: print("\t is\t",-1.e9/decay," ns") self.transitionMatrix.append(transitionVector) np.array(self.transitionMatrix) self.transitionMatrix = np.transpose(self.transitionMatrix) def drawSpectra(self): self.fig, = plt.subplots(1, 1,figsize=(16, 5)) lineWavelength = [] lineStrength = [] lineName = [] i = 0 while i<len(self.levelLabel): j = 0 while j<len(self.levelLabel): if (i!=j): wavelength = self.atom.getTransitionWavelength(\ self.levelLabel[i][0],\ self.levelLabel[i][1],self.levelLabel[i][2], self.levelLabel[j][0],\ self.levelLabel[j][1],self.levelLabel[j][2]) intensity = self.atom.getTransitionRate(self.levelLabel[i][0],\ self.levelLabel[i][1],self.levelLabel[i][2],\ self.levelLabel[j][0],\ self.levelLabel[j][1],self.levelLabel[j][2]) lineWavelength.append(abs(wavelength)*1.e9) lineStrength.append(abs(intensity)) lineName.append(printStateString(self.levelLabel[i][0],\ self.levelLabel[i][1],\ self.levelLabel[i][2])+\ " -> "+ printStateString(self.levelLabel[j][0],\ self.levelLabel[j][1],\ self.levelLabel[j][2])) j = j+1 i = i+1 self.spectraX = np.copy(lineWavelength) self.spectraY = np.copy(lineStrength) self.spectraLine = np.copy(lineName) def drawSpectraConvoluted(self,lowerWavelength, higherWavelength,points,gamma): wavelengths = linspace(lowerWavelength,higherWavelength,points) spectra = np.zeros(points) i = 0 while i<len(wavelengths): value = 0 j = 0 while j<len(self.spectraX): value = value + self.spectraY[j]*gamma/\ ((self.spectraX[j]-wavelengths[i])**2+gamma**2) j = j+1 spectra[i] = value i = i+1,spectra,"g-") def showSpectra(self,saveInFile="",showTransitionPoints=True): if showTransitionPoints:,self.spectraY,"ro",picker=5)"Wavelength (nm)")"Intensity (arb.un)") self.fig.subplots_adjust(right=0.95,left=0.1),600) self.fig.canvas.mpl_connect('pick_event', self.onpick3) if (saveInFile != ""): self.fig.savefig(saveInFile)
[docs] def drawLevels(self): """ Draws a level diagram plot """ self.fig, = plt.subplots(1, 1,figsize=(9.0, 11.5)) i = 0 while i<len(self.listX):[self.listX[i]-self.width,self.listX[i]+self.width], \ [self.listY[i],self.listY[i]],"b-",picker=4) if (i<len(self.populations) and (self.populations[i]>1e-3)):[self.listX[i]],[self.listY[i]],"ro",alpha=self.populations[i]) i = i+1
[docs] def showPlot(self): """ Shows a level diagram plot """"Energy (eV)"),self.lTo+0.5) # X AXIS majorLocator = MultipleLocator(1) tickNames = [" "] for l in xrange(self.lFrom,self.lTo+1): tickNames.append(printStateLetter(l)) tickNum = len( self.fig.canvas.draw() self.fig.canvas.mpl_connect('pick_event', self.onpick2)
def findState(self,x,y): distance = 100000000.0 state=[0,0,0] i = 0 while i<len(self.listX): dx = self.listX[i]-x dy = self.listY[i]-y dist = sqrt(dx*dx+dy*dy) if (dist<distance): distance = dist state = self.levelLabel[i] i = i+1 return state def findStateNo(self,state): # returns no of the given state in the basis i = 0 while i<len(self.levelLabel): if (self.levelLabel[i][0] == state[0])and\ (self.levelLabel[i][1] == state[1])and\ (abs(self.levelLabel[i][2] - state[2])<0.01): return i i = i+1 print("Error: requested state ") print(state) print("could not be found!") return -1 def findLine(self,x,y): distance = 1.e19 line="" i = 0 while i<len(self.spectraLine): dx = self.spectraX[i]-x dy = self.spectraY[i]-y dist = sqrt(dx*dx+dy*dy) if (dist<distance): distance = dist line = self.spectraLine[i] i = i+1 return line def onpick2(self,event): if isinstance(event.artist, matplotlib.lines.Line2D): thisline = event.artist xdata = thisline.get_xdata() ydata = thisline.get_ydata() state = self.findState((xdata[0]+xdata[0])/2., ydata[0]) if (self.state1[0]==0 or (state[1]== self.state2[1])): self.state1 = state[0],state[1],state[2])+" -> ") self.state2=[-1,-1,-1] else: title = "" if (state[1] != self.state1[1]) and (state[1]!= self.state2[1]): title = printStateString(self.state1[0],\ self.state1[1],\ self.state1[2])+\ " -> "+\ printStateString(state[0],state[1],state[2])+" " title = title+(" %.2f nm (%.3f GHz)" % \ (self.atom.getTransitionWavelength(self.state1[0],\ self.state1[1],\ self.state1[2],\ state[0],state[1],\ state[2])*1e9,\ self.atom.getTransitionFrequency(self.state1[0],\ self.state1[1],\ self.state1[2],\ state[0],\ state[1],\ state[2])*1e-9)) self.state1=[0,0,0] self.state2[1] = state[1] event.canvas.draw() def onpick3(self,event): if isinstance(event.artist, Line2D): thisline = event.artist xdata = thisline.get_xdata() ydata = thisline.get_ydata() ind = event.ind print(ind[0]) line = self.findLine(xdata[ind][0], ydata[ind][0]) event.canvas.draw()
[docs]def printState(n,l,j): """ Prints state spectroscopic label for numeric :math:`n`, :math:`l`, :math:`s` label of the state Args: n (int): principal quantum number l (int): orbital angular momentum j (float): total angular momentum """ print(n," ",printStateLetter(l),(" %.0d/2" % (j*2)))
[docs]def printStateString(n,l,j): """ Returns state spectroscopic label for numeric :math:`n`, :math:`l`, :math:`s` label of the state Args: n (int): principal quantum number l (int): orbital angular momentum j (float): total angular momentum Returns: string: label for the state in standard spectroscopic notation """ return str(n)+" "+printStateLetter(l)+(" %.0d/2" % (j*2))