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,

"""

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.linalg import eigh
from numpy.ma 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
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 _. 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:
..  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 = 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

: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.

: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).

:obj:eFieldList, :obj:y, :obj:diagonalise
"""

# pointers towards figure
self.fig = 0
self.ax = 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

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==n) and (abs(s-l)<0.1) and (abs(s-j)<0.1) and\
(abs(s-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()

self.mat1[ii][ii] = self.atom.getEnergy(states[ii],\
states[ii],states[ii])\
* C_e/C_h*1e-9 \
+ self.atom.getZeemanEnergyShift(
states[ii],
states[ii],
states[ii],
self.Bz) / C_h * 1.0e-9

for jj in xrange(ii+1,dimension):
coupling = self._eFieldCouplingDivE(states[ii]\
,states[ii],\
states[ii],mj,\
states[jj],\
states[jj],\
states[jj],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)

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):
if progressOutput: print("Finding driving field coupling...")
# get first what was the state we are calculating coupling with
state1 = drivingFromState
n1 = int(round(state1))
l1 = int(round(state1))
j1 = state1
m1 = state1
q = state1

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]-l1))==1)and\
(int(abs(self.basisStates[i]-j1))<=1) and\
(int(abs(self.basisStates[i]-m1-q))==0):
state2 = self.basisStates[i]
n2 = int(state2)
l2 = int(state2)
j2 = state2
m2 = state2
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.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
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 = datetime.datetime.now().strftime(fmt)

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.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)),
self.drivingFromState))

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', \
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', \
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', \
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',
"""
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

if progressOutput:
print("plotting...")

originalState = self.basisStates[self.indexOfCoupledState]
n = originalState
l = originalState
j = originalState

existingPlot = False
if (self.fig == 0 or not addToExistingPlot):
if (self.fig != 0): plt.close()
self.fig, self.ax = 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:
self.ax.scatter(eFieldList/100.,y*0.03336,s=1,color="k",picker=5)
else:
cm = rvb
cNorm  = 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)
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.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:
self.ax.scatter(eFieldList/100.,y,\
s=1,color="k",picker=5) # in GHz
else:
cm = rvb
cNorm  = 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)
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.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+10
lowery = self.atom.getEnergy(n,l,j)*C_e/C_h*1e-9*0.03336-10
self.ax.set_ylabel("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
self.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.)
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:
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 = 0
self.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")
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.ax.plot([self.eFieldList[i]/100.],\
[self.y[i][j]*scaleFactor],"bs",\
linewidth=0,zorder=3)

self.ax.set_title(("[%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): value+= "+" value = value+ ("%.2f" % stateVector[i])+\ self._addState(*self.basisStates[stateVector[i]]) totalContribution += abs(stateVector[i])**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

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):
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
l = originalState
j = originalState
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, 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,\
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.e3," MHz cm^2 / V^2 ")

y_fit = []
for val in xOriginalState:
y_fit.append(polarizabilityFit(val,popt,popt))
y_fit = np.array(y_fit)

if showPlot:
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 = xOriginalState
self.fitY = yOriginalState
self.fittedCurveY = y_fit

return popt*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
self.ax = 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<=lTo and state>=self.nFrom:
self.listX.append(state)
self.listY.append(self.atom.getEnergy(state,state,state))
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-state2)==1)and\
(abs(state1-state2)<=1.01)
if (dipoleAllowed):
# decay to this state
rate = self.atom.getTransitionRate(state2,state2,state2,\
state1,state1,state1,\
temperature=environmentTemperature)

transitionVector.append(rate)

# decay from this state
rate = self.atom.getTransitionRate(state1,state1,state1,\
state2,state2,state2,\
temperature=environmentTemperature)

decay = decay-rate
else:
transitionVector.append(0.0)

transitionVector[i] = decay
if printDecays:
print("Decay time of ")
printState(state1, state1, state1)
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, self.ax = 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],\
self.levelLabel[i],self.levelLabel[i],
self.levelLabel[j],\
self.levelLabel[j],self.levelLabel[j])

intensity = self.atom.getTransitionRate(self.levelLabel[i],\
self.levelLabel[i],self.levelLabel[i],\
self.levelLabel[j],\
self.levelLabel[j],self.levelLabel[j])

lineWavelength.append(abs(wavelength)*1.e9)
lineStrength.append(abs(intensity))
lineName.append(printStateString(self.levelLabel[i],\
self.levelLabel[i],\
self.levelLabel[i])+\
" -> "+
printStateString(self.levelLabel[j],\
self.levelLabel[j],\
self.levelLabel[j]))

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
self.ax.plot(wavelengths,spectra,"g-")

def showSpectra(self,saveInFile="",showTransitionPoints=True):
if showTransitionPoints:
self.ax.plot(self.spectraX,self.spectraY,"ro",picker=5)
self.ax.set_xlabel("Wavelength (nm)")
self.ax.set_ylabel("Intensity (arb.un)")
#self.ax.set_xlim(300,600)
self.fig.canvas.mpl_connect('pick_event', self.onpick3)
if (saveInFile != ""):
self.fig.savefig(saveInFile)
plt.show()

[docs]    def drawLevels(self):
"""
Draws a level diagram plot
"""
self.fig, self.ax = plt.subplots(1, 1,figsize=(9.0, 11.5))

i = 0
while i<len(self.listX):
self.ax.plot([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.ax.plot([self.listX[i]],[self.listY[i]],"ro",alpha=self.populations[i])

i = i+1

[docs]    def showPlot(self):
"""
Shows a level diagram plot
"""
self.ax.set_ylabel("Energy (eV)")
self.ax.set_xlim(-0.5+self.lFrom,self.lTo+0.5)

# X AXIS
majorLocator   = MultipleLocator(1)

self.ax.xaxis.set_major_locator(majorLocator)
tickNames = [" "]
for l in xrange(self.lFrom,self.lTo+1):
tickNames.append(printStateLetter(l))
tickNum = len(self.ax.get_xticklabels())

self.fig.canvas.draw()
self.ax.set_xticklabels(tickNames)
self.fig.canvas.mpl_connect('pick_event', self.onpick2)
plt.show()

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] == state)and\
(self.levelLabel[i] == state)and\
(abs(self.levelLabel[i] - state)<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+xdata)/2., ydata)
if (self.state1==0 or (state== self.state2)):
self.state1 = state
self.ax.set_title(printStateString(state,state,state)+" -> ")
self.state2=[-1,-1,-1]
else:
title = ""
if (state != self.state1) and (state!= self.state2):
title = printStateString(self.state1,\
self.state1,\
self.state1)+\
" -> "+\
printStateString(state,state,state)+" "
title = title+(" %.2f nm (%.3f GHz)" % \
(self.atom.getTransitionWavelength(self.state1,\
self.state1,\
self.state1,\
state,state,\
state)*1e9,\
self.atom.getTransitionFrequency(self.state1,\
self.state1,\
self.state1,\
state,\
state,\
state)*1e-9))
self.ax.set_title(title)
self.state1=[0,0,0]

self.state2 = state
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)

line = self.findLine(xdata[ind], ydata[ind])
self.ax.set_title(line)
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))