"""
Steady-state solver for analytical Doppler averaging
Code in this module implements analytical doppler averaging techniques described in
Omar Nagib and Thad G. Walker, *Exact steady state of perturbed open quantum systems*,
Phys. Rev. Research **7** 033076 (2025) https://doi.org/10.1103/kgsg-3npp
This solver computes the doppler averaged steady state of a sensor.
In 1 spatial dimension problems, it averages analytically.
In 2 or 3 spatial dimension problems, it averages one axis analytically
and the remaining numerically.
"""
import warnings
import numpy as np
from importlib.metadata import version
from copy import deepcopy
from typing import Optional, Set
from scipy.special import erfcx
from .sensor import Sensor
from .sensor_utils import _hamiltonian_term, generate_eom, make_real, _squeeze_dims
from .sensor_solution import Solution
from .doppler_utils import doppler_classes, doppler_mesh, apply_doppler_weights, MeshMethod, get_doppler_equations
from .slicing.slicing import matrix_slice, get_slice_num_hybrid
from .exceptions import RydiquleError, RydiquleWarning, PopulationNotConservedWarning
def _doppler_eigvec_array(lamdas: np.ndarray, rtol: float = 1e-5, atol: float = 1e-9) -> np.ndarray:
"""
Helper function for vectorizing the analytical summation of eigenvalues.
Note that because rydiqule provides L_v as containing sqrt(2)*sigma_v,
that term is rolled into the definition of lamdas here.
Parameters
----------
lamdas: numpy.ndarray
Array of complex eigenvalues
rtol: float, optional
Relative tolerance parameter for checking 0-eigenvalues when calculating the doppler prefactor.
Passed to :external+numpy:func:`~numpy.isclose`.
Defaults to 1e-5.
atol: float, optional
Absolute tolerance parameter for checking 0-eigenvalues when calculating the doppler prefactor.
Passed to :external+numpy:func:`~numpy.isclose`.
Defaults to 1e-9.
Returns
-------
numpy.ndarray
Analytical Doppler broadening function evaluated at the given eigenvalues.
Values are complex.
"""
doppler_array = np.ones_like(lamdas) # for 0-eigenvalues, set value to 1
idx = np.where(~np.isclose(lamdas, 0, rtol, atol))
if len(idx[0]) > 0:
# split the calculation into two lines for readability, using the complementary error function to avoid
# catastrophic cancellation when |lambda|~0
p1 = np.sqrt(np.pi)/(np.sqrt(-1*lamdas[idx]**2))
p2 = erfcx(-1*np.sqrt(-1*lamdas[idx]**2)/(lamdas[idx]**2))
doppler_array[idx] = p1*p2
return doppler_array
def _get_rho0(L0: np.ndarray) -> np.ndarray:
"""
Helper function for computing rho0 (null vector) of a stack of equations of motion
via the inverse power method.
Parameters
----------
L0: np.ndarray
Stack of steady state Liouvillians
Returns
-------
numpy.ndarray
Stack of steady state (vectorized) density matrices
"""
stack_shape = L0.shape[0:-2]
n = int(np.sqrt(L0.shape[-1]))
rho0 = np.random.rand(*stack_shape, n**2)[..., np.newaxis]
rho0 /= np.linalg.norm(rho0, axis=-1, keepdims=True)
I = np.eye(n**2)
converged_flags = np.zeros(stack_shape, dtype=bool)
# Compute rho0 by finding the null vector of L0 via the shifted inverse power method
for iteration in range(50):
if np.all(converged_flags):
break
remaining_flags_index = np.where(~converged_flags)
current_L0 = L0[remaining_flags_index]
current_shifted_L0 = current_L0 + 1e-14 * I
current_rho0 = rho0[remaining_flags_index]
z = np.linalg.solve(current_shifted_L0, current_rho0) # compute (L0 + 1e-14)^-1 * rho0
rho0_new = z / (np.linalg.norm(z, axis=-2, keepdims=True) + np.finfo(float).eps)
# Estimate magnitude of eigenvalues by Rayleigh quotient
L0rho0 = np.einsum('...ij,...j->...i', current_L0, rho0_new[..., 0])
numerator = np.sum(rho0_new[..., 0] * L0rho0, axis=-1)
denominator = np.sum(rho0_new[..., 0] * rho0_new[..., 0], axis=-1)
current_eigenvalues = np.abs(numerator / denominator)
rho_0_diff = np.linalg.norm(current_rho0 - rho0_new, axis=-1)
converged_flags[remaining_flags_index] = (rho_0_diff < 1e-15).any(axis=-1) | (current_eigenvalues < 1e-14)
rho0[remaining_flags_index] = rho0_new
rho0 = rho0.squeeze(axis=-1)
rho0 *= np.sign(rho0[...,0])[...,None] # remove arbitrary sign from null-vector so all pops are positive
pops = np.sum(rho0[...,::n+1], axis=-1) # calculate trace of each vector
rho0 /= pops[..., None] # normalize vectors by trace
return rho0
[docs]
def solve_doppler_analytic(sensor: Sensor, doppler_mesh_method: Optional[MeshMethod] = None,
analytic_axis: Optional[int] = None, n_slices: Optional[int] = None, rtol: float = 1e-5,
atol: float = 1e-9) -> Solution:
"""
Solves a sensor in steady state in the presence of doppler broadening,
with one dimension analytically averaged.
If the broadening is 1 dimensional, this function will solve analytically.
If the broadening is 2 or 3 dimensional,
this function will average analytically over the specified axis
and numerically over the remaining axes.
This function uses the method outlined in Ref [1]_ for the analytic dimension.
This solver is considered more accurate than :func:`~.solve_steady_state`
since it replaces direct sampling and solving of the velocity classes
with a few tensor inversions and calculation of the numerical prefactor.
This also leads to faster solves,
approximately dictated by the ratio of samples along the doppler axis
relative to the other parameter dimensions.
Additionally, in sensors with 2 or 3 dimensional doppler broadening,
this solver effectively reduces the dimension to 1 or 2, respectively,
leading to faster solves.
If the sensor contains couplings with `time_dependence`, this solver will add those
couplings at their :math:`t=0` value to the steady-state hamiltonian to solve.
If insufficient system memory is available to solve the system in a single call,
system is broken into "slices" of manageable memory footprint which are solved individually.
This slicing behavior does not affect the result.
Parameters
----------
sensor : :class:`~.Sensor`
The sensor for which the solution will be calculated.
It must define 1 and only 1 dimension of doppler shifts
(ie one or more couplings with `kvec` with non-zero values on the same dimension).
doppler_mesh_method: dict, optional
If not `None`, should be a dictionary of meshing parameters to be passed
to :func:`~.doppler_classes`. See :func:`~.doppler_classes` for more
information on supported methods and arguments. If `None, uses the
default doppler meshing. Default is `None`.
analytic_axis: int, optional
Specifies over which axis the solver will average analytically.
Defaults to the first nonzero axis.
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`, uses the minimum
number of slices to solve the system without a memory error. Detailed information about
slicing behavior can be found in :func:`~.slicing.slicing.matrix_slice`. Default is `None`.
rtol : float, optional
Relative tolerance parameter for checking 0-eigenvalues when calculating the doppler prefactor.
Passed to :external+numpy:func:`~numpy.isclose`.
Defaults to 1e-5.
atol : float, optional
Absolute tolerance parameter for checking 0-eigenvalues when calculating the doppler prefactor.
Passed to :external+numpy:func:`~numpy.isclose`.
Defaults to 1e-9.
Returns
-------
:class:`~.Solution`
An object containing the solution and related information.
References
----------
.. [1] Omar Nagib and Thad G. Walker,
Exact steady state of perturbed open quantum systems,
Phys. Rev. Research **7** 033076 (2025)
https://doi.org/10.1103/kgsg-3npp
"""
spatial_dim = sensor.spatial_dim()
if spatial_dim == 0:
raise RydiquleError("System must have at least 1 spatial dimension of Doppler shifts")
# Create a sorted array of nonzero axes
kvecs = sensor.get_value_dictionary('kvec')
nonzero: Set[int] = set()
for _, v in kvecs.items():
current_nonzero_axes = np.where(~np.isclose(v, 0.0))[0]
nonzero.update(current_nonzero_axes)
nonzero_axes = np.array(sorted(list(nonzero)))
# Set the default analytic axis to the first nonzero axis
if analytic_axis is None:
analytic_axis: int = nonzero_axes[0]
# Account for analytic axis set to a zero axis, e.g., kvec of (1,1,0) and inputted analytic axis of 2
elif analytic_axis not in nonzero_axes:
raise RydiquleError(f"analytic_axis={analytic_axis} has no doppler shifts, "
+ f" valid options are: {nonzero_axes}")
# Remap the physical axes to the reduced problem space, e.g., kvec of (0,0,1) and inputted analytic axis of 2
analytic_axis = np.where(nonzero_axes == analytic_axis)[0][0]
n = sensor.basis_size
# 1. Get base Liouvillian (L0) and Doppler shifts
L0, dummy_const = generate_eom(sensor.get_time_hamiltonian(t=0), sensor.decoherence_matrix(),
remove_ground_state=False, real_eom=True)
all_shifts = sensor.get_doppler_shifts()
# 2. Compute perturbation and velocity mesh
numeric_axes = [ax for ax in range(spatial_dim) if ax != analytic_axis]
num_numeric_dims = len(numeric_axes)
dopp_classes = doppler_classes(method = doppler_mesh_method)
dopp_velocities, dopp_volumes = doppler_mesh(dopp_classes, num_numeric_dims)
analytic_shift = np.take(all_shifts, analytic_axis, axis=0)
numeric_shifts = np.delete(all_shifts, analytic_axis, axis=0)
L_pert_complex = _hamiltonian_term(analytic_shift.squeeze())
L_pert, _ = make_real(L_pert_complex, dummy_const, ground_removed=False)
n_vel_points = len(dopp_classes)
numeric_doppler_shape = (n_vel_points,) * num_numeric_dims
param_stack_shape = L0.shape[:-2]
# 3. Compute number of slices and loop over each slice computing rho0, L0m, the analytic integral,
# and the numeric weighting/summing
n_slices, out_sol_shape = get_slice_num_hybrid(n, param_stack_shape, numeric_doppler_shape,
n_slices=n_slices)
if n_slices > 1:
print(f"Breaking parameter stack into {n_slices} slices...")
sols = np.zeros(out_sol_shape, dtype = np.complex128)
for i, (idx, L0_slice) in enumerate(matrix_slice(L0, n_slices=n_slices)):
if n_slices > 1:
print(f"Solving slice {i+1}/{n_slices}", end='\r')
if num_numeric_dims == 0:
L_base_slice = L0_slice
else:
L_base_slice = get_doppler_equations(L0_slice, numeric_shifts, dopp_velocities, ground_removed=False)
rho0_slice = _get_rho0(L_base_slice)
vec1 = np.eye(n).flatten()
L0m_slice = (np.linalg.inv(L_base_slice + rho0_slice[..., np.newaxis] * vec1[np.newaxis, :])
- rho0_slice[..., np.newaxis] * vec1[np.newaxis, :])
lamdas, r_eigvecs = np.linalg.eig(L0m_slice @ L_pert)
prefix = _doppler_eigvec_array(lamdas)
suffix = np.linalg.solve(r_eigvecs, rho0_slice[..., np.newaxis]).squeeze(axis=-1)
rho_dopp_slice = np.einsum('...j,...ij,...j->...i', prefix, r_eigvecs, suffix)
sols_weighted = apply_doppler_weights(rho_dopp_slice, dopp_velocities, dopp_volumes)
axes_to_sum = tuple(range(num_numeric_dims))
sols_slice = np.sum(sols_weighted, axis=axes_to_sum)
sols[idx] = sols_slice
# 4. Postprocess solution
imag_tol = 10000
sols_real = np.real_if_close(sols, tol=imag_tol)
if np.iscomplexobj(sols_real):
rho_dopp_imag = np.abs(sols.imag)
count = np.count_nonzero(rho_dopp_imag > np.finfo(float).eps*imag_tol)
warnings.warn('Doppler-averaged solution has complex parts outside of tolerance, solution is suspect. ' +
f'{count:d} of {sols.size:d} elements larger than cutoff {np.finfo(float).eps*imag_tol:.3g}. ' +
f'Max Abs(Imag): {sols.max():.3g}, Std Abs(Imag): {np.std(sols):.3g}',
RydiquleWarning)
pops_dopp = np.sum(sols_real[...,::n+1], axis=-1)
if not np.isclose(pops_dopp, 1.0, rtol=.01).all():
warnings.warn('Doppler-averaged solution has populations not conserved, solution is suspect. ' +
f'{np.count_nonzero(~np.isclose(pops_dopp, 1.0)):d} of {pops_dopp.size:d} have non-unit trace. ' +
f'Min trace is {pops_dopp.min():f}',
PopulationNotConservedWarning)
# 5. Package into a solution object
solution = Solution()
solution.rho = sols_real[...,1:]
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 = (sensor.axis_labels() + ["density_matrix"])
solution.axis_values = ([val for _,_,val,_ in sensor.variable_parameters()]
+ [sensor.dm_basis()])
solution.dm_basis = sensor.dm_basis()
solution.rq_version = version("rydiqule")
return solution