From a1cd30886d1f712a4a5de518c67bf3b1beae8cc7 Mon Sep 17 00:00:00 2001 From: lhoenen Date: Wed, 16 Apr 2025 18:47:42 +0200 Subject: [PATCH] Uploaded initial "working" version of PyLCP simulation. This is not well structured yet. Nor is it well tested or commented. --- Dy_atom.py | 62 +++++ PlotStuff.py | 256 ++++++++++++++++++ common.py | 679 +++++++++++++++++++++++++++++++++++++++++++++++ custom_rateeq.py | 264 ++++++++++++++++++ 4 files changed, 1261 insertions(+) create mode 100644 Dy_atom.py create mode 100644 PlotStuff.py create mode 100644 common.py create mode 100644 custom_rateeq.py diff --git a/Dy_atom.py b/Dy_atom.py new file mode 100644 index 0000000..8de62b1 --- /dev/null +++ b/Dy_atom.py @@ -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) \ No newline at end of file diff --git a/PlotStuff.py b/PlotStuff.py new file mode 100644 index 0000000..4c28f99 --- /dev/null +++ b/PlotStuff.py @@ -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() + diff --git a/common.py b/common.py new file mode 100644 index 0000000..a206edd --- /dev/null +++ b/common.py @@ -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) + diff --git a/custom_rateeq.py b/custom_rateeq.py new file mode 100644 index 0000000..f3133b2 --- /dev/null +++ b/custom_rateeq.py @@ -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 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 \ No newline at end of file