lhoenen/Trapping_Simulation
1
0
Fork 0
Code Issues Pull Requests Actions Packages repo.project_board Releases Wiki Activity

Uploaded initial "working" version of PyLCP simulation.

Browse Source 
This is not well structured yet. Nor is it well tested or commented.
This commit is contained in:
lhoenen 2025-04-16 18:47:42 +02:00
commit a1cd30886d
4 changed files with 1261 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

62
Dy_atom.py Normal file
View File

@ -0,0 +1,62 @@
from pylcp.atom import atom as base_atom
from pylcp.atom import state
import numpy as np
import scipy.constants as cts
class DyAtom(base_atom):
def __init__(self, species):
# Prepare to add in some useful electronic states:
self.state = []
if species == "161Dy" or species == "Dy161":
self.I = 5/2 # nuclear spin
self.gI = -0.1922 # nuclear magnetic moment
self.mass = 160.93*cts.value('atomic mass constant')
# Ground state:
self.state.append(state(n=6, L=6, J=8, S=2, lam=np.inf, tau=np.inf,
gJ=1.2415867, Ahfs=-116.2322e6,
Bhfs =1091.5748e6, Chfs=0))
#Blue Transition:
self.state.append(state(n=6, L=7, J=9, S=2, lam=421.291e-9,
tau=4.94e-9, gJ=1.22, Ahfs=-86.9e6,
Bhfs=1747.4e6, Chfs=0))
elif species == "163Dy" or species == "Dy163":
self.I = 5/2 # nuclear spin
self.gI = 0.269#-0.192 # nuclear magnetic moment
self.mass = 162.93*cts.value('atomic mass constant')
# Ground state:
self.state.append(state(n=6, L=6, J=8, S=2, lam=np.inf, tau=np.inf,
gJ=1.2415867, Ahfs=162.7543e6,
Bhfs =1153.8684e6, Chfs=0))
#Blue Transition:
self.state.append(state(n=6, L=7, J=9, S=2, lam=421.291e-9,
tau=4.94e-9, gJ=1.22, Ahfs=121.62e6,
Bhfs=1844.9e6, Chfs=0))
elif species == "164Dy" or species == "Dy164":
self.I = 0 # nuclear spin
self.gI = 0 # nuclear magnetic moment
self.mass = 163.929*cts.value('atomic mass constant')
# Ground state:
self.state.append(state(n=6, L=6, J=8, S=2, lam=np.inf, tau=np.inf,
gJ=1.25))
#Blue Transition:
self.state.append(state(n=6, L=7, J=9, S=2, lam=421.291e-9,
tau=4.94e-9, gJ=1.22))
#Red Transition:
self.state.append(state(n=6, J=9, S=2, gJ=1.29, lam=626e-9, tau=1200e-9))
else:
# Fallback to normal atom behavior
super().__init__(species)
# Take the states and make transitions:
base_atom._atom__make_transitions(self)

256
PlotStuff.py Normal file
View File

@ -0,0 +1,256 @@
import pickle
import os
import matplotlib.pyplot as plt
from pprint import pprint
import numpy as np
from matplotlib.ticker import ScalarFormatter
import common
from datetime import date
"""
These functions plot stuff. One can but does not have to use them.
"""
def plot_state_densities_and_force_field_with_trajectories(Isotope, file_name, date_str = date.today().isoformat()):
folder_name = "Data\\"+date_str
sols_folder = folder_name+'\Solutions'
other_data_folder = folder_name+"\OtherData\\"
other_data_list = []
for filename in os.listdir(other_data_folder):
if Isotope in filename and file_name in filename and filename.endswith(".npz"):
other_data_list.append(np.load(other_data_folder+filename))
print(other_data_folder+filename)
#other_data_list.append(np.load("Z:\Users\Lenny\Lenny\ (Master)\Simulations\ and\ other\ Code\AtomLightSimulations\CleanedUpSimulation\Data\\2025-04-07\OtherData\data_dir_161into2DMOT.npz"))
if len(other_data_list) > 1:
print("There were multiple 'Other Data' Files in the directory,"
"that belong to you Isotope. The list contains:\n", other_data_list)
elif len(other_data_list) == 0:
print("Couldn't find a forcefield to plot. List seems to be empty.")
other_data_dir = other_data_list[0]
Sols = {}
for filename in os.listdir(sols_folder):
if Isotope in filename and file_name in filename and filename.endswith(".pkl"):
key = os.path.splitext(filename)[0] # Remove .pkl from filename
full_path = os.path.join(sols_folder, filename)
with open(full_path, 'rb') as f:
Sols[key] = pickle.load(f)
basis_g = other_data_dir['basis_g']
basis_e = other_data_dir['basis_e']
x0 = other_data_dir['x0']/100
t0 = other_data_dir['t0']
for names in Sols:
sol = Sols[names]
data = sol.N
num_states = data.shape[0]
num_timesteps = data.shape[1]
v0 = sol.v[0,0]
# Normalize (optional)
#data = data / data.max()
# Assume your data is in `data` with shape (timesteps, num_states)
data_g = np.transpose(data[:len(basis_g), :]) # First g states
data_e = np.transpose(data[len(basis_g):, :]) # Last e states
fig, axes = plt.subplots(1, 2, figsize=(14, 6), sharey=True)
# Plot System 1
im1 = axes[0].imshow(data_g, aspect='auto', origin='lower',
cmap='viridis', interpolation='none')
axes[0].set_title('J = 8')
axes[0].set_xlabel('State Index')
axes[0].set_ylabel('position x[m]')
state_labels_str = [f"{int(s[0])},{int(s[1])}" for s in basis_g[::]]
axes[0].set_xticks(np.arange(len(basis_g)))
axes[0].set_xticklabels(state_labels_str, rotation=90, fontsize=8)
indices = np.linspace(0, len(sol.t) - 1, 40, dtype=int) #write 40 labels
position_label = sol.r[0]*x0
pretty_position_label = [f"{x:.3f}" for x in position_label[indices]]
axes[0].set_yticks(np.arange(len(sol.t))[indices])
axes[0].set_yticklabels(pretty_position_label, rotation=0, fontsize=8)
fig.colorbar(im1, ax=axes[0], label='State Density')
# Plot System 2
im2 = axes[1].imshow(data_e, aspect='auto', origin='lower', cmap='viridis', interpolation='none')
axes[1].set_title('J = 9')
axes[1].set_xlabel('State Index')
state_labels_str = [f"{int(s[0])},{int(s[1])}" for s in basis_e[::-1]]
axes[1].set_xticks(np.arange(len(basis_e)))
axes[1].set_xticklabels(state_labels_str, rotation=90, fontsize=8)
fig.colorbar(im2, ax=axes[1], label='State Density')
fig.suptitle("Density of mF states at v_0="+str(v0*(x0/t0))+"mps")
plt.tight_layout(pad=2.0)
plotname = "StateDensityDy"+Isotope+"at"+str(round(v0*(x0/t0)))+"mps"
filedir = common.get_data_filepath("Plots", plotname, ".png")
plt.show()
#plt.savefig(filedir, dpi=200)
plt.close()
Fx_RE = other_data_dir['acceleration_profile']
x = other_data_dir['x']
v = other_data_dir['v']
X,V = np.meshgrid(x, v)
formatter = ScalarFormatter(useMathText=True)
formatter.set_powerlimits((-2, 2)) # Forces scientific notation outside this range
fig, ax0 = plt.subplots(nrows=1, ncols=1)
im0 = ax0.pcolormesh(X*x0, V*x0/t0, np.nan_to_num(Fx_RE))#, cmap="hot")
cb0 = plt.colorbar(im0)
cb0.ax.yaxis.set_major_formatter(formatter)
cb0.set_label('$a [m/s^2]$')
ax0.set_xlabel('$x[m]$')
ax0.set_ylabel('$v[m/s]$')
ax0.set_title("Acceleration from rate equaitons")
for names in Sols:
sol = Sols[names]
ax0.plot(sol.r[0]*x0, sol.v[0]*x0/t0, "w")
ax0.set_xlim(x[0]*x0, x[-1]*x0)
plt.tight_layout()
plotname = "PhaseSpaceDy"+Isotope
filedir = common.get_data_filepath("Plots", plotname, ".png")
plt.show()
#plt.savefig(filedir, dpi=200)
plt.close()
def plot_force_field_and_trajectories(Isotope, file_name, date_str = date.today().isoformat()):
folder_name = "Data\\"+date_str
sols_folder = folder_name+'\Solutions'
other_data_folder = folder_name+"\OtherData\\"
other_data_list = []
for filename in os.listdir(other_data_folder):
if Isotope in filename and file_name in filename and filename.endswith(".npz"):
other_data_list.append(np.load(other_data_folder+filename))
print(other_data_folder+filename)
#other_data_list.append(np.load("Z:\Users\Lenny\Lenny\ (Master)\Simulations\ and\ other\ Code\AtomLightSimulations\CleanedUpSimulation\Data\\2025-04-07\OtherData\data_dir_161into2DMOT.npz"))
if len(other_data_list) > 1:
print("There were multiple 'Other Data' Files in the directory,"
"that belong to you Isotope. The list contains:\n", other_data_list)
elif len(other_data_list) == 0:
print("Couldn't find a forcefield to plot. List seems to be empty.")
other_data_dir = other_data_list[0]
#other_data_dir = np.load("Data\\2025-04-14\OtherData\data_dir_Dy161RijIntoPushBeam0Gauss0DetPol-1.npz")
Sols = {}
for filename in os.listdir(sols_folder):
if Isotope in filename and file_name in filename and filename.endswith(".pkl"):
key = os.path.splitext(filename)[0] # Remove .pkl from filename
full_path = os.path.join(sols_folder, filename)
with open(full_path, 'rb') as f:
Sols[key] = pickle.load(f)
basis_g = other_data_dir['basis_g']
basis_e = other_data_dir['basis_e']
x0 = other_data_dir['x0']/100
t0 = other_data_dir['t0']
Fx_RE = other_data_dir['acceleration_profile']
x = other_data_dir['x']
v = other_data_dir['v']
X,V = np.meshgrid(x, v)
formatter = ScalarFormatter(useMathText=True)
formatter.set_powerlimits((-2, 2)) # Forces scientific notation outside this range
fig, ax0 = plt.subplots(nrows=1, ncols=1)
im0 = ax0.pcolormesh(X*x0, V*x0/t0, np.nan_to_num(Fx_RE))
cb0 = plt.colorbar(im0)
cb0.ax.yaxis.set_major_formatter(formatter)
cb0.set_label('$a [m/s^2]$')
ax0.set_xlabel('$x[m]$')
ax0.set_ylabel('$v[m/s]$')
ax0.set_title("Acceleration from rate equaitons")
for names in Sols:
sol = Sols[names]
ax0.plot(sol.r[0]*x0, sol.v[0]*x0/t0, "w")
ax0.set_xlim(x[0]*x0, x[-1]*x0)
plt.tight_layout()
plotname = "PhaseSpaceDy"+Isotope
filedir = common.get_data_filepath("Plots", plotname, ".png")
plt.show()
#plt.savefig(filedir, dpi=200)
plt.close()
def plot_force_field(Isotope, file_name, date_str = date.today().isoformat()):
folder_name = "Data\\"+date_str
sols_folder = folder_name+'\Solutions'
other_data_folder = folder_name+"\OtherData\\"
other_data_list = []
for filename in os.listdir(other_data_folder):
if file_name in filename and Isotope in filename and filename.endswith(".npz"):
other_data_list.append(np.load(other_data_folder+filename))
print(other_data_folder+filename)
#other_data_list.append(np.load("Z:\Users\Lenny\Lenny\ (Master)\Simulations\ and\ other\ Code\AtomLightSimulations\CleanedUpSimulation\Data\\2025-04-07\OtherData\data_dir_161into2DMOT.npz"))
if len(other_data_list) > 1:
print("There were multiple 'Other Data' Files in the directory,"
"that belong to you Isotope. The first entry will be chosen. The list contains:\n", other_data_list)
elif len(other_data_list) == 0:
print("Couldnt read a file in 'Other Data'. Check that you chose the correct date, Isotope and filename.")
other_data_dir = other_data_list[0]
basis_g = other_data_dir['basis_g']
basis_e = other_data_dir['basis_e']
x0 = other_data_dir['x0']/100
t0 = other_data_dir['t0']
Fx_RE = other_data_dir['acceleration_profile']
x = other_data_dir['x']
v = other_data_dir['v']
X,V = np.meshgrid(x, v)
formatter = ScalarFormatter(useMathText=True)
formatter.set_powerlimits((-2, 2)) # Forces scientific notation outside this range
fig, ax0 = plt.subplots(nrows=1, ncols=1)
im0 = ax0.pcolormesh(X*x0, V*x0/t0, np.nan_to_num(Fx_RE))
cb0 = plt.colorbar(im0)
cb0.ax.yaxis.set_major_formatter(formatter)
cb0.set_label('$a [m/s^2]$')
ax0.set_xlabel('$x[m]$')
ax0.set_ylabel('$v[m/s]$')
ax0.set_title("Acceleration from rate equaitons")
plt.tight_layout()
plotname = "PhaseSpaceDy"+Isotope
filedir = common.get_data_filepath("Plots", plotname, ".png")
plt.show()
#plt.savefig(filedir, dpi=200)
plt.close()

679
common.py Normal file
View File

@ -0,0 +1,679 @@
import numpy as np
import matplotlib.pyplot as plt
import scipy.constants as cts
from pprint import pprint
import os
from datetime import date
import pylcp
import pickle
import custom_rateeq
import Dy_atom
from sympy import Rational
from pylcp.common import progressBar
"""
These are some functions common to all applications. These will be used for each run. Some of them are just for bookkeeping (checked_save() for example),
some of them are setting up the simulations for each run (init_hamiltonian...() for example) and some define specific simulation runs.
These will have to be extended for each case (Zeeman Slower / MOT / Mollasses etc.)
"""
def init_hamiltonian_singleF(F0, F1, gF0, gF1, mass, muB=1, gamma=1, k=1, returnBasis=False):
"""
Create Hamiltonian for a two-level sytem of states F0 and F1
Parameters
----------
F0 : int or float
Total angular momentum F of the ground state
F1 : int or float
Total angular momentum F of the excited state
gF0 : float
(F-)Landé g-factor of the ground state
gF1 : float
(F-)Landé g-factor of the excited state
mass : float
(unitless) Mass of the atom
muB : float
(unitless) Bohr magneton (default: 1)
gamma : float
(unitless) Natural linewidth of the excited state (default: 1)
k : float
(unitless) Wavenumber of the driving laser field (default: 1)
returnBasis : bool, optional
Whether to return the basis states used for the ground and excited
state Hamiltonians (default: False)
Returns
-------
Ham : pylcp.hamiltonian
Hamiltonian for the 2 level system, ready to use in simulations
E_g : np.ndarray
Energies (eigenvalues) of the ground state
E_e : np.ndarray
Energies (eigenvalues) of the excited state
basis_g : list (only if returnBasis is True)
Basis states used for ground state Hamiltonian
basis_e : list (only if returnBasis is True)
Basis states used for excited state Hamiltonian
"""
H_g, mu_q_g, basis_g = pylcp.hamiltonians.singleF(F=F0, gF=gF0, muB=muB, return_basis=returnBasis)
H_e, mu_q_e, basis_e = pylcp.hamiltonians.singleF(F=F1, gF=gF1, muB=muB, return_basis=returnBasis)
d_ijq = pylcp.hamiltonians.dqij_two_bare_hyperfine(F0, F1)
Ham = pylcp.hamiltonian(H_g, H_e, mu_q_g, mu_q_e, d_ijq, mass = mass, muB=muB, gamma=gamma, k=k)
E_g = np.unique(np.diagonal(H_g))
E_e = np.unique(np.diagonal(H_e))
if returnBasis:
return Ham, E_g, E_e, basis_g, basis_e
else:
return Ham, E_g, E_e
def init_hamiltonian_hyperfine_coupled(Isotope, mass, t0, muB=1, gamma=1, k=1, state0=0, state1=1, returnBasis=False):
"""
Create Hamiltonian for a two-level atom including hyperfine splitting
Parameters
----------
Isotope : object
Atom object from PyLCP.atom class with nuclear magnetic moment, mass, states and transitions
mass : float
(unitless) Mass of the atom
t0 : float, optional
time unit, typically 1/Gamma
muB : float, optional
(unitless) Bohr magneton (default: 1)
gamma : float, optional
(unitless) Natural linewidth of the excited state (default: 1)
k : float, optional
(unitless) Wavenumber of the driving laser field (default: 1)
state0 : int, optional
Index of the ground state in Isotope.state (default: 0)
state1 : int, optional
Index of the excited state in Isotope.state (default: 1)
returnStates : bool, optional
Whether to return the basis states used in the Hamiltonians (default: False)
Returns
-------
Hamiltonian : pylcp.hamiltonian
Hamiltonian for the two-level hyperfine system
E_g : np.ndarray
Energies of the ground state manifold
E_e : np.ndarray
Energies of the excited state manifold
basis_g : list (only if returnStates is True)
Basis states used for ground state Hamiltonian
basis_e : list (only if returnStates is True)
Basis states used for excited state Hamiltonian
"""
H_g, mu_q_g, basis_g = pylcp.hamiltonians.hyperfine_coupled(
Isotope.state[state0].J, Isotope.I, Isotope.state[state0].gJ, Isotope.gI,
Ahfs=Isotope.state[state0].Ahfs/Isotope.state[1].gammaHz,
Bhfs=Isotope.state[state0].Bhfs/Isotope.state[1].gammaHz, Chfs=0, muB=muB, return_basis=True)
H_e, mu_q_e, basis_e = pylcp.hamiltonians.hyperfine_coupled(
Isotope.state[state1].J, Isotope.I, Isotope.state[state1].gJ, Isotope.gI,
Ahfs=Isotope.state[state1].Ahfs/Isotope.state[1].gammaHz,
Bhfs=Isotope.state[state1].Bhfs/Isotope.state[1].gammaHz, Chfs=0, muB=muB, return_basis= True)
dijq = pylcp.hamiltonians.dqij_two_hyperfine_manifolds(
Isotope.state[state0].J, Isotope.state[state1].J, Isotope.I)
E_g = np.unique(np.diagonal(H_g))
E_e = np.unique(np.diagonal(H_e))
Hamiltonian = pylcp.hamiltonian(H_g, H_e, mu_q_g, mu_q_e, dijq, mass = mass, muB=muB, gamma=gamma, k=k)
if returnBasis:
return Hamiltonian, E_g, E_e, basis_g, basis_e
else:
return Hamiltonian, E_g, E_e
def init_LaserBeams2DMOT(detuning, beam_waist, aperture, k, saturation):
laserBeams = pylcp.fields.laserBeams(
[
{
"kvec": k * np.array([+1, +1, 0.0])/np.sqrt(2),
"s": saturation,
"pol": -1,
"delta": detuning ,
"wb": beam_waist,
"rs": aperture,
},
{
"kvec": k * np.array([-1, -1, 0.])/np.sqrt(2),
"s": saturation,
"pol": -1,
"delta": detuning,
"wb": beam_waist,
"rs": aperture,
},
{
"kvec": k * np.array([1, -1, 0.0])/np.sqrt(2),
"s": saturation,
"pol": 1,
"delta": detuning,
"wb": beam_waist,
"rs": aperture,
},
{
"kvec": k * np.array([-1, 1,0.])/np.sqrt(2),
"s": saturation,
"pol": 1,
"delta": detuning,
"wb": beam_waist,
"rs": aperture,
}
],
beam_type=pylcp.fields.clippedGaussianBeam,
)
return laserBeams
def init_LaserBeamSingleZeeman(detuning, saturation, polarisation, beam_waist, aperture, k=1):
laserBeam = pylcp.fields.laserBeams(
[
{
"kvec": k * np.array([-1., 0., 0.0]),
"s": saturation,
"pol": polarisation,
"delta": detuning ,
"wb": beam_waist,
"rs": aperture,
}
],
beam_type=pylcp.fields.clippedGaussianBeam,
)
return laserBeam
def init_1DMOTBeams(detuning, saturation, polarisation, beam_waist, aperture, k=1):
laserBeam = pylcp.fields.laserBeams(
[
{
"kvec": k * np.array([-1., 0., 0.0]),
"s": saturation,
"pol": polarisation,
"delta": detuning ,
"wb": beam_waist,
"rs": aperture,
},
{
"kvec": k * np.array([1., 0., 0.0]),
"s": saturation,
"pol": polarisation,
"delta": detuning ,
"wb": beam_waist,
"rs": aperture,
}
],
beam_type=pylcp.fields.clippedGaussianBeam,
)
return laserBeam
def get_data_filepath(subfolder, base_name, extension):
"""
Returns a unique file path inside the correct daily Data folder structure.
If a file already exists under the same name an index will be appended to avoid overwriting files.
Parameters:
- subfolder: one of ['Plots', 'OtherData', 'Solutions']
- base_name: the desired base file name (without extension)
- extension: file extension including dot, e.g. '.npz'
Returns:
- full_path: a unique full path like Data/2025-03-17/OtherData/RateEqu_1.npz
"""
assert subfolder in ['Plots', 'OtherData', 'Solutions'], "Invalid subfolder name!"
# Create today's date string
date_str = date.today().isoformat()
# Build full folder path
base_dir = os.path.join(".\Data", date_str)
subfolder_path = os.path.join(base_dir, subfolder)
# Create directories if needed
os.makedirs(subfolder_path, exist_ok=True)
# Generate filename
filename = f"{base_name}{extension}"
full_path = os.path.join(subfolder_path, filename)
# Make filename unique
counter = 1
while os.path.exists(full_path):
filename = f"{base_name}_{counter}{extension}"
full_path = os.path.join(subfolder_path, filename)
counter += 1
return full_path
def get_data_filepath_and_check(subfolder, base_name, extension):
"""
Returns a unique file path inside the correct daily Data folder structure.
Checks if a file already excists and if yes waits for a user input.
User can choose to replace existing or tp add an index to the path.
Parameters:
- subfolder: one of ['Plots', 'OtherData', 'Solutions']
- base_name: the desired base file name (without extension)
- extension: file extension including dot, e.g. '.npz'
Returns:
- full_path: a unique full path like Data/2025-03-17/OtherData/RateEqu_1.npz
"""
assert subfolder in ['Plots', 'OtherData', 'Solutions'], "Invalid subfolder name!"
# Create today's date string
date_str = date.today().isoformat()
# Build full folder path
base_dir = os.path.join(".\Data", date_str)
subfolder_path = os.path.join(base_dir, subfolder)
# Create directories if needed
os.makedirs(subfolder_path, exist_ok=True)
# Generate filename
filename = f"{base_name}{extension}"
full_path = os.path.join(subfolder_path, filename)
if os.path.exists(full_path):
print(f"Warning: A file named '{filename}' already exists in the directory.")
user_input = input("Do you want to replace the existing file? (yes/no): ").strip().lower()
if user_input == 'yes':
print("You chose to overwrite a file.")
return True, full_path # Return True to indicate overwrite permission
else:
# Make filename unique
counter = 1
while os.path.exists(full_path):
filename = f"{base_name}_{counter}{extension}"
full_path = os.path.join(subfolder_path, filename)
counter += 1
return False, full_path
else:
return False, full_path # As the user did not explicitly agree to overwrite, return False.
def checked_save(check, path, data):
if check:
# Overwrite the file if check is True
if path.endswith(".npz"):
np.savez(path, **data)
print(f"File '{os.path.basename(path)}' has been saved (overwrite allowed).")
elif path.endswith(".pkl"):
with open(path, 'wb') as f:
pickle.dump(data, f)
print(f"File '{os.path.basename(path)}' has been saved (overwrite allowed).")
else:
print("File type couldnt be recognised. Can only be '.pkl' or '.npz'.")
else:
# Ensure the file does not exist before saving
if not os.path.exists(path):
if path.endswith(".npz"):
np.savez(path, **data)
print(f"File saved as '{os.path.basename(path)}'.")
elif path.endswith(".pkl"):
with open(path, 'wb') as f:
pickle.dump(data, f)
print(f"File saved as '{os.path.basename(path)}'.")
else:
print("File type couldnt be recognised. Can only be '.pkl' or '.npz'.")
else:
print(f"Error: Cannot save '{os.path.basename(path)}' because overwrite is not allowed.")
def calc_force_profile_Single_Beam(Isotope_name, params, resolution_x, resolution_v, basename):
check1, path1 = get_data_filepath_and_check('OtherData', 'data_dir_'+Isotope_name+basename, '.npz')
pylcp.rateeq.evolve_motion = custom_rateeq.evolve_motion #replace the evolve motion funciton with custom
Isotope = Dy_atom.DyAtom(Isotope_name)
k_lab = Isotope.transition[0].k # 1/cm not 2*pi/cm
gamma_lab = Isotope.state[1].gammaHz # Hz This is 32.2MHz, not 2*pi*32.2MHz
# Define "natural units"
x0 = 1/(k_lab) # cm
t0 = 1/gamma_lab # s
muB = 1
k = k_lab*x0
gamma = gamma_lab*t0
mass = Isotope.mass*(x0/100)**2/(cts.h * t0)
det = params["det"]
s = params["s"]
pol = params["polarisation"]
field_mag = params["field_mag_Gauss_cm"]*1e-4
field_type = params["field_type"]
waist = params["waist_m"]/(x0/100)
aperture = params["aperture_m"]/(x0/100)
trap_length = params["trap_length_m"] #m
zrange = trap_length/(x0/100)
velocity = params["max_velocity"] # m/s
vrange = velocity/((x0/100)/t0)
v0s = params["initial_velocities_m"]/((x0/100)/t0)
Ham, E_g, E_e , basis_g, basis_e = init_hamiltonian_hyperfine_coupled(Isotope, mass, t0, muB=muB, gamma=gamma, k=k, returnBasis=True)
basis_g = np.transpose(basis_g)
basis_e = np.transpose(basis_e)
if Isotope_name == "Dy161":
det_L = (E_e[0] - E_g[0]) + det
elif Isotope_name == "Dy163":
det_L = (E_e[-1] - E_g[-1]) + det
elif Isotope_name == "Dy164":
det_L = det
else:
raise(ValueError(f"Couldn't recognise Isotope name {Isotope_name}"))
LaserBeams = init_LaserBeamSingleZeeman(detuning=det_L, saturation=s, polarisation=pol, beam_waist=waist, aperture=aperture, k=k)
if field_type == "2DMOT":
alpha_2DMOT = field_mag * cts.value('Bohr magneton') * t0 * x0 / cts.h #field mag is given as Gauss/cm
magField = pylcp.fields.magField(lambda R: [alpha_2DMOT*R[1],alpha_2DMOT*R[0],0])
elif field_type == "Gradient":
alpha_gradient = field_mag * cts.value('Bohr magneton') * t0 * x0 / cts.h #field mag is given as Gauss/cm
magField = pylcp.fields.magField(lambda R: [alpha_gradient*R[0],0,0])
elif field_type == "Offset":
alpha_offset = field_mag * cts.value('Bohr magneton') * t0 / cts.h #field mag is given as Gauss/cm
magField = pylcp.fields.constantMagneticField([alpha_offset,0,0])
else:
raise ValueError("Error: Unrecognised field type. Has to be either '2DMOT' or 'Gradient' or 'Offset'.")
x = np.linspace(-zrange, zrange, resolution_x)
v = np.linspace(-vrange, vrange, resolution_v)
X, V = np.meshgrid(x, v)
Rvec = np.array([X, np.zeros(X.shape), np.zeros(X.shape)])
Vvec = np.array([V, np.zeros(V.shape), np.zeros(V.shape)])
print("Calculating rate equations:")
Rateeq = pylcp.rateeq(LaserBeams, magField, Ham, include_mag_forces=True)
print("Completed. \nCalculating Force Profile:")
Rateeq.set_initial_pop(np.concatenate([np.ones(len(basis_g))/len(basis_g), np.zeros(len(basis_e))]))
accel_profile = np.full((resolution_v, resolution_x), np.nan)
progress = progressBar()
count = 0
num = resolution_v*resolution_x
for i_v in range(resolution_v):
for i_x in range(resolution_x):
count += 1
r_vec = np.array([X[i_v, i_x], 0.0, 0.0])
v_vec = np.array([V[i_v, i_x], 0.0, 0.0])
_, R_ij = Rateeq.construct_evolution_matrix(r_vec, v_vec)
magnitude = -1*np.sum(*R_ij['g->e'])/t0*cts.h*k_lab*100/Isotope.mass
accel_profile[i_v, i_x] = magnitude
progress.update(count/num)
data_dict = {
'basis_g': basis_g,
'basis_e': basis_e,
'x0': x0,
't0': t0,
'x': x,
'v': v,
'acceleration_profile': accel_profile,
'magField': field_mag
}
checked_save(check1, path1, data_dict)
def calc_force_field_and_trajectories_Single_Beam(Isotope_name, params, resolution_x, resolution_v, t_max, basename):
check1, path1 = get_data_filepath_and_check('OtherData', 'data_dir_'+Isotope_name+basename, '.npz')
Isotope = Dy_atom.DyAtom(Isotope_name)
k_lab = Isotope.transition[0].k # 1/cm not 2*pi/cm
gamma_lab = Isotope.state[1].gammaHz # Hz This is 32.2MHz, not 2*pi*32.2MHz
# Define "natural units"
x0 = 1/(k_lab) # cm
t0 = 1/gamma_lab # s
muB = 1
k = k_lab*x0
gamma = gamma_lab*t0
mass = Isotope.mass*(x0/100)**2/(cts.h * t0)
det = params["det"]
s = params["s"]
pol = params["polarisation"]
field_mag = params["field_mag_Gauss_cm"]*1e-4
field_type = params["field_type"]
waist = params["waist_m"]/(x0/100)
aperture = params["aperture_m"]/(x0/100)
trap_length = params["trap_length_m"] #m
zrange = trap_length/(x0/100)
velocity = params["max_velocity"] # m/s
vrange = velocity/((x0/100)/t0)
v0s = params["initial_velocities_m"]/((x0/100)/t0)
pylcp.rateeq.evolve_motion = custom_rateeq.evolve_motion #replace the evolve motion funciton with custom
Ham, E_g, E_e , basis_g, basis_e = init_hamiltonian_hyperfine_coupled(Isotope, mass, t0, muB=muB, gamma=gamma, k=k, returnBasis=True)
basis_g = np.transpose(basis_g)
basis_e = np.transpose(basis_e)
if Isotope_name == "Dy164":
det_L = det
elif Isotope_name == "Dy163":
det_L = (E_e[-1] - E_g[-1]) + det
elif Isotope_name == "Dy161":
det_L = (E_e[0] - E_g[0]) + det
else:
raise ValueError(f"Isotope name '{Isotope_name}' not recognised.")
LaserBeams = init_LaserBeamSingleZeeman(detuning=det_L, saturation=s, polarisation=pol, beam_waist=waist, aperture=aperture, k=k)
if field_type == "2DMOT":
alpha_2DMOT = field_mag * cts.value('Bohr magneton') * x0 * t0 / cts.h #field mag is given as Gauss/cm
magField = pylcp.fields.magField(lambda R: [-alpha_2DMOT*R[1],-alpha_2DMOT*R[0],0])
elif field_type == "Gradient":
alpha_gradient = field_mag * cts.value('Bohr magneton') * x0 * t0 / cts.h #field mag is given as Gauss/cm
magField = pylcp.fields.magField(lambda R: [alpha_gradient*R[0],0,0])
elif field_type == "Offset":
alpha_offset = field_mag * cts.value('Bohr magneton') * t0 / cts.h #field mag is given as Gauss/cm
magField = pylcp.fields.constantMagneticField([alpha_offset,0,0])
else:
raise ValueError("Error: Unrecognised field type. Has to be either '2DMOT' or 'Gradient' or 'Offset'.")
x = np.linspace(-zrange, zrange, resolution_x)
v = np.linspace(-vrange, vrange, resolution_v)
X, V = np.meshgrid(x, v)
Rvec = np.array([X, np.zeros(X.shape), np.zeros(X.shape)])
Vvec = np.array([V, np.zeros(V.shape), np.zeros(V.shape)])
print("Calculating rate equations:")
Rateeq = pylcp.rateeq(LaserBeams, magField, Ham, include_mag_forces=True)
print("Completed. \nCalculating Force Profile:")
Rateeq.set_initial_pop(np.concatenate([np.ones(len(basis_g))/len(basis_g), np.zeros(len(basis_e))]))
accel_profile = np.full((resolution_v, resolution_x), np.nan)
progress = progressBar()
count = 0
num = resolution_v*resolution_x
for i_v in range(resolution_v):
for i_x in range(resolution_x):
count += 1
r_vec = np.array([X[i_v, i_x], 0.0, 0.0])
v_vec = np.array([V[i_v, i_x], 0.0, 0.0])
_, R_ij = Rateeq.construct_evolution_matrix(r_vec, v_vec)
magnitude = -1*np.sum(*R_ij['g->e'])/t0*cts.h*k_lab*100/Isotope.mass
accel_profile[i_v, i_x] = magnitude
progress.update(count/num)
data_dict = {
'basis_g': basis_g,
'basis_e': basis_e,
'x0': x0,
't0': t0,
'x': x,
'v': v,
'acceleration_profile': accel_profile
}
checked_save(check1, path1, data_dict)
initial_position = params["initial_position_m"]/(x0/100)#x[int(len(x)*(1-(aperture/zrange)))]
boundary = params["boundary_m"]/(x0/100)#initial_position*1.2
for i, v0 in enumerate(v0s):
check2, path2 = get_data_filepath_and_check('Solutions', 'sol'+Isotope_name+basename+str(round(v0*(x0/(100*t0))))+'mps', '.pkl')
Rateeq.set_initial_position_and_velocity(np.array([initial_position, 0., 0.]), np.array([v0, 0, 0]))
if isinstance(Rateeq, pylcp.rateeq):
Rateeq.set_initial_pop(np.ones(len(basis_e)+len(basis_g))/(len(basis_e)+len(basis_g)))
print("Calculating Trajectory:", i+1, "out of ", len(v0s))
Rateeq.evolve_motion([0., t_max/t0], max_step=1e-4/t0, boundary=boundary, progress_bar=True)
sol = Rateeq.sol
checked_save(check2, path2, sol)
def calc_force_field_and_trajectories_angled_2DMOT(Isotope_name, params, resolution_x, resolution_v, t_max, basename):
check1, path1 = get_data_filepath_and_check('OtherData', 'data_dir_'+Isotope_name+basename, '.npz')
pylcp.rateeq.evolve_motion = custom_rateeq.evolve_motion #replace the evolve motion funciton with custom
Isotope = Dy_atom.DyAtom(Isotope_name)
k_lab = Isotope.transition[0].k # 1/cm not 2*pi/cm
gamma_lab = Isotope.state[1].gammaHz # Hz This is 32.2MHz, not 2*pi*32.2MHz
# Define "natural units"
x0 = 1/(k_lab) # cm
t0 = 1/gamma_lab # s
muB = 1
k = k_lab*x0
gamma = gamma_lab*t0
mass = Isotope.mass*(x0/100)**2/(cts.h * t0)
det = params["det"]
s = params["s"]
field_mag = params["field_mag_Gauss_cm"]*1e-4
field_type = params["field_type"]
waist = params["waist_m"]/(x0/100)
aperture = params["aperture_m"]/(x0/100)
trap_length = params["trap_length_m"] #m
zrange = trap_length/(x0/100)
velocity = params["max_velocity"] # m/s
vrange = velocity/((x0/100)/t0)
v0s = params["initial_velocities_m"]/((x0/100)/t0)
Ham, E_g, E_e , basis_g, basis_e = init_hamiltonian_hyperfine_coupled(Isotope, mass, t0, muB=muB, gamma=gamma, k=k, returnBasis=True)
basis_g = np.transpose(basis_g)
basis_e = np.transpose(basis_e)
if Isotope_name == "Dy164":
det_L = det
elif Isotope_name == "Dy163":
det_L = (E_e[-1] - E_g[-1]) + det
elif Isotope_name == "Dy161":
det_L = (E_e[0] - E_g[0]) + det
else:
raise ValueError(f"Isotope name '{Isotope_name}' not recognised.")
LaserBeams = init_LaserBeams2DMOT(detuning=det_L, saturation=s, beam_waist=waist, aperture=aperture, k=k)
if field_type == "2DMOT":
alpha_2DMOT = field_mag * cts.value('Bohr magneton') * x0 * t0 / cts.h #field mag is given as Gauss/cm
magField = pylcp.fields.magField(lambda R: [-alpha_2DMOT*R[1],-alpha_2DMOT*R[0],0])
else:
raise ValueError("Error: Unrecognised field type. Field Type has to be '2DMOT'.")
x = np.linspace(-zrange, zrange, resolution_x)
v = np.linspace(-vrange, vrange, resolution_v)
X, V = np.meshgrid(x, v)
Rvec = np.array([X, np.zeros(X.shape), np.zeros(X.shape)])
Vvec = np.array([V, np.zeros(V.shape), np.zeros(V.shape)])
print("Calculating rate equations:")
Rateeq = pylcp.rateeq(LaserBeams, magField, Ham, include_mag_forces=True)
print("Completed. \nCalculating Force Profile:")
Rateeq.set_initial_pop(np.concatenate([np.ones(len(basis_g))/len(basis_g), np.zeros(len(basis_e))]))
accel_profile = np.full((resolution_v, resolution_x), np.nan)
progress = progressBar()
count = 0
num = resolution_v*resolution_x
for i_v in range(resolution_v):
for i_x in range(resolution_x):
count += 1
r_vec = np.array([X[i_v, i_x], 0.0, 0.0])
v_vec = np.array([V[i_v, i_x], 0.0, 0.0])
_, R_ij = Rateeq.construct_evolution_matrix(r_vec, v_vec)
#Rij contains an entry for each laser beam. So we want to sum all contributtions of all beams into one array:
summed_array = R_ij['g->e'].sum(axis=0)
#And the magnitude of the force will be all contributions of all transiotions summed, times some factors.
magnitude = -1*np.sum(summed_array)/t0*cts.h*k_lab*100/Isotope.mass
accel_profile[i_v, i_x] = magnitude
progress.update(count/num)
data_dict = {
'basis_g': basis_g,
'basis_e': basis_e,
'x0': x0,
't0': t0,
'x': x,
'v': v,
'acceleration_profile': accel_profile
}
checked_save(check1, path1, data_dict)
initial_position = params["initial_position_m"]/(x0/100)
if "boundary_m" in params:
boundary = params["boundary_m"]/(x0/100)
else:
boundary = None
if "MOTvelocity_ms" in params:
MOTvelocity = params["MOTvelocity_ms"]/((x0/100)/t0)
else:
MOTvelocity = None
if "MOTradius_m" in params:
MOTradius = params["MOTradius_m"]/(x0/100)
else:
MOTradius = None
for i, v0 in enumerate(v0s):
check2, path2 = get_data_filepath_and_check('Solutions', 'sol'+Isotope_name+basename+str(round(v0*(x0/(100*t0))))+'mps', '.pkl')
Rateeq.set_initial_position_and_velocity(np.array([initial_position, 0., 0.]), np.array([v0, 0, 0]))
if isinstance(Rateeq, pylcp.rateeq):
Rateeq.set_initial_pop(np.ones(len(basis_e)+len(basis_g))/(len(basis_e)+len(basis_g)))
print("Calculating Trajectory:", i+1, "out of ", len(v0s))
Rateeq.evolve_motion([0., t_max/t0], max_step=1e-4/t0, boundary=boundary, MOTradius=MOTradius, MOTvelocity=MOTvelocity, progress_bar=True)
sol = Rateeq.sol
checked_save(check2, path2, sol)

264
custom_rateeq.py Normal file
View File

@ -0,0 +1,264 @@
import numpy as np
from scipy.integrate import solve_ivp
from pylcp.common import (progressBar, random_vector)
from pylcp.integration_tools import solve_ivp_random
from scipy.interpolate import interp1d
"""
Here, the function "rateeq" in the rateeq module of pylcp is adjusted. Typically, there will be no need to change anything here.
Goal:
Be able to interrupt the calculation of a trajectory in a more useful way.
Implementation:
Set a max boundary and a radius in phase space that are both able
to interrupt the calculation. If the atom leaves a usefull region or is trapped
in the MOT, the calculation will be stopped. Thus, a long timespan can be specified
and the calculation will stop automatically, when reasonable.
"""
def evolve_motion(self, t_span, freeze_axis=[False, False, False],
random_recoil=False, random_force=False,
max_scatter_probability=0.1, progress_bar=False,
record_force=False, rng=np.random.default_rng(),
boundary=None, MOTradius=None, MOTvelocity=None, **kwargs):
"""
Evolve the populations :math:`N` and the motion of the atom in time.
This function evolves the rate equations, moving the atom through space,
given the instantaneous force, for some period of time.
Parameters
----------
t_span : list or array_like
A two element list or array that specify the initial and final time
of integration.
freeze_axis : list of boolean
Freeze atomic motion along the specified axis.
Default: [False, False, False]
random_recoil : boolean
Allow the atom to randomly recoil from scattering events.
Default: False
random_force : boolean
Rather than calculating the force using the rateeq.force() method,
use the calculated scattering rates from each of the laser beam
(combined with the instantaneous populations) to randomly add photon
absorption events that cause the atom to recoil randomly from the
laser beam(s).
Default: False
max_scatter_probability : float
When undergoing random recoils and/or force, this sets the maximum
time step such that the maximum scattering probability is less than
or equal to this number during the next time step. Default: 0.1
progress_bar : boolean
If true, show a progress bar as the calculation proceeds.
Default: False
record_force : boolean
If true, record the instantaneous force and store in the solution.
Default: False
rng : numpy.random.Generator()
A properly-seeded random number generator. Default: calls
``numpy.random.default.rng()``
boundary : float
Define a boundary of the simulation. This is the maximal position
the atoms can reach before the calculation is interrupted by an event.
If boundary = 0, no interruption will happen. The calculation is run
until t_spn[1] is reached. sry for the stupid coding. (default: 0)
MOTradius : float
Same as boundary, but specifies a small radius of "MOT region".
When the atoms reach this radius and have a velocity smaller than
MOTvelocity, the caculation will be interrupted by an event.
If 0, no interruption. (default: 0)
MOTvelocity : float
See above. (default : 0)
**kwargs :
Additional keyword arguments get passed to solve_ivp_random, which
is what actually does the integration.
Returns
-------
sol : OdeSolution
Bunch object that contains the following fields:
* t: integration times found by solve_ivp
* N: population vs. time
* v: atomic velocity
* r: atomic position
It contains other important elements, which can be discerned from
scipy's solve_ivp documentation.
"""
free_axes = np.bitwise_not(freeze_axis)
if progress_bar:
progress = progressBar()
if record_force:
ts = []
Fs = []
def motion(t, y):
N = y[:-6]
v = y[-6:-3]
r = y[-3:]
Rev, Rijl = self.construct_evolution_matrix(r, v, t)
if not random_force:
if record_force:
F = self.force(r, t, N, return_details=True)
ts.append(t)
Fs.append(F)
F = F[0]
else:
F = self.force(r, t, N, return_details=False)
dydt = np.concatenate((Rev @ N,
F*free_axes/self.hamiltonian.mass+
self.constant_accel,
v))
else:
dydt = np.concatenate((Rev @ N,
self.constant_accel,
v))
if np.any(np.isnan(dydt)):
raise ValueError('Enountered a NaN!')
if progress_bar:
progress.update(t/t_span[-1])
return dydt
def random_force_func(t, y, dt):
total_P = 0
num_of_scatters = 0
# Go over all available keys:
for key in self.laserBeams:
# Extract the pumping rate from each laser:
Rl = np.sum(self.Rijl[key], axis=(1,2))
# Calculate the probability to scatter a photon from the laser:
P = Rl*dt
# Roll the dice N times, where $N=\sum(lasers)
dice = rng.random(len(P))
# Give them kicks!
for ii in np.arange(len(Rl))[dice<P]:
num_of_scatters += 1
y[-6:-3] += self.laserBeams[key].beam_vector[ii].kvec(y[-3:], t)/\
self.hamiltonian.mass
# Can branch to a differe, lower state, but let's ignore that
# for the moment.
y[-6:-3] += self.recoil_velocity[key]*(random_vector(rng, free_axes))
total_P += np.sum(P)
# Calculate a new maximum dt to make sure we evolve while not
# exceeding dt max:
new_dt_max = (max_scatter_probability/total_P)*dt
return (num_of_scatters, new_dt_max)
def random_recoil_func(t, y, dt):
num_of_scatters = 0
total_P = 0.
# Go over each block in the Hamiltonian and compute the decay:
for key in self.decay_rates:
P = dt*self.decay_rates[key]*y[self.decay_N_indices[key]]
# Roll the dice N times, where $N=\sum(n_i)
dice = rng.random(len(P))
# For any random number that is lower than P_i, add a
# recoil velocity.
for ii in range(np.sum(dice<P)):
num_of_scatters += 1
y[-6:-3] += self.recoil_velocity[key]*(random_vector(rng, free_axes)+
random_vector(rng, free_axes))
# Save the total probability of a scatter:
total_P += np.sum(P)
# Calculate a new maximum dt to make sure we evolve while not
# exceeding dt max:
new_dt_max = (max_scatter_probability/total_P)*dt
return (num_of_scatters, new_dt_max)
eventList=[]
# Define the event function to stop the simulation when r > boundary
def boundary_event(boundary):
def event(t, y):
return 0 if y[-3]**2+y[-2]**2+y[-1]**2 > boundary**2 else 1 # Trigger event when x > boundary
return event
# Define the event function to stop the simulation when position < 0.01 and velocity < 1
def MOT_event(MOTradius, MOTvelocity):
def event(t, y):
radius = y[-3]**2+y[-2]**2+y[-1]**2
velocity = y[-6]**2+y[-5]**2+y[-4]**2
return radius - MOTradius**2 if velocity < MOTvelocity**2 else 1
return event
if boundary:
if type(boundary) == float:
boundary_event_func = boundary_event(boundary)
boundary_event_func.terminal = True # Stop the integration on condition
boundary_event_func.direction = -1 # trigger when decreasing
eventList.append(boundary_event_func)
else:
raise(TypeError(f"Boundary is type {type(boundary)}. Expected float."))
if MOTradius and MOTvelocity:
if type(MOTradius) == float and type(MOTvelocity) == float:
MOT_event_func = MOT_event(MOTradius, MOTvelocity)
MOT_event_func.terminal = True # Stop the integration on condition
MOT_event_func.direction = -1 # Trigger when decreasing
eventList.append(MOT_event_func)
else:
raise(TypeError(f"Motvelocity is type {type(MOTvelocity)} and MOTradius is type {type(MOTradius)}. Expected float for both."))
y0 = np.concatenate((self.N0, self.v0, self.r0))
if random_force:
self.sol = solve_ivp_random(motion, random_force_func, t_span, y0,
initial_max_step=max_scatter_probability,
**kwargs)
elif random_recoil:
self.sol = solve_ivp_random(motion, random_recoil_func, t_span, y0,
initial_max_step=max_scatter_probability,
**kwargs)
else:
self.sol = solve_ivp(motion, t_span, y0, events=eventList, **kwargs)
if progress_bar:
# Just in case the solve_ivp_random terminated due to an event.
progress.update(1.)
# Rearrange the solution:
self.sol.N = self.sol.y[:-6]
self.sol.v = self.sol.y[-6:-3]
self.sol.r = self.sol.y[-3:]
if record_force:
f = interp1d(ts[:-1], np.array([f[0] for f in Fs[:-1]]).T)
self.sol.F = f(self.sol.t)
f = interp1d(ts[:-1], np.array([f[2] for f in Fs[:-1]]).T)
self.sol.fmag = f(self.sol.t)
self.sol.f = {}
for key in Fs[0][1]:
f = interp1d(ts[:-1], np.array([f[1][key] for f in Fs[:-1]]).T)
self.sol.f[key] = f(self.sol.t)
self.sol.f[key] = np.swapaxes(self.sol.f[key], 0, 1)
del self.sol.y
return self.sol