922 lines
38 KiB
Python
922 lines
38 KiB
Python
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)
|