#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Tools for solving the rate equations.
"""
import numpy as np
import copy
from scipy.optimize import minimize
from scipy.integrate import solve_ivp
from inspect import signature
from .fields import laserBeams, magField
from .common import (progressBar, random_vector, spherical_dot,
cart2spherical, spherical2cart, base_force_profile)
from .governingeq import governingeq
from .integration_tools import solve_ivp_random
from scipy.interpolate import interp1d
#@numba.vectorize([numba.float64(numba.complex128),numba.float32(numba.complex64)])
def abs2(x):
return x.real**2 + x.imag**2
[docs]class force_profile(base_force_profile):
"""
Rate equation force profile
The force profile object stores all of the calculated quantities created by
the rateeq.generate_force_profile() method. It has following attributes:
Attributes
----------
R : array_like, shape (3, ...)
Positions at which the force profile was calculated.
V : array_like, shape (3, ...)
Velocities at which the force profile was calculated.
F : array_like, shape (3, ...)
Total equilibrium force at position R and velocity V.
f_mag : array_like, shape (3, ...)
Magnetic force at position R and velocity V.
f : dictionary of array_like
The forces due to each laser, indexed by the
manifold the laser addresses. The dictionary is keyed by the transition
driven, and individual lasers are in the same order as in the
pylcp.laserBeams object used to create the governing equation.
Neq : array_like
Equilibrium population found.
Rijl : dictionary of array_like
The pumping rates of each laser, indexed by the
manifold the laser addresses. The dictionary is keyed by the transition
driven, and individual lasers are in the same order as in the
pylcp.laserBeams object used to create the governing equation.
"""
def __init__(self, R, V, laserBeams, hamiltonian):
super().__init__(R, V, laserBeams, hamiltonian)
# Add in the specific instance of the Rijl:
self.Rijl = {}
for key in laserBeams:
self.Rijl[key] = np.zeros(
self.R[0].shape+(len(laserBeams[key].beam_vector),
hamiltonian.blocks[hamiltonian.laser_keys[key]].n,
hamiltonian.blocks[hamiltonian.laser_keys[key]].m)
)
def store_data(self, ind, Neq, F, F_laser, F_mag, Rijl):
super().store_data(ind, Neq, F, F_laser, F_mag)
for key in Rijl:
self.Rijl[key][ind] = Rijl[key]
[docs]class rateeq(governingeq):
"""
The rate equations
This class constructs the rate equations from the given laser
beams, magnetic field, and hamiltonian.
Parameters
----------
laserBeams : dictionary of pylcp.laserBeams, pylcp.laserBeams, or list of pylcp.laserBeam
The laserBeams that will be used in constructing the optical Bloch
equations. which transitions in the block diagonal hamiltonian. It can
be any of the following:
* A dictionary of pylcp.laserBeams: if this is the case, the keys of
the dictionary should match available :math:`d^{nm}` matrices
in the pylcp.hamiltonian object. The key structure should be
`n->m`.
* pylcp.laserBeams: a single set of laser beams is assumed to
address the transition `g->e`.
* a list of pylcp.laserBeam: automatically promoted to a
pylcp.laserBeams object assumed to address the transtion `g->e`.
magField : pylcp.magField or callable
The function or object that defines the magnetic field.
hamiltonian : pylcp.hamiltonian
The internal hamiltonian of the particle.
a : array_like, shape (3,), optional
A default acceleraiton to apply to the particle's motion, usually
gravity. Default: [0., 0., 0.]
include_mag_forces : boolean
Optional flag to inculde magnetic forces in the force calculation.
Default: True
r0 : array_like, shape (3,), optional
Initial position. Default: [0., 0., 0.]
v0 : array_like, shape (3,), optional
Initial velocity. Default: [0., 0., 0.]
"""
def __init__(self, laserBeams, magField, hamitlonian,
a=np.array([0., 0., 0.]), include_mag_forces=True,
svd_eps=1e-10, r0=np.array([0., 0., 0.]),
v0=np.array([0., 0., 0.])):
# First step is to save the imported laserBeams, magField, and
# hamiltonian.
super().__init__(laserBeams, magField, hamitlonian,
a=a, r0=r0, v0=v0)
self.include_mag_forces = include_mag_forces
self.svd_eps = svd_eps
# Check function signatures for any time dependence:
self.tdepend = {}
self.tdepend['B'] = False
"""self.tdepend['pol'] = False
self.tdepend['kvec'] = False
self.tdepend['det'] = False
self.tdepend['intensity'] = False"""
"""if 't' in str(signature(self.magField)):
self.tdepend['B'] = True
for key in self.laserBeams:
for beam in self.laserBeams[key]:
if not beam.pol_sig is None and 't' in beam.pol_sig:
self.tdepend['pol'] = True
if not beam.kvec_sig is None and 't' in beam.kvec_sig:
self.tdepend['kvec'] = True
if not beam.delta_sig is None and 't' in beam.delta_sig:
self.tdepend['det'] = True
if not beam.intensity_sig is None and 't' in beam.intensity_sig:
self.tdepend['intensity'] = True"""
# Set up two dictionaries that are useful for both random forces and
# random recoils:
self.decay_rates = {}
self.decay_N_indices = {}
# If the matrix is diagonal, we get to do something cheeky. Let's just
# construct the decay part of the evolution once:
if np.all(self.hamiltonian.diagonal):
self._calc_decay_comp_of_Rev(self.hamiltonian)
# Save the recoil velocity for the relevant transitions:
self.recoil_velocity = {}
for key in self.hamiltonian.laser_keys:
self.recoil_velocity[key] = \
self.hamiltonian.blocks[self.hamiltonian.laser_keys[key]].parameters['k']\
/self.hamiltonian.mass
# A dictionary to store the pumping rates.
self.Rijl = {}
# Set up a dictionary to store the profiles.
self.profile = {}
def _calc_decay_comp_of_Rev(self, rotated_ham):
"""
Constructs the decay portion of the evolution matrix.
Parameters
----------
rotated_ham: pylcp.hamiltonian object
The diagonalized hamiltonian with rotated d_q matrices
"""
self.Rev_decay = np.zeros((self.hamiltonian.n, self.hamiltonian.n))
# Go through each of the blocks and calculate the decay rate
# contributions to the rate equations. This would be simpler with
# rotated_ham.make_full_matrices(), but I would imagine it would
# be significantly slower.
# Let's use the laser_keys dictionary to go through the non-zero d_q:
for key in self.hamiltonian.laser_keys:
# The index of the current d_q matrix:
ind = rotated_ham.laser_keys[key]
# Grab the block of interest:
d_q_block = rotated_ham.blocks[ind]
# The offset index for the lower states:
noff = int(np.sum(rotated_ham.ns[:ind[0]]))
# The offset index for the higher states:
moff = int(np.sum(rotated_ham.ns[:ind[1]]))
# The number of lower states:
n = rotated_ham.ns[ind[0]]
# The number of higher states:
m = rotated_ham.ns[ind[1]]
# Calculate the decay of the states attached to this block:
gamma = d_q_block.parameters['gamma']
k = d_q_block.parameters['k']
# Let's see if we can avoid a for loop here:
# for jj in range(rotated_ham.ns[ll]):
# self.Rev_decay[n+jj, n+jj] -= gamma*\
# np.sum(np.sum(abs2(
# other_block.matrix[:, :, jj]
# )))
# Save the decay rates out of the excited state:
self.decay_rates[key] = gamma*np.sum(abs2(
d_q_block.matrix[:, :, :]
), axis=(0,1))
# Save the indices for the excited states of this d_q block
# for the random_recoil function:
self.decay_N_indices[key] = np.arange(moff, moff+m)
# Add (more accurately, subtract) these decays to the evolution matrix:
self.Rev_decay[(np.arange(moff, moff+m), np.arange(moff, moff+m))] -= \
self.decay_rates[key]
# Now calculate the decays into the lower state connected by this
# d_q:
self.Rev_decay[noff:noff+n, moff:moff+m] += \
gamma*np.sum(abs2(d_q_block.matrix), axis=0)
return self.Rev_decay
def _calc_pumping_rates(self, r, v, t, Bhat):
"""
This method calculates the pumping rates for each laser beam:
"""
for key in self.laserBeams:
# Extract the relevant d_q matrix:
ind = self.hamiltonian.rotated_hamiltonian.laser_keys[key]
d_q = self.hamiltonian.rotated_hamiltonian.blocks[ind].matrix
gamma = self.hamiltonian.blocks[ind].parameters['gamma']
# Extract the energies:
E1 = np.diag(self.hamiltonian.rotated_hamiltonian.blocks[ind[0],ind[0]].matrix)
E2 = np.diag(self.hamiltonian.rotated_hamiltonian.blocks[ind[1],ind[1]].matrix)
E2, E1 = np.meshgrid(E2, E1)
# Initialize the pumping matrix:
self.Rijl[key] = np.zeros((len(self.laserBeams[key].beam_vector),) +
d_q.shape[1:])
# Grab the laser parameters:
kvecs = self.laserBeams[key].kvec(r, t)
intensities = self.laserBeams[key].intensity(r, t)
deltas = self.laserBeams[key].delta(t)
projs = self.laserBeams[key].project_pol(Bhat, R=r, t=t)
# Loop through each laser beam driving this transition:
for ll, (kvec, intensity, proj, delta) in enumerate(zip(kvecs, intensities, projs, deltas)):
fijq = np.abs(d_q[0]*proj[2] + d_q[1]*proj[1] +d_q[2]*proj[0])**2
# Finally, calculate the scattering rate the polarization
# onto the appropriate basis:
self.Rijl[key][ll] = gamma*intensity/2*\
fijq/(1 + 4*(-(E2 - E1) + delta - np.dot(kvec, v))**2/gamma**2)
def _add_pumping_rates_to_Rev(self):
# Now add the pumping rates into the rate equation propogation matrix:
for key in self.laserBeams:
ind = self.hamiltonian.rotated_hamiltonian.laser_keys[key]
n_off = sum(self.hamiltonian.rotated_hamiltonian.ns[:ind[0]])
n = self.hamiltonian.rotated_hamiltonian.ns[ind[0]]
m_off = sum(self.hamiltonian.rotated_hamiltonian.ns[:ind[1]])
m = self.hamiltonian.rotated_hamiltonian.ns[ind[1]]
# Sum over all lasers:
Rij = np.sum(self.Rijl[key], axis=0)
# Add the off diagonal components:
self.Rev[n_off:n_off+n, m_off:(m_off+m)] += Rij
self.Rev[m_off:(m_off+m), n_off:n_off+n] += Rij.T
# Add the diagonal components.
self.Rev[(np.arange(n_off, n_off+n), np.arange(n_off, n_off+n))] -= np.sum(Rij, axis=1)
self.Rev[(np.arange(m_off, m_off+m), np.arange(m_off, m_off+m))] -= np.sum(Rij, axis=0)
[docs] def construct_evolution_matrix(self, r, v, t=0., default_axis=np.array([0., 0., 1.])):
"""
Constructs the evolution matrix at a given position and time.
Parameters
----------
r : array_like, shape (3,)
Position at which to calculate the equilibrium population
v : array_like, shape (3,)
Velocity at which to calculate the equilibrium population
t : float
Time at which to calculate the equilibrium population
"""
if self.tdepend['B']:
B = self.magField.Field(r, t)
else:
B = self.magField.Field(r)
# Calculate its magnitude:
Bmag = np.linalg.norm(B, axis=0)
# Calculate the Bhat direction:
if Bmag > 1e-10:
Bhat = B/Bmag
else:
Bhat = default_axis
# Diagonalize the hamiltonian at this location:
self.hamiltonian.diag_static_field(Bmag)
# Re-initialize the evolution matrix:
self.Rev = np.zeros((self.hamiltonian.n, self.hamiltonian.n))
if not np.all(self.hamiltonian.diagonal):
# Reconstruct the decay matrix to match this new field.
self.Rev_decay = self._calc_decay_comp_of_Rev(
self.hamiltonian.rotated_hamiltonian
)
self.Rev += self.Rev_decay
# Recalculate the pumping rates:
Bhat = self._calc_pumping_rates(r, v, t, Bhat)
# Add the pumping rates to the evolution matrix:
self._add_pumping_rates_to_Rev()
return self.Rev, self.Rijl
[docs] def equilibrium_populations(self, r, v, t, **kwargs):
"""
Returns the equilibrium population as determined by the rate equations
This method uses singular matrix decomposition to find the equilibrium
state of the rate equations at a given position, velocity, and time.
Parameters
----------
r : array_like, shape (3,)
Position at which to calculate the equilibrium population
v : array_like, shape (3,)
Velocity at which to calculate the equilibrium population
t : float
Time at which to calculate the equilibrium population
return_details : boolean, optional
In addition to the equilibrium populations, return the full
population evolution matrix and the scattering rates for each of the
lasers
Returns
-------
Neq : array_like
Equilibrium population vector
Rev : array_like
If return details is True, the evolution matrix for the state
populations.
Rijl : dictionary of array_like
If return details is True, the scattering rates for each laser and
each combination of states between the manifolds specified by the
dictionary's index.
"""
return_details = kwargs.pop('return_details', False)
Rev, Rijl = self.construct_evolution_matrix(r, v, t, **kwargs)
# Find the singular values:
U, S, VH = np.linalg.svd(Rev)
Neq = np.compress(S <= self.svd_eps, VH, axis=0).T
Neq /= np.sum(Neq)
if Neq.shape[1] > 1:
Neq = np.nan*Neq[:, 0]
#raise ValueError("more than one equilbrium state found")
# The operations above return a column vector, this collapses the column
# vector into an array:
Neq = Neq.flatten()
if return_details:
return (Neq, Rev, Rijl)
else:
return Neq
[docs] def force(self, r, t, N, return_details=True):
"""
Calculates the instantaneous force
Parameters
----------
r : array_like
Position at which to calculate the force
t : float
Time at which to calculate the force
N : array_like
Relative state populations
return_details : boolean, optional
If True, returns the forces from each laser and the magnetic forces.
Returns
-------
F : array_like
total equilibrium force experienced by the atom
F_laser : dictionary of array_like
If return_details is True, the forces due to each laser, indexed
by the manifold the laser addresses. The dictionary is keyed by
the transition driven, and individual lasers are in the same order
as in the pylcp.laserBeams object used to create the governing
equation.
F_mag : array_like
If return_details is True, the forces due to the magnetic field.
"""
F = np.zeros((3,))
f = {}
for key in self.laserBeams:
f[key] = np.zeros((3, len(self.laserBeams[key].beam_vector)))
ind = self.hamiltonian.laser_keys[key]
n_off = sum(self.hamiltonian.ns[:ind[0]])
n = self.hamiltonian.ns[ind[0]]
m_off = sum(self.hamiltonian.ns[:ind[1]])
m = self.hamiltonian.ns[ind[1]]
Ne, Ng = np.meshgrid(N[m_off:(m_off+m)], N[n_off:(n_off+n)], )
for ll, beam in enumerate(self.laserBeams[key].beam_vector):
kvec = beam.kvec(r, t)
f[key][:, ll] += kvec*np.sum(self.Rijl[key][ll]*(Ng - Ne), axis=(0,1))
F += np.sum(f[key], axis=1)
fmag = np.array([0., 0., 0.])
if self.include_mag_forces:
gradBmag = self.magField.gradFieldMag(r)
for ii, block in enumerate(np.diag(self.hamiltonian.blocks)):
ind1 = int(np.sum(self.hamiltonian.ns[:ii]))
ind2 = int(np.sum(self.hamiltonian.ns[:ii+1]))
if self.hamiltonian.diagonal[ii]:
if isinstance(block, tuple):
fmag += np.sum(np.real(
block[1].matrix[1] @ N[ind1:ind2]
))*gradBmag
elif isinstance(block, self.hamiltonian.vector_block):
fmag += np.sum(np.real(
block.matrix[1] @ N[ind1:ind2]
))*gradBmag
else:
if isinstance(block, tuple):
fmag += np.sum(np.real(
self.hamiltonian.U[ii].T @ block[1].matrix[1] @
self.hamiltonian.U[ii]) @ N[ind1:ind2])*gradBmag
elif isinstance(block, self.hamiltonian.vector_block):
fmag += np.sum(np.real(
self.hamiltonian.U[ii].T @ block.matrix[1] @
self.hamiltonian.U[ii]) @ N[ind1:ind2])*gradBmag
F += fmag
if return_details:
return F, f, fmag
else:
return F
[docs] def set_initial_pop(self, N0):
"""
Sets the initial populations
Parameters
----------
N0 : array_like
The initial state population vector :math:`N_0`. It must have
:math:`n` elements, where :math:`n` is the total number of states
in the system.
"""
if len(N0) != self.hamiltonian.n:
raise ValueError('Npop has only %d entries for %d states.' %
(len(N0), self.hamiltonian.n))
if np.any(np.isnan(N0)) or np.any(np.isinf(N0)):
raise ValueError('Npop has NaNs or Infs!')
self.N0 = N0
[docs] def set_initial_pop_from_equilibrium(self):
"""
Sets the initial populations based on the equilibrium population at
the initial position and velocity and time t=0.
"""
self.N0 = self.equilibrium_populations(self.r0, self.v0, t=0.)
[docs] def evolve_populations(self, t_span, **kwargs):
"""
Evolve the state population in time.
This function integrates the rate equations to determine how the
populations evolve in time. Any initial velocity is kept constant.
It is analogous to obe.evolve_density().
Parameters
----------
t_span : list or array_like
A two element list or array that specify the initial and final time
of integration.
**kwargs :
Additional keyword arguments get passed to solve_ivp, which is
what actually does the integration.
Returns
-------
sol : OdeSolution
Bunch object that contains the following fields:
* t: integration times found by solve_ivp
* rho: density matrix
* v: atomic velocity (constant)
* r: atomic position
It contains other important elements, which can be discerned from
scipy's solve_ivp documentation.
"""
if any([self.tdepend[key] for key in self.tdepend.keys()]):
raise NotImplementedError('Time dependence not yet implemented.')
else:
Rev, Rijl = self.construct_evolution_matrix(self.r0, self.v0)
self.sol = solve_ivp(lambda t, N: Rev @ N, t_span, self.N0, **kwargs)
return self.sol
[docs] 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(),
**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()``
**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)
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, **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
[docs] def find_equilibrium_force(self, return_details=False, **kwargs):
"""
Finds the equilibrium force at the initial position
This method works by finding the equilibrium population through the
rateeq.equilibrium_population() function, then calculating the resulting
force.
Parameters
----------
return_details : boolean, optional
If true, returns the forces from each laser and the scattering rate
matrix. Default: False
kwargs :
Any additional keyword arguments to be passed to
equilibrium_populations()
Returns
-------
F : array_like
total equilibrium force experienced by the atom
F_laser : array_like
If return_details is True, the forces due to each laser.
Neq : array_like
If return_details is True, the equilibrium populations.
Rijl : dictionary of array_like
If return details is True, the scattering rates for each laser and
each combination of states between the manifolds specified by the
dictionary's index.
F_mag : array_like
If return_details is True, the forces due to the magnetic field.
ii : int
Number of iterations needed to converge.
"""
if any([self.tdepend[key] for key in self.tdepend.keys()]):
raise NotImplementedError('Time dependence not yet implemented.')
else:
N_eq, Rev, Rijl = self.equilibrium_populations(
self.r0, self.v0, t=0., return_details=True, **kwargs
)
F_eq, f_eq, f_mag = self.force(self.r0, 0., N_eq)
if return_details:
return F_eq, f_eq, N_eq, Rijl, f_mag
else:
return F_eq
[docs] def generate_force_profile(self, R, V, name=None, progress_bar=False,
**kwargs):
"""
Map out the equilibrium force vs. position and velocity
Parameters
----------
R : array_like, shape(3, ...)
Position vector. First dimension of the array must be length 3, and
corresponds to :math:`x`, :math:`y`, and :math:`z` components,
repsectively.
V : array_like, shape(3, ...)
Velocity vector. First dimension of the array must be length 3, and
corresponds to :math:`v_x`, :math:`v_y`, and :math:`v_z` components,
repsectively.
name : str, optional
Name for the profile. Stored in profile dictionary in this object.
If None, uses the next integer, cast as a string, (i.e., '0') as
the name.
progress_bar : boolean, optional
Displays a progress bar as the proceeds. Default: False
kwargs :
Any additional keyword arguments to be passed to
rateeq.find_equilibrium_force()
Returns
-------
profile : pylcp.rateeq.force_profile
Resulting force profile.
"""
if not name:
name = '{0:d}'.format(len(self.profile))
self.profile[name] = force_profile(R, V, self.laserBeams,
self.hamiltonian)
it = np.nditer([R[0], R[1], R[2], V[0], V[1], V[2]],
flags=['refs_ok', 'multi_index'],
op_flags=[['readonly'], ['readonly'], ['readonly'],
['readonly'], ['readonly'], ['readonly']])
if progress_bar:
progress = progressBar()
for (x, y, z, vx, vy, vz) in it:
# Construct the rate equations:
r = np.array([x, y, z])
v = np.array([vx, vy, vz])
# Update position (use multi_index), populations, and forces.
self.set_initial_position_and_velocity(r, v)
try:
F, f, Neq, Rijl, f_mag = self.find_equilibrium_force(
return_details=True, **kwargs
)
except:
raise ValueError(
"Unable to find solution at " +\
"r=({0:0.2f},{1:0.2f},{2:0.2f})".format(x, y, z) + " and " +
"v=({0:0.2f},{1:0.2f},{2:0.2f})".format(vx, vy, vz)
)
if progress_bar:
progress.update((it.iterindex+1)/it.itersize)
self.profile[name].store_data(it.multi_index, Neq, F, f, f_mag, Rijl)
