189 lines
7.6 KiB
Python
189 lines
7.6 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,find_peaks
|
|
from tqdm import tqdm
|
|
|
|
import fewfermions.analysis.units as si
|
|
from fewfermions.simulate.traps.twod.trap import DoubleTweezer
|
|
|
|
def plot_solutions(trap,n_levels,left_cutoff,right_cutoff,n_pot_steps=1000,display_plot=-1,state_mult=1e4,plot=True,ret_results=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)
|
|
"""
|
|
|
|
# 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(left_cutoff)
|
|
|
|
if plot:
|
|
z_np = np.linspace(left_cutoff, right_cutoff, num=n_pot_steps)
|
|
ax: plt.Axes
|
|
fig, ax = plt.subplots(figsize=(5, 5))
|
|
ax.plot(z_np / si.um, potential(z_np) / const.h / si.kHz,color="cornflowerblue" ,marker="None")
|
|
ax.set_title(f"{np.sum(bound_states)} bound solutions, tweezer distance: {trap.subs(trap.distance_tweezers)/si.um}um")
|
|
ax.set_xlabel(r"x ($\mathrm{\mu m}$)")
|
|
ax.set_ylabel(r"E / h (kHz)")
|
|
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,
|
|
color="cornflowerblue"
|
|
)
|
|
|
|
count = 0
|
|
for i, bound in enumerate(bound_states):
|
|
if not bound:
|
|
continue
|
|
energy = energies[i]
|
|
state = states[i]
|
|
ax.plot(
|
|
z_np / si.um,
|
|
np.where(
|
|
(energy > potential(z_np)),
|
|
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 *state_mult, marker="None",label=f"state {count}")#, c="k")
|
|
count += 1
|
|
|
|
plt.legend()
|
|
|
|
|
|
if ret_results:
|
|
return np.sum(bound_states), energies[bound_states], states[bound_states], potential
|
|
|
|
|
|
def loop_distances(trap,distances,n_levels=10,n_pot_steps=2000):
|
|
"""
|
|
returns unsorted arrays for energies, states and the potential for a given array of tweezer distances
|
|
"""
|
|
left_cutoff = -0.5*np.max(distances)-3*np.max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])
|
|
right_cutoff = 0.5*np.max(distances)+3*np.max([float(trap.subs(trap.waist_tweezer1)),float(trap.subs(trap.waist_tweezer2))])
|
|
x_np = np.linspace(left_cutoff,right_cutoff,n_pot_steps)
|
|
|
|
energies = np.zeros((len(distances),n_levels))
|
|
states = np.zeros((len(distances),n_levels,n_pot_steps))
|
|
potential = np.zeros((len(distances),n_pot_steps))
|
|
|
|
#diagonalise hamiltonian for all distances
|
|
for i, dist in enumerate(distances):
|
|
trap[trap.distance_tweezers] = dist
|
|
num, ener, state, pot = plot_solutions(trap,n_levels,np.min(x_np),np.max(x_np),display_plot=0,state_mult=450,n_pot_steps=n_pot_steps,plot=False,ret_results=True)
|
|
|
|
# fill in arrays
|
|
energies[i,:len(ener)] = ener
|
|
states[i,:len(ener)] = state
|
|
potential[i] = pot(x_np)
|
|
|
|
return energies, states, potential
|
|
|
|
def find_crossovers_topstate(energies):
|
|
"""
|
|
returns the indices(distance) of all crossovers (with calculated and not calculated states) of the largest calculated state
|
|
"""
|
|
arr = np.abs(np.gradient(np.gradient(energies[:,-1])))
|
|
index = find_peaks(arr,width=(1,4))[0]
|
|
return index
|
|
|
|
def find_crossovers(energies,mult_min_energy=0.01):
|
|
"""
|
|
returns the indices(distance) of all crossovers between two calculated states
|
|
(very dependent on mult_min_energy)
|
|
"""
|
|
n_levels = len(energies[0,:])
|
|
index = np.array([],dtype=int)
|
|
swap_index =np.empty((0,2),dtype=int)
|
|
|
|
for i in range(n_levels):
|
|
for j in range(i+1,n_levels):
|
|
diff = np.abs(energies[:,i] - energies[:,j])
|
|
mask = np.where(diff < mult_min_energy*np.max(diff),diff,mult_min_energy*np.max(diff))
|
|
peaks = find_peaks(-mask)[0]
|
|
index = np.append(index,peaks)
|
|
for k in range(len(peaks)):
|
|
swap_index = np.append(swap_index,[[i,j]],axis=0)
|
|
|
|
return index, swap_index
|
|
|
|
def swapped_loop_distance(distances, energies, states, potentials):
|
|
"""
|
|
returns the results of loop_distance() with energies and states matching to individual states
|
|
"""
|
|
|
|
index_top = find_crossovers_topstate(energies)
|
|
index, swap_index = find_crossovers(energies)
|
|
|
|
new_energies = np.full((energies.shape[0],energies.shape[1]+len(index_top)),np.nan)
|
|
new_states = np.full((states.shape[0],states.shape[1]+len(index_top),states.shape[2]),np.nan)
|
|
|
|
swapped_index = np.arange(0,energies.shape[1])
|
|
for i, dist in enumerate(distances):
|
|
if np.any(i == index):
|
|
j = np.where(index == i)
|
|
swapped_index[swap_index[j]] = swapped_index[np.roll(swap_index[j],1)]
|
|
#print(f"crossover of states {swap_index[j]} at {dist/si.um:.2f}um")
|
|
|
|
elif np.any(i == index_top):
|
|
#print(f"crossover of top state at {dist/si.um:.2f}um")
|
|
swapped_index[-1] = np.max(swapped_index)+1
|
|
|
|
|
|
new_energies[i,swapped_index] = energies[i]
|
|
new_states[i,swapped_index] = states[i]
|
|
|
|
return new_energies, new_states, potentials, index_top, index, swap_index
|
|
|
|
def find_ass_tweezer(new_energies,new_states,return_deltaE=True):
|
|
"""
|
|
Sorts energies and states into those coming from left and right tweezer.
|
|
|
|
Returns the minimum energy difference of the groundstate of the shallow tweezer to any of the other tweezers eigenstates.
|
|
Also returns the occupation number when the groundstates of both tweezers were initially occupied (when return_deltaE=True)
|
|
"""
|
|
#seperate energies and states corresponding to the initial tweezer they belonged to
|
|
mask_right = np.where(np.sum(new_states[0,:,:int(new_states.shape[2]/2)]**2,axis=1) < 0.5,True, False)
|
|
mask_left = np.logical_not(mask_right)
|
|
|
|
energies_right = new_energies[:,mask_right]
|
|
energies_left = new_energies[:,mask_left]
|
|
|
|
states_right = new_states[:,mask_right,:]
|
|
states_left = new_states[:,mask_left,:]
|
|
|
|
#identify GS of shallower tweezer and calculated minimal energy difference to other states
|
|
if mask_left[np.nanargmin(new_energies[0])]:
|
|
energy_GS = energies_right[:,0]
|
|
state_GS = states_right[:,0]
|
|
#smallest energy difference
|
|
delta_E_min = np.nanmin(np.abs(energy_GS[:,np.newaxis] - energies_left))
|
|
#final occupation if two GS are involved
|
|
occ_numbers = np.array([0,np.searchsorted(energies_left[-1],energy_GS[-1])])
|
|
else:
|
|
energy_GS = energies_left[:,0]
|
|
state_GS = states_left[:,0]
|
|
#smallest energy difference
|
|
delta_E_min = np.nanmin(np.abs(energy_GS[:,np.newaxis] - energies_right))
|
|
#final occupation if two GS are involved
|
|
occ_numbers = np.array([0,np.searchsorted(energies_right[-1],energy_GS[-1])])
|
|
|
|
if return_deltaE:
|
|
return energies_left, energies_right, states_left, states_right, delta_E_min, occ_numbers
|
|
else:
|
|
return energies_left, energies_right, states_left, states_right |