Added functionality to compute the modulated single Gaussian beam and did some other cleanup and refactoring.

calculateDipoleTrapPotential.py

@@ -36,18 +36,23 @@ def find_nearest(array, value):
     idx = (np.abs(array - value)).argmin()
     return idx
 
-def arccos_modulation(mod_amp, n_points):
+def modulation_function(mod_amp, n_points, func = 'arccos'):
-    phi = np.linspace(0, 2*np.pi, n_points)
+    if func == 'arccos':
-    first_half  = 2/np.pi * np.arccos(phi/np.pi-1) - 1
+        phi         = np.linspace(0, 2*np.pi, n_points)
-    second_half = np.flip(first_half)
+        first_half  = 2/np.pi * np.arccos(phi/np.pi-1) - 1
-    return mod_amp * np.concatenate((first_half, second_half))
+        second_half = np.flip(first_half)
+        mod_func    = mod_amp * np.concatenate((first_half, second_half))
+        dx          = (max(mod_func) - min(mod_func))/(2*n_points)
+        return dx, mod_func
+    else:
+        return None
 
 #####################################################################
 # BEAM PARAMETERS                                                   #
 #####################################################################
 
-# Rayleigh range
+# Rayleigh length
-def z_R(w_0:np.ndarray, lamb:float)->np.ndarray:
+def z_R(w_0, lamb)->np.ndarray:
     return np.pi*w_0**2/lamb
 
 # Beam Radius
@@ -145,14 +150,17 @@ def harmonic_potential(pos, v, xoffset, yoffset, m = 164*u.u):
     U_Harmonic = ((0.5 * m * (2 * np.pi * v*u.Hz)**2 * (pos*u.um - xoffset*u.um)**2)/ac.k_B).to(u.uK) + yoffset*u.uK
     return U_Harmonic.value
 
-# def modulated_single_gaussian_beam_potential(positions: "np.ndarray|u.quantity.Quantity", waists: "np.ndarray|u.quantity.Quantity", alpha:"float|u.quantity.Quantity", P:"float|u.quantity.Quantity"=1, wavelength:"float|u.quantity.Quantity"=1.064*u.um)->np.ndarray:
+def modulated_single_gaussian_beam_potential(positions: "np.ndarray|u.quantity.Quantity", waists: "np.ndarray|u.quantity.Quantity", alpha:"float|u.quantity.Quantity", P:"float|u.quantity.Quantity"=1, wavelength:"float|u.quantity.Quantity"=1.064*u.um, mod_amp:"float|u.quantity.Quantity"=1)->np.ndarray:
-#     mod_amp = 0.5 * waists[0]
+    mod_amp = mod_amp * waists[0]
-#     n_points = int(len(positions[0,:])/2)
+    n_points = int(len(positions[0,:])/2)
-#     x_mod = arccos_modulation(mod_amp, n_points)
+    dx, x_mod = modulation_function(mod_amp, n_points, func = 'arccos')
-#     A = 2*P/(np.pi*w(positions[1,:], waists[0], wavelength)*w(positions[1,:], waists[1], wavelength))
+    A = 2*P/(np.pi*w(positions[1,:], waists[0], wavelength)*w(positions[1,:], waists[1], wavelength))
-#     U_tilde = (1 / (2 * ac.eps0 * ac.c)) * alpha * (4 * np.pi * ac.eps0 * ac.a0**3)
+    U_tilde = (1 / (2 * ac.eps0 * ac.c)) * alpha * (4 * np.pi * ac.eps0 * ac.a0**3)
-#     U = - U_tilde * A * np.trapz(np.exp(-2 * (np.subtract(x_mod, positions[0,:])/w(positions[1,:], waists[0], wavelength))**2))
+    dU = np.zeros(2*n_points)
-#     return U
+    for i in range(len(x_mod)):
+        dU = np.vstack((dU, np.exp(-2 * (np.subtract(x_mod[i], positions[0,:])/w(positions[1,:], waists[0], wavelength))**2)))
+    U = - U_tilde * A * 1/(2*mod_amp) * np.trapz(dU, dx = dx, axis = 0)
+    return U
 
 #####################################################################
 # COMPUTE/EXTRACT TRAP POTENTIAL AND PARAMETERS                     #
@@ -352,11 +360,11 @@ def plotPotential(Positions, ComputedPotentials, axis, Params = [], listToIterat
             elif i % 2 != 0:
                 plt.plot(Positions[axis], ComputedPotentials[i][axis], label = 'Effective Trap Depth = ' + str(round(EffectiveTrapDepthInKelvin.value, 2)) + ' ' + str(EffectiveTrapDepthInKelvin.unit) + '; ' + generate_label(v, dv))
     if axis == 0:
-        dir = 'X'
+        dir = 'X - Horizontal'
     elif axis == 1:
-        dir = 'Y'
+        dir = 'Y - Propagation'
     else:
-        dir = 'Z'
+        dir = 'Z - Vertical'
     
     plt.ylim(top = 0)
     plt.xlabel(dir + ' Direction (um)', fontsize= 12, fontweight='bold')
@@ -368,6 +376,96 @@ def plotPotential(Positions, ComputedPotentials, axis, Params = [], listToIterat
         plt.savefig('pot_' + dir + '.png')
     plt.show()
 
+def plotIntensityProfileAndPotentials(Power, waists, alpha, wavelength, options):
+        w_x = waists[0]
+        w_z = waists[1]
+        extent = options['extent']
+        modulation = options['modulation']
+        
+        if not modulation:
+            extent = 50
+            x_Positions      = np.arange(-extent, extent, 1)*u.um 
+            y_Positions      = np.arange(-extent, extent, 1)*u.um
+            z_Positions      = np.arange(-extent, extent, 1)*u.um 
+            
+            idx = np.where(y_Positions==0)[0][0]
+
+            alpha = Polarizability
+            wavelength = 1.064*u.um
+            
+            xm,ym,zm = np.meshgrid(x_Positions, y_Positions, z_Positions, sparse=True, indexing='ij')
+
+            ## Single Gaussian Beam
+            A = 2*Power/(np.pi*w(ym, w_x, wavelength)*w(ym, w_z, wavelength))
+            I = A * np.exp(-2 * ((xm/w(ym, w_x, wavelength))**2 + (zm/w(ym, w_z, wavelength))**2))
+            I = np.transpose(I.to(u.MW/(u.cm*u.cm)))
+            
+            U_tilde = (1 / (2 * ac.eps0 * ac.c)) * alpha * (4 * np.pi * ac.eps0 * ac.a0**3)
+            U = - U_tilde * I
+            U = (U/ac.k_B).to(u.uK)
+
+            fig = plt.figure(figsize=(12, 6))
+            ax = fig.add_subplot(121)
+            ax.set_title('Intensity Profile ($MW/cm^2$)\n Aspect Ratio = %.2f' %(w_x/w_z))
+            im = plt.imshow(I[:,idx,:].value, cmap="coolwarm", extent=[np.min(x_Positions.value), np.max(x_Positions.value), np.min(z_Positions.value), np.max(z_Positions.value)])
+            plt.xlabel('X - Horizontal (um)', fontsize= 12, fontweight='bold')
+            plt.ylabel('Z - Vertical (um)', fontsize= 12, fontweight='bold')
+            ax.set_aspect('equal')
+            fig.colorbar(im, fraction=0.046, pad=0.04, orientation='vertical')
+                
+            bx = fig.add_subplot(122)
+            bx.set_title('Trap Potential')
+            plt.plot(x_Positions, U[np.where(x_Positions==0)[0][0], idx, :], label = 'X - Horizontal')
+            plt.plot(z_Positions, U[:, idx, np.where(z_Positions==0)[0][0]], label = 'Z - Vertical')
+            plt.ylim(top = 0)
+            plt.xlabel('Extent (um)', fontsize= 12, fontweight='bold')
+            plt.ylabel('Depth (uK)', fontsize= 12, fontweight='bold')
+            plt.tight_layout()
+            plt.grid(visible=1)
+            plt.legend(prop={'size': 12, 'weight': 'bold'})
+            plt.show()    
+        else:
+            mod_amp = options['modulation_amplitude']
+            x_Positions      = np.arange(-extent, extent, 1)*u.um 
+            y_Positions      = np.arange(-extent, extent, 1)*u.um
+            z_Positions      = np.arange(-extent, extent, 1)*u.um 
+
+            mod_amp   = mod_amp * w_x
+            n_points  = int(len(x_Positions)/2)
+            dx, xmod_Positions = modulation_function(mod_amp, n_points, func = 'arccos')
+            
+            idx = np.where(y_Positions==0)[0][0]
+
+            xm,ym,zm,xmodm = np.meshgrid(x_Positions, y_Positions, z_Positions, xmod_Positions, sparse=True, indexing='ij')
+            A = 2*Power/(np.pi*w(0*u.um , w_x, wavelength)*w(0*u.um , w_z, wavelength))
+            intensity_profile = A * 1/(2*mod_amp) * np.trapz(np.exp(-2 * (((xmodm - xm)/w(ym, w_x, wavelength))**2 + (zm/w(ym, w_z, wavelength))**2)), dx = dx, axis = -1)
+            I = np.transpose(intensity_profile[:, idx, :].to(u.MW/(u.cm*u.cm)))
+            
+            U_tilde = (1 / (2 * ac.eps0 * ac.c)) * alpha * (4 * np.pi * ac.eps0 * ac.a0**3)
+            U = - U_tilde * I
+            U = (U/ac.k_B).to(u.uK)
+
+            fig = plt.figure(figsize=(12, 6))
+            ax = fig.add_subplot(121)
+            ax.set_title('Intensity Profile ($MW/cm^2$)')
+            im = plt.imshow(I.value, cmap="coolwarm", extent=[np.min(x_Positions.value), np.max(x_Positions.value), np.min(z_Positions.value), np.max(z_Positions.value)])
+            plt.xlabel('X - Horizontal (um)', fontsize= 12, fontweight='bold')
+            plt.ylabel('Z - Vertical (um)', fontsize= 12, fontweight='bold')
+            ax.set_aspect('equal')
+            fig.colorbar(im, fraction=0.046, pad=0.04, orientation='vertical')
+            
+            bx = fig.add_subplot(122)
+            bx.set_title('Trap Potential')
+            plt.plot(x_Positions, U[np.where(x_Positions==0)[0][0], :], label = 'X - Horizontal')
+            plt.plot(z_Positions, U[:, np.where(z_Positions==0)[0][0]], label = 'Z - Vertical')
+            plt.ylim(top = 0)
+            plt.xlabel('Extent (um)', fontsize= 12, fontweight='bold')
+            plt.ylabel('Depth (uK)', fontsize= 12, fontweight='bold')
+            plt.tight_layout()
+            plt.grid(visible=1)
+            plt.legend(prop={'size': 12, 'weight': 'bold'})
+            plt.show()
+
 def plotScatteringLengths():
     BField = np.arange(0, 2.59, 1e-3) * u.G
     a_s_array = np.zeros(len(BField)) * ac.a0
@@ -397,8 +495,13 @@ if __name__ == '__main__':
 
     Power = 40*u.W
     Polarizability = 184.4 # in a.u, most precise measured value of Dy polarizability
+    Wavelength = 1.064*u.um
     w_x, w_z = 27.5*u.um, 33.8*u.um # Beam Waists in the x and y directions
 
+    # Power = 11*u.W
+    # Polarizability = 184.4 # in a.u, most precise measured value of Dy polarizability
+    # w_x, w_z = 54.0*u.um, 54.0*u.um # Beam Waists in the x and y directions
+    
     options = {
         'axis': 0, # axis referenced to the beam along which you want the dipole trap potential
         'extent': 3e2, # range of spatial coordinates in one direction to calculate trap potential over
@@ -412,24 +515,55 @@ if __name__ == '__main__':
         'disp_foci': 3 * z_R(w_0 = np.asarray([30]), lamb = 1.064)[0]*u.um # difference in position of the foci along the propagation direction (Astigmatism)
     }
     
-    ComputedPotentials = [] 
+    """Plot ideal and trap potential resulting for given parameters only"""
-    Params = []  
+    # ComputedPotentials = [] 
+    # Params = []  
 
-    Positions, IdealTrappingPotential, TrappingPotential, TrapDepthsInKelvin, CalculatedTrapFrequencies, ExtractedTrapFrequencies = computeTrapPotential(w_x, w_z, Power, Polarizability, options)
+    # Positions, IdealTrappingPotential, TrappingPotential, TrapDepthsInKelvin, CalculatedTrapFrequencies, ExtractedTrapFrequencies = computeTrapPotential(w_x, w_z, Power, Polarizability, options)
-    ComputedPotentials.append(IdealTrappingPotential)
+    # ComputedPotentials.append(IdealTrappingPotential)
-    ComputedPotentials.append(TrappingPotential)
+    # ComputedPotentials.append(TrappingPotential)
-    Params.append([TrapDepthsInKelvin, CalculatedTrapFrequencies, ExtractedTrapFrequencies])
+    # Params.append([TrapDepthsInKelvin, CalculatedTrapFrequencies, ExtractedTrapFrequencies])
 
-    ComputedPotentials = np.asarray(ComputedPotentials)
+    # ComputedPotentials = np.asarray(ComputedPotentials)
-    plotPotential(Positions, ComputedPotentials, options['axis'], Params)
+    # plotPotential(Positions, ComputedPotentials, options['axis'], Params)
 
+    """Plot harmonic fit for trap potential resulting for given parameters only"""
+    # v, dv, popt, pcov = extractTrapFrequency(Positions, TrappingPotential, options['axis'])
+    # plotHarmonicFit(Positions, TrappingPotential, TrapDepthsInKelvin, options['axis'], popt, pcov)
+
+    """Plot trap potential resulting for given parameters (with one parameter being a list of values and the potential being computed for each of these values) only"""
+    # ComputedPotentials = [] 
+    # Params = []  
+    # Power = [10, 20, 25, 30, 35, 40]*u.W # Single Beam Power
+    # for p in Power: 
+    #     Positions, IdealTrappingPotential, TrappingPotential, TrapDepthsInKelvin, CalculatedTrapFrequencies, ExtractedTrapFrequencies = computeTrapPotential(w_x, w_z, p, Polarizability, options)
+    #     ComputedPotentials.append(IdealTrappingPotential)
+    #     ComputedPotentials.append(TrappingPotential)
+    #     Params.append([TrapDepthsInKelvin, CalculatedTrapFrequencies, ExtractedTrapFrequencies])
+        
+    # ComputedPotentials = np.asarray(ComputedPotentials)
+    # plotPotential(Positions, ComputedPotentials, options['axis'], Params)
+
+    """Plot transverse intensity profile and trap potential resulting for given parameters only"""
+    # options = {
+    #     'extent': 70, # range of spatial coordinates in one direction to calculate trap potential over
+    #     'modulation': True,
+    #     'modulation_amplitude': 4.37
+    # }
+
+    # plotIntensityProfileAndPotentials(Power, [w_x, w_z], Polarizability, Wavelength, options)
+
+    
+    """Calculate relevant parameters for evaporative cooling"""
     AtomNumber = 1.13 * 1e7
-    Temperature = 30 * u.uK
+    Temperature = 100 * u.uK
     BField = 2.1 * u.G
+    modulation = False
     
-    aspect_ratio = 3.67
+    if modulation:
-    init_ar = w_x / w_z
+        aspect_ratio = 3.67
-    w_x = w_x * (aspect_ratio / init_ar)
+        init_ar = w_x / w_z
+        w_x = w_x * (aspect_ratio / init_ar)
 
     n = particleDensity(w_x, w_z, Power, Polarizability, N = AtomNumber, T = Temperature, m = 164*u.u).decompose().to(u.cm**(-3))
     Gamma_elastic = calculateElasticCollisionRate(w_x, w_z, Power, Polarizability, N = AtomNumber, T = Temperature, B = BField)
@@ -451,15 +585,4 @@ if __name__ == '__main__':
     
     # plotScatteringLengths()
     
-    # v, dv, popt, pcov = extractTrapFrequency(Positions, TrappingPotential, options['axis'])
-    # plotHarmonicFit(Positions, TrappingPotential, TrapDepthsInKelvin, options['axis'], popt, pcov)
     
-    # Power = [10, 20, 25, 30, 35, 40]*u.W # Single Beam Power
-    # for p in Power: 
-    #     Positions, IdealTrappingPotential, TrappingPotential, TrapDepthsInKelvin, CalculatedTrapFrequencies, ExtractedTrapFrequencies = computeTrapPotential(w_x, w_z, p, Polarizability, options)
-    #     ComputedPotentials.append(IdealTrappingPotential)
-    #     ComputedPotentials.append(TrappingPotential)
-    #     Params.append([TrapDepthsInKelvin, CalculatedTrapFrequencies, ExtractedTrapFrequencies])
-        
-    # ComputedPotentials = np.asarray(ComputedPotentials)
-    # plotPotential(Positions, ComputedPotentials, options['axis'], Params)