Added harmonic fit to extract trap frequency of potential, corrected few bugs.

calculateDipoleTrapPotential.py

@@ -1,18 +1,20 @@
 import math
 import numpy as np
 import matplotlib.pyplot as plt
+from scipy.optimize import curve_fit
 from astropy import units as u, constants as ac
 
+def orderOfMagnitude(number):
+    return math.floor(math.log(number, 10))
+
 def rotation_matrix(axis, theta):
     """
     Return the rotation matrix associated with counterclockwise rotation about
     the given axis by theta radians.
-
     In 2-D it is just,
         thetaInRadians = np.radians(theta)
         c, s = np.cos(thetaInRadians), np.sin(thetaInRadians)
         R = np.array(((c, -s), (s, c)))
-
     In 3-D, one way to do it is use the Euler-Rodrigues Formula as is done here   
     """
     axis = np.asarray(axis)
@@ -39,7 +41,7 @@ def trap_depth(w_1:"float|u.quantity.Quantity", w_2:"float|u.quantity.Quantity",
 def gravitational_potential(positions: "np.ndarray|u.quantity.Quantity", m:"float|u.quantity.Quantity"):
     return m * ac.g0 * positions
 
-def single_gaussian_beam_potential(positions: "np.ndarray|u.quantity.Quantity", waists: "np.ndarray|u.quantity.Quantity", P:"float|u.quantity.Quantity"=1, wavelength:"float|u.quantity.Quantity"=1.064*u.um, alpha:"float|u.quantity.Quantity"=184.4)->np.ndarray:
+def 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:
     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 = - U_tilde * A * np.exp(-2 * ((positions[0,:]/w(positions[1,:], waists[0], wavelength))**2 + (positions[2,:]/w(positions[1,:], waists[1], wavelength))**2)) 
@@ -48,44 +50,79 @@ def single_gaussian_beam_potential(positions: "np.ndarray|u.quantity.Quantity",
 def single_gaussian_beam_potential_harmonic_approximation(positions: "np.ndarray|u.quantity.Quantity", waists: "np.ndarray|u.quantity.Quantity", depth:"float|u.quantity.Quantity"=1, wavelength:"float|u.quantity.Quantity"=1.064*u.um)->np.ndarray:
     U = - depth * (1 - (2 * (positions[0,:]/waists[0])**2) - (2 * (positions[2,:]/waists[1])**2) - (0.5 * positions[1,:]**2 * np.sum(np.reciprocal(z_R(waists, wavelength)))**2))
     return U
+    
+def harmonic_potential(pos, v, offset):
+    U_Harmonic = ((0.5 * 164*u.u * (2 * np.pi * v*u.Hz)**2 * (pos*u.um)**2)/ac.k_B).to(u.uK) + offset*u.uK
+    return U_Harmonic.value
+
+def extractTrapFrequency(Positions, TrappingPotential, TrapDepthInKelvin, axis):
+    tmp_pos = Positions[axis, :]
+    center_idx = np.where(tmp_pos == 0)[0][0]
+    lb = int(round(center_idx - len(tmp_pos)/20, 1))
+    ub = int(round(center_idx + len(tmp_pos)/20, 1))
+    xdata = tmp_pos[lb:ub]
+    tmp_pot = TrappingPotential[axis]
+    Potential = tmp_pot[lb:ub]
+    p0=[1e3, -TrapDepthInKelvin.value]
+    popt, pcov = curve_fit(harmonic_potential, xdata, Potential, p0)
+    v = popt[0]
+    dv = pcov[0][0]**0.5 
+    return v, dv, popt, pcov
+
+def plotHarmonicFit(Positions, TrappingPotential, TrapDepthInKelvin, axis, popt, pcov):
+    v = popt[0]
+    dv = pcov[0][0]**0.5
+    happrox = harmonic_potential(Positions[axis, :].value, *popt)
+    plt.figure()
+    plt.plot(Positions[axis, :].value, happrox, '-r', label = '\u03BD = %.1f \u00B1 %.2f Hz'% tuple([v,dv]))
+    plt.plot(Positions[axis, :], TrappingPotential[axis], 'ob', label = 'Gaussian Potential')
+    plt.xlabel('Distance (um)', fontsize= 12, fontweight='bold')
+    plt.ylabel('Trap Potential (uK)', fontsize= 12, fontweight='bold')
+    plt.ylim([-TrapDepthInKelvin.value, max(TrappingPotential[axis].value)])
+    plt.tight_layout()
+    plt.grid(visible=1)
+    plt.legend(prop={'size': 12, 'weight': 'bold'})
+    plt.show()
 
 def plotPotential(Positions, Powers, ComputedPotentials, axis, TrapDepthLabels):
 
     ## plot of the measured parameter vs. scan  parameter
     plt.figure(figsize=(9, 7))
     for i in range(np.size(ComputedPotentials, 0)):
-        plt.plot(Positions[axis], ComputedPotentials[i][axis], label = 'P = ' + str(Powers[i]) + ' W; ' + TrapDepthLabels[i]) 
+        v, dv, popt, pcov = extractTrapFrequency(Positions, ComputedPotentials[i], TrapDepthInKelvin, axis)
+        unit = 'Hz'
+        if v <= 0:
+            v        = np.nan
+            dv       = np.nan
+            unit = 'Hz'
+        elif v > 0 and orderOfMagnitude(v) > 2:
+            v    = v / 1e3 # in kHz
+            dv   = dv / 1e3 # in kHz
+            unit = 'kHz'
+        tf_label = '\u03BD = %.1f \u00B1 %.2f %s'% tuple([v,dv,unit])
+        plt.plot(Positions[axis], ComputedPotentials[i][axis], label = 'P = ' + str(Powers[i]) + ' W; ' + TrapDepthLabels[i] + '; ' + tf_label) 
     if axis == 0:
         dir = 'X'
     elif axis == 1:
         dir = 'Y'
     else:
         dir = 'Z'
-    # maxPotentialValue   = max(ComputedPotentials.flatten())
-    # minPotentialValue   = min(ComputedPotentials.flatten())
-    # PotentialValueRange = maxPotentialValue - minPotentialValue
-    # upperlimit = 5
-    # if maxPotentialValue > 0:
-    #     upperlimit = maxPotentialValue
-    # lowerlimit = min(ComputedPotentials.flatten()) - PotentialValueRange/6
-    # plt.ylim([lowerlimit, upperlimit])
     plt.xlabel(dir + ' Direction (um)', fontsize= 12, fontweight='bold')
     plt.ylabel('Trap Potential (uK)', fontsize= 12, fontweight='bold')
     plt.tight_layout()
     plt.grid(visible=1)
     plt.legend(prop={'size': 12, 'weight': 'bold'})
-    plt.tight_layout()
+    plt.show()
-    # plt.show()
+    # plt.savefig('pot_' + dir + '.png')
-    plt.savefig('pot_' + dir + '.png')
 
 if __name__ == '__main__':
 
     # Powers = [0.1, 0.5, 2]
     # Powers = [5, 10, 20, 30, 40]
     Powers = [40]
-    Polarizability = 160 # in a.u., should we use alpha = 136 or 160 or 184.4?
+    Polarizability = 184.4 # in a.u, most precise measured value of Dy polarizability
-    # w_x, w_z = 34*u.um, 27.5*u.um # Beam Waists in the x and y directions
+    w_x, w_z = 34*u.um, 27.5*u.um # Beam Waists in the x and y directions
-    w_x, w_z = 35*u.um, 35*u.um # Beam Waists in the x and y directions
+    # w_x, w_z = 70*u.um, 70*u.um # Beam Waists in the x and y directions
     # w_x, w_z = 20.5*u.um, 20.5*u.um
 
     axis = 1 # axis referenced to the beam along which you want the dipole trap potential
@@ -95,11 +132,11 @@ if __name__ == '__main__':
     ComputedPotentials = []
     TrapDepthLabels = []
 
-    gravity = False
+    gravity = True
     astigmatism = False
 
     tilt_gravity = True
-    theta = 0 # in degrees
+    theta = 1 # in degrees
     tilt_axis = [1, 0, 0] # lab space coordinates are rotated about x-axis in reference frame of beam
 
     for p in Powers: 
@@ -111,10 +148,13 @@ if __name__ == '__main__':
         TrapDepthLabels.append("Trap Depth = " + str(round(TrapDepthInKelvin.value, 2)) + " " + str(TrapDepthInKelvin.unit))
 
         projection_axis = np.array([0, 1, 0]) # default
+
         if axis == 0:
             projection_axis = np.array([1, 0, 0]) # radial direction (X-axis)
+        
         elif axis == 1:
             projection_axis = np.array([0, 1, 0]) # propagation direction (Y-axis)
+        
         elif axis == 2:
             projection_axis = np.array([0, 0, 1]) # vertical direction (Z-axis)
         
@@ -138,14 +178,19 @@ if __name__ == '__main__':
             gravity_axis_positions   = np.vstack((x_Positions, y_Positions, z_Positions)) * gravity_axis[:, np.newaxis]
             TrappingPotential        = single_gaussian_beam_potential(Positions, np.asarray([w_x.value, w_z.value])*u.um, P = Power, alpha = Polarizability) + gravitational_potential(gravity_axis_positions, m)
             TrappingPotential        = (TrappingPotential/ac.k_B).to(u.uK)
+            
+        elif not gravity and astigmatism:
+            # Influence of Astigmatism
+            pass
         
+        else:
+            # Influence of Gravity and Astigmatism
+            pass
+        
+        # v, dv, popt, pcov = extractTrapFrequency(Positions, TrappingPotential, TrapDepthInKelvin, axis)
+        # plotHarmonicFit(Positions, TrappingPotential, TrapDepthInKelvin, axis, popt, pcov)
+
         ComputedPotentials.append(TrappingPotential)
     
     ComputedPotentials = np.asarray(ComputedPotentials)
-    plotPotential(Positions, Powers, ComputedPotentials, axis, TrapDepthLabels)
+    plotPotential(Positions, Powers, ComputedPotentials, axis, TrapDepthLabels)
-
-# Influence of Astigmatism
-
-# TrappingPotential = single_gaussian_beam_potential_harmonic_approximation(Positions, np.asarray([w_x.value, w_z.value])*u.um, depth = TrapDepth)
-# TrappingPotential = (TrappingPotential/ac.k_B).to(u.uK)
-