def test1():
### 1. Fig 3 of Generalised treatment ... Rotondaro JOSAB 2015 paper
### Normal Faraday spectrum
import time
d = np.arange(-10000,10000,10)
#Voigt
p_dict = {'Bfield':300,'rb85frac':1,'Btheta':0,'lcell':75e-3,'T':58,'Dline':'D2','Elem':'Cs'}
#timing:
st = time.clock()
TF = get_spectra2(d,[1,0,0],p_dict,outputs=['Iy'])
et = time.clock() - st
print(('E-field - Elapsed time (s):', et))
#check vs old elecsus
from elecsus.libs import spectra as old_spec
st = time.clock()
TF_old = old_spec.get_spectra(d,p_dict,outputs=['Iy'])
et = time.clock() - st
print(('Old elecsus - Elapsed time (s):', et))
index = 0 # Iy
fig = plt.figure("Faraday comparison")
ax1 = fig.add_subplot(111)
ax1.plot(d,TF[index],'r',lw=2,label='Faraday')
ax1.plot(d,TF_old[0],'k--',lw=2,label='Vanilla ElecSus')
ax1.legend(loc=0)
ax1.set_xlabel('Detuning (MHz)')
ax1.set_ylabel('Transmission')
plt.show()
def test4():
"""
4. Fig 8 of Rotondaro paper
Arbitrary Filter - optimised
"""
print('This takes a while to compute - be patient!')
d = np.linspace(-15000, 15000, 300)
#Voigt
#p_dict = {'Bfield':700,'rb85frac':1,'Btheta':90*np.pi/180,'lcell':75e-3,'T':84,'Dline':'D2','Elem':'Cs'}
p_dict = {
'Bfield': 1000,
'rb85frac': 1,
'Btheta': 88 * np.pi / 180,
'Bphi': 00 * np.pi / 180,
'lcell': 75e-3,
'T': 93,
'Dline': 'D2',
'Elem': 'Cs'
}
pol = np.array([1.0, 0.0, 0.0])
TVx = get_spectra(d,
pol,
p_dict,
outputs=['I_M45', 'I_P45', 'Ix', 'Iy', 'S0', 'Iz'])
fig2 = plt.figure()
ax1a = fig2.add_subplot(411)
ax2a = fig2.add_subplot(412, sharex=ax1a)
ax3a = fig2.add_subplot(413, sharex=ax1a)
ax4a = fig2.add_subplot(414, sharex=ax1a)
ax1a.plot(d, TVx[0], 'r', lw=2, label=r'$I_{-45}$')
ax2a.plot(d, TVx[1], 'b', lw=2, label=r'$I_{+45}$')
ax3a.plot(d, TVx[2], 'r', lw=2, label=r'$I_x$')
ax4a.plot(d, TVx[3], 'b', lw=2, label=r'$I_y$')
ax4a.plot(d, TVx[0] + TVx[1], 'r:', lw=3.5, label=r'$I_{+45}+I_{-45}$')
ax4a.plot(d, TVx[2] + TVx[3], 'k:', lw=2.5, label=r'$I_x + I_y$')
ax4a.plot(d, TVx[4], 'g--', lw=1.5, label='$S_0$')
# ax4a.plot(d,TVx[5],'c--',lw=2.5,label='$I_z$')
ax4a.set_xlabel('Detuning (MHz)')
ax1a.set_ylabel('I -45')
ax2a.set_ylabel('I +45')
ax3a.set_ylabel('Ix')
ax4a.set_ylabel('Iy')
ax4a.set_xlim(d[0], d[-1] + 3000)
ax4a.legend(loc=0)
plt.show()
def test1():
### 1. Fig 3 of Generalised treatment ... Rotondaro JOSAB 2015 paper
### Normal Faraday spectrum
import os
import time
if hasattr(time, 'process_time'):
from time import process_time as timing # Python 3.3
elif os.name == 'posix':
from time import time as timing #Timing for linux or apple
else:
from time import clock as timing #Timing for windows
d = np.arange(-10000, 10000, 10)
#Voigt
p_dict = {
'Bfield': 300,
'rb85frac': 1,
'Btheta': 0,
'lcell': 75e-3,
'T': 58,
'Dline': 'D2',
'Elem': 'Cs'
}
#timing:
st = timing()
TF = get_spectra2(d, [1, 0, 0], p_dict, outputs=['Iy'])
et = timing() - st
print(('E-field - Elapsed time (s):', et))
#check vs old elecsus
from elecsus.libs import spectra as old_spec
st = timing()
TF_old = old_spec.get_spectra(d, p_dict, outputs=['Iy'])
et = timing() - st
print(('Old elecsus - Elapsed time (s):', et))
index = 0 # Iy
fig = plt.figure("Faraday comparison")
ax1 = fig.add_subplot(111)
ax1.plot(d, TF[index], 'r', lw=2, label='Faraday')
ax1.plot(d, TF_old[0], 'k--', lw=2, label='Vanilla ElecSus')
ax1.legend(loc=0)
ax1.set_xlabel('Detuning (MHz)')
ax1.set_ylabel('Transmission')
plt.show()
def test2():
"""
2. Fig 4/5 of Rotondaro paper
Voigt Filter
"""
d = np.linspace(-15000, 15000, 2500)
#Voigt
## 700 G, 84 C, Cs, 75mm
p_dict = {
'Bfield': 700,
'Btheta': 90 * np.pi / 180,
'lcell': 75e-3,
'T': 84,
'Dline': 'D2',
'Elem': 'Cs'
}
pol = 1. / np.sqrt(2) * np.array([1.0, 1.0, 0.0])
TVx = get_spectra(d,
pol,
p_dict,
outputs=['I_M45', 'I_P45', 'Ix', 'Iy', 'S0', 'Iz'])
fig2 = plt.figure()
ax1a = fig2.add_subplot(411)
ax2a = fig2.add_subplot(412, sharex=ax1a)
ax3a = fig2.add_subplot(413, sharex=ax1a)
ax4a = fig2.add_subplot(414, sharex=ax1a)
ax1a.plot(d, TVx[0], 'r', lw=2, label=r'$I_{-45}$')
ax2a.plot(d, TVx[1], 'b', lw=2, label=r'$I_{+45}$')
ax3a.plot(d, TVx[2], 'r', lw=2, label=r'$I_x$')
ax4a.plot(d, TVx[3], 'b', lw=2, label=r'$I_y$')
ax4a.plot(d, TVx[0] + TVx[1], 'r:', lw=3.5, label=r'$I_{+45}+I_{-45}$')
ax4a.plot(d, TVx[2] + TVx[3], 'k:', lw=2.5, label=r'$I_x + I_y$')
ax4a.plot(d, TVx[4], 'g--', lw=1.5, label='$S_0$')
# ax4a.plot(d,TVx[5],'c--',lw=2.5,label='$I_z$')
ax4a.set_xlabel('Detuning (MHz)')
ax1a.set_ylabel('I -45')
ax2a.set_ylabel('I +45')
ax3a.set_ylabel('Ix')
ax4a.set_ylabel('Iy')
ax4a.set_xlim(d[0], d[-1] + 3000)
ax4a.legend(loc=0)
plt.show()
def test1():
"""
1. Fig 3 of Generalised treatment ... Rotondaro JOSAB 2015 paper
Normal Faraday spectrum
"""
d = np.arange(-10000, 10000, 10) # MHz
#Voigt
p_dict = {
'Bfield': 300,
'rb85frac': 1,
'Btheta': 0,
'lcell': 75e-3,
'T': 58,
'Dline': 'D2',
'Elem': 'Cs'
}
#timing:
st = time.clock()
TF = get_spectra(d, [1, 0, 0], p_dict, outputs=['Iy'])
et = time.clock() - st
print(('E-field - Elapsed time (s):', et))
'''
#check vs old elecsus
from elecsus_v2.libs import spectra as old_spec
st = time.clock()
TF_old = old_spec.get_spectra(d,p_dict,outputs=['Iy'])
et = time.clock() - st
print 'Old elecsus - Elapsed time (s):', et
'''
fig = plt.figure("Faraday comparison")
ax1 = fig.add_subplot(111)
ax1.plot(d, TF[0], 'r', lw=2, label='Faraday')
#ax1.plot(d,TF_old[0],'k--',lw=2,label='Vanilla ElecSus')
#ax1.legend(loc=0)
ax1.set_xlabel('Detuning (MHz)')
ax1.set_ylabel('Transmission')
plt.show()
def test2():
"""
2. Fig 4/5 of Rotondaro paper
Voigt Filter
"""
d = np.linspace(-15000,15000,2500)
#Voigt
## 700 G, 84 C, Cs, 75mm
p_dict = {'Bfield':700,'Btheta':90*np.pi/180,'lcell':75e-3,'T':84,'Dline':'D2','Elem':'Cs'}
pol = 1./np.sqrt(2)*np.array([1.0,1.0,0.0])
TVx = get_spectra(d,pol,p_dict,outputs=['I_M45','I_P45','Ix','Iy','S0','Iz'])
fig2 = plt.figure()
ax1a = fig2.add_subplot(411)
ax2a = fig2.add_subplot(412,sharex=ax1a)
ax3a = fig2.add_subplot(413,sharex=ax1a)
ax4a = fig2.add_subplot(414,sharex=ax1a)
ax1a.plot(d,TVx[0],'r',lw=2,label=r'$I_{-45}$')
ax2a.plot(d,TVx[1],'b',lw=2,label=r'$I_{+45}$')
ax3a.plot(d,TVx[2],'r',lw=2,label=r'$I_x$')
ax4a.plot(d,TVx[3],'b',lw=2,label=r'$I_y$')
ax4a.plot(d,TVx[0]+TVx[1],'r:',lw=3.5,label=r'$I_{+45}+I_{-45}$')
ax4a.plot(d,TVx[2]+TVx[3],'k:',lw=2.5,label=r'$I_x + I_y$')
ax4a.plot(d,TVx[4],'g--',lw=1.5,label='$S_0$')
# ax4a.plot(d,TVx[5],'c--',lw=2.5,label='$I_z$')
ax4a.set_xlabel('Detuning (MHz)')
ax1a.set_ylabel('I -45')
ax2a.set_ylabel('I +45')
ax3a.set_ylabel('Ix')
ax4a.set_ylabel('Iy')
ax4a.set_xlim(d[0],d[-1]+3000)
ax4a.legend(loc=0)
plt.show()
def test4():
"""
4. Fig 8 of Rotondaro paper
Arbitrary Filter - optimised
"""
print('This takes a while to compute - be patient!')
d = np.linspace(-15000,15000,300)
#Voigt
#p_dict = {'Bfield':700,'rb85frac':1,'Btheta':90*np.pi/180,'lcell':75e-3,'T':84,'Dline':'D2','Elem':'Cs'}
p_dict = {'Bfield':1000,'rb85frac':1,'Btheta':88*np.pi/180,'Bphi':00*np.pi/180,'lcell':75e-3,'T':93,'Dline':'D2','Elem':'Cs'}
pol = np.array([1.0,0.0,0.0])
TVx = get_spectra(d,pol,p_dict,outputs=['I_M45','I_P45','Ix','Iy','S0','Iz'])
fig2 = plt.figure()
ax1a = fig2.add_subplot(411)
ax2a = fig2.add_subplot(412,sharex=ax1a)
ax3a = fig2.add_subplot(413,sharex=ax1a)
ax4a = fig2.add_subplot(414,sharex=ax1a)
ax1a.plot(d,TVx[0],'r',lw=2,label=r'$I_{-45}$')
ax2a.plot(d,TVx[1],'b',lw=2,label=r'$I_{+45}$')
ax3a.plot(d,TVx[2],'r',lw=2,label=r'$I_x$')
ax4a.plot(d,TVx[3],'b',lw=2,label=r'$I_y$')
ax4a.plot(d,TVx[0]+TVx[1],'r:',lw=3.5,label=r'$I_{+45}+I_{-45}$')
ax4a.plot(d,TVx[2]+TVx[3],'k:',lw=2.5,label=r'$I_x + I_y$')
ax4a.plot(d,TVx[4],'g--',lw=1.5,label='$S_0$')
# ax4a.plot(d,TVx[5],'c--',lw=2.5,label='$I_z$')
ax4a.set_xlabel('Detuning (MHz)')
ax1a.set_ylabel('I -45')
ax2a.set_ylabel('I +45')
ax3a.set_ylabel('Ix')
ax4a.set_ylabel('Iy')
ax4a.set_xlim(d[0],d[-1]+3000)
ax4a.legend(loc=0)
plt.show()
def big_diagram(BFIELD=1000, output='S0'):
"""
Main code to plot 'big' diagram with the following components:
- Theoretical absorption spectrum (top panel)
- Breit Rabi diagram for 0 to specified B-field (left)
- Energy levels for ground and excited states (bottom panel)
- Arrows for each transition, underneath the corresponding part of the spectrum
"""
##
## First part - calculate the absorption spectrum
##
# Define the detuning axis based on what the magnetic field strength is (in GHz)
# Values for BFIELD should be given in Gauss (1 G = 1e-4 T)
Dmax = max(6, 5 + (BFIELD / 1e4 * 3 * mu_B))
det_range = np.linspace(-Dmax, Dmax, int(3e4))
# Input parameters to calculate the spectrum
Bfield = BFIELD #alias
ELEM = 'Rb'
DLINE = 'D2'
RB85FRAC = 0.0 # Pure Rb87
LCELL = 1e-3
TEMP = 100 # C ~ 373K
# Voigt, horizontal polarisation
pol = [1, 0, 0]
p_dict = {
'T': TEMP,
'lcell': LCELL,
'Elem': ELEM,
'rb85frac': RB85FRAC,
'Dline': DLINE,
'Bfield': BFIELD,
'Btheta': 90 * np.pi / 180,
'Bphi': 45 * np.pi / 180,
'BoltzmannFactor': True
}
[S0, S1, S2, S3] = get_spectra(det_range * 1e3,
pol,
p_dict,
outputs=['S0', 'S1', 'S2', 'S3'])
lenergy87, lstrength87, ltransno87, lgl87, lel87, \
renergy87, rstrength87, rtransno87, rgl87, rel87, \
zenergy87, zstrength87, ztransno87, zgl87, zel87 = calc_chi_energies([1], p_dict)
##
## Second part - calculate the Breit-Rabi diagram
##
BreitRabiVals = np.linspace(0, BFIELD, 2000)
BreitRabiVals = np.append(BreitRabiVals, BreitRabiVals[-1])
Bstep = BreitRabiVals[1] - BreitRabiVals[0]
# Calculate Zeeman-shifted energy levels in parallel (uses multiprocessing module)
po = Pool()
res = po.map_async(eval_energies, ((
"Rb87",
"D2",
BreitRabiVals[k],
) for k in range(len(BreitRabiVals))))
energies = res.get()
gnd_energies = np.zeros((len(energies[0][0]), len(BreitRabiVals)))
exc_energies = np.zeros((len(energies[0][1]), len(BreitRabiVals)))
for jj, energyB in enumerate(energies):
gnd_energies[:, jj] = energyB[0]
exc_energies[:, jj] = energyB[1]
po.close()
po.join()
# Energies at largest B-field value
final_gnd_energies, final_exc_energies = eval_energies(
("Rb87", "D2", BreitRabiVals[-1]))
##
## Third part - calculate state decomposition
##
## Below values are for Rb-87. **Change for other atoms**.
I = 3.0 / 2
L = 0
S = 1.0 / 2
J = 1.0 / 2
output_states = AM_StateDecomp(I, L, S, J, atom='Rb', B=BFIELD / 1e4)
print('\nState decomposition at B = ', BFIELD / 1e4)
print(output_states)
##
## Fourth part - arrange the plot panels
##
fig = plt.figure("Big diagram at " + str(BFIELD / 1e4) + ' T',
facecolor=None,
figsize=(12, 8))
plt.clf()
# Subplot arrangement
xBR = 2
xspec = 6
yBRe = 3
yBRg = 5
yspec = 4
xx = xBR + xspec
yy = yBRe + yBRg + yspec
ax_spec = plt.subplot2grid((yy, xx), (0, xBR),
colspan=xspec,
rowspan=yspec)
ax_excBR = plt.subplot2grid((yy, xx), (yspec, 0),
colspan=xBR,
rowspan=yBRe)
ax_gndBR = plt.subplot2grid((yy, xx), (yspec + yBRe, 0),
colspan=xBR,
rowspan=yBRg,
sharex=ax_excBR)
ax_eLev = plt.subplot2grid((yy, xx), (yspec, xBR),
colspan=xspec,
rowspan=yBRe,
sharex=ax_spec,
sharey=ax_excBR)
ax_gLev = plt.subplot2grid((yy, xx), (yspec + yBRe, xBR),
colspan=xspec,
rowspan=yBRg,
sharex=ax_spec,
sharey=ax_gndBR)
# Turn off axes for eLev and gLev axes
for ax in [ax_eLev, ax_gLev]:
ax.set_frame_on(False)
for parameter in [
ax.get_xticklabels(),
ax.get_yticklabels(),
ax.get_xticklines(),
ax.get_yticklines()
]:
plt.setp(parameter, visible=False)
plt.setp(ax_excBR.get_xticklabels(), visible=False)
ax_excBR.spines['right'].set_color('none')
ax_gndBR.spines['right'].set_color('none')
ax_gndBR.spines['top'].set_color('none')
ax_excBR.spines['top'].set_color('none')
ax_excBR.spines['bottom'].set_color('none')
ax_gndBR.xaxis.set_ticks_position('bottom')
ax_excBR.xaxis.set_ticks_position('none')
ax_excBR.tick_params(axis='y', left=True, right=False)
ax_gndBR.tick_params(axis='y', left=True, right=False)
# axis labels
ax_spec.set_xlabel('Detuning (GHz)')
ax_spec.xaxis.set_label_position('top')
ax_spec.tick_params(axis='x',
bottom=True,
top=True,
labelbottom=False,
labeltop=True)
ax_excBR.set_ylabel('$5P_{3/2}$ energy (GHz)')
ax_gndBR.set_ylabel('$5S_{1/2}$ energy (GHz)')
ax_gndBR.set_xlabel('Magnetic Field (T)')
fig.subplots_adjust(left=0.07,
right=0.98,
top=0.93,
bottom=0.085,
hspace=0.34,
wspace=0)
#Ghost axes for actually plotting the Breit-Rabi data
eleft = ax_excBR.get_position().extents[0:2]
eright = ax_eLev.get_position().extents[2:]
gleft = ax_gndBR.get_position().extents[0:2]
gright = ax_gLev.get_position().extents[2:]
ax_e_bound = np.append(eleft, eright - eleft)
ax_g_bound = np.append(gleft, gright - gleft)
print('\nAxes bounds for B-R diagram:')
print(ax_e_bound)
print(ax_g_bound)
ax_e = fig.add_axes(ax_e_bound, frameon=False, facecolor=None)
ax_g = fig.add_axes(ax_g_bound, frameon=False, facecolor=None)
ax_g.set_xticks([])
ax_g.set_yticks([])
ax_e.set_xticks([])
ax_e.set_yticks([])
##
## Fifth part - Add the data to the figure
##
# Edit last magnetic field value
BreitRabiVals[-1] = BreitRabiVals[-2] * ((xspec + xBR) / xBR)
print('\nMagnetic field values (Breit-Rabi diagram)')
print(BreitRabiVals)
if output == 'S0':
ax_spec.set_ylabel('Transmission, $S_{0}$')
ax_spec.plot(det_range, S0.real, lw=2, color=d_black)
elif output == 'S1':
ax_spec.set_ylabel('$S_{1}$')
ax_spec.plot(det_range, S1.real, lw=2, color=d_black)
elif output == 'S2':
ax_spec.set_ylabel('$S_{2}$')
ax_spec.plot(det_range, S2.real, lw=2, color=d_black)
elif output == 'S3':
ax_spec.set_ylabel('$S_{3}$')
ax_spec.plot(det_range, S3.real, lw=2, color=d_black)
#convert to GHz from MHz
exc_energies /= 1e3
gnd_energies /= 1e3
final_exc_energies /= 1e3
final_gnd_energies /= 1e3
for energy in exc_energies[int(len(final_exc_energies) / 3):]:
ax_e.plot(BreitRabiVals / 1e4, energy, color=d_black, lw=1)
for energy in gnd_energies:
ax_g.plot(BreitRabiVals / 1e4, energy, color=d_black, lw=1.5)
ax_excBR.set_xlim(0, (Bfield + 10 * Bstep) / 1e4)
for ax in [ax_g, ax_e]:
ax.set_ylim(ax.get_ylim()[0] * 1.15, ax.get_ylim()[1] * 1.15)
ax.set_xlim(BreitRabiVals[0] / 1e4, BreitRabiVals[-1] / 1e4)
ax_excBR.set_ylim(ax_e.get_ylim())
ax_gndBR.set_ylim(ax_g.get_ylim())
ax_spec.set_xlim(det_range[0], det_range[-1])
ax_spec.set_ylim(ax_spec.get_ylim()[0], 1.01)
##
## Sixth part - Add arrows for each transition
##
print('Sigma minus transitions:')
print(sorted(lenergy87))
print('Sigma plus transitions:')
print(sorted(renergy87))
print('Pi transitions:')
print(sorted(zenergy87))
for energy in lenergy87:
ax_spec.axvline(energy / 1e3, color=d_purple, lw=1.5)
for energy in renergy87:
ax_spec.axvline(energy / 1e3, color=d_blue, lw=1.5)
for energy in zenergy87:
ax_spec.axvline(energy / 1e3,
color=d_olive,
lw=1.5,
linestyle='dashed')
# Coordinates for arrows - sigma minus transitions (purple)
xy1s = zip(lenergy87 / 1e3, lgl87 / 1e3)
xy2s = zip(lenergy87 / 1e3, lel87 / 1e3)
ecol = d_purple
fcol = 0.5 * (np.array(d_lightpurple) + np.array(d_purple))
alpha = 0.9
#styles = ['solid','solid','solid','solid','dashed','dashed','dashed','dashed']
for xy1, xy2, strength in zip(xy1s, xy2s, lstrength87):
#if (xy1[0] > 15) or (xy1[0]<-15):
coordsA = 'data'
coordsB = 'data'
con = ConnectionPatch(xy1,
xy2,
coordsA,
coordsB,
arrowstyle="simple",
shrinkB=0,
axesA=ax_gLev,
axesB=ax_eLev,
mutation_scale=25,
ec=ecol,
fc=fcol,
lw=1.25,
alpha=alpha)
ax_gLev.add_artist(con)
# Coordinates for arrows - sigma plus transitions (blue)
xy1s = zip(renergy87 / 1e3, rgl87 / 1e3)
xy2s = zip(renergy87 / 1e3, rel87 / 1e3)
ecol = d_blue
fcol = 0.5 * (np.array(d_midblue) + np.array(d_blue))
alpha = 0.9
#styles = ['solid','solid','solid','solid','dashed','dashed','dashed','dashed']
for xy1, xy2, strength in zip(xy1s, xy2s, rstrength87):
#if (xy1[0] > 15) or (xy1[0]<-15):
coordsA = 'data'
coordsB = 'data'
con = ConnectionPatch(xy1,
xy2,
coordsA,
coordsB,
arrowstyle="simple",
shrinkB=0,
axesA=ax_gLev,
axesB=ax_eLev,
mutation_scale=25,
ec=ecol,
fc=fcol,
lw=1.25,
alpha=alpha)
ax_gLev.add_artist(con)
# Coordinates for arrows - pi transitions (olive)
xy1s = zip(zenergy87 / 1e3, zgl87 / 1e3)
xy2s = zip(zenergy87 / 1e3, zel87 / 1e3)
ecol = d_darkolive
fcol = d_olive #darkyellow#olive #(0.16,0.85,0.16)
alpha = 0.6
#styles = ['solid','solid','solid','solid','dashed','dashed','dashed','dashed']
for xy1, xy2, strength in zip(xy1s, xy2s, zstrength87):
#if (xy1[0] < 15) and (xy1[0]>-15):
coordsA = 'data'
coordsB = 'data'
con = ConnectionPatch(xy1,
xy2,
coordsA,
coordsB,
arrowstyle="simple",
shrinkB=0,
axesA=ax_gLev,
axesB=ax_eLev,
mutation_scale=25,
ec=ecol,
fc=fcol,
lw=1.25,
alpha=alpha)
ax_gLev.add_artist(con)
# Add B-field info to plot - top left
fig.text(0.1,
0.78 - 0.03,
'L = ' + str(LCELL * 1e3) + ' mm',
size=18,
ha='center')
fig.text(0.1,
0.82 - 0.03,
r'T = ' + str(TEMP) + ' $^{\circ}$C',
size=18,
ha='center')
fig.text(0.1,
0.86 - 0.03,
'B = ' + str(Bfield / 1e4) + ' T',
size=18,
ha='center')
fig.text(0.1, 0.90 - 0.03, str(DLINE) + ' Line', size=18, ha='center')
fig.text(0.1, 0.94 - 0.03, '$^{87}$Rb', size=18, ha='center')
##
## Finally - show the plot and save the figure
##
ax_spec.set_xlim(-Dmax, Dmax)
# fig.savefig('./BR_plot_'+str(Bfield)+str(output)'.pdf',dpi=300)
# fig.savefig('./BR_plot_'+str(Bfield)+str(output)'.png',dpi=300)
plt.show()
print('--- End of calculations ---')
return fig
def field_gradient_fit(fit=True):
"""
Fit theory curve to experimental data using the non-uniform magnetic field model above (fieldgrad_fitfn).
Uses lmfit module for fitting. Differential evolution is recommended to find global optimum parameters
(leastsq method may work depending on choice of intiial parameters, but not guaranteed). Fitting takes a while,
around 10-15 minutes depending on computer speed.
The 'fit' keyword argument is used to either perform the fitting (True) or load in best-fit parameters from previous
calculations (False).
This method will reproduce figure 8 of the ElecSus paper completely.
For this the file S0_Bgradient.csv from the 'ElecSusTestData\Field Gradient' GitHub repository is required to be
in the same directory as this file.
"""
## initial guess params for fit
sep = 92.
offset_adj = 0.
temperature = 17.7
# experimental data; load and crop
d_expt, S0_expt = np.loadtxt('./S0_Bgradient.csv',delimiter=',').T
d_expt = d_expt[::200]
S0_expt = S0_expt[::200]
#######################################
if fit:
# Fitting takes a while..!
x = np.array(d_expt)
y = np.array(S0_expt)
p_dict = {'sep':sep, 'offset_adj':offset_adj, 'temperature':temperature}
# Have a quick look at the guess parameter curve to see if it's close
S0_trial = fieldgrad_fitfn(d_expt, sep, offset_adj, temperature)
plt.plot(d_expt,S0_trial)
plt.plot(d_expt,S0_expt,'.')
plt.title('Curve based on initial parameters:')
plt.show() # halts the program here until plot window is closed...
model = lm.Model(fieldgrad_fitfn)
params = model.make_params(**p_dict)
# Turn off all parameters varying by default, unless specified in p_dict_bools
allkeys = params.valuesdict()
for key in allkeys:
params[key].vary = False
p_dict_bools = {'sep':True, 'offset_adj':True, 'temperature':True}
p_dict_bounds = {'sep': [75., 120.], 'offset_adj':[-10,10], 'temperature':[15,25]} # sensible boundaries based on lab conditions
# Turn on fitting parameters as specified in p_dict_bools, and set boundaries
for key in p_dict_bools:
params[key].vary = p_dict_bools[key]
if p_dict_bounds is not None:
if key in p_dict_bounds:
params[key].min = p_dict_bounds[key][0]
params[key].max = p_dict_bounds[key][1]
# need to use global solver - there are local minima in parameter space
#method = 'leastsq'
method = 'differential_evolution'
# run the fit
result = model.fit(y, x=x, params=params, method=method)
print((result.fit_report()))
S0_thy = result.best_fit
# make quick residual plot to look for any remaining structure
result.plot_residuals()
# save parameters to pickled file so we don't have to re-run the fit every time we want to look at the plot
fit_params = result.best_values
pickle.dump(fit_params, open('./bgradient_fitparams.pkl','wb'))
else:
# Load in previously calculated fit parameters
fit_params = pickle.load(open('./bgradient_fitparams.pkl','rb'))
# extract parameters from dictionary for convenience
sep = fit_params['sep']
offset_adj = fit_params['offset_adj']
temperature= fit_params['temperature']
# detuning axis for theory arrays
d = np.linspace(-10,10,1000)
S0, S1 = fieldgrad_fitfn(d, sep, offset_adj, temperature,return_S1=True) # evaluated for best-fit parameters
# scale parameters back into metres, from mm
sep *= 1e-3
offset_adj *= 1e-3
# plot magnetic field with optimised parameters
z0 = sep/2
LCELL = 75e-3
offset = z0 - LCELL/2 - offset_adj
zs = np.linspace(-60e-3,120e-3,1400)
Bf = tophat_profile((zs-z0+offset),z0)
# Calculate average magnetic field across cell
zs_cell = np.linspace(0,LCELL,1000)
Bf2 = tophat_profile(zs_cell-z0+offset,z0)
B_avg = Bf2.mean() * 1e4 # in Gauss
# Compare against single pass with average magnetic field - i.e. no gradient
p_dict = {'rb85frac':72.17,'Btheta':0,'Bphi':0,'lcell':LCELL,'T':temperature,'Dline':'D2','Elem':'Rb'}
p_dict['Bfield'] = B_avg
print(('Average field,', B_avg))
print(('Min / Max field across cell: ', Bf2.min()*1e4, Bf2.max()*1e4))
# calculate specra with average field
pol_in = 1./np.sqrt(2) * np.array([1,1,0])
[Iy_avg, S1_avg, S0_avg] = get_spectra(d*1e3,pol_in,p_dict,outputs=['Iy','S1','S0'])
# also calculate zero-field spectra for comparison (olive shading in fig 8)
p_dict['Bfield'] = 0
[Iy_0, S1_0, S0_0] = get_spectra(d*1e3,pol_in,p_dict,outputs=['Iy','S1','S0'])
# arbitrary scaling of zero-field spectra to fit nicely on the same axes limits as the other data.
S0_0 = S0_0 * 0.5 + 0.5
# Set up figure panels with subplot2grid
fig2 = plt.figure("Figure 8 of ElecSus paper",figsize=(5,6))
yy = 3
xx = 9
axM = plt.subplot2grid((yy,xx),(0,1),colspan=7)
ax = plt.subplot2grid((yy,xx),(1,0), colspan=xx)
ax2 = plt.subplot2grid((yy,xx),(2,0),colspan=xx,sharex=ax)
# Plot magnetic field profile
axM.plot(zs*1e3,-Bf,color=d_blue)
# format axes for this sub-plot
axM.xaxis.set_label_position('top')
axM.tick_params(axis='x',bottom=True,top=True,labelbottom=False,labeltop=True)
axM.set_xlabel('Axial position, $z$ (mm)')
axM.set_ylabel('$B_z(z)$ (T)')
axM.set_yticks([-.4000,-.2000,0,.2000,.4000,.6000])
axM.set_ylim(-0.1,0.45)
axM.set_xlim(-50,120)
# magnet length parameters (for shading on axes)
L1 = 13.3e-3
L2 = 9.5e-3
# Add shading for cell and axial extent of magnets
axM.axvspan(1e3*(-offset),1e3*(-offset-L1-L2),alpha=0.5,color=d_midblue)
axM.axvspan(1e3*(sep-offset),1e3*(sep-offset+L1+L2),alpha=0.5,color=d_midblue)
axM.axvspan(0*1e3,LCELL*1e3,alpha=0.35, color=d_purple)
# Add S0 spectral data
ax.fill_between(d,1,S0_0,label=r'Zero field',color=d_olive,alpha=0.25)
ax.plot(d,S0_avg,label=r'Uniform, $\langle B_z(z) \rangle$',color=2.5*np.array(d_black),linestyle='dashed')
ax.plot(d,S0,label='Gradient, $B_z(z)$',color=d_blue)
ax.plot(d_expt,S0_expt,'.', ms=4, label='Gradient, $B_z(z)$',color=d_purple)
# Add S1 curves
ax2.plot(d,S1_avg,label=r'Uniform, ($\langle B_z(z) \rangle$)',color=2.5*np.array(d_black),linestyle='dashed')
ax2.plot(d,S1,label='Gradient, ($B_z(z)$)',color=d_blue)
# Format axes
ax2.set_xlabel('Detuning (GHz)')
ax.set_ylabel('$S_0$')
ax2.set_ylabel('$S_1$')
plt.setp(ax.get_xticklabels(),visible=False)
ax.set_xlim(-10,10)
# Scale to full figure canvas
plt.tight_layout()
# Save figure images
fig2.savefig('./field_gradient_fit.png')
fig2.savefig('./field_gradient_fit.pdf')
# Show figure interactively
plt.show()
def fieldgrad_fitfn(x, sep, offset_adj, temperature, n_segments=25,return_S1=False, verbose=False):
"""
Fit function to calculate transmission spectra in Faraday geometry, with a magnetic field gradient
The magnetic field is from a pair of Nd top-hat magnets placed next to a 75 mm reference cell,
calculated using the tophat_profile() method above
The calculation is based on splitting the cell into multiple (n_segments) elements, and propagating the electric field through each element separately with a local value of the magnetic field. Convergence tests should be run to ensure that a suitable number of elements are used.
Floating parameters are the separation between the magnets, their position relative to the vapour cell, and the cell temperature.
sep and offset_adj are scaled to mm so that all fit parameters are around the same order of magnitude to prevent potential issues with levenburg-marquardt fitting routines
Inputs:
x : 1D-array, detuning axis for the calculation, in GHz
sep : Separation between the two top-hat magnets (in mm)
offset_adj : Adjust the position between the magnets and the vapour cell (in mm)
temperature : cell temperature in C
Options:
n_segments : number of segments cell is split into for calculation
return_S1 : If True, returns both S0 and S1 spectra (need False to run fit)
verbose : If True, give more information about present cell segments
Outputs:
S0 : 1D numpy array, Transmission after the cell
S1 : Only if return_S1 is True. 1D numpy array, Faraday rotation signal after the cell
"""
# so we can see what the fit is doing...
print((sep, offset_adj, temperature))
sep *= 1e-3 # convert to m from mm
offset_adj *= 1e-3 # convert to m from mm
# cell length
LCELL = 75e-3
#magnetic field angles
BTHETA = 0
BPHI = 0
#element, Dline and isotopic abundance
RB85FRAC = 72.17
DLINE = 'D2'
ELEM = 'Rb'
# cell starts at z=0, ends at z=LCELL. Offset needed to adjust the position of the magnets such that
# they are initially centred at the middle of the cell. This offset is adjustable through offset_adj to adjust the relative position of cell and magnet pair.
z0 = sep/2
offset = z0 - LCELL/2 - offset_adj
#parameter dictionary
p_dict = {'Bfield':0,'rb85frac':RB85FRAC,'Btheta':BTHETA*np.pi/180,'Bphi':00*np.pi/180,'lcell':LCELL,'T':temperature,'Dline':DLINE,'Elem':ELEM}
# input polarisation
pol_in = 1./np.sqrt(2) * np.array([1,1,0]) # from expt
I_in = 1.
# number of parts to split the cell into (see convergence testing method below)
#n_segments = 25
seg_length = LCELL/n_segments
# centre position of each segment
edges = np.linspace(0,LCELL,n_segments+1)
seg_centres = (edges[:-1] + edges[1:])/2
# calculate field at centre of each segment
#seg_fields = ringMagnetBfield(seg_centres,z0,M_ID,M_OD,M_T)
seg_fields = tophat_profile(seg_centres-z0+offset,z0)
# Evolve electric field over each segment
for i in range(n_segments):
if verbose: print(('Cell segment ', i+1, '/', n_segments))
# average magnetic field in the segment
p_dict['Bfield'] = seg_fields[i] * 1e4 # convert to Gauss from Tesla
p_dict['lcell'] = seg_length
[E_out] = get_spectra(x*1e3,pol_in,p_dict,outputs=['E_out'])
# input field for the next iteration is the output field of this iteration
pol_in = E_out
# Calculate Stokes parameters...
Ey = np.array(JM.VertPol_xy * E_out[:2])
Iy = (Ey * Ey.conjugate()).sum(axis=0) / I_in
Ex = np.array(JM.HorizPol_xy * E_out[:2])
Ix = (Ex * Ex.conjugate()).sum(axis=0) / I_in
S0 = Ix + Iy
#Other Stokes parameters...
S1 = Ix - Iy
if return_S1:
return S0, S1
else:
return S0
