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

617 lines
22 KiB
Python

import matplotlib.pyplot as plt
import numpy as np
import sympy as sp
from IPython.display import Math, display
from matplotlib.axes import Axes
from scipy import constants as const
from scipy.integrate import quad
from scipy.optimize import root_scalar
from scipy.signal import argrelmax,argrelmin
from tqdm import tqdm
import fewfermions.analysis.units as si
from fewfermions.simulate.traps.twod.trap import PancakeTrap
from fewfermions.style import FIGS_PATH, setup
colors, colors_alpha = setup()
def plot_solutions(trap,n_levels,left_cutoff,right_cutoff,n_pot_steps=1000,display_plot=-1,state_mult=1e4):
"""Plot the potential and solutions for a given trap object
(diplay_plot=-1 for all, 0,...,n for individual ones and -2 for none)
"""
x, y, z = trap.x, trap.y, trap.z
axial_width = trap.get_tweezer_rayleigh()
pot_ax = trap.subs(trap.get_potential())
pot_diff_ax = sp.diff(pot_ax, trap.z)
pot_diff2_ax = sp.diff(pot_diff_ax, trap.z)
pot_diff3_ax = sp.diff(pot_diff2_ax, trap.z)
pot_diff_ax_numpy = sp.lambdify(trap.z, pot_diff_ax.subs({x: 0, y: 0}))
pot_diff2_ax_numpy = sp.lambdify(trap.z, pot_diff2_ax.subs({x: 0, y: 0}))
pot_diff3_ax_numpy = sp.lambdify(trap.z, pot_diff3_ax.subs({x: 0, y: 0}))
barrier = root_scalar(
pot_diff_ax_numpy,
x0=1.5 * float(trap.subs(axial_width)),
fprime=pot_diff2_ax_numpy,
xtol=1e-18,
fprime2=pot_diff3_ax_numpy,
).root
minimum = root_scalar(
pot_diff_ax_numpy,
x0=0,
fprime=pot_diff2_ax_numpy,
xtol=1e-28,
fprime2=pot_diff3_ax_numpy,
).root
# Solve the hamiltonian numerically in axial direction
energies, states, potential, coords = trap.nstationary_solution(
trap.z, (left_cutoff, right_cutoff), n_pot_steps, k=n_levels
)
# States that are below the potential barrier
bound_states = energies < potential(barrier)
# Density of states is larger on the left than on the right
# Likely that the state in question is a true bound state
true_bound_states = np.logical_and(
bound_states,
np.sum(states[:, coords[z] < barrier] ** 2, axis=1)
> np.sum(states[:, coords[z] > barrier] ** 2, axis=1),
)
transmission_probability = np.full_like(energies, np.nan, dtype=float)
for j, energy in enumerate(energies):
if not true_bound_states[j]:
continue
intersect_end = root_scalar(
lambda x: potential(x) - energy,
bracket=(barrier, 3 * float(trap.subs(axial_width))),
).root
intersect_start = root_scalar(
lambda x: potential(x) - energy,
bracket=(minimum, barrier),
).root
barrier_interval = np.logical_and(
coords[z] > intersect_start, coords[z] < intersect_end
)
s = quad(
lambda x: np.sqrt(
2
* float(trap.subs(trap.m))
* np.clip(potential(x) - energy, a_min=0, a_max=None)
)
/ const.hbar,
intersect_start,
intersect_end,
)
transmission_probability[j] = sp.exp(-2 * s[0])
lifetime = (
const.h / (transmission_probability * np.abs(energies - potential(minimum)))
)
print(f"{lifetime[true_bound_states]} s")
z_np = np.linspace(left_cutoff, right_cutoff, num=n_pot_steps)
ax: plt.Axes
fig, ax = plt.subplots(figsize=(2.5, 2.5))
# ax.set_title("Axial")
abs_min = np.min(potential(z_np))
print(abs_min)
ax.fill_between(
z_np / si.um,
potential(z_np) / const.h / si.kHz,
abs_min / const.h / si.kHz,
fc=colors_alpha["red"],
alpha=0.5,
)
# ax2 = ax.twinx()
count = 0
for i, bound in enumerate(true_bound_states):
if not bound:
continue
energy = energies[i]
state = states[i]
ax.plot(
z_np / si.um,
np.where(
(energy > potential(z_np)) & (z_np < barrier),
energy / const.h / si.kHz,
np.nan,
),
c="k",
lw=0.5,
marker="None",
)
if display_plot == -1 or display_plot == count:
ax.plot(z_np/si.um, state**2 *state_mult, marker="None", c="k")
count += 1
ax.plot(z_np / si.um, potential(z_np) / const.h / si.kHz, marker="None")
ax.vlines(barrier/ si.um,np.min(potential(z_np) / const.h/ si.kHz),np.max(potential(z_np) / const.h/ si.kHz))
ax.vlines(minimum/ si.um,np.min(potential(z_np) / const.h/ si.kHz),np.max(potential(z_np) / const.h/ si.kHz))
ax.set_title(f"{np.sum(true_bound_states)} bound solutions, power: {trap.subs(trap.power_tweezer)/si.mW}mW")
ax.set_xlabel(r"z ($\mathrm{\mu m}$)")
ax.set_ylabel(r"E / h (kHz)")
def plot_solutions_ax(trap,n_levels,left_cutoff,right_cutoff,n_pot_steps=1000,display_plot=-1,state_mult=1e4):
"""Plot the potential and solutions in the z-direction for a given trap object with no spilling
(diplay_plot=-1 for all, 0,...,n for individual ones and -2 for none)
"""
#turn off magnetic gradient and gravity
trap[trap.g] = 0
initial_gradient = trap.grad_z
trap[trap.grad_z] = 0
trap
x, y, z = trap.x, trap.y, trap.z
axial_width = trap.get_tweezer_rayleigh()
pot_ax = trap.subs(trap.get_potential())
pot_diff_ax = sp.diff(pot_ax, trap.z)
pot_diff2_ax = sp.diff(pot_diff_ax, trap.z)
pot_diff3_ax = sp.diff(pot_diff2_ax, trap.z)
pot_diff_ax_numpy = sp.lambdify(trap.z, pot_diff_ax.subs({x: 0, y: 0}))
pot_diff2_ax_numpy = sp.lambdify(trap.z, pot_diff2_ax.subs({x: 0, y: 0}))
pot_diff3_ax_numpy = sp.lambdify(trap.z, pot_diff3_ax.subs({x: 0, y: 0}))
"""
barrier = root_scalar(
pot_diff_ax_numpy,
x0=1.5 * float(trap.subs(axial_width)),
fprime=pot_diff2_ax_numpy,
xtol=1e-18,
fprime2=pot_diff3_ax_numpy,
).root
minimum = root_scalar(
pot_diff_ax_numpy,
x0=0,
fprime=pot_diff2_ax_numpy,
xtol=1e-28,
fprime2=pot_diff3_ax_numpy,
).root
"""
# Solve the hamiltonian numerically in axial direction
energies, states, potential, coords = trap.nstationary_solution(
trap.z, (left_cutoff, right_cutoff), n_pot_steps, k=n_levels
)
# States that are below the potential barrier
#bound_states = energies < potential(barrier)
bound_states = energies < potential(left_cutoff)
# Density of states is larger on the left than on the right
# Likely that the state in question is a true bound state
"""true_bound_states = np.logical_and(
bound_states,
np.sum(states[:, coords[x] < barrier] ** 2, axis=1)
> np.sum(states[:, coords[x] > barrier] ** 2, axis=1),
)"""
true_bound_states = bound_states
"""
transmission_probability = np.full_like(energies, np.nan, dtype=float)
for j, energy in enumerate(energies):
if not true_bound_states[j]:
continue
intersect_end = root_scalar(
lambda x: potential(x) - energy,
bracket=(barrier, 3 * float(trap.subs(axial_width))),
).root
intersect_start = root_scalar(
lambda x: potential(x) - energy,
bracket=(minimum, barrier),
).root
barrier_interval = np.logical_and(
coords[x] > intersect_start, coords[x] < intersect_end
)
s = quad(
lambda x: np.sqrt(
2
* float(trap.subs(trap.m))
* np.clip(potential(x) - energy, a_min=0, a_max=None)
)
/ const.hbar,
intersect_start,
intersect_end,
)
transmission_probability[j] = sp.exp(-2 * s[0])
lifetime = (
const.h / (transmission_probability * np.abs(energies - potential(minimum)))
)
print(f"{lifetime[true_bound_states]} s")"""
z_np = np.linspace(left_cutoff, right_cutoff, num=n_pot_steps)
ax: plt.Axes
fig, ax = plt.subplots(figsize=(2.5, 2.5))
# ax.set_title("Axial")
abs_min = np.min(potential(z_np))
print(abs_min)
ax.fill_between(
z_np / si.um,
potential(z_np) / const.h / si.kHz,
abs_min / const.h / si.kHz,
fc=colors_alpha["red"],
alpha=0.5,
)
# ax2 = ax.twinx()
count = 0
for i, bound in enumerate(true_bound_states):
if not bound:
continue
energy = energies[i]
print((np.max(potential(z_np))-energy)/ const.h / si.kHz)
state = states[i]
ax.plot(
z_np / si.um,
np.where(
(energy > potential(z_np)), #& (z_np < barrier),
energy / const.h / si.kHz,
np.nan,
),
c="k",
lw=0.5,
marker="None",
)
if display_plot == -1 or display_plot == count:
ax.plot(z_np/si.um, state**2 *state_mult, marker="None", c="k")
count += 1
ax.plot(z_np / si.um, potential(z_np) / const.h / si.kHz, marker="None")
#ax.vlines(barrier/ si.um,np.min(potential(z_np) / const.h/ si.kHz),np.max(potential(z_np) / const.h/ si.kHz))
#ax.vlines(minimum/ si.um,np.min(potential(z_np) / const.h/ si.kHz),np.max(potential(z_np) / const.h/ si.kHz))
#ax.set_title(f"{np.sum(true_bound_states)} bound solutions, power: {trap.subs(trap.power_tweezer)/si.mW}mW")
ax.set_xlabel(r"z ($\mathrm{\mu m}$)")
ax.set_ylabel(r"E / h (kHz)")
#turn on magnetic gradient and gravity
trap[trap.g] = const.g
trap[trap.grad_z] = initial_gradient
def plot_solutions_rad(trap,n_levels,left_cutoff,right_cutoff,n_pot_steps=1000,display_plot=-1,state_mult=1e4):
"""Plot the potential and solutions in the x-direction for a given trap object
(diplay_plot=-1 for all, 0,...,n for individual ones and -2 for none)
"""
x, y, z = trap.x, trap.y, trap.z
axial_width = trap.get_tweezer_rayleigh()
pot_ax = trap.subs(trap.get_potential())
pot_diff_ax = sp.diff(pot_ax, trap.x)
pot_diff2_ax = sp.diff(pot_diff_ax, trap.x)
pot_diff3_ax = sp.diff(pot_diff2_ax, trap.x)
pot_diff_ax_numpy = sp.lambdify(trap.x, pot_diff_ax.subs({x: 0, y: 0}))
pot_diff2_ax_numpy = sp.lambdify(trap.x, pot_diff2_ax.subs({x: 0, y: 0}))
pot_diff3_ax_numpy = sp.lambdify(trap.x, pot_diff3_ax.subs({x: 0, y: 0}))
"""
barrier = root_scalar(
pot_diff_ax_numpy,
x0=1.5 * float(trap.subs(axial_width)),
fprime=pot_diff2_ax_numpy,
xtol=1e-18,
fprime2=pot_diff3_ax_numpy,
).root
minimum = root_scalar(
pot_diff_ax_numpy,
x0=0,
fprime=pot_diff2_ax_numpy,
xtol=1e-28,
fprime2=pot_diff3_ax_numpy,
).root
"""
# Solve the hamiltonian numerically in axial direction
energies, states, potential, coords = trap.nstationary_solution(
trap.x, (left_cutoff, right_cutoff), n_pot_steps, k=n_levels
)
# States that are below the potential barrier
#bound_states = energies < potential(barrier)
bound_states = energies < potential(left_cutoff)
# Density of states is larger on the left than on the right
# Likely that the state in question is a true bound state
"""true_bound_states = np.logical_and(
bound_states,
np.sum(states[:, coords[x] < barrier] ** 2, axis=1)
> np.sum(states[:, coords[x] > barrier] ** 2, axis=1),
)"""
true_bound_states = bound_states
"""
transmission_probability = np.full_like(energies, np.nan, dtype=float)
for j, energy in enumerate(energies):
if not true_bound_states[j]:
continue
intersect_end = root_scalar(
lambda x: potential(x) - energy,
bracket=(barrier, 3 * float(trap.subs(axial_width))),
).root
intersect_start = root_scalar(
lambda x: potential(x) - energy,
bracket=(minimum, barrier),
).root
barrier_interval = np.logical_and(
coords[x] > intersect_start, coords[x] < intersect_end
)
s = quad(
lambda x: np.sqrt(
2
* float(trap.subs(trap.m))
* np.clip(potential(x) - energy, a_min=0, a_max=None)
)
/ const.hbar,
intersect_start,
intersect_end,
)
transmission_probability[j] = sp.exp(-2 * s[0])
lifetime = (
const.h / (transmission_probability * np.abs(energies - potential(minimum)))
)
print(f"{lifetime[true_bound_states]} s")"""
z_np = np.linspace(left_cutoff, right_cutoff, num=n_pot_steps)
ax: plt.Axes
fig, ax = plt.subplots(figsize=(2.5, 2.5))
# ax.set_title("Axial")
abs_min = np.min(potential(z_np))
print(abs_min)
ax.fill_between(
z_np / si.um,
potential(z_np) / const.h / si.kHz,
abs_min / const.h / si.kHz,
fc=colors_alpha["red"],
alpha=0.5,
)
# ax2 = ax.twinx()
count = 0
for i, bound in enumerate(true_bound_states):
if not bound:
continue
energy = energies[i]
print((np.max(potential(z_np))-energy)/ const.h / si.kHz)
state = states[i]
ax.plot(
z_np / si.um,
np.where(
(energy > potential(z_np)), #& (z_np < barrier),
energy / const.h / si.kHz,
np.nan,
),
c="k",
lw=0.5,
marker="None",
)
if display_plot == -1 or display_plot == count:
ax.plot(z_np/si.um, state**2 *state_mult, marker="None", c="k")
count += 1
ax.plot(z_np / si.um, potential(z_np) / const.h / si.kHz, marker="None")
#ax.vlines(barrier/ si.um,np.min(potential(z_np) / const.h/ si.kHz),np.max(potential(z_np) / const.h/ si.kHz))
#ax.vlines(minimum/ si.um,np.min(potential(z_np) / const.h/ si.kHz),np.max(potential(z_np) / const.h/ si.kHz))
#ax.set_title(f"{np.sum(true_bound_states)} bound solutions, power: {trap.subs(trap.power_tweezer)/si.mW}mW")
ax.set_xlabel(r"z ($\mathrm{\mu m}$)")
ax.set_ylabel(r"E / h (kHz)")
def plot_occupation(trap,n_levels,left_cutoff,right_cutoff,t_spill=25*si.ms,n_spill_steps=100,power_fac_down=0.4,power_fac_up=0.7,n_pot_steps=1000, results=False):
"""
Plotting the occupation of the tweezer when varying the tweezer power.
"""
x, y, z = trap.x, trap.y, trap.z
axial_width = trap.get_tweezer_rayleigh()
initial_power = trap[trap.power_tweezer]
spill_power_factor = np.linspace(power_fac_up, power_fac_down, num=n_spill_steps)
powers = trap[trap.power_tweezer] * spill_power_factor
atom_number = np.zeros_like(powers,dtype=float)
# Resolution of the potential when solving numerically
#n_pot_steps = 1000
for i, power in enumerate(tqdm(powers)):
trap[trap.power_tweezer] = power
# Solve the hamiltonian numerically in axial direction
energies, states, potential, coords = trap.nstationary_solution(
trap.z, (left_cutoff, right_cutoff), n_pot_steps, k=n_levels
)
# Determine the potential and its derivatives
pot_ax = trap.subs(trap.get_potential())
pot_diff_ax = sp.diff(pot_ax, trap.z)
pot_diff2_ax = sp.diff(pot_diff_ax, trap.z)
pot_diff3_ax = sp.diff(pot_diff2_ax, trap.z)
pot_diff_ax_numpy = sp.lambdify(trap.z, pot_diff_ax.subs({x: 0, y: 0}))
pot_diff2_ax_numpy = sp.lambdify(trap.z, pot_diff2_ax.subs({x: 0, y: 0}))
pot_diff3_ax_numpy = sp.lambdify(trap.z, pot_diff3_ax.subs({x: 0, y: 0}))
barrier = root_scalar(
pot_diff_ax_numpy,
x0=1.5 * float(trap.subs(axial_width)),
fprime=pot_diff2_ax_numpy,
xtol=1e-28,
fprime2=pot_diff3_ax_numpy,
).root
minimum = root_scalar(
pot_diff_ax_numpy,
x0=0,
fprime=pot_diff2_ax_numpy,
xtol=1e-28,
fprime2=pot_diff3_ax_numpy,
).root
# States that are below the potential barrier
bound_states = energies < potential(barrier)
# Density of states is larger on the left than on the right
# Likely that the state in question is a true bound state
true_bound_states = np.logical_and(
bound_states,
np.sum(states[:, coords[z] < barrier] ** 2, axis=1)
> np.sum(states[:, coords[z] > barrier] ** 2, axis=1),
)
if t_spill == 0:
atom_number[i] = np.sum(true_bound_states)
continue
transmission_probability = np.full_like(energies, np.nan, dtype=float)
for j, energy in enumerate(energies):
if not true_bound_states[j]:
continue
intersect_end = root_scalar(
lambda x: potential(x) - energy,
bracket=(barrier, 3 * float(trap.subs(axial_width))),
).root
intersect_start = root_scalar(
lambda x: potential(x) - energy,
bracket=(minimum, barrier),
).root
barrier_interval = np.logical_and(
coords[z] > intersect_start, coords[z] < intersect_end
)
s = quad(
lambda x: np.sqrt(
2
* float(trap.subs(trap.m))
* np.clip(potential(x) - energy, a_min=0, a_max=None)
)
/ const.hbar,
intersect_start,
intersect_end,
)
transmission_probability[j] = sp.exp(-2 * s[0])
tunneling_rate = (
transmission_probability * np.abs(energies - potential(minimum)) / const.h
)
atom_number[i] = np.sum(np.exp(-t_spill * tunneling_rate[true_bound_states]))
ax: plt.Axes
fig: plt.Figure
fig, ax = plt.subplots(figsize=(2.5, 2.5))
ax.set_title(f"initial power:{initial_power/si.mW}mW, B':{trap[trap.grad_z]/si.G*si.cm}G/cm, w_0:{trap[trap.waist_tweezer]/si.um}um, wvl:{trap[trap.wvl]/si.nm}nm")
ax.set_xlabel("rel. tweezer power (a.u.)")
ax.set_ylabel("atom number")
ax.plot(spill_power_factor, atom_number, marker="None")
#ax.fill_between(spill_power_factor, atom_number, fc=colors_alpha["red"], alpha=0.5)
if results:
return spill_power_factor, atom_number
def plot_occupation_grad(trap,n_levels,left_cutoff,right_cutoff,t_spill=25*si.ms,n_spill_steps=100,grad_fac_down=0.4,grad_fac_up=0.7,n_pot_steps=1000):
"""
Plotting the occupation of the tweezer when varying the magnetic gradient.
"""
x, y, z = trap.x, trap.y, trap.z
axial_width = trap.get_tweezer_rayleigh()
initial_power = trap[trap.power_tweezer]
initial_grad = trap[trap.grad_z]
spill_grad_factor = np.linspace(grad_fac_up, grad_fac_down, num=n_spill_steps)
grads = trap[trap.grad_z] * spill_grad_factor
atom_number = np.zeros_like(grads)
# Resolution of the potential when solving numerically
#n_pot_steps = 1000
for i, grad in enumerate(tqdm(grads)):
trap[trap.grad_z] = grad
# Solve the hamiltonian numerically in axial direction
energies, states, potential, coords = trap.nstationary_solution(
trap.z, (left_cutoff, right_cutoff), n_pot_steps, k=n_levels
)
# Determine the potential and its derivatives
pot_ax = trap.subs(trap.get_potential())
pot_diff_ax = sp.diff(pot_ax, trap.z)
pot_diff2_ax = sp.diff(pot_diff_ax, trap.z)
pot_diff3_ax = sp.diff(pot_diff2_ax, trap.z)
pot_diff_ax_numpy = sp.lambdify(trap.z, pot_diff_ax.subs({x: 0, y: 0}))
pot_diff2_ax_numpy = sp.lambdify(trap.z, pot_diff2_ax.subs({x: 0, y: 0}))
pot_diff3_ax_numpy = sp.lambdify(trap.z, pot_diff3_ax.subs({x: 0, y: 0}))
barrier = root_scalar(
pot_diff_ax_numpy,
x0=1.5 * float(trap.subs(axial_width)),
fprime=pot_diff2_ax_numpy,
xtol=1e-28,
fprime2=pot_diff3_ax_numpy,
).root
minimum = root_scalar(
pot_diff_ax_numpy,
x0=0,
fprime=pot_diff2_ax_numpy,
xtol=1e-28,
fprime2=pot_diff3_ax_numpy,
).root
# States that are below the potential barrier
bound_states = energies < potential(barrier)
# Density of states is larger on the left than on the right
# Likely that the state in question is a true bound state
true_bound_states = np.logical_and(
bound_states,
np.sum(states[:, coords[z] < barrier] ** 2, axis=1)
> np.sum(states[:, coords[z] > barrier] ** 2, axis=1),
)
if t_spill == 0:
atom_number[i] = np.sum(true_bound_states)
continue
transmission_probability = np.full_like(energies, np.nan, dtype=float)
for j, energy in enumerate(energies):
if not true_bound_states[j]:
continue
intersect_end = root_scalar(
lambda x: potential(x) - energy,
bracket=(barrier, 3 * float(trap.subs(axial_width))),
).root
intersect_start = root_scalar(
lambda x: potential(x) - energy,
bracket=(minimum, barrier),
).root
barrier_interval = np.logical_and(
coords[z] > intersect_start, coords[z] < intersect_end
)
s = quad(
lambda x: np.sqrt(
2
* float(trap.subs(trap.m))
* np.clip(potential(x) - energy, a_min=0, a_max=None)
)
/ const.hbar,
intersect_start,
intersect_end,
)
transmission_probability[j] = sp.exp(-2 * s[0])
tunneling_rate = (
transmission_probability * np.abs(energies - potential(minimum)) / const.h
)
atom_number[i] = np.sum(np.exp(-t_spill * tunneling_rate[true_bound_states]))
ax: plt.Axes
fig: plt.Figure
fig, ax = plt.subplots(figsize=(2.5, 2.5))
ax.plot(spill_grad_factor, atom_number, marker="None")
ax.fill_between(spill_grad_factor, atom_number, fc=colors_alpha["red"], alpha=0.5)
ax.set_title(f"power:{initial_power/si.mW}mW, initial B':{initial_grad/si.G*si.cm}G/cm, w_0:{trap[trap.waist_tweezer]/si.um}um, wvl:{trap[trap.wvl]/si.nm}nm")
ax.set_xlabel("rel. gradient (a.u.)")
ax.set_ylabel("atom number")