def plotStarkMap(calc,units=1,xlim=[],ylim=[],filename=""):
originalState = calc.basisStates[calc.indexOfCoupledState]
n = originalState[0]
l = originalState[1]
j = originalState[2]
ax = webPlot()
x = []
y = []
yState = []
ax.xlabel = "E field (V/cm)"
coeff = 1.0
ax.ylabel = "Energy/h (GHz)"
if (units==1 ):
## in cm^{-1}
coeff = 0.03336 # conversion factor from GHz to cm^{-1}
ax.ylabel = "Energy/(h c) (cm^{-1})"
if (ylim == []):
ylim = [calc.atom.getEnergy(n,l,j)*elemCharge/h*1e-9*coeff-10,\
calc.atom.getEnergy(n,l,j)*elemCharge/h*1e-9*coeff+10]
for br in xrange(len(calc.y)):
for i in xrange(len(calc.y[br])):
yt = calc.y[br][i]*coeff
if (yt<ylim[1] and ylim[0]<yt):
x.append(calc.eFieldList[i])
y.append(yt)
yState.append(calc.highlight[br][i])
yState = np.array(yState)
sortOrder = yState.argsort(kind='heapsort')
x = np.array(x)
y = np.array(y)
x = x[sortOrder]
y = y[sortOrder]
yState = yState[sortOrder]
ct = "|< %s | \mu > |^2" % printStateString(n,l,j)
ax.scatter(x/100.,y,c=yState,cmin=0,cmax=1,ctitle=ct)
if (xlim==[]):
xlim = [min(x)/100.,max(x)/100.]
ax.printPlot(xlim=xlim,ylim=ylim,filename=filename,name="starkdiv1",\
height=600)
return 0
def plotInteractionLevels(calc,xlim=[],ylim=[],filename=""):
ax = webPlot()
ax.xlabel = "R (\mu m)"
ax.ylabel = "\Delta E (GHz)"
if (calc.drivingFromState[0] == 0):
# colouring is based on the contribution of the original pair state here
ct = r"|< %s %.1f , %s %.1f | \mu > |^2$" % \
(printStateString(calc.n, calc.l,calc.j),\
calc.m1,\
printStateString(calc.nn, calc.ll,calc.jj),\
calc.m1)
else:
# colouring is based on the coupling to different states
ct = "\Omega_\mu/\Omega"
x=[]
y=[]
yState=[]
for br in xrange(len(calc.y)):
for i in xrange(len(calc.y[br])):
x.append(calc.r[i])
y.append(calc.y[br][i])
yState.append(calc.highlight[br][i])
yState = np.array(yState)
sortOrder = yState.argsort(kind='heapsort')
x = np.array(x)
y = np.array(y)
x = x[sortOrder]
y = y[sortOrder]
yState = yState[sortOrder]
ax.scatter(x,y,c=yState,cmin=0,cmax=1,ctitle=ct)
ax.printPlot(xlim=xlim,ylim=ylim,filename=filename,name="levelintdiv")
return
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
self.ax.set_title(
printStateString(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.ax.set_title(title)
self.state1 = [0, 0, 0]
self.state2[1] = state[1]
event.canvas.draw()
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][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 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 = 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] < 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)
