549 lines
21 KiB
Python
549 lines
21 KiB
Python
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
import sympy as sp
|
|
from scipy import constants as const
|
|
from scipy.integrate import quad
|
|
from scipy.optimize import root_scalar
|
|
from tqdm import tqdm
|
|
|
|
import trap_units as si
|
|
|
|
|
|
def plot_solutions(trap,n_levels,left_cutoff,right_cutoff,n_pot_steps=1000,display_plot=-1,state_mult=1e4,plot=True,ret_num=False):
|
|
"""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),
|
|
)
|
|
|
|
if plot:
|
|
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))
|
|
ax.fill_between(
|
|
z_np / si.um,
|
|
potential(z_np) / const.h / si.kHz,
|
|
abs_min / const.h / si.kHz,
|
|
alpha=0.5,
|
|
)
|
|
|
|
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)")
|
|
|
|
if ret_num:
|
|
return np.sum(true_bound_states)
|
|
|
|
def solutions_power(trap,power,n_levels,left_cutoff,right_cutoff):
|
|
trap[trap.power_tweezer] = power
|
|
return plot_solutions(trap,n_levels,left_cutoff,right_cutoff,plot=False, ret_num=True)
|
|
|
|
def root_power(trap,num_atoms,bracket,n_levels,left_cutoff,right_cutoff):
|
|
f = lambda x: solutions_power(trap,x,n_levels,left_cutoff,right_cutoff) - num_atoms
|
|
return root_scalar(f,bracket=bracket).root
|
|
|
|
def sweep_power(trap, gradient,n_levels,left_cutoff,right_cutoff,t_spill=25*si.ms,n_pot_steps=1000,max_atoms=10,n_plotpoints=100):
|
|
"""
|
|
Sweeps over power and gets the power for 0 atom boundary and the precision level
|
|
"""
|
|
|
|
x, y, z = trap.x, trap.y, trap.z
|
|
zr = trap.get_tweezer_rayleigh()
|
|
|
|
trap[trap.grad_z] = gradient
|
|
initial_power = float(4/3/np.sqrt(3)*trap.subs(zr) * np.pi* trap.subs(trap.m * trap.g + trap.mu_b * trap.grad_z) * trap.subs(trap.waist_tweezer**2/trap.a))
|
|
final_power = root_power(trap,max_atoms,[initial_power,initial_power*5],n_levels,left_cutoff,right_cutoff)
|
|
powers = np.linspace(initial_power,final_power,n_plotpoints)
|
|
atom_number = np.zeros_like(powers,dtype=float)
|
|
|
|
|
|
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(zr)),
|
|
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_num = np.sum(true_bound_states)
|
|
atom_number = np.append(atom_number,atom_num)
|
|
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(zr))),
|
|
).root
|
|
intersect_start = root_scalar(
|
|
lambda x: potential(x) - energy,
|
|
bracket=(minimum, barrier),
|
|
).root
|
|
|
|
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]))
|
|
|
|
return powers, atom_number
|
|
|
|
def sweep_power_old(trap, gradient,dp,n_levels,left_cutoff,right_cutoff,t_spill=25*si.ms,max_spill_steps=200,n_pot_steps=1000,max_atoms=10):
|
|
"""
|
|
Sweeps over power and gets the power for 0 atom boundary and the precision level
|
|
"""
|
|
|
|
x, y, z = trap.x, trap.y, trap.z
|
|
zr = trap.get_tweezer_rayleigh()
|
|
|
|
trap[trap.grad_z] = gradient
|
|
initial_power = 4/3/np.sqrt(3)*trap.subs(zr) * np.pi* trap.subs(trap.m * trap.g + trap.mu_b * trap.grad_z) * trap.subs(trap.waist_tweezer**2/trap.a)
|
|
trap[trap.power_tweezer] = initial_power
|
|
powers = np.array([initial_power],dtype=float)
|
|
atom_number = np.array([0.0],dtype=float)
|
|
|
|
i = 0
|
|
pbar = tqdm(disable=True)
|
|
while atom_number[i] <max_atoms and i<max_spill_steps:
|
|
#print(i)
|
|
trap[trap.power_tweezer] = initial_power + (i+1)*dp
|
|
powers = np.append(powers, initial_power + (i+1)*dp)
|
|
|
|
# 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(zr)),
|
|
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_num = np.sum(true_bound_states)
|
|
atom_number = np.append(atom_number,atom_num)
|
|
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(zr))),
|
|
).root
|
|
intersect_start = root_scalar(
|
|
lambda x: potential(x) - energy,
|
|
bracket=(minimum, barrier),
|
|
).root
|
|
|
|
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_num = np.sum(np.exp(-t_spill * tunneling_rate[true_bound_states]))
|
|
atom_number = np.append(atom_number,atom_num)
|
|
i += 1
|
|
pbar.update(1)
|
|
|
|
pbar.close()
|
|
if i == max_spill_steps:
|
|
print(f"{max_atoms} atoms not reached")
|
|
return powers,atom_number
|
|
raise Exception(f"{max_atoms} atoms not reached")
|
|
|
|
return powers, atom_number
|
|
|
|
def solutions_gradient(trap,gradient,n_levels,left_cutoff,right_cutoff):
|
|
trap[trap.grad_z] = gradient
|
|
return plot_solutions(trap,n_levels,left_cutoff,right_cutoff,plot=False, ret_num=True)
|
|
|
|
def root_gradient(trap,num_atoms,bracket,n_levels,left_cutoff,right_cutoff):
|
|
f = lambda x: solutions_gradient(trap,x,n_levels,left_cutoff,right_cutoff) - num_atoms
|
|
return root_scalar(f,bracket=bracket).root
|
|
|
|
def sweep_gradient(trap, power,n_levels,left_cutoff,right_cutoff,t_spill=25*si.ms,n_pot_steps=1000,max_atoms=10,n_plotpoints=100):
|
|
"""
|
|
Sweeps over gradient and gets the gradient for 0 atom boundary and the precision level
|
|
"""
|
|
|
|
x, y, z = trap.x, trap.y, trap.z
|
|
zr = trap.get_tweezer_rayleigh()
|
|
|
|
trap[trap.power_tweezer] = power
|
|
initial_grad = float(trap.subs(1/trap.mu_b * (3*np.sqrt(3)/4 * power * trap.a/np.pi/trap.waist_tweezer**2/zr - trap.m * trap.g)))
|
|
final_grad = root_gradient(trap,max_atoms,[-2.89*si.G/si.cm,initial_grad],n_levels,left_cutoff,right_cutoff)
|
|
gradients = np.linspace(initial_grad,final_grad,n_plotpoints,dtype=float)
|
|
atom_number = np.zeros_like(gradients,dtype=float)
|
|
|
|
for i, gradient in enumerate(tqdm(gradients)):
|
|
#print(i)
|
|
trap[trap.grad_z] = gradient
|
|
|
|
# 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(zr)),
|
|
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_num = np.sum(true_bound_states)
|
|
atom_number = np.append(atom_number,atom_num)
|
|
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(zr))),
|
|
).root
|
|
intersect_start = root_scalar(
|
|
lambda x: potential(x) - energy,
|
|
bracket=(minimum, barrier),
|
|
).root
|
|
|
|
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]))
|
|
|
|
return gradients, atom_number
|
|
|
|
def sweep_gradient_old(trap, power,dgrad,n_levels,left_cutoff,right_cutoff,t_spill=25*si.ms,max_spill_steps=200,n_pot_steps=1000,max_atoms=10):
|
|
"""
|
|
Sweeps over gradient and gets the gradient for 0 atom boundary and the precision level
|
|
"""
|
|
|
|
x, y, z = trap.x, trap.y, trap.z
|
|
zr = trap.get_tweezer_rayleigh()
|
|
|
|
trap[trap.power_tweezer] = power
|
|
initial_grad = float(trap.subs(1/trap.mu_b * (3*np.sqrt(3)/4 * power * trap.a/np.pi/trap.waist_tweezer**2/zr - trap.m * trap.g)))
|
|
trap[trap.grad_z] = initial_grad
|
|
gradients = np.array([initial_grad],dtype=float)
|
|
atom_number = np.array([0.0],dtype=float)
|
|
|
|
i = 0
|
|
pbar = tqdm(disable=True)
|
|
while atom_number[i] <max_atoms and i<max_spill_steps:
|
|
#print(i)
|
|
trap[trap.grad_z] = initial_grad - (i+1)*dgrad
|
|
gradients = np.append(gradients, initial_grad - (i+1)*dgrad)
|
|
|
|
# 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(zr)),
|
|
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_num = np.sum(true_bound_states)
|
|
atom_number = np.append(atom_number,atom_num)
|
|
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(zr))),
|
|
).root
|
|
intersect_start = root_scalar(
|
|
lambda x: potential(x) - energy,
|
|
bracket=(minimum, barrier),
|
|
).root
|
|
|
|
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_num = np.sum(np.exp(-t_spill * tunneling_rate[true_bound_states]))
|
|
atom_number = np.append(atom_number,atom_num)
|
|
i += 1
|
|
pbar.update(1)
|
|
|
|
pbar.close()
|
|
if i == max_spill_steps:
|
|
print(f"{max_atoms} atoms not reached")
|
|
return gradients,atom_number
|
|
raise Exception(f"{max_atoms} atoms not reached")
|
|
|
|
return gradients, atom_number
|
|
|
|
def calculate_stepsize(powers, atom_number,plot=False,max_atoms=10,cutoff=0.1):
|
|
"""
|
|
Given an array of powers (or gradients) and atom_numbers, returns the average length of steps.
|
|
The step length is the area where atom number is within cutoff percent of an integer.
|
|
"""
|
|
|
|
bool_array = np.logical_and(np.abs(atom_number - np.round(atom_number,0)) < cutoff, atom_number > cutoff) #find points where atomnumber is close to integer
|
|
bool_array = np.logical_and(bool_array, atom_number < (max_atoms-0.5))
|
|
if plot:
|
|
plt.plot(powers,atom_number)
|
|
plt.plot(powers[bool_array],atom_number[bool_array],"b.")
|
|
|
|
diff = np.diff(bool_array.astype(int)) # Convert to int to use np.diff
|
|
start_indices = np.where(diff == 1)[0] + 1 # Indices where blobs start
|
|
end_indices = np.where(diff == -1)[0] + 1 # Indices where blobs end
|
|
|
|
# Special case: handle if the array starts or ends with a True
|
|
if bool_array[0]:
|
|
start_indices = np.insert(start_indices, 0, 0) # Add the start of the array
|
|
if bool_array[-1]:
|
|
end_indices = np.append(end_indices, len(bool_array)) # Add the end of the array
|
|
|
|
#step length
|
|
step_len = np.abs(np.mean(np.diff(powers)))
|
|
|
|
# Blob lengths
|
|
blob_lengths = float(np.mean((end_indices - start_indices)*step_len))
|
|
|
|
# Results
|
|
return blob_lengths |