diff --git a/calculateDipoleTrapPotential.py b/calculateDipoleTrapPotential.py
index f35bfe9..6600848 100644
--- a/calculateDipoleTrapPotential.py
+++ b/calculateDipoleTrapPotential.py
@@ -57,10 +57,19 @@ def single_gaussian_beam_potential_harmonic_approximation(positions: "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
+def harmonic_potential(pos, v, offset, m = 164*u.u):
+    U_Harmonic = ((0.5 * m * (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 calculateTrapFrequency(w_x, w_z, Power, Polarizability, m = 164*u.u, dir = 'x'):
+    TrapDepth = trap_depth(w_x, w_z, Power, alpha=Polarizability)
+    TrapFrequency = np.nan
+    if dir == 'x':
+        TrapFrequency = ((1/(2 * np.pi)) * np.sqrt(4 * TrapDepth / (m*w_x**2))).decompose()
+    elif dir == 'z':
+        TrapFrequency = ((1/(2 * np.pi)) * np.sqrt(4 * TrapDepth/ (m*w_z**2))).decompose()
+    return round(TrapFrequency.value, 2)*u.Hz
+
 def extractTrapFrequency(Positions, TrappingPotential, TrapDepthInKelvin, axis):
     tmp_pos = Positions[axis, :]
     center_idx = np.where(tmp_pos == 0)[0][0]
@@ -107,11 +116,14 @@ def plotPotential(Positions, Powers, ComputedPotentials, axis, TrapDepthLabels):
             dv   = dv / 1e3 # in kHz
             unit = 'kHz'
         tf_label = '\u03BD = %.1f \u00B1 %.2f %s'% tuple([v,dv,unit])
-        if i % 2 == 0 and j < len(Powers):
-            plt.plot(Positions[axis], ComputedPotentials[i][axis], '--',label = 'P = ' + str(Powers[j]) + ' W; ' + TrapDepthLabels[j] + '; ' + tf_label) 
-        elif i % 2 != 0 and j < len(Powers):
-            plt.plot(Positions[axis], ComputedPotentials[i][axis], label = 'P = ' + str(Powers[j]) + ' W; ' + tf_label) 
-            j = j + 1
+        if np.size(ComputedPotentials, 0) == len(Powers):
+            plt.plot(Positions[axis], ComputedPotentials[i][axis], label = 'P = ' + str(Powers[i]) + ' W; ' + TrapDepthLabels[i] + '; ' + tf_label) 
+        else:
+            if i % 2 == 0 and j < len(Powers):
+                plt.plot(Positions[axis], ComputedPotentials[i][axis], '--',label = 'P = ' + str(Powers[j]) + ' W; ' + TrapDepthLabels[j] + '; ' + tf_label) 
+            elif i % 2 != 0 and j < len(Powers):
+                plt.plot(Positions[axis], ComputedPotentials[i][axis], label = 'P = ' + str(Powers[j]) + ' W; ' + tf_label) 
+                j = j + 1
     if axis == 0:
         dir = 'X'
     elif axis == 1:
@@ -129,7 +141,7 @@ def plotPotential(Positions, Powers, ComputedPotentials, axis, TrapDepthLabels):
 if __name__ == '__main__':
 
     # Powers = [0.1, 0.5, 2]
-    # Powers = [5, 10, 20, 30, 40]
+    # Powers = [5, 20, 40]
     Powers = [40]
     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
@@ -143,14 +155,14 @@ if __name__ == '__main__':
     ComputedPotentials = []
     TrapDepthLabels = []
 
-    gravity = False
+    gravity = True
     astigmatism = True
 
     tilt_gravity = True
-    theta = 1 # in degrees
+    theta = 5 # in degrees
     tilt_axis = [1, 0, 0] # lab space coordinates are rotated about x-axis in reference frame of beam
 
-    disp_foci = 1.5 * z_R(w_0 = np.asarray([30]), lamb = 1.064)[0]*u.um # difference in position of the foci along the propagation direction (Astigmatism)
+    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)
 
     for p in Powers: 
 
@@ -214,11 +226,13 @@ if __name__ == '__main__':
         else:
             TrappingPotential = IdealTrappingPotential
 
+        # v_x = calculateTrapFrequency(w_x, w_z, Power, Polarizability, dir = 'x')
+        # v_z = calculateTrapFrequency(w_x, w_z, Power, Polarizability, dir = 'z')
+        
         # v, dv, popt, pcov = extractTrapFrequency(Positions, TrappingPotential, TrapDepthInKelvin, axis)
         # plotHarmonicFit(Positions, TrappingPotential, TrapDepthInKelvin, axis, popt, pcov)
-
+        
         ComputedPotentials.append(TrappingPotential)
     
-    # print(np.shape(ComputedPotentials))
     ComputedPotentials = np.asarray(ComputedPotentials)
     plotPotential(Positions, Powers, ComputedPotentials, axis, TrapDepthLabels)
\ No newline at end of file