556 lines
		
	
	
		
			21 KiB
		
	
	
	
		
			Python
		
	
	
	
	
	
			
		
		
	
	
			556 lines
		
	
	
		
			21 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=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,
 | |
|             fc=colors_alpha["red"],
 | |
|             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 |