LennartNaeve_code/spilling_code/harmonic_approximation/ODT.py
2025-04-25 20:52:11 +02:00

187 lines
8.9 KiB
Python

import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
from scipy.signal import find_peaks
from astropy import units as u, constants as ac
class ODP_1D:
"""
1D optical dipole trap in prop_dir direction with gravity in z direction and
magnetic field in mag_dir direction
"""
def __init__(self, options):
self.prop_dir = options["prop_dir"]
self.wavelength = options["wavelength"]
self.alpha = options["alpha"]*(4 * np.pi * ac.eps0 * ac.a0**3)/(2 * ac.eps0 * ac.c) # in atomic units
self.m = options["dy_mass"]
self.mu = options["dy_mag_moment"]
self.power = options["power"] # laser power
self.w_0 = options["w_0"] # beam waist in transverse direction
self.gravity = options["gravity"] # boolean for gravity in z direction
self.mag_dir = options["mag_dir"] # direction of magnetic field and gradient
self.mag_grad = options["mag_grad"] # magnetic gradient
self.mag_offset = options["mag_offset"] # magnetic field offset
self.mJ = options["mJ"] #magnetic quantum number
self.ref_ind = options["ref_ind"] #refractive index
#self.z_r = (np.pi * self.w_0**2 * self.ref_ind / self.wavelength).to(u.um) #rayleigh range
#self.nu_r = (np.sqrt(2*self.power*self.alpha/np.pi/self.m) * 2/self.w_0**2).to(u.Hz)/(2*np.pi)
#self.nu_z = (np.sqrt(4*self.power*self.alpha/np.pi**3/self.m) * self.wavelength/self.w_0**3/self.ref_ind).to(u.Hz)/(2*np.pi)
@property
def z_r(self):
return (np.pi * self.w_0**2 * self.ref_ind / self.wavelength).to(u.um) #rayleigh range
@property
def nu_r(self):
return (np.sqrt(2*self.power*self.alpha/np.pi/self.m) * 2/self.w_0**2).to(u.Hz)/(2*np.pi)
@property
def nu_z(self):
return (np.sqrt(4*self.power*self.alpha/np.pi**3/self.m) * self.wavelength/self.w_0**3/self.ref_ind).to(u.Hz)/(2*np.pi)
def get_axial_potential(self,z):
"""
input: z(array) positions to evaluate potential with unit of length
output: V(z)(array) potential energies at z
"""
beam_pot = (-self.alpha * (2*self.power/np.pi/self.w_0**2) / (1+(z/self.z_r)**2)).to(u.J)
#check if magnetic field is along axial direction
if self.mag_dir == self.prop_dir:
mag_grad_pot = (self.mu * self.mag_grad * z).to(u.J)
mag_offset_pot = (self.mu * self.mag_offset * self.mJ).to(u.J)
else:
mag_grad_pot = 0
mag_offset_pot = 0
#check if gravity is on and axial direction is along z
if self.gravity and (self.prop_dir == "v"):
gravity_pot = (self.m * ac.g0 * z).to(u.J)
else:
gravity_pot = 0
return beam_pot + gravity_pot + mag_grad_pot + mag_offset_pot
def get_radial_potential(self,r):
"""
input: r(array) positions to evaluate potential with unit of length
output: V(r)(array) potential energies at z
"""
beam_pot = (-self.alpha * (2*self.power/np.pi/self.w_0**2) * np.exp(-2*r**2/self.w_0**2))
#check if magnetic field is along radial direction
if self.mag_dir != self.prop_dir:
mag_grad_pot = (self.mu * self.mag_grad * r).to(u.J)
mag_offset_pot = (self.mu * self.mag_offset * self.mJ).to(u.J)
else:
mag_grad_pot = 0
mag_offset_pot = 0
#check if gravity is on and axial direction is horizontal
if self.gravity and (self.prop_dir == "h"):
gravity_pot = (self.m * ac.g0 * r).to(u.J)
else:
gravity_pot = 0
return beam_pot + gravity_pot + mag_grad_pot + mag_offset_pot
def harm_pot(self, z, nu, V_0, z_0):
return V_0 + 0.5*self.m.to(u.kg).value*(2*np.pi*nu)**2 * (z-z_0)**2
def get_axial_occ_analytic(self,z,plot=False):
"""
Finds the number of occupied energy states that are bound in the potential well,
using the analytic harmonic approximation of the gaussian beam.
"""
pot = self.get_axial_potential(z)
if plot:
plt.plot(z,(pot/ac.k_B).to(u.uK), label='actual potential')
#find local minimum and maximum of potential
try:
E_max = pot[find_peaks(pot)[0]]
except:
raise Exception("potential maximum not found")
try:
ind_min = find_peaks(-pot)[0]
E_min = pot[ind_min]
z_min = z[ind_min][0]
except:
raise Exception("potential minimum not found")
delta_E = E_max - E_min
#get eigenenergies of harmonic oscillator as temperatures
n_max = int(np.ceil((delta_E/(ac.hbar*2*np.pi*self.nu_z)-0.5).value))
n_temps = np.arange(n_max)
temps = ((ac.hbar*2*np.pi*self.nu_z*(n_temps +0.5)+E_min)/ac.k_B).to(u.uK).value
#temps = np.array([((ac.hbar*2*np.pi*self.nu_z*(n +0.5)+E_min)/ac.k_B).to(u.uK).value for n in range(n_max)])
#plot results
if plot:
#plot energy levels
plt.hlines(temps, np.min(z).to(u.um).value, np.max(z).to(u.um).value, linewidth=0.1,color="black", label="harmonic eigenenergies")
#plot harmonic approximation
#harm_approx = self.harm_pot(z.to(u.m).value, self.nu_z.to(u.Hz).value,E_min.to(u.J).value,z_min.to(u.m).value)*u.J
z_plot = np.array([zpt for zpt in z.to(u.m).value if self.harm_pot(zpt, self.nu_z.to(u.Hz).value,E_min.to(u.J).value,z_min.to(u.m).value)*u.J < np.max(pot)])*u.m
plt.plot(z_plot.to(u.um).value,(self.harm_pot(z_plot.to(u.m).value, self.nu_z.to(u.Hz).value,E_min.to(u.J).value,z_min.to(u.m).value)*u.J/ac.k_B).to(u.uK),linestyle="--", label="harmonic approximation")
plt.xlabel("z position [um]")
plt.ylabel("trap depth [uK]")
plt.title(f"power={self.power}, w_0={self.w_0}, n={n_max}, nu={self.nu_z}")
plt.legend()
plt.grid()
plt.show()
return n_max
def get_axial_occ(self, z,fit_fac=1,p0=[1e8,-1.5e-27,0],plot=True,plot_fac=1):
"""
#input: z(array) positions
fit_fac(float) factor to make the interval for fitting the parabola bigger
p0(list) initial guess for the fit
plot(bool) enable/disable plotting
plot_fac(float) factor to make the plot interval for parabola bigger
#output: n_max(int) number of modes that fit into the trap
(nu(float) eigenfrequency)
"""
#get potential
pot = self.get_axial_potential(z)
if plot:
plt.plot(z,(pot/ac.k_B).to(u.uK), label='actual potential')
#get z values and potential in region to be fitted
#z_min = z[find_peaks(-pot)[0]][0]
#interval=(z_min - self.w_0, z_min + self.w_0)
#z_fit = np.array([zpt for zpt in z.to(u.m).value if interval[0].to(u.m).value< zpt <interval[1].to(u.m).value])
#z_fit = np.linspace(z_min - self.w_0, z_min + self.w_0,100).to(u.m).value
z_fit = np.linspace((-self.w_0*fit_fac),(+self.w_0*fit_fac),100).to(u.m).value
pot_fit = self.get_axial_potential(z_fit*u.m).to(u.J).value
#perform fit
try:
popt, pcov = curve_fit(self.harm_pot, z_fit,pot_fit,p0=p0)
nu = popt[0]*u.Hz
print(nu)
V_0 = popt[1]*u.J
#get eigenenergies of harmonic oscillator as temperatures
#find local maximum
E_max = np.abs(pot[find_peaks(pot)[0]] - V_0)
n_max = int(np.ceil((E_max/(ac.hbar*2*np.pi*nu)-0.5).value))
temps = np.array([((ac.hbar*2*np.pi*nu*(n +0.5)+V_0)/ac.k_B).to(u.uK).value for n in range(n_max)])
#print(f"{n_max} modes of trap frequncy {nu} were found in potential well.")
#plot results
if plot:
plt.hlines(temps, np.min(z).to(u.um).value, np.max(z).to(u.um).value, linewidth=0.1,color="black", label="harmonic eigenenergies")
z_plot = z_fit*u.m*plot_fac #np.linspace(interval[0],interval[1],100)
plt.plot(z_plot.to(u.um).value,(self.harm_pot(z_plot.to(u.m).value, *popt)*u.J/ac.k_B).to(u.uK),linestyle="--", label="harmonic approximation")
plt.xlabel("z position [um]")
plt.ylabel("trap depth [uK]")
plt.title(f"power={self.power}, w_0={self.w_0}, n={n_max}")
plt.legend()
plt.grid()
plt.show()
except:
raise Exception("fit was unsuccessful :(")
#plot potential for trouble shooting
if plot:
plt.xlabel("z position [um]")
plt.ylabel("trap depth [uK]")
plt.title(f"power={self.power}, w_0={self.w_0}, n={n_max}")
plt.legend()
plt.grid()
plt.show()
return n_max
########################