Analysis_Calibration.py

'''
authored collaboratively by the authors of PUBLICATION_LINK

            ###############################################
            ### Rydberg Atom Electric Field Calibration ###
            ###############################################

This script walks through finding the electric field experienced by the atoms.            

It is assumed you have access to the following variables:

Oscilloscope data [Volts] in np.arrays 'probe' and 'LIA' with an optical frequency array 'optical_freqs'
Spectrum Analyzer data [dBm] from the waveguide in np.array 'SA_dBms' with frequencies 'SA_freqs'
Applied RF frequency [Hz] as float 'LO'
'''

#################
##### DUMMY ##### 
#################
# I'll delete this before publication, but this is to check for errors when using the variables below
probe = np.array([])
LIA = np.array([])
optical_freqs = np.array([])
SA_dBms = np.array([])
SA_freqs = np.array([])
LO = 0.
freq_shifts = np.array([])
measured_dBms = np.array([])
#################
#################
#################

import matplotlib.pyplot as plt
plt.rcParams['font.family'] = 'serif' # style choice
import numpy as np
import scipy.constants as cons
    
    # varibales of convenience
size = 2**6 # useful to scale image text
figsize = (8,4.9) # usefult to scale image size

################################
### CHOOSE WHAT TO SHOW/SAVE ###
################################
show_plots = True
save_plots = False

if show_plots: # recreates version of Figure 2 from PUBLICATION_LINK    
    fig,ax=plt.subplots(2, figsize = figsize, sharex = True, gridspec_kw={'hspace': 0}) 
        # plot data
    ax[0].plot(optical_freqs,probe)
    ax[1].plot(optical_freqs,LIA)    
        # labels
    ax[-1].set_xlabel(r'$\Delta_{C} (MHz)$', fontsize = size)
    ax[0].set_ylabel("Photodiode Output (V)", fontsize = size)
    ax[1].set_ylabel("Lock-in Amplifier Output (V)", fontsize = size)
        # axes
    ax[0].tick_params(axis = 'x', which = 'both', bottom = True, top = False, labelbottom = False)
    ax[0].tick_params(axis = 'y', labelsize = size)
    ax[0].minorticks_on()            
    ax[0].tick_params(which = 'major', length = 12)    
        # axes
    ax[1].tick_params(axis = 'x', which = 'both', bottom = True, top = False, labelbottom = True, labelsize = size)
    ax[1].tick_params(axis = 'y', labelsize = size)
    ax[1].minorticks_on()        
    ax[1].tick_params(which = 'major', length = 12)
        # figure layout
    fig.suptitle('Raw Oscilloscope Data', fontsize = size)
    fig.tight_layout()
    
    if save_plots: # save plot
        fig.savefig('Oscilloscope_data.png')

###################################
### GET SPECTRUMA ANALYZER PEAK ###
###################################
    # find the Local Oscillator peak dBm
wanted_freq_idx = np.argmin([abs(x-LO) for x in SA_freqs])    
SA_dBm = SA_dBms[wanted_freq_idx]

if show_plots:        
    fig,ax=plt.subplots(figsize = figsize)
        # plot data
    ax.plot(SA_freqs,SA_dBms)
    ax.scatter(SA_freqs[wanted_freq_idx], SA_dBm, marker = '*', s = 2**8, color = 'red')   
        # labels
    ax.set_xlabel('Frequency (Hz)', fontsize = size)
    ax.set_ylabel('RF Power (dBm)', fontsize = size)
        # axes
    ax.tick_params(axis = 'both', labelsize = size)
    ax.minorticks_on()            
    ax.tick_params(which = 'major', length = 12)
        # figure layout
    fig.suptitle('Spectrum Analyzer Data', fontsize = size)
    fig.tight_layout()
    
    if save_plots:
        fig.savefig('SpectrumAnalyzer_data.png')

###################################
### GET OPTICAL FREQUENCY SHIFT ###
###################################
    # this defines a boundary between the minimum and maximum of the LIA output
if (np.argmax(LIA) > np.argmin(LIA)):
    boundary = ((optical_freqs > optical_freqs[np.argmax(LIA)]) & (optical_freqs < optical_freqs[np.argmin(LIA)]))
elif (np.argmax(LIA) < np.argmin(LIA)):
    boundary = ((optical_freqs > optical_freqs[np.argmin(LIA)]) & (optical_freqs < optical_freqs[np.argmax(LIA)]))
LIA = LIA[boundary]
optical_freqs = optical_freqs[boundary]
    # find the frequency shift, which is the zero crossing of the LIA
freqShift = optical_freqs[np.argmin(np.abs(LIA))]

if show_plots:
    fig,ax=plt.subplots(figsize = figsize)
        # plot data
    ax.plot(optical_freqs, LIA)
    ax.scatter(optical_freqs[np.argmin(np.abs(LIA))], \
        LIA[np.argmin(np.abs(LIA))], marker="*", color="red", s=2**8, label = 'zero crossing')
        # labels
    ax.set_xlabel(r'$\Delta_{C} (MHz)$', fontsize = size)
    ax.set_ylabel("Lock-in Amplifier Output (V)", fontsize = size)    
        # axes
    ax.tick_params(axis = 'both', labelsize = size)
    ax.minorticks_on()            
    ax.tick_params(which = 'major', length = 12)
        # legend
    ax.legend(loc = 'upper right', fontsize = size//2)
        # figure layout
    fig.suptitle('Cropped Lock-in Amplifer Output', fontsize = size)
    fig.tight_layout()
    
    if save_plots:
        fig.savefig('Cropped_LIA.png')

'''
##################        
### SAVE DATA ###
##################
Save the following data to a file:
    SA_dBm
    freqShift

Now one can convert the optical frequency shift of the EIT feature to an electric field via Eq. 1 in PUBLICATION_LINK
'''

#################################################
### CONVERT FREQUENCY SHIFT TO ELECTRIC FIELD ###
#################################################
# define your atom and Rydberg state n, j, l
# this uses the Alkali Rydberg Calculator package: https://arc-alkali-rydberg-calculator.readthedocs.io/en/latest/
from arc import *
atom = Rubidium() # example                    
n, l, j = [55, 3, 7/2] # example 55F_7/2 Rydberg state
dp = DynamicPolarizability(atom, n, l, j)
dp.defineBasis(max(0, n-30), min(n+30, 100)) # beyond n +/- 30, alpha0 changes minutely
alpha0 = dp.getPolarizability(cons.c/LO)[0]

'''
Assume you have np.arrays of:
    Measured applied RF powers 'measured_dBms' populated with each SA_dBm from above
    Measured frequency shifts 'freq_shifts' populated with each freq_shift from above
'''
    # get applied RF power
measured_RF_mW = 10 ** (measured_dBms / 10) # mW
    # get Efield
freq_shifts = abs(freq_shifts - max(freq_shifts)) # peak will shift towards negative coupler frequency
Efields = np.sqrt((2/alpha0)*freq_shifts) # RMS                    

#####################################
### GET ATOMIC CALIBRATION FACTOR ###
#####################################
fit = np.polyfit(x = np.sqrt(measured_RF_mW), y = Efields, deg = 1, cov=True)
slope = fit[0][0]; intercept = fit[0][1]; linear_cov = fit[1] # save these should you wish

if show_plots: # recreates version of Figure 3 from PUBLICATION_LINK
    fig,ax = plt.subplots(figsize = figsize)
        # plot data
    x = np.sqrt(measured_RF_mW) # convenience
    ax.scatter(x,Efields, s=2**8)
    ax.plot(x,slope*x+intercept, linestyle = "--", color = "red", linewidth = 2**3)                                            
        # labels
    ax.set_xlabel(r'Applied LO Power ($\sqrt{mW}$)', fontsize = size)
    ax.set_ylabel('Electric Field (V/m)',fontsize = size)
        # axes
    ax.tick_params(axis = 'x', labelsize = size)
    ax.tick_params(axis = 'both', labelsize = size)
    ax.minorticks_on()            
    ax.tick_params(which = 'major', length = 12)
        # figure layout
    fig.suptitle('Atomic Calibration', fontsize = size)
    fig.tight_layout()

    if save_plots:
        fig.savefig('Calibration_Factor.png')