"""
Solvers for time domain analysis with an arbitrary RF field
"""
import numpy as np
from importlib.metadata import version
from copy import deepcopy
from .sensor_utils import _hamiltonian_term, generate_eom, make_real, TimeFunc, check_positive_semi_definite, _squeeze_dims
from .sensor_solution import Solution
from .sensor import Sensor # only needing for type hinting
from .slicing.slicing import matrix_slice, get_slice_num_t
from .exceptions import RydiquleError
from rydiqule.solvers import solve_steady_state
from rydiqule.doppler_utils import *
from rydiqule.stack_solvers.scipy_solver import scipy_solve
from rydiqule.stack_solvers.cyrk_solver import cyrk_solve
from typing import Optional, Tuple, List, Callable, Literal, Union, Dict
stack_solvers: Dict[str, Callable] = {"scipy": scipy_solve,
"cyrk": cyrk_solve,
}
solver_type = Union[Callable, Literal['scipy', 'cyrk']]
[docs]
def solve_time(sensor: Sensor, end_time: float, num_pts: int,
init_cond: Optional[np.ndarray] = None, doppler: bool = False,
doppler_mesh_method: Optional[MeshMethod] = None, sum_doppler: bool = True,
weight_doppler: bool = True,
n_slices: Union[int, None] = None, solver: solver_type = "scipy",
**kwargs) -> Solution:
"""
Solves the response of the optical sensor in the time domain given the its time-dependent inputs
If insufficient system memory is available to solve the system all at once,
system is broken into "slices" of manageable memory footprint which are solved individually.
This slicing behavior does not affect the result.
All couplings that include a "time_dependence" argument will be solved in the time domain.
A number of solver backends work with rydiqule, but the default `"scipy"` ivp solver is the
is recommended backend in almost all cases, as it is the most fully-featured and
documented. Advanced users have the ability to define their own solver backends
by creating a function that follows the call signature for rydiqule timesolver
backends. Additional arguments to the solver backend can be supplied with `**kwargs`.
Parameters
-----------
sensor : :class:`~.Sensor`
The sensor object representing the atomic/laser arrangement of the system.
end_time : float
Amount of time, in microseconds, for which to simulate the system
num_pts : int
The number of points along the range `(0, end_time)` for which
the solution is evaluated. This does not affect the number of function
evaluations during the solve, rather the spacing of the points in
the reported solution.
init_cond : numpy.ndarray or `None`, optional
Density matrix representing the initial state of the system.
If specified, the shape should be either `(n)` in the case of a single
initial condition for all parameter values, or should be of shape `(*l, n)`
matching the output shape of a steady state solve if the initial condition
may be different for different combinations of parameters.
If `None`, will solve the problem in the steady state with all time-dependent fields
at their :math:`t=0` value and use the solution as the initial condition.
Other possible manual options might include a matrix populated by zeros representing the
entire population in the ground state. Defaults to `None`.
doppler : bool, optional
Whether to account for doppler shift among moving atoms in
the gas. If True, the solver will implicitly define a velocity distribution
for particles in the cell, solve the problem for each velocity class,
and return a weighted average of the results. Note that solving in this
manner carries a substantial performance penalty, as each doppler velocity class
is solved as its own problem. If solved with doppler, only axis specified by a `"kvec"'
argument in one of the sensor couplings will be average over. The time solver currently
supports doppler averaging in any number of spatial dimensions, up to the limit of
3 imposed by the macroscopic physical world. Defaults to `False`.
doppler_mesh_method : dict, optional
Dictionary that controls the doppler meshing method. Exact details of this are
found in the documentation of :func:`~doppler_classes`. Ignored if
`doppler=False`. Default is `None`.
sum_doppler : bool, optional
Whether to average over doppler classes after the solve
is complete. Setting to false will not perform the sum, allowing viewing
of the weighted results of the solve for each doppler class. Ignored
if `doppler=False`. Default is `True`.
weight_doppler : bool
Whether to apply weights to doppler solution to perform
averaging. If `False`, will **not** apply weights or perform a doppler_average,
regardless of the value of `sum_doppler`. Changing from default intended
only for internal use. Ignored if `doppler=False` or `sum_doppler=False`.
Default is `True`.
n_slices : int or None, optional
How many sets of equations to break the full equations into.
The actual number of slices will be the largest between this value and the minimum
number of slices to solve the system without a memory error. If `None`, solver uses the
minimum number of slices required to solve without a `memoryError`. Defaults to None.
solver : {"scipy", "cyrk"} or callable
The backend solver used to solve the ivp generated by the sensor.
All string values correspond to backend solvers built in to rydiqule.
Valid string values are:
- "scipy": Solves equations with :func:`scipy:scipy.integrate.solve_ivp`.
The default, most stable, and well-supported option.
- "cyrk": Solves jit-compiled equations with a cython compiled RK solver from `CyRK`.
Due to some jit compilation, only faster for moderate length problems
(ie problems with a moderate number of required time steps).
Additionally, can be specified with a callable that matches
rydiqule's time-solver convention,
enabling using a custom solver backend.
.. note::
Unless otherwise noted, backends other than scipy are considered experimental.
Issues with their use are considered features not fully implemented rather than
bugs.
**kwargs: Additional keyword arguments passed to the backend solver.
See documentation of the relevant solver (i.e. :func:`scipy:scipy.integrate.solve_ivp`)
for details and supported arguments.
Returns
-------
:class:`~.Solution`
An object containing the solution and related information.
Timesolver-specific defined attributes are `t` and `init_cond`,
corresponding respectively to the times at which the solution is sampled and
the initial conditions used for the solve.
Examples
--------
A basic solve for a 3-level system would have a "density matrix" solution of size 8 (3^2-1).
Here we use a trivial time dependence for demonstration purposes, but in practice the time
dependence is likely more complicated. Below the most basic use of `solve_time` is demonstrated
>>> s = rq.Sensor(3)
>>> td = lambda t: 1
>>> s.add_coupling((0,1), detuning = 1, rabi_frequency=1)
>>> s.add_coupling((1,2), detuning = 2, rabi_frequency=2, time_dependence=td)
>>> s.add_transit_broadening(0.1)
>>> end_time = 10 #microseconds
>>> n_pts = 1000 #interpolated points in solution
>>> sol = rq.solve_time(s, end_time, n_pts)
>>> print(type(sol))
<class 'rydiqule.sensor_solution.Solution'>
>>> print(type(sol.rho))
<class 'numpy.ndarray'>
>>> print(sol.rho.shape)
(1000, 8)
Defining an array-like parameter will automatically calculate the density matrix solution
for every value. Here we use 11 values, resulting in 11 density matrices. The `axis_labels`
attribute of the solution can clarify which axes are which.
>>> s = rq.Sensor(3)
>>> td = lambda t: 1
>>> det = np.linspace(-1,1,11)
>>> s.add_coupling((0,1), detuning = det, rabi_frequency=1)
>>> s.add_coupling((1,2), detuning = 2, rabi_frequency=2, time_dependence=td)
>>> s.add_transit_broadening(0.1)
>>> sol = rq.solve_time(s, end_time, n_pts)
>>> print(sol.rho.shape)
(11, 1000, 8)
>>> print(sol.axis_labels)
['(0,1)_detuning', 'time', 'density_matrix']
As expected, multiple axes of scanned parameters are handled the same way as they are in the
steady-state case, with the expected additions from the time solver.
>>> s = rq.Sensor(3)
>>> td = lambda t: 1
>>> det = np.linspace(-1,1,11)
>>> s.add_coupling((0,1), detuning = det, rabi_frequency=1)
>>> s.add_coupling((1,2), detuning = det, rabi_frequency=2, time_dependence=td)
>>> s.add_transit_broadening(0.1)
>>> sol = rq.solve_time(s, end_time, n_pts)
>>> print(sol.rho.shape)
(11, 11, 1000, 8)
>>> print(sol.axis_labels)
['(0,1)_detuning', '(1,2)_detuning', 'time', 'density_matrix']
If the solve uses doppler broadening, all doppler classes will be computed a weighted average
will be taken over the doppler axis, and the shape of the solution will not change. While this
is the desired behavior in most situations, the `sum_doppler` argument can be used to override
this behavior, leave (or more) solution axis corresponding to different doppler classes.
>>> s = rq.Sensor(3, vP=1)
>>> td = lambda t: 1
>>> det = np.linspace(-1,1,11)
>>> s.add_coupling((0,1), detuning = det, rabi_frequency=1, kvec=(4,0,0))
>>> s.add_coupling((1,2), detuning = 2, rabi_frequency=2, kvec=(-4,0,0), time_dependence=td)
>>> s.add_transit_broadening(0.1)
>>> end_time = 10 #microseconds
>>> n_pts = 1000 #interpolated points in solution
>>> sol_dop = rq.solve_time(s, end_time, n_pts, doppler=True)
>>> sol_dop_nosum = rq.solve_time(s, end_time, n_pts, doppler=True, sum_doppler=False)
>>> print(sol_dop.rho.shape)
(11, 1000, 8)
>>> print(sol_dop_nosum.rho.shape)
(561, 11, 1000, 8)
>>> print(sol_dop.axis_labels)
['(0,1)_detuning', 'time', 'density_matrix']
>>> print(sol_dop_nosum.axis_labels)
['doppler_0', '(0,1)_detuning', 'time', 'density_matrix']
"""
if len(sensor.couplings_with("time_dependence")) == 0:
raise RydiquleError("At least one time-dependent coupling is required")
if isinstance(solver, str):
try:
solver = stack_solvers[solver]
except KeyError as err:
raise RydiquleError(
f"{solver} is not a built-in solver."
f" Supported built-in solvers are {list(stack_solvers.keys())}") from err
if not callable(solver):
raise RydiquleError("Solvers must be callable functions")
solution = Solution()
# relevant sensor-related quantities
stack_shape = sensor._stack_shape()
basis_size = sensor.basis_size
spatial_dim = sensor.spatial_dim()
# set up time parameters
time_range = (0.0, end_time)
t_eval = np.linspace(*time_range, num=num_pts, dtype=np.float64)
# initialize doppler-related quantities
doppler_axis_shape: Tuple[int, ...] = ()
dop_classes = None
doppler_shifts = None
out_doppler_axes: Tuple[slice, ...] = ()
doppler_axes: Tuple[slice, ...] = ()
# update doppler-related values
if doppler:
dop_classes = doppler_classes(method=doppler_mesh_method)
doppler_shifts = sensor.get_doppler_shifts()
doppler_axis_shape = tuple(len(dop_classes) for _ in range(spatial_dim))
doppler_axes = tuple(slice(None) for _ in range(spatial_dim))
if not sum_doppler:
out_doppler_axes = doppler_axes
if init_cond is None:
init_cond = solve_steady_state(
sensor, doppler=doppler, weight_doppler=False, sum_doppler=False,
doppler_mesh_method=doppler_mesh_method
).rho
else:
# check that user provided init_cond are physical
check_positive_semi_definite(init_cond)
# use available memory to figure out how to slice the hamiltonian
n_slices, out_sol_shape = get_slice_num_t(basis_size, stack_shape, doppler_axis_shape,
num_pts, sum_doppler, weight_doppler, n_slices)
if n_slices > 1:
print(f"Breaking equations of motion into {n_slices} sets of equations...")
# allocate arrays
hamiltonians = sensor.get_hamiltonian()
hamiltonians_time, hamiltonians_time_i = sensor.get_time_hamiltonian_components()
n_time = len(hamiltonians_time)
time_functions = sensor.get_time_dependence()
gamma = sensor.decoherence_matrix()
sols = np.zeros(out_sol_shape, dtype="float64")
n_slices_true = sum(1 for _ in matrix_slice(gamma, n_slices=n_slices))
for i, (idx, H, G, *time_hams) in enumerate(matrix_slice(hamiltonians, gamma,
*hamiltonians_time,
*hamiltonians_time_i,
n_slices=n_slices)):
if n_slices_true > 1:
print(f"Solving equation slice {i+1}/{n_slices_true}", end='\r')
Ht = np.array(time_hams[:n_time])
Ht_i = np.array(time_hams[n_time:])
full_idx = (*out_doppler_axes, *idx, slice(None))
sols[full_idx] = _solve_hamiltonian_stack(H, Ht, Ht_i,
G, time_functions, t_eval,
init_cond[(*doppler_axes, *idx)], solver,
doppler=doppler,
dop_classes=dop_classes, sum_doppler=sum_doppler,
weight_doppler=weight_doppler,
doppler_shifts=doppler_shifts,
spatial_dim=spatial_dim, **kwargs)
# save results to Solution object
solution.rho = sols
# specific to observable calculations
solution._eta = sensor.eta
solution._kappa = sensor.kappa
solution._cell_length = sensor.cell_length
solution._beam_area = sensor.beam_area
solution._probe_freq = sensor.probe_freq
solution._probe_tuple = sensor.probe_tuple
sensor._expand_dims()
solution.couplings = deepcopy(sensor.couplings)
_squeeze_dims(sensor.couplings)
solution.axis_labels = ([f'doppler_{i:d}' for i in range(spatial_dim) if not sum_doppler]
+ sensor.axis_labels()
+ ["time", "density_matrix"])
solution.axis_values = ([dop_classes for i in range(spatial_dim) if not sum_doppler]
+ [val for _,_,val,_ in sensor.variable_parameters()]
+ [t_eval, sensor.dm_basis()])
solution.dm_basis = sensor.dm_basis()
solution.rq_version = version("rydiqule")
solution.doppler_classes = dop_classes
# time solver specific
solution.t = t_eval
solution.init_cond = init_cond
return solution
def _solve_hamiltonian_stack(hamiltonians_base: np.ndarray, hamiltonians_time: np.ndarray,
hamiltonians_time_i: np.ndarray, gamma_matrix: np.ndarray,
time_functions: List[TimeFunc], t_eval: np.ndarray,
init_cond: np.ndarray, solver, doppler: bool = False,
dop_classes: Optional[np.ndarray] = None, sum_doppler: bool = True,
weight_doppler: bool = True,
doppler_shifts: Optional[np.ndarray] = None,
spatial_dim: int = 0, **kwargs
) -> np.ndarray:
"""
Internal functions which solve the equations of a given hamiltonian stack
with the given parameters.
"""
eom_base, const = generate_eom(hamiltonians_base, gamma_matrix,
remove_ground_state=True, real_eom=True)
eom_time_r, const_r = generate_eom_time(hamiltonians_time)
eom_time_i, const_i = generate_eom_time(hamiltonians_time_i)
if doppler:
assert dop_classes is not None and doppler_shifts is not None
dop_velocities, dop_volumes = doppler_mesh(dop_classes, spatial_dim)
eom_base = get_doppler_equations(eom_base, doppler_shifts, dop_velocities)
sols = solve_eom_stack(eom_base, const, eom_time_r, const_r,
eom_time_i, const_i, time_functions,
t_eval, init_cond, solver, **kwargs)
if weight_doppler:
sols_weighted = apply_doppler_weights(sols, dop_velocities, dop_volumes)
if sum_doppler:
sum_axes = tuple(np.arange(spatial_dim))
sols_final = np.sum(sols_weighted, axis=sum_axes)
else:
sols_final = sols_weighted
else:
sols_final = sols
else:
sols_final = solve_eom_stack(eom_base, const, eom_time_r, const_r,
eom_time_i, const_i, time_functions, t_eval,
init_cond, solver, **kwargs)
return sols_final.swapaxes(-1,-2)
[docs]
def solve_eom_stack(eoms_base: np.ndarray, const: np.ndarray,
eom_time_r: np.ndarray, const_r: np.ndarray,
eom_time_i: np.ndarray, const_i: np.ndarray,
time_inputs: List[TimeFunc],
t_eval: np.ndarray, init_cond: np.ndarray, solver, **kwargs
) -> np.ndarray:
"""
Solve a stack of equations of motion with shape `(*l, n, n)` in the time domain.
Companion function to :func:`~.timesolvers.solve_time`, but can be invoked on its
for equations already formatted.
Parameters
----------
eoms_base : numpy.ndarray
Array of shape `(*l, n, n)` representing the part of
equations of motion of the system which do not respond to external fields.
const : numpy.ndarray
constant term of shape (n,) added in differential equations. Typically
generated by :func:`~.sensor_utils.generate_eom`.
eoms_time_r : list[numpy.ndarray]
list of arrays of shape `(basis_size^2-1, basis_size^2-1)`
representing the parts of the OBEs with a real-valued time-dependence. In the solver,
this array will be multiplied by a time-dependent rabi frequency. Typically a matrix of
mostly zeros, with non-zero terms corresponding to a particular time-dependent coupling
const_r : numpy.ndarray
Constant term of shape (n,) added in a real time-dependent portion of
differential equations. Typically generated by :func:`~.sensor_utils.generate_eom_time`.
eoms_time_i : numpy.ndarray
list of arrays of shape `(basis_size^2-1, basis_size^2-1)`
representing the parts of the OBEs with an imaginary-valued time-dependence. In the solver,
this array will be multiplied by a time-dependent rabi frequency.
const_i : numpy.ndarray
constant term of shape (n,) added in an imaginary time-dependent portion of
differential equations. Typically generated by :func:`~.sensor_utils.generate_eom_time`.
t_eval : numpy.ndarray
1-D array of times, in microseconds, at which to
evaluate the solution. Does not affect evaluations in the solve.
time_inputs : list[function float->float]
List of functions which represent
the rabi frequency of a field as a function of time. list length should
be identical to the length of obes_time. In the solver, the *i* th time
input will be evaluated at time *t* and multiplied by the *i* th entry of
obes_time.
time_range tuple(float):
Pair of values represent the start and end time, in microseconds, of the simulation.
init_cond : numpy.ndarray or `None`, optional
Density matrix representing the initial state of the system.
If specified, the shape should be either `(n)` in the case of a single
initial condition for all parameter values, or should be of shape `(*l, n)`
matching the output shape of a steady state solve if the initial condition
may be different for different combinations of parameters.
If `None`, will solve the problem in the steady state with all time-dependent fields
"off" and use the solution as the initial condition for the time behavior. Other
possible manual options might include a matrix populated by zeros representing the
entire population in the ground state. Defaults to `None`.
Returns
-------
numpy.ndarray
Flattened solution array corresponding to time points.
"""
stack_shape = eoms_base.shape[:-2]
solution_shape = stack_shape + (eoms_base.shape[-1], len(t_eval))
if init_cond.shape == eoms_base.shape[-1:]:
init_cond = np.broadcast_to(init_cond, solution_shape[:-1])
elif init_cond.shape != solution_shape[:-1]:
msg = f"""Initial condition shape {init_cond.shape} does not match expected
solution shape {solution_shape[:-1]}"""
raise RydiquleError(msg)
solutions = solver(eoms_base, const, eom_time_r, const_r, eom_time_i, const_i,
time_inputs, t_eval, init_cond, **kwargs)
return solutions
[docs]
def generate_eom_time(hamiltonians_time: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
"""
Generates the Optical Bloch Equations for just the rf terms.
Uses the convention of the hamiltonian_rf return of the get_time_hamiltonian function.
The equations of motion returned are assumed to be used in conjunction with an electric field.
Parameters
----------
hamiltonians_time : numpy.ndarray
A matrix of shape (basis_size, basis_size), where the off-diagonal
terms (i,j) are the dipole matrix elements in e a_b of the transition
coupling state i to state j.
Returns
-------
numpy.ndarray: Part of the Optical Bloch Equations corresponding to time_dependent couplings.
To produce equations to solve, these values must be multiplied by an electric
field in V/m.
numpy.ndarray: Constant term of the time-dependent portion of the equations
of motion. Same units as the equations themselves.
"""
eom_time = _hamiltonian_term(hamiltonians_time)
eqns, const = remove_ground(eom_time)
eom_time, const = make_real(eqns,const)
return eom_time.real, const