import warnings from collections.abc import Callable from functools import cache from itertools import permutations from numbers import Number from typing import Any import datetime import pickle import numpy as np import numpy.typing as npt import sympy as sp from scipy import constants as const from scipy.sparse import save_npz from scipy.optimize import minimize_scalar from sympy import LT, default_sort_key, topological_sort from sympy.core.numbers import Zero import trap_units as si from trap_numerics import (nstationary_solution, nstationary_solution_export, discrete_hamiltonian, functional_hamiltonian) import trap_potential as pot import trap_properties as p from trap_typing import float_or_array Numeric = int | float | complex | np.number | sp.Number ExprNum = sp.Expr | Numeric class PancakeTrap: #: Spacial coordinates x, y, z = sp.symbols("x, y, z", real=True, finite=True) #: Offset from the waist location in z d_0 = sp.symbols("d_0", real=True, finite=True) #: Beam waist in x-direction (collimated). waist_x0 = sp.Symbol("W_{x0}", real=True, positive=True, finite=True, nonzero=True) #: Beam waist in z-direction at the focus. waist_z0 = sp.Symbol("W_{z0}", real=True, positive=True, finite=True, nonzero=True) #: Waist of the tweezer (at the focus) waist_tweezer = sp.Symbol( "W_t", real=True, positive=True, finite=True, nonzero=True ) #: Power per beam. Total power is 2x. power = sp.Symbol("P", real=True, positive=True, finite=True, nonzero=True) #: Power of the tweezer beam power_tweezer = sp.Symbol( "P_t", real=True, positive=True, finite=True, nonzero=True ) #magnetic offset field B_x, B_y, B_z = sp.symbols( "B_x, B_y, B_z", real=True, finite=True ) #: Magnetic field gradients. grad_r, grad_z = sp.symbols( "\\nabla{B_x}, \\nabla{B_z}", real=True, finite=True, nonzero=True ) #: Real part of polarizability :math:`\alpha` a = sp.Symbol("a", real=True, finite=True, nonzero=True, positive=True) #: Pancake spacing d = sp.Symbol("d", real=True, positive=True, finite=True, nonzero=True) #: Bohr magneton mu_b = sp.Symbol("mu_b", real=True, positive=True, finite=True, nonzero=True) #: Wavelength wvl = sp.Symbol("lambda", real=True, positive=True, finite=True, nonzero=True) #: Angle of the incident beams theta = sp.Symbol("theta", real=True, positive=True, finite=True, nonzero=True) #: Mass of the trapped atoms m = sp.Symbol("m", real=True, positive=True, finite=True, nonzero=True) #: contact-interaction scattering length of atoms a_s = sp.Symbol("a_s", real=True, positive=True, finite=True, nonzero=True) #: Trap frequency in each spacial direction omega_x, omega_y, omega_z = sp.symbols( "omega_x, omega_y, omega_z", real=True, positive=True, finite=True, nonzero=True ) #: Trap frequency of the tweezer omega_r_tweezer, omega_ax_tweezer = sp.symbols( "omega_t_r, omega_t_ax", real=True, positive=True, finite=True, nonzero=True ) #: Resonance frequency omega_0 = sp.Symbol("omega_0", real=True, positive=True, finite=True, nonzero=True) #: Laser frequency omega_l = sp.Symbol("omega_l", real=True, positive=True, finite=True, nonzero=True) #: Decay rate gamma gamma = sp.Symbol("Gamma", real=True, positive=True, finite=True, nonzero=True) #: Angle of the magnetic field gradient beta = sp.Symbol("beta", real=True, positive=True, finite=True) #: gravitational acceleration (set 0 if not needed) g = sp.Symbol("g", real=True, positive=True, finite=True) #: Default substitutions applied using :meth:`.subs` _defaults: dict[str, Any] = { #"a": ( # (3 * sp.pi * const.c**2) # / (2 * omega_0**3) # * (gamma / (omega_0 - omega_l) + gamma / (omega_0 + omega_l)) #), "d": wvl / (2 * sp.sin(theta)), "omega_l": 2 * np.pi * const.c / wvl, #"omega_0": 2 * np.pi * const.c / (626 * si.nm), "d_0": 0, "wvl": 532 * si.nm, #1064 * si.nm, "theta": 0, #np.deg2rad(7.35), "beta": 0, "gamma": 2 * np.pi * 135 * const.kilo, #2 * np.pi * 5.8724 * const.mega, "grad_r": 0.0066 * grad_z, # see labbook 2023-11-21 "grad_z": 85 * si.G / si.cm, "B_x": 0*si.G, "B_y": 0*si.G, "B_z": 1e-14*si.G, #apply small field to avoid zero "m": 164 * const.value("atomic mass constant"), #6.0151228 * const.value("atomic mass constant"), "a_s": 85* const.value("Bohr radius"), "mu_b": 9.93 * const.value("Bohr magneton" ), #using mu_b as mu for the atom "g": const.g, } def __init__(self, **values: Numeric): self._values: dict[sp.Symbol, Numeric] = { self._get_key_symbol(key): value for key, value in self._defaults.items() } for key, value in values.items(): self[key] = value self._omega_solution = {} def subs(self, expr: sp.Expr, symbolic_only: bool = False) -> sp.Expr: # Only substitute symbols that have a value of type Expr, # i.e. they are expressed through other symbols. symbolic_values: list[tuple[sp.Symbol, sp.Expr | Numeric]] = [ (key, value) for key, value in self._values.items() if isinstance(value, sp.Expr) or value == 0 ] # Find all edges where an expression contains the symbol of # a different subsitution. edges = [ (i, j) for i, j in permutations(symbolic_values, 2) if isinstance(i[1], sp.Expr) and i[1].has(j[0]) ] # Sort the substitutions topologically such that sorted_symbolic_values = topological_sort( [symbolic_values, edges], default_sort_key ) expr_symbolic = expr.subs(sorted_symbolic_values) if symbolic_only: return expr_symbolic else: # Also substitute all non-symbolic values return expr_symbolic.subs( [ (key, value) for key, value in self._values.items() if not isinstance(value, sp.Expr) ] ) def __setitem__(self, key: str | sp.Symbol, value): self._values[self._get_key_symbol(key)] = value def __getitem__(self, item: str | sp.Symbol): return self._values[self._get_key_symbol(item)] def _get_key_symbol(self, key: str | sp.Symbol) -> sp.Symbol: if isinstance(key, sp.Symbol): return key if not isinstance(key, str): raise ValueError( f"Expected type 'str' or 'sympy.Symbol' but got {type(key)!s}" ) if not hasattr(self, key): raise ValueError(f"Unknown attribute '{key!s}'") key_symbol = getattr(self, key) if not isinstance(getattr(self, key), sp.Symbol): raise ValueError(f"Attribute '{key!s}' is not a symbol") return key_symbol def get_tweezer_rayleigh(self) -> sp.Expr: return sp.pi * self.waist_tweezer**2 / self.wvl def get_tweezer_intensity(self) -> sp.Expr: rayleigh_length = self.get_tweezer_rayleigh() waist = self.waist_tweezer * sp.sqrt(1 + (self.z / rayleigh_length) ** 2) i_0 = 2 * self.power_tweezer / (sp.pi * self.waist_tweezer**2) r = sp.sqrt(self.x**2 + self.y**2) return i_0 * (self.waist_tweezer / waist) ** 2 * sp.exp(-2 * r**2 / (waist**2)) def get_tweezer_potential(self) -> sp.Expr: return -self.a * self.get_tweezer_intensity() def get_omega_tweezer(self, omega: sp.Symbol) -> sp.Expr: """Get tweezer trap frequency""" return self.get_omega(omega, with_tweezer=True, with_2dtrap=False) def get_omega_r_tweezer(self) -> sp.Expr: return self.get_omega_tweezer(self.omega_r_tweezer) def get_omega_ax_tweezer(self) -> sp.Expr: return self.get_omega_tweezer(self.omega_ax_tweezer) def get_beam_waist_z(self, length: sp.Symbol | Number) -> sp.Expr: """The beam is focused in z-direction, leading to a waist that depends on the distance :math:`l` from the lens. Depending on which beam (upper or lower), the distance from the lens :math:`l` is given by :math:`y \\cos(\\theta) \\mp z \\sin(\\theta)` :param length: Distance :math:`l` from the lens. """ z_r = sp.pi * self.waist_z0**2 / self.wvl return self.waist_z0 * sp.sqrt(1 + ((length - self.d_0) / z_r) ** 2) def get_upper_beam_intensity(self) -> sp.Expr: """Intensity distribution of the upper beam""" length = self.y * sp.cos(self.theta) - self.z * sp.sin(self.theta) waist_z = self.get_beam_waist_z(length) # Intensity of th0e beam at the origin i_0 = 2 * self.power / (sp.pi * self.waist_x0 * self.waist_z0) # Turn off automatic formatting to preserve readability # fmt: off return i_0 * self.waist_z0 / waist_z * sp.exp( - 2 * self.x ** 2 / self.waist_x0 ** 2 - 2 * (self.z * sp.cos(self.theta) + self.y * sp.sin(self.theta)) ** 2 / waist_z ** 2 ) # fmt: on def get_lower_beam_intensity(self) -> sp.Expr: """Intensity distribution of the upper beam""" length = self.y * sp.cos(self.theta) + self.z * sp.sin(self.theta) waist_z = self.get_beam_waist_z(length) # Intensity of the beam at the origin i_0 = 2 * self.power / (sp.pi * self.waist_x0 * self.waist_z0) # Turn off automatic formatting to preserve readability # fmt: off return i_0 * self.waist_z0 / waist_z * sp.exp( - 2 * self.x ** 2 / self.waist_x0 ** 2 - 2 * (self.z * sp.cos(self.theta) - self.y * sp.sin(self.theta)) ** 2 / waist_z ** 2 ) # fmt: on def get_total_intensity( self, with_tweezer: bool = True, with_2dtrap: bool = True ) -> sp.Expr: """Total intensity distribution of both beams""" i_1 = self.get_upper_beam_intensity() i_2 = self.get_lower_beam_intensity() i_tweezer = int(with_tweezer) * self.get_tweezer_intensity() i_2dtrap = int(with_2dtrap) * ( i_1 + i_2 + 2 * sp.sqrt(i_1 * i_2) * sp.cos(2 * sp.pi * self.z / self.d) ) return i_2dtrap + i_tweezer def get_potential( self, with_tweezer: bool = True, with_2dtrap: bool = True, with_grad: bool = True, apply_zero_offset: bool = True, ) -> sp.Expr: """Get the potential of the 2d trap. The total potential consists of the laser induced optical dipole potential and the magnetic field potential. """ grad_r_coordinate = self.x * sp.cos(self.beta) - self.y * sp.sin(self.beta) pot = -self.a * self.get_total_intensity( with_tweezer=with_tweezer, with_2dtrap=with_2dtrap ) - int(with_grad) * self.mu_b * ( self.grad_r * grad_r_coordinate + self.grad_z * self.z ) - self.m * self.g * self.z offset = int(apply_zero_offset) * pot.subs({self.x: 0, self.y: 0, self.z: 0}) return pot - offset @cache def get_omega( self, omega: sp.Symbol, with_tweezer: bool = True, with_2dtrap: bool = True, ) -> sp.Expr: """Calculate trap frequency for one spacial direction. Uses :func:`sympy.series`, equates the 2nd order term with the model of a harmonic oscillator and solves for omega. :param omega: Which spacial trap frequency to calculate. One of :attr:`.omega_x`, :attr:`.omega_y`, or :attr:`.omega_z`. :param force: Force re-evaluation of the trap frequency. The analytic solution contains many symbols and is thus rather lengthy. To save time, the result is cached in the :attr:`._omega_solution` attribute. :param with_tweezer: Whether or not to include tweezer potential in the calculation. :param with_2dtrap: Whether or not to include 2D trap potential in the calculation. :return: Trap frequency as a sympy expression. """ others: set[sp.Symbol] = {self.x, self.y, self.z} variable: sp.Symbol if omega == self.omega_x or omega == self.omega_r_tweezer: variable = self.x elif omega == self.omega_y: variable = self.y elif omega == self.omega_z or omega == self.omega_ax_tweezer: variable = self.z else: raise ValueError(f"Unknown omega '{omega!r}'") others.remove(variable) series = self.subs( self.get_total_intensity( with_tweezer=with_tweezer, with_2dtrap=with_2dtrap ).subs({symbol: 0 for symbol in others}), # Only substitute symbols with other symbols symbolic_only=True, ).series(x=variable, x0=0, n=3) quad_term = sp.LT(series.removeO(), gens=(variable,)) equation = sp.Eq(-self.a * quad_term, self.m * omega**2 * variable**2 / 2) solution = sp.solve(equation, omega) if len(solution) == 0: raise RuntimeError(f"Unable to solve for '{omega}'") return solution[0] def get_omega_x( self, with_tweezer: bool = True, with_2dtrap: bool = True ) -> sp.Expr: """Calculate trap frequency in x. See :meth:`.get_omega`. """ return self.get_omega( self.omega_x, with_tweezer=with_tweezer, with_2dtrap=with_2dtrap ) def get_omega_y( self, with_tweezer: bool = True, with_2dtrap: bool = True ) -> sp.Expr: """Calculate trap frequency in y. See :meth:`.get_omega`. """ return self.get_omega( self.omega_y, with_tweezer=with_tweezer, with_2dtrap=with_2dtrap ) def get_omega_z( self, with_tweezer: bool = True, with_2dtrap: bool = True ) -> sp.Expr: """Calculate trap frequency in z. See :meth:`.get_omega`. """ return self.get_omega( self.omega_z, with_tweezer=with_tweezer, with_2dtrap=with_2dtrap ) def waist_z_to_waist_y(self) -> sp.Expr: """Calculate waist in y from :math:`W_z`""" return self.waist_z0 / sp.sin(self.theta) def waist_z_to_waist_ax(self) -> sp.Expr: """Calculates the axial waist from :math:`W_z`""" return self.waist_z0 / sp.cos(self.theta) def nstationary_solution( self, dims: tuple[sp.Symbol, ...], extend: tuple[tuple[ExprNum, ExprNum], ...] | tuple[ExprNum, ExprNum], size: tuple[int, ...], k: int, method="sparse", export=False ): """Numeric stationary solution""" if isinstance(dims, sp.Expr): dims = (dims,) ndims = len(dims) if not (np.shape(extend) != (2,) or np.shape(extend) != (ndims, 2)): raise ValueError( f"Shape of parameter 'extend' has to be either (2,) or ({ndims}, 2), " f"but is {np.shape(extend)}." ) if not (isinstance(size, int) or len(size) == ndims): raise ValueError( "Parameter 'size' has to be either of type 'int' or " f"a tuple of length {ndims}." ) # Make sure 'size' is always a tuple, even if all dimensions # have the same size. if isinstance(size, int): size = (size,) * ndims others: set[sp.Symbol] = {self.x, self.y, self.z} coordinates: dict[sp.Symbol, npt.NDArray] = {} for i, dim in enumerate(dims): if dim not in others: raise ValueError(f"Symbols '{dim}' unsupported or used multiple times") others.remove(dim) # Check if a single extend is used for all dimensions. This # boolean is later used to determine if the extend should # be overwritten. This avoids re-evaluating the same extend # multiple times. single_extend = np.shape(extend) == (2,) # The extend is casted to a list to make the items mutable dim_extend: list[ExprNum, ExprNum] = ( list(extend) if single_extend else list(extend[i]) ) # Make sure that all limits in the extend are numbers that # numpy can work with. for j, lim in enumerate(dim_extend): if isinstance(lim, sp.Expr): new_lim = self.subs(lim).evalf() if not isinstance(new_lim, (float, sp.Float)): raise ValueError( f"Unable to convert '{lim}' to float (using substitutions)" ) dim_extend[j] = float(new_lim) if single_extend: # To avoid re-evaluation of the symbolic extends, fix # the extends to the calculated value extend = tuple(dim_extend) # cast list to tuple again coordinates[dim] = np.linspace(*dim_extend, num=size[i]) # Calculate size of the volume elements """ coords = np.array(list(coordinates.values())) dvol: npt.NDArray = np.mean(coords[:, 1:] - coords[:, :-1], axis=1) assert dvol.shape == (ndims,)""" dvol = np.array([]) for coord in coordinates.values(): dvol = np.append(dvol,np.mean(np.diff(coord))) # Create potential array pot = self.subs(self.get_potential(apply_zero_offset=False).subs({k: 0 for k in others})) pot_numpy = sp.lambdify(tuple(coordinates.keys()), pot, cse=True) if ndims > 1: meshgrid = np.meshgrid(*coordinates.values(),indexing="ij") potential = pot_numpy(*meshgrid) else: potential = pot_numpy(*coordinates.values()) # Perform the diagonalization of the Hamiltonian energies, states = nstationary_solution( tuple(dvol), potential, float(self.subs(self.m).evalf()), k=k, method=method) states = states.reshape((k, *size)) #export hamiltonian and parameters if export: #check if we are trying to solve DoubleTweezer and save corr. parameters current_time = datetime.datetime.now() #add to overview the key parameters of this run if isinstance(self,DoubleTweezer): distance_tweezers = float(self.subs(self.distance_tweezers).evalf()) power1 = float(self.subs(self.power_tweezer1).evalf()) power2 = float(self.subs(self.power_tweezer2).evalf()) waist1 = float(self.subs(self.waist_tweezer1).evalf()) waist2 = float(self.subs(self.waist_tweezer2).evalf()) wvl = float(self.subs(self.wvl).evalf()) omega_x1, omega_x2 = self.get_both_omega(self.x) omega_x1 = float(self.subs(omega_x1).evalf()) omega_x2 = float(self.subs(omega_x2).evalf()) with open('data/overview.txt', 'a') as f: f.write(f'{current_time.strftime("%Y-%m-%d_%H-%M-%S")}; {size}; {wvl}; {power1}; {power2}; {waist1}; {waist2}; {distance_tweezers}; {omega_x1/2/np.pi}; {omega_x2/2/np.pi} \n') elif isinstance(self,TwoSiteLattice): lattice_spacing_x = float(self.subs(self.lattice_spacing_x).evalf()) lattice_spacing_y = float(self.subs(self.lattice_spacing_y).evalf()) lattice_spacing_z = float(self.subs(self.lattice_spacing_z).evalf()) omega_x1, omega_x2 = self.get_both_omega(self.x) omega_z1, omega_z2 = self.get_both_omega(self.z) aspect_ratio = float(self.subs(sp.sqrt(omega_x1/omega_z1))) wvl = float(self.subs(self.wvl).evalf()) with open('data/overview.txt', 'a') as f: f.write(f'{current_time.strftime("%Y-%m-%d_%H-%M-%S")}; {size}; {wvl}; {lattice_spacing_x}; {lattice_spacing_y}; {lattice_spacing_z}; {aspect_ratio}; TwoSiteLattice\n') #save results as npz x3D,y3D,z3D = np.meshgrid(coordinates[self.x],coordinates[self.y],coordinates[self.z],indexing="ij") np.savez(f"data/{current_time.strftime("%Y-%m-%d_%H-%M-%S")}_results.npz", energies=energies, states=states, dvol=dvol, k=k, size=size, extend=extend, pot=pot_numpy(x3D,y3D,z3D), x=coordinates[self.x],y=coordinates[self.y],z=coordinates[self.z]) #save hamiltonian as sparse matrix save_npz(f"data/{current_time.strftime("%Y-%m-%d_%H-%M-%S")}_hamiltonian.npz", discrete_hamiltonian(tuple(dvol), potential, float(self.subs(self.m).evalf()))) #save trap object with open(f"data/{current_time.strftime("%Y-%m-%d_%H-%M-%S")}_trap.npz", 'wb') as file: pickle.dump(self, file) print(f"files saved with ...._{current_time.strftime("%Y-%m-%d_%H-%M-%S")}") return energies, states, pot_numpy, coordinates def nstationary_solution_export( self,path, dims: tuple[sp.Symbol, ...], extend: tuple[tuple[ExprNum, ExprNum], ...] | tuple[ExprNum, ExprNum], size: tuple[int, ...], k: int, ): """Numeric stationary solution""" if isinstance(dims, sp.Expr): dims = (dims,) ndims = len(dims) if not (np.shape(extend) != (2,) or np.shape(extend) != (ndims, 2)): raise ValueError( f"Shape of parameter 'extend' has to be either (2,) or ({ndims}, 2), " f"but is {np.shape(extend)}." ) if not (isinstance(size, int) or len(size) == ndims): raise ValueError( "Parameter 'size' has to be either of type 'int' or " f"a tuple of length {ndims}." ) # Make sure 'size' is always a tuple, even if all dimensions # have the same size. if isinstance(size, int): size = (size,) * ndims others: set[sp.Symbol] = {self.x, self.y, self.z} coordinates: dict[sp.Symbol, npt.NDArray] = {} for i, dim in enumerate(dims): if dim not in others: raise ValueError(f"Symbols '{dim}' unsupported or used multiple times") others.remove(dim) # Check if a single extend is used for all dimensions. This # boolean is later used to determine if the extend should # be overwritten. This avoids re-evaluating the same extend # multiple times. single_extend = np.shape(extend) == (2,) # The extend is casted to a list to make the items mutable dim_extend: list[ExprNum, ExprNum] = ( list(extend) if single_extend else list(extend[i]) ) # Make sure that all limits in the extend are numbers that # numpy can work with. for j, lim in enumerate(dim_extend): if isinstance(lim, sp.Expr): new_lim = self.subs(lim).evalf() if not isinstance(new_lim, (float, sp.Float)): raise ValueError( f"Unable to convert '{lim}' to float (using substitutions)" ) dim_extend[j] = float(new_lim) if single_extend: # To avoid re-evaluation of the symbolic extends, fix # the extends to the calculated value extend = tuple(dim_extend) # cast list to tuple again coordinates[dim] = np.linspace(*dim_extend, num=size[i]) # Calculate size of the volume elements """ coords = np.array(list(coordinates.values())) dvol: npt.NDArray = np.mean(coords[:, 1:] - coords[:, :-1], axis=1) assert dvol.shape == (ndims,)""" dvol = np.array([]) for coord in coordinates.values(): dvol = np.append(dvol,np.mean(np.diff(coord))) # Create potential array pot = self.subs(self.get_potential(apply_zero_offset=False).subs({k: 0 for k in others})) pot_numpy = sp.lambdify(tuple(coordinates.keys()), pot, cse=True) if ndims > 1: meshgrid = np.meshgrid(*coordinates.values(),indexing="ij") potential = pot_numpy(*meshgrid) else: potential = pot_numpy(*coordinates.values()) # Perform the diagonalization of the Hamiltonian e_min = nstationary_solution_export(path, tuple(dvol), potential, float(self.subs(self.m).evalf()), k=k) x3D,y3D,z3D = np.meshgrid(coordinates[self.x],coordinates[self.y],coordinates[self.z],indexing="ij") #check if we are trying to solve DoubleTweezer and save corr. parameters if isinstance(self,DoubleTweezer): np.savez(path + r"\parameters.npz",e_min=e_min, k=k, size=size, pot=pot_numpy(x3D,y3D,z3D), x=coordinates[self.x],y=coordinates[self.y],z=coordinates[self.z], mass=self.m, mu=self.mu_b,distance_tweezers=self.distance_tweezers, power1=self.power_tweezer1, power2=self.power_tweezer2, waist1=self.waist_tweezer1, waist2=self.waist_tweezer2) else: np.savez(path + r"\parameters.npz",e_min=e_min, k=k, size=size, pot=pot_numpy(x3D,y3D,z3D), x=coordinates[self.x],y=coordinates[self.y],z=coordinates[self.z], mass=self.m, mu=self.mu_b) print("files saved under " + path) class TrueHarmonicTweezer(PancakeTrap): """Pancake Trap but with a truly harmonic tweezer (i.e. quadratic in z and r). Note that spilling is not possible, as the gradient simply shifts the center of the trap. """ def get_tweezer_intensity(self) -> sp.Expr: i_xz = super().get_tweezer_intensity().subs({self.y: 0}) series_x = i_xz.subs({self.z: 0}).series(x=self.x, x0=0, n=3).removeO() series_z = i_xz.subs({self.x: 0}).series(x=self.z, x0=0, n=3).removeO() series_r = series_x.subs(self.x, sp.sqrt(self.x**2 + self.y**2)) quad_term_z = sp.LT(series_z, gens=(self.z,)) return series_r + quad_term_z class DoubleTweezer(PancakeTrap): """ Pancake trap but with two tweezers. Each waist and power, and their spacing can be adjusted. """ def __init__(self,**values): self.power_tweezer1 = sp.Symbol("P_t1", real=True, positive=True, finite=True, nonzero=True) self.power_tweezer2 = sp.Symbol("P_t2", real=True, positive=True, finite=True, nonzero=True) self.waist_tweezer1 = sp.Symbol("W_t1", real=True, positive=True, finite=True, nonzero=True) self.waist_tweezer2 = sp.Symbol("W_t2", real=True, positive=True, finite=True, nonzero=True) self.distance_tweezers = sp.Symbol("d_t", real=True, positive=True, finite=True, nonzero=True) super().__init__(**values) def get_tweezer_rayleigh1(self): return sp.pi * self.waist_tweezer1**2 / self.wvl def get_tweezer_rayleigh2(self): return sp.pi * self.waist_tweezer2**2 / self.wvl def get_tweezer_intensity1(self,x_offset) -> sp.Expr: rayleigh_length = self.get_tweezer_rayleigh1() waist = self.waist_tweezer1 * sp.sqrt(1 + (self.z / rayleigh_length) ** 2) i_0 = 2 * self.power_tweezer1 / (sp.pi * self.waist_tweezer1**2) r = sp.sqrt((self.x-x_offset)**2 + self.y**2) return i_0 * (self.waist_tweezer1 / waist) ** 2 * sp.exp(-2 * r**2 / (waist**2)) def get_tweezer_intensity2(self,x_offset) -> sp.Expr: rayleigh_length = self.get_tweezer_rayleigh2() waist = self.waist_tweezer2 * sp.sqrt(1 + (self.z / rayleigh_length) ** 2) i_0 = 2 * self.power_tweezer2 / (sp.pi * self.waist_tweezer2**2) r = sp.sqrt((self.x-x_offset)**2 + self.y**2) return i_0 * (self.waist_tweezer2 / waist) ** 2 * sp.exp(-2 * r**2 / (waist**2)) def get_tweezer_intensity(self) -> sp.Expr: return self.get_tweezer_intensity1(-self.distance_tweezers/2) + self.get_tweezer_intensity2(self.distance_tweezers/2) def get_both_omega(self, dim: sp.Symbol): """ Returns omega_1, omega_2 the frequency in direction dim (sympy coord) of the left(1) and right(2) tweezer """ V = self.subs(self.get_potential(apply_zero_offset=False)) a = float(self.subs(self.distance_tweezers)) #find minima of potential def V_func(x): return float(V.subs(self.x,x).subs(self.y,0).subs(self.z,0)) x_right = minimize_scalar(V_func,bracket=[0,a]).x x_left = minimize_scalar(V_func,bracket=[-a,0]).x #catch case where both potentials have already merged tunneling_dist = abs(x_right-x_left) if tunneling_dist < 1e-20: raise Exception("potential has only one minmum") #trapping frequencies through second derivative V_prime = sp.diff(V, dim) V_double_prime = sp.diff(V_prime, dim) omega_1 = sp.sqrt(V_double_prime.subs({self.x:x_left, self.y:0, self.z:0})/self.m) omega_2 = sp.sqrt(V_double_prime.subs({self.x:x_left, self.y:0, self.z:0})/self.m) return omega_1, omega_2 class TwoSiteLattice(PancakeTrap): """ Trap with lattice potential that has just two sites """ def __init__(self,**values): self.lattice_depth_x = sp.Symbol("V0_x", real=True, positive=True, finite=True, nonzero=True) self.lattice_depth_y = sp.Symbol("V0_y", real=True, positive=True, finite=True, nonzero=True) self.lattice_depth_z = sp.Symbol("V0_z", real=True, positive=True, finite=True, nonzero=True) #lattice spacings are normaly half the wavelength in that direction self.lattice_spacing_x = sp.Symbol("d_x", real=True, positive=True, finite=True, nonzero=True) self.lattice_spacing_y = sp.Symbol("d_y", real=True, positive=True, finite=True, nonzero=True) self.lattice_spacing_z = sp.Symbol("d_z", real=True, positive=True, finite=True, nonzero=True) super().__init__(**values) def get_potential( self, with_tweezer: bool = True, with_2dtrap: bool = True, with_lattice: bool = True, with_grad: bool = True, apply_zero_offset: bool = True, ) -> sp.Expr: """Get the potential of the 2d trap. The total potential consists of the laser induced optical dipole potential and the magnetic field potential. """ a = self.lattice_spacing_x b = self.lattice_spacing_y c = self.lattice_spacing_z #use sine in x direction to have two wells around the origin pot_lattice = (self.lattice_depth_x* sp.sin(np.pi/a*self.x)**2 + self.lattice_depth_y* sp.cos(np.pi/b*self.y)**2 + self.lattice_depth_z* sp.cos(np.pi/c*self.z)**2) grad_r_coordinate = self.x * sp.cos(self.beta) - self.y * sp.sin(self.beta) pot = -self.a * self.get_total_intensity( with_tweezer=with_tweezer, with_2dtrap=with_2dtrap ) - int(with_grad) * self.mu_b * ( self.grad_r * grad_r_coordinate + self.grad_z * self.z ) - int(with_lattice)* (pot_lattice ) - self.m * self.g * self.z offset = int(apply_zero_offset) * pot.subs({self.x: 0, self.y: 0, self.z: 0}) return pot - offset def get_both_omega(self, dim: sp.Symbol): """ Returns omega_1, omega_2 the frequency in direction dim (sympy coord) of the left(1) and right(2) tweezer """ V = self.subs(self.get_potential(apply_zero_offset=False)) a = float(self.subs(self.lattice_spacing_x)) #find minima of potential x_right = a/2 x_left = -a/2 #catch case where both potentials have already merged tunneling_dist = abs(x_right-x_left) if tunneling_dist < 1e-20: raise Exception("potential has only one minmum") #trapping frequencies through second derivative V_prime = sp.diff(V, dim) V_double_prime = sp.diff(V_prime, dim) omega_1 = sp.sqrt(V_double_prime.subs({self.x:x_left, self.y:0, self.z:0})/self.m) omega_2 = sp.sqrt(V_double_prime.subs({self.x:x_left, self.y:0, self.z:0})/self.m) return omega_1, omega_2 def _optional_a( a: float_or_array | None, wavelength: float_or_array | None ) -> float_or_array: if a is None: if wavelength is None: raise ValueError("Either 'a' or 'wavelength' has to be specified.") a = get_a(wavelength) elif wavelength is not None: warnings.warn( "Both 'a' and 'wavelength' were provided. " "Keyword argument 'wavelength' will be ignored.", RuntimeWarning, stacklevel=2, ) return a def get_d( waist_r: float_or_array | None = None, waist_z: float_or_array | None = None, wavelength: float_or_array = p.wavelength_trap, theta: float_or_array | None = p.theta, # in rad ) -> float_or_array: """Get fringe spacing (in z-direction) for the 2D trap pancakes. If the waists are specified, the angle theta is ignored as it can be substituted using the waists. Either both waists or the angle have to be specified. :param waist_r: :param waist_z: :param wavelength: :param theta: :return: """ if waist_r is not None and waist_z is not None: return wavelength / (2 * waist_z) * waist_r elif theta is not None: return wavelength / (2 * np.cos(theta)) else: raise ValueError( "Either 'waist_r' and 'waist_z', or 'theta' have to be specified." ) def get_a(wavelength: float_or_array) -> float_or_array: """Get real part of the polarizability (Re{alpha}) for a :param wavelength: Laser wavelength :return: """ omega_laser = 2 * np.pi * const.c / wavelength return ( (3 * np.pi * const.c**2) / (2 * p.omega_0**3) * (p.gamma / (p.omega_0 - omega_laser) + p.gamma / (p.omega_0 + omega_laser)) ) def get_power_radial_gaussian( omega: float_or_array, waist: float_or_array, a: float_or_array | None = None, wavelength: float_or_array | None = None, ) -> float_or_array: """Calculate the required power for a radial gaussian beam with a given desired trap frequency and waist. :param omega: Desired trap frequency to compensate. :param waist: Waist of the repulsion beam. :param a: Pre-factor 'a' for the trap depth. :param wavelength: Optional to calculate 'a'. :return: Power of the repulsion beam. """ a = _optional_a(a, wavelength) return np.abs(omega**2 * np.pi * waist**4 * p.mass_lithium6 / (8 * a)) def get_omega_r_trap( power: float_or_array, waist_r: float_or_array, waist_z: float_or_array, a: float_or_array | None = None, wavelength: float_or_array | None = None, ) -> float_or_array: """ :param power: Power of the trapping beam. :param waist_r: Radial waist, denoted :math:`W_{0x}` in Ram-Janik's thesis. Optional. :param waist_z: Waist in z-direction, denoted :math:`W_{0z}`. Optional. :param a: Pre-factor 'a' for the trap depth. :param wavelength: Optional to calculate 'a'. :return: Trapping frequency trap :math:`\\omega_{x}`. """ a = _optional_a(a, wavelength) return np.sqrt(32 * a * power / (np.pi * p.mass_lithium6 * waist_r**3 * waist_z)) def get_omega_z_trap( power_trap: float_or_array, waist_r: float_or_array, waist_z: float_or_array, theta: float_or_array = p.theta, # in rad wavelength: float_or_array = p.wavelength_trap, ) -> float_or_array: a = get_a(wavelength) return np.sqrt( 16 * a * power_trap / (np.pi * p.mass_lithium6 * waist_r * waist_z) * ( 2 * np.sin(theta) ** 2 / waist_z**2 + (wavelength * np.sin(theta) / np.pi) ** 2 * (1 / waist_r**4 + 1 / waist_z**4) + (np.pi / get_d(waist_r, waist_z, wavelength)) ** 2 ) ) def get_power_omega_r( omega_r: float_or_array, waist_r: float_or_array = p.waist_trap_r, waist_z: float_or_array = p.waist_trap_z, a: float_or_array | None = None, wavelength: float_or_array | None = None, ) -> float_or_array: """Calculate the power of the trap based upon the radial properties. :param omega_r: Radial trapping frequency. :param waist_r: Radial waist, denoted :math:`W_{0x}` in Ram-Janik's thesis. Optional. :param waist_z: Waist in z-direction, denoted :math:`W_{0z}`. Optional. :param a: Pre-factor 'a' for the trap depth. :param wavelength: Optional to calculate 'a'. :return: Power of the trapping beam :math:`P_{1,2}`. """ intensity = pot.TotalIntensity() series_x = intensity.subs(pot.z, 0).subs(pot.y, 0).series(x=pot.x, x0=0, n=3) a = _optional_a(a, wavelength) return omega_r**2 * np.pi * waist_r**3 * waist_z * p.mass_lithium6 / (32 * a) def get_power_omega_z( omega_z: float_or_array, waist_r: float_or_array = p.waist_trap_r, waist_z: float_or_array = p.waist_trap_z, a: float_or_array | None = None, wavelength: float_or_array | None = None, ) -> float_or_array: """Calculate the power of the trap based upon the axial properties. :param omega_z: Axial trapping frequency. :param waist_r: Radial waist, denoted :math:`W_{0x}` in Ram-Janik's thesis. Optional. :param waist_z: Waist in z-direction, denoted :math:`W_{0z}`. Optional. :param a: Pre-factor 'a' for the trap depth. :param wavelength: Optional to calculate 'a'. :return: Power of the trapping beam :math:`P_{1,2}`. """ a = _optional_a(a, wavelength) return def gaussian_beam_1d( r: float_or_array, power: float_or_array, waist: float_or_array ) -> float_or_array: """Gaussian beam profile in 1d. :param r: Radial coordinate :param power: Power of the beam :param waist: Waist """ return 2 * power / (np.pi * waist**2) * np.exp(-2 * r**2 / waist**2)