264 lines
10 KiB
Python
264 lines
10 KiB
Python
import numpy as np
|
|
from scipy.integrate import solve_ivp
|
|
from pylcp.common import (progressBar, random_vector)
|
|
from pylcp.integration_tools import solve_ivp_random
|
|
from scipy.interpolate import interp1d
|
|
|
|
"""
|
|
Here, the function "rateeq" in the rateeq module of pylcp is adjusted. Typically, there will be no need to change anything here.
|
|
|
|
Goal:
|
|
Be able to interrupt the calculation of a trajectory in a more useful way.
|
|
|
|
Implementation:
|
|
Set a max boundary and a radius in phase space that are both able
|
|
to interrupt the calculation. If the atom leaves a usefull region or is trapped
|
|
in the MOT, the calculation will be stopped. Thus, a long timespan can be specified
|
|
and the calculation will stop automatically, when reasonable.
|
|
"""
|
|
|
|
def evolve_motion(self, t_span, freeze_axis=[False, False, False],
|
|
random_recoil=False, random_force=False,
|
|
max_scatter_probability=0.1, progress_bar=False,
|
|
record_force=False, rng=np.random.default_rng(),
|
|
boundary=None, MOTradius=None, MOTvelocity=None, **kwargs):
|
|
"""
|
|
Evolve the populations :math:`N` and the motion of the atom in time.
|
|
|
|
This function evolves the rate equations, moving the atom through space,
|
|
given the instantaneous force, for some period of time.
|
|
|
|
Parameters
|
|
----------
|
|
t_span : list or array_like
|
|
A two element list or array that specify the initial and final time
|
|
of integration.
|
|
freeze_axis : list of boolean
|
|
Freeze atomic motion along the specified axis.
|
|
Default: [False, False, False]
|
|
random_recoil : boolean
|
|
Allow the atom to randomly recoil from scattering events.
|
|
Default: False
|
|
random_force : boolean
|
|
Rather than calculating the force using the rateeq.force() method,
|
|
use the calculated scattering rates from each of the laser beam
|
|
(combined with the instantaneous populations) to randomly add photon
|
|
absorption events that cause the atom to recoil randomly from the
|
|
laser beam(s).
|
|
Default: False
|
|
max_scatter_probability : float
|
|
When undergoing random recoils and/or force, this sets the maximum
|
|
time step such that the maximum scattering probability is less than
|
|
or equal to this number during the next time step. Default: 0.1
|
|
progress_bar : boolean
|
|
If true, show a progress bar as the calculation proceeds.
|
|
Default: False
|
|
record_force : boolean
|
|
If true, record the instantaneous force and store in the solution.
|
|
Default: False
|
|
rng : numpy.random.Generator()
|
|
A properly-seeded random number generator. Default: calls
|
|
``numpy.random.default.rng()``
|
|
boundary : float
|
|
Define a boundary of the simulation. This is the maximal position
|
|
the atoms can reach before the calculation is interrupted by an event.
|
|
If boundary = 0, no interruption will happen. The calculation is run
|
|
until t_spn[1] is reached. sry for the stupid coding. (default: 0)
|
|
MOTradius : float
|
|
Same as boundary, but specifies a small radius of "MOT region".
|
|
When the atoms reach this radius and have a velocity smaller than
|
|
MOTvelocity, the caculation will be interrupted by an event.
|
|
If 0, no interruption. (default: 0)
|
|
MOTvelocity : float
|
|
See above. (default : 0)
|
|
**kwargs :
|
|
Additional keyword arguments get passed to solve_ivp_random, which
|
|
is what actually does the integration.
|
|
|
|
Returns
|
|
-------
|
|
sol : OdeSolution
|
|
Bunch object that contains the following fields:
|
|
|
|
* t: integration times found by solve_ivp
|
|
* N: population vs. time
|
|
* v: atomic velocity
|
|
* r: atomic position
|
|
|
|
It contains other important elements, which can be discerned from
|
|
scipy's solve_ivp documentation.
|
|
"""
|
|
free_axes = np.bitwise_not(freeze_axis)
|
|
|
|
if progress_bar:
|
|
progress = progressBar()
|
|
|
|
if record_force:
|
|
ts = []
|
|
Fs = []
|
|
|
|
def motion(t, y):
|
|
N = y[:-6]
|
|
v = y[-6:-3]
|
|
r = y[-3:]
|
|
|
|
Rev, Rijl = self.construct_evolution_matrix(r, v, t)
|
|
if not random_force:
|
|
if record_force:
|
|
F = self.force(r, t, N, return_details=True)
|
|
|
|
ts.append(t)
|
|
Fs.append(F)
|
|
|
|
F = F[0]
|
|
else:
|
|
F = self.force(r, t, N, return_details=False)
|
|
|
|
dydt = np.concatenate((Rev @ N,
|
|
F*free_axes/self.hamiltonian.mass+
|
|
self.constant_accel,
|
|
v))
|
|
else:
|
|
dydt = np.concatenate((Rev @ N,
|
|
self.constant_accel,
|
|
v))
|
|
|
|
if np.any(np.isnan(dydt)):
|
|
raise ValueError('Enountered a NaN!')
|
|
|
|
if progress_bar:
|
|
progress.update(t/t_span[-1])
|
|
|
|
return dydt
|
|
|
|
def random_force_func(t, y, dt):
|
|
total_P = 0
|
|
num_of_scatters = 0
|
|
|
|
# Go over all available keys:
|
|
for key in self.laserBeams:
|
|
# Extract the pumping rate from each laser:
|
|
Rl = np.sum(self.Rijl[key], axis=(1,2))
|
|
|
|
# Calculate the probability to scatter a photon from the laser:
|
|
P = Rl*dt
|
|
|
|
# Roll the dice N times, where $N=\sum(lasers)
|
|
dice = rng.random(len(P))
|
|
|
|
# Give them kicks!
|
|
for ii in np.arange(len(Rl))[dice<P]:
|
|
num_of_scatters += 1
|
|
y[-6:-3] += self.laserBeams[key].beam_vector[ii].kvec(y[-3:], t)/\
|
|
self.hamiltonian.mass
|
|
# Can branch to a differe, lower state, but let's ignore that
|
|
# for the moment.
|
|
y[-6:-3] += self.recoil_velocity[key]*(random_vector(rng, free_axes))
|
|
|
|
total_P += np.sum(P)
|
|
|
|
# Calculate a new maximum dt to make sure we evolve while not
|
|
# exceeding dt max:
|
|
new_dt_max = (max_scatter_probability/total_P)*dt
|
|
|
|
return (num_of_scatters, new_dt_max)
|
|
|
|
|
|
def random_recoil_func(t, y, dt):
|
|
num_of_scatters = 0
|
|
total_P = 0.
|
|
|
|
# Go over each block in the Hamiltonian and compute the decay:
|
|
for key in self.decay_rates:
|
|
P = dt*self.decay_rates[key]*y[self.decay_N_indices[key]]
|
|
|
|
# Roll the dice N times, where $N=\sum(n_i)
|
|
dice = rng.random(len(P))
|
|
|
|
# For any random number that is lower than P_i, add a
|
|
# recoil velocity.
|
|
for ii in range(np.sum(dice<P)):
|
|
num_of_scatters += 1
|
|
y[-6:-3] += self.recoil_velocity[key]*(random_vector(rng, free_axes)+
|
|
random_vector(rng, free_axes))
|
|
|
|
# Save the total probability of a scatter:
|
|
total_P += np.sum(P)
|
|
|
|
# Calculate a new maximum dt to make sure we evolve while not
|
|
# exceeding dt max:
|
|
new_dt_max = (max_scatter_probability/total_P)*dt
|
|
|
|
return (num_of_scatters, new_dt_max)
|
|
|
|
eventList=[]
|
|
|
|
# Define the event function to stop the simulation when r > boundary
|
|
def boundary_event(boundary):
|
|
def event(t, y):
|
|
return 0 if y[-3]**2+y[-2]**2+y[-1]**2 > boundary**2 else 1 # Trigger event when x > boundary
|
|
return event
|
|
|
|
# Define the event function to stop the simulation when position < 0.01 and velocity < 1
|
|
def MOT_event(MOTradius, MOTvelocity):
|
|
def event(t, y):
|
|
radius = y[-3]**2+y[-2]**2+y[-1]**2
|
|
velocity = y[-6]**2+y[-5]**2+y[-4]**2
|
|
return radius - MOTradius**2 if velocity < MOTvelocity**2 else 1
|
|
return event
|
|
|
|
if boundary:
|
|
if type(boundary) == float:
|
|
boundary_event_func = boundary_event(boundary)
|
|
boundary_event_func.terminal = True # Stop the integration on condition
|
|
boundary_event_func.direction = -1 # trigger when decreasing
|
|
eventList.append(boundary_event_func)
|
|
else:
|
|
raise(TypeError(f"Boundary is type {type(boundary)}. Expected float."))
|
|
|
|
if MOTradius and MOTvelocity:
|
|
if type(MOTradius) == float and type(MOTvelocity) == float:
|
|
MOT_event_func = MOT_event(MOTradius, MOTvelocity)
|
|
MOT_event_func.terminal = True # Stop the integration on condition
|
|
MOT_event_func.direction = -1 # Trigger when decreasing
|
|
eventList.append(MOT_event_func)
|
|
else:
|
|
raise(TypeError(f"Motvelocity is type {type(MOTvelocity)} and MOTradius is type {type(MOTradius)}. Expected float for both."))
|
|
|
|
y0 = np.concatenate((self.N0, self.v0, self.r0))
|
|
if random_force:
|
|
self.sol = solve_ivp_random(motion, random_force_func, t_span, y0,
|
|
initial_max_step=max_scatter_probability,
|
|
**kwargs)
|
|
elif random_recoil:
|
|
self.sol = solve_ivp_random(motion, random_recoil_func, t_span, y0,
|
|
initial_max_step=max_scatter_probability,
|
|
**kwargs)
|
|
else:
|
|
self.sol = solve_ivp(motion, t_span, y0, events=eventList, **kwargs)
|
|
|
|
if progress_bar:
|
|
# Just in case the solve_ivp_random terminated due to an event.
|
|
progress.update(1.)
|
|
|
|
# Rearrange the solution:
|
|
self.sol.N = self.sol.y[:-6]
|
|
self.sol.v = self.sol.y[-6:-3]
|
|
self.sol.r = self.sol.y[-3:]
|
|
|
|
if record_force:
|
|
f = interp1d(ts[:-1], np.array([f[0] for f in Fs[:-1]]).T)
|
|
self.sol.F = f(self.sol.t)
|
|
|
|
f = interp1d(ts[:-1], np.array([f[2] for f in Fs[:-1]]).T)
|
|
self.sol.fmag = f(self.sol.t)
|
|
|
|
self.sol.f = {}
|
|
for key in Fs[0][1]:
|
|
f = interp1d(ts[:-1], np.array([f[1][key] for f in Fs[:-1]]).T)
|
|
self.sol.f[key] = f(self.sol.t)
|
|
self.sol.f[key] = np.swapaxes(self.sol.f[key], 0, 1)
|
|
|
|
del self.sol.y
|
|
|
|
return self.sol |