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