"""Python solver entry points."""
from __future__ import annotations
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from numbers import Number
from typing import Any, cast
import numpy as np
from numpy.typing import NDArray
from scipy import sparse as sp # type: ignore[import-untyped]
from scipy.linalg import expm as _dense_expm # type: ignore[import-untyped]
from ._julia_bridge import JuliaBridgeUnavailable, load_backend
from .observables import (
StateObservable,
_state_observable_spec,
evaluate_state_observables,
)
from .operators import Operator
from .result import Options, Result
from .timedep import InterpolatedCoefficient, TimeDependentHamiltonian
Array = NDArray[np.complex128]
FloatArray = NDArray[np.float64]
@dataclass(frozen=True)
class _MCStateObservableEntry:
name: str
operator_offset: int | None = None
constant: float | None = None
@dataclass(frozen=True)
class _MCStateObservablePlan:
operators: tuple[Array, ...]
entries: tuple[_MCStateObservableEntry, ...]
[docs]
def mesolve(
H: Operator | TimeDependentHamiltonian,
rho0: Array,
tlist: Sequence[float],
*,
c_ops: Sequence[Operator] | None = None,
e_ops: Sequence[Operator] | None = None,
state_observables: Mapping[str, StateObservable] | None = None,
options: Options | None = None,
) -> Result:
"""Solve a Lindblad master equation.
Supports time-independent Hamiltonians and ``TimeDependentHamiltonian``
objects with scalar coefficients.
"""
opts = options or Options()
times = np.asarray(tlist, dtype=np.float64)
rho0_array = np.asarray(rho0, dtype=np.complex128)
H_data = H.base.data if isinstance(H, TimeDependentHamiltonian) else H.data
c_arrays = [op.data for op in c_ops or []]
e_arrays = [op.data for op in e_ops or []]
needs_states = bool(opts.save_states or state_observables)
_validate_mesolve_inputs(H, rho0_array, times, c_arrays, e_arrays)
try:
backend = load_backend()
except JuliaBridgeUnavailable as exc:
msg = "mesolve requires the Julia backend; run `oqs setup-julia` first."
raise NotImplementedError(msg) from exc
if isinstance(H, TimeDependentHamiltonian):
raw = backend.mesolve_time_dependent(
_matrix_payload(backend, H_data),
[
_matrix_payload(
backend,
term.operator.data,
)
for term in H.terms
],
[_coefficient_payload(backend, term.coefficient) for term in H.terms],
rho0_array,
times,
[_matrix_payload(backend, array) for array in c_arrays],
[_matrix_payload(backend, array) for array in e_arrays],
float(opts.rtol),
float(opts.atol),
needs_states,
bool(opts.compute_entropy),
)
else:
raw = backend.mesolve(
_matrix_payload(backend, H_data),
rho0_array,
times,
[_matrix_payload(backend, array) for array in c_arrays],
[_matrix_payload(backend, array) for array in e_arrays],
float(opts.rtol),
float(opts.atol),
needs_states,
str(opts.method),
int(opts.krylov_dim),
bool(opts.compute_entropy),
)
raw_times = np.asarray(_field(raw, "times"), dtype=np.float64)
raw_expect = np.asarray(_field(raw, "expect"), dtype=np.complex128)
raw_entropy = (
np.asarray(_field(raw, "entropy"), dtype=np.float64)
if opts.compute_entropy
else None
)
raw_states = np.asarray(_field(raw, "states"), dtype=np.complex128)
stats = _to_python_dict(_field(raw, "solver_stats"))
expects = [raw_expect[idx, :].copy() for idx in range(raw_expect.shape[0])]
state_series = _density_states(raw_states, len(raw_times)) if needs_states else []
states = state_series if opts.save_states else None
observed = (
evaluate_state_observables(state_series, state_observables)
if state_observables
else {}
)
if observed:
stats["state_observables"] = list(observed)
return Result(
times=raw_times,
states=states,
expect=expects,
expect_std=[],
expect_stderr=[],
state_observables=observed,
entropy=raw_entropy,
solver_stats=stats,
)
[docs]
def mcsolve(
H: Operator,
psi0: Array,
tlist: Sequence[float],
*,
c_ops: Sequence[Operator] | None = None,
e_ops: Sequence[Operator] | None = None,
n_traj: int | None = None,
state_observables: Mapping[str, StateObservable] | None = None,
options: Options | None = None,
) -> Result:
"""Run Monte Carlo wave-function trajectories.
Built-in ``state_observables`` created by :func:`openquantumsim.state_metrics`
are aggregated by the backend when they can be represented as linear
expectation values or pure-trajectory constants. Arbitrary Python callbacks
remain available through :func:`single_trajectory` or deterministic
:func:`mesolve` runs with saved state access.
"""
opts = options or Options()
times = np.asarray(tlist, dtype=np.float64)
psi0_array = np.asarray(psi0, dtype=np.complex128).reshape(-1)
c_arrays = [op.data for op in c_ops or []]
e_arrays = [op.data for op in e_ops or []]
state_observable_plan = _mcsolve_state_observable_plan(
H.shape[0],
state_observables,
)
backend_e_arrays = [*e_arrays, *state_observable_plan.operators]
trajectory_count = int(n_traj if n_traj is not None else opts.n_traj)
seed = 0 if opts.seed is None else int(opts.seed)
_validate_mcsolve_inputs(
H,
psi0_array,
times,
c_arrays,
backend_e_arrays,
trajectory_count,
opts.max_step,
opts.n_jobs,
opts.checkpoint_every,
)
try:
backend = load_backend()
except JuliaBridgeUnavailable as exc:
msg = "mcsolve requires the Julia backend; run `oqs setup-julia` first."
raise NotImplementedError(msg) from exc
raw = backend.mcsolve(
_matrix_payload(backend, H.data),
psi0_array,
times,
[_matrix_payload(backend, array) for array in c_arrays],
[_matrix_payload(backend, array) for array in backend_e_arrays],
trajectory_count,
seed,
float(opts.max_step),
int(opts.n_jobs),
"" if opts.checkpoint_file is None else str(opts.checkpoint_file),
int(opts.checkpoint_every),
bool(opts.progress),
)
raw_times = np.asarray(_field(raw, "times"), dtype=np.float64)
raw_expect = np.asarray(_field(raw, "expect"), dtype=np.complex128)
raw_expect_std = np.asarray(_field(raw, "expect_std"), dtype=np.float64)
raw_expect_stderr = np.asarray(_field(raw, "expect_stderr"), dtype=np.float64)
raw_entropy = np.asarray(_field(raw, "entropy"), dtype=np.float64)
stats = _to_python_dict(_field(raw, "solver_stats"))
n_user_e_ops = len(e_arrays)
expects = [raw_expect[idx, :].copy() for idx in range(n_user_e_ops)]
expect_std = [
raw_expect_std[idx, :].copy() for idx in range(n_user_e_ops)
]
expect_stderr = [
raw_expect_stderr[idx, :].copy() for idx in range(n_user_e_ops)
]
observed, observed_std, observed_stderr = _mcsolve_state_observable_results(
state_observable_plan,
raw_expect,
raw_expect_std,
raw_expect_stderr,
n_user_e_ops,
len(raw_times),
)
if observed:
stats["state_observables"] = list(observed)
stats["state_observables_backend"] = True
return Result(
times=raw_times,
states=None,
expect=expects,
expect_std=expect_std,
expect_stderr=expect_stderr,
state_observables=observed,
state_observables_std=observed_std,
state_observables_stderr=observed_stderr,
entropy=raw_entropy,
solver_stats=stats,
)
[docs]
def single_trajectory(
H: Operator,
psi0: Array,
tlist: Sequence[float],
*,
c_ops: Sequence[Operator] | None = None,
e_ops: Sequence[Operator] | None = None,
state_observables: Mapping[str, StateObservable] | None = None,
options: Options | None = None,
) -> Result:
"""Run one Monte Carlo wave-function trajectory.
Set ``Options(save_states=True)`` to return the ket at each requested time.
"""
opts = options or Options()
times = np.asarray(tlist, dtype=np.float64)
psi0_array = np.asarray(psi0, dtype=np.complex128).reshape(-1)
c_arrays = [op.data for op in c_ops or []]
e_arrays = [op.data for op in e_ops or []]
seed = 0 if opts.seed is None else int(opts.seed)
needs_states = bool(opts.save_states or state_observables)
_validate_mcsolve_inputs(
H,
psi0_array,
times,
c_arrays,
e_arrays,
1,
opts.max_step,
1,
1,
)
try:
backend = load_backend()
except JuliaBridgeUnavailable as exc:
msg = (
"single_trajectory requires the Julia backend; "
"run `oqs setup-julia` first."
)
raise NotImplementedError(msg) from exc
raw = backend.single_trajectory(
_matrix_payload(backend, H.data),
psi0_array,
times,
[_matrix_payload(backend, array) for array in c_arrays],
[_matrix_payload(backend, array) for array in e_arrays],
seed,
float(opts.max_step),
needs_states,
)
raw_times = np.asarray(_field(raw, "times"), dtype=np.float64)
raw_expect = np.asarray(_field(raw, "expect"), dtype=np.complex128)
raw_states = np.asarray(_field(raw, "states"), dtype=np.complex128)
expects = [raw_expect[idx, :].copy() for idx in range(raw_expect.shape[0])]
state_series = _ket_states(raw_states, len(raw_times)) if needs_states else []
states = state_series if opts.save_states else None
observed = (
evaluate_state_observables(state_series, state_observables)
if state_observables
else {}
)
return Result(
times=raw_times,
states=states,
expect=expects,
expect_std=[],
expect_stderr=[],
state_observables=observed,
entropy=np.zeros_like(raw_times),
solver_stats={
"nsteps": max(len(raw_times) - 1, 0),
"retcode": "Success",
"n_traj": 1,
"max_step": float(opts.max_step),
"seed": seed,
"save_states": bool(opts.save_states),
"state_observables": list(observed),
},
)
[docs]
def steadystate(
H: Operator,
c_ops: Sequence[Operator],
*,
method: str = "iterative-gmres",
options: Options | None = None,
) -> Array:
"""Compute a Lindblad steady state."""
opts = options or Options()
c_arrays = [op.data for op in c_ops]
if H.shape[0] != H.shape[1]:
msg = "H must be square."
raise ValueError(msg)
for op in c_arrays:
if op.shape != H.shape:
msg = "collapse operators must have the same shape as H."
raise ValueError(msg)
try:
backend = load_backend()
except JuliaBridgeUnavailable as exc:
msg = "steadystate requires the Julia backend; run `oqs setup-julia` first."
raise NotImplementedError(msg) from exc
raw = backend.steadystate(
_matrix_payload(backend, H.data),
[_matrix_payload(backend, array) for array in c_arrays],
method=str(method),
rtol=float(opts.rtol),
krylov_dim=int(opts.krylov_dim),
)
return np.asarray(raw, dtype=np.complex128)
[docs]
def lindblad_superoperator(
H: Operator,
c_ops: Sequence[Operator] | None = None,
) -> Array:
"""Return the dense Lindblad superoperator matrix in column-stacked form.
The returned matrix ``L`` acts on ``vec(rho)`` (column-stacked, matching
NumPy's Fortran-order reshape) as ``vec(L(rho)) = L @ vec(rho)``.
"""
h_dim = H.dim
h_dense = np.asarray(H.to_numpy(), dtype=np.complex128)
identity = np.eye(h_dim, dtype=np.complex128)
generator = -1j * (np.kron(identity, h_dense) - np.kron(h_dense.T, identity))
for c_op in c_ops or []:
c_dense = np.asarray(c_op.to_numpy(), dtype=np.complex128)
cdc = c_dense.conj().T @ c_dense
generator += np.kron(c_dense.conj(), c_dense)
generator -= 0.5 * (np.kron(identity, cdc) + np.kron(cdc.T, identity))
return generator
[docs]
def apply_superoperator(superoperator: Array, rho: Array) -> Array:
"""Apply a column-stacked superoperator to a density matrix."""
matrix = np.asarray(superoperator, dtype=np.complex128)
state = np.asarray(rho, dtype=np.complex128)
if state.ndim == 1:
state = np.asarray(np.outer(state, state.conj()), dtype=np.complex128)
if state.ndim != 2 or state.shape[0] != state.shape[1]:
msg = "rho must be a square density matrix."
raise ValueError(msg)
dim = state.shape[0]
if matrix.shape != (dim * dim, dim * dim):
msg = "superoperator dimension does not match rho."
raise ValueError(msg)
out = matrix @ state.reshape(-1, order="F")
return np.asarray(out.reshape((dim, dim), order="F"), dtype=np.complex128)
[docs]
def brmesolve(
H: Operator,
rho0: Array,
tlist: Sequence[float],
*,
a_ops: Sequence[tuple[Operator, Callable[[float], float]]],
e_ops: Sequence[Operator] | None = None,
options: Options | None = None,
) -> Result:
"""Solve the secular Bloch-Redfield master equation.
``a_ops`` pairs each bath-coupling operator with a spectrum callable
``gamma(omega)`` giving the dissipator weight at Bohr frequency
``omega``. A constant weight ``g`` is equivalent to a Lindblad collapse
operator ``sqrt(g) * A``; a zero-temperature bath returns ``g`` for
``omega > 0`` and ``0`` for ``omega < 0``.
Transitions sharing a Bohr frequency are treated as a single secular
channel. The Hamiltonian must be Hermitian.
"""
opts = options or Options()
times = np.asarray(tlist, dtype=np.float64)
rho0_array = np.asarray(rho0, dtype=np.complex128)
a_pairs = [(op, spectrum) for op, spectrum in a_ops]
a_arrays = [op.to_numpy() for op, _ in a_pairs]
spectra = [spectrum for _, spectrum in a_pairs]
e_arrays = [op.to_numpy() for op in e_ops or []]
_validate_mesolve_inputs(H, rho0_array, times, a_arrays, e_arrays)
try:
backend = load_backend()
except JuliaBridgeUnavailable as exc:
msg = "brmesolve requires the Julia backend; run `oqs setup-julia` first."
raise NotImplementedError(msg) from exc
raw = backend.brmesolve(
_matrix_payload(backend, H.data),
rho0_array,
times,
[_matrix_payload(backend, array) for array in a_arrays],
list(spectra),
[_matrix_payload(backend, array) for array in e_arrays],
float(opts.rtol),
float(opts.atol),
bool(opts.save_states),
str(opts.method),
int(opts.krylov_dim),
bool(opts.compute_entropy),
)
raw_times = np.asarray(_field(raw, "times"), dtype=np.float64)
raw_expect = np.asarray(_field(raw, "expect"), dtype=np.complex128)
raw_entropy = (
np.asarray(_field(raw, "entropy"), dtype=np.float64)
if opts.compute_entropy
else None
)
raw_states = np.asarray(_field(raw, "states"), dtype=np.complex128)
stats = _to_python_dict(_field(raw, "solver_stats"))
needs_states = bool(opts.save_states)
expects = [raw_expect[idx, :].copy() for idx in range(raw_expect.shape[0])]
state_series = _density_states(raw_states, len(raw_times)) if needs_states else []
return Result(
times=raw_times,
states=state_series if needs_states else None,
expect=expects,
expect_std=[],
expect_stderr=[],
state_observables={},
entropy=raw_entropy,
solver_stats=stats,
)
[docs]
def propagator(
H: Operator,
tlist: Sequence[float],
*,
c_ops: Sequence[Operator] | None = None,
) -> list[Operator]:
"""Propagators for a time-independent system at each time in ``tlist``.
Without collapse operators this returns the unitary propagators
``U(t) = exp(-i H t)``. With collapse operators it returns the
propagator superoperators ``S(t) = exp(L t)`` built from the Lindblad
generator (see :func:`lindblad_superoperator`); apply them to a density
matrix with :func:`apply_superoperator`.
"""
times = np.asarray(tlist, dtype=np.float64)
if times.ndim != 1 or times.size == 0:
msg = "tlist must be a non-empty 1-D sequence."
raise ValueError(msg)
if np.any(np.diff(times) < 0):
msg = "tlist must be increasing."
raise ValueError(msg)
h_dense = np.asarray(H.to_numpy(), dtype=np.complex128)
if c_ops:
generator = lindblad_superoperator(H, c_ops)
return [
Operator(
_dense_expm(generator * float(time)),
H.space,
f"S({time})",
)
for time in times
]
return [
Operator(
_dense_expm(-1j * h_dense * float(time)),
H.space,
f"U({time})",
)
for time in times
]
def _validate_mesolve_inputs(
H: Operator | TimeDependentHamiltonian,
rho0: Array,
times: FloatArray,
c_ops: Sequence[Any],
e_ops: Sequence[Any],
) -> None:
if H.shape[0] != H.shape[1]:
msg = "H must be square."
raise ValueError(msg)
if rho0.shape != H.shape:
msg = "rho0 must have the same shape as H."
raise ValueError(msg)
if times.ndim != 1 or len(times) == 0:
msg = "tlist must be a non-empty one-dimensional sequence."
raise ValueError(msg)
if np.any(np.diff(times) < 0):
msg = "tlist must be sorted in ascending order."
raise ValueError(msg)
for op in [*c_ops, *e_ops]:
if op.shape != H.shape:
msg = "collapse and expectation operators must have the same shape as H."
raise ValueError(msg)
def _coefficient_payload(backend: Any, coefficient: object) -> object:
if isinstance(coefficient, InterpolatedCoefficient):
return backend.InterpolatedCoefficient(coefficient.times, coefficient.values)
if isinstance(coefficient, Number):
return complex(cast(Any, coefficient))
if callable(coefficient):
return coefficient
msg = f"unsupported coefficient type: {type(coefficient).__name__}"
raise TypeError(msg)
def _matrix_payload(backend: Any, matrix_like: object) -> object:
"""Return a Julia-friendly dense or sparse matrix payload."""
if sp.issparse(matrix_like):
matrix = sp.csc_matrix(matrix_like, dtype=np.complex128)
matrix.eliminate_zeros()
return _sparse_matrix_payload(backend, matrix)
array = np.asarray(matrix_like, dtype=np.complex128)
if _should_send_sparse(array):
matrix = sp.csc_matrix(array)
matrix.eliminate_zeros()
return _sparse_matrix_payload(backend, matrix)
return array
def _sparse_matrix_payload(backend: Any, matrix: Any) -> object:
colptr = np.asarray(matrix.indptr, dtype=np.int64) + 1
rowval = np.asarray(matrix.indices, dtype=np.int64) + 1
nzval = np.asarray(matrix.data, dtype=np.complex128)
return backend.sparse_from_csc(
int(matrix.shape[0]),
int(matrix.shape[1]),
colptr,
rowval,
nzval,
)
def _should_send_sparse(array: Array) -> bool:
if array.ndim != 2 or array.size < 64:
return False
nnz = int(np.count_nonzero(array))
return 4 * nnz <= int(array.size)
def _mcsolve_state_observable_plan(
dim: int,
state_observables: Mapping[str, object] | None,
) -> _MCStateObservablePlan:
if not state_observables:
return _MCStateObservablePlan(operators=(), entries=())
operators: list[Array] = []
entries: list[_MCStateObservableEntry] = []
for name, observable in state_observables.items():
_validate_state_observable_name(name)
if not callable(observable):
msg = f"state observable {name!r} must be callable."
raise TypeError(msg)
spec = _state_observable_spec(observable)
if spec is None:
msg = (
"mcsolve can aggregate built-in state observables created by "
"openquantumsim.state_metrics and the named observable helpers. "
"Arbitrary Python callbacks require mesolve or single_trajectory "
"with saved states."
)
raise NotImplementedError(msg)
kind = spec.get("kind")
if kind == "pure_constant":
entries.append(
_MCStateObservableEntry(
name=name,
constant=float(cast(float, spec["value"])),
),
)
elif kind == "population":
index = int(cast(int, spec["index"]))
operators.append(_population_projector(dim, index))
entries.append(
_MCStateObservableEntry(
name=name,
operator_offset=len(operators) - 1,
),
)
elif kind == "bloch":
index = int(cast(int, spec["index"]))
operators.append(_bloch_operator(dim, index))
entries.append(
_MCStateObservableEntry(
name=name,
operator_offset=len(operators) - 1,
),
)
elif kind == "fidelity":
reference = np.asarray(spec["reference"], dtype=np.complex128)
operators.append(_pure_reference_projector(reference, dim))
entries.append(
_MCStateObservableEntry(
name=name,
operator_offset=len(operators) - 1,
),
)
elif kind == "unsupported":
reason = str(spec.get("reason", "the metric is nonlinear"))
msg = (
f"mcsolve cannot aggregate state observable {name!r}: {reason}. "
"Use single_trajectory with save_states=True for per-trajectory "
"nonlinear callbacks."
)
raise NotImplementedError(msg)
else:
msg = f"unsupported built-in state observable kind for mcsolve: {kind!r}."
raise NotImplementedError(msg)
return _MCStateObservablePlan(operators=tuple(operators), entries=tuple(entries))
def _mcsolve_state_observable_results(
plan: _MCStateObservablePlan,
raw_expect: Array,
raw_expect_std: FloatArray,
raw_expect_stderr: FloatArray,
operator_offset: int,
n_times: int,
) -> tuple[dict[str, Array], dict[str, FloatArray], dict[str, FloatArray]]:
observed: dict[str, Array] = {}
observed_std: dict[str, FloatArray] = {}
observed_stderr: dict[str, FloatArray] = {}
for entry in plan.entries:
if entry.operator_offset is None:
value = complex(0.0 if entry.constant is None else entry.constant)
observed[entry.name] = np.full(n_times, value, dtype=np.complex128)
observed_std[entry.name] = np.zeros(n_times, dtype=np.float64)
observed_stderr[entry.name] = np.zeros(n_times, dtype=np.float64)
continue
idx = operator_offset + entry.operator_offset
observed[entry.name] = raw_expect[idx, :].copy()
observed_std[entry.name] = raw_expect_std[idx, :].copy()
observed_stderr[entry.name] = raw_expect_stderr[idx, :].copy()
return observed, observed_std, observed_stderr
def _population_projector(dim: int, index: int) -> Array:
if index < 0 or index >= dim:
msg = f"population index {index} is outside Hilbert dimension {dim}."
raise ValueError(msg)
projector = np.zeros((dim, dim), dtype=np.complex128)
projector[index, index] = 1.0
return projector
def _bloch_operator(dim: int, index: int) -> Array:
if dim != 2:
msg = "Bloch-vector state observables require a two-level mcsolve state."
raise ValueError(msg)
if index == 0:
return np.array([[0.0, 1.0], [1.0, 0.0]], dtype=np.complex128)
if index == 1:
return np.array([[0.0, -1.0j], [1.0j, 0.0]], dtype=np.complex128)
if index == 2:
return np.array([[1.0, 0.0], [0.0, -1.0]], dtype=np.complex128)
msg = f"Bloch-vector component index must be 0, 1, or 2; got {index}."
raise ValueError(msg)
def _pure_reference_projector(reference: Array, dim: int) -> Array:
array = np.asarray(reference, dtype=np.complex128)
if array.ndim == 1:
if array.shape != (dim,):
msg = "fidelity reference ket length must match H."
raise ValueError(msg)
norm = np.linalg.norm(array)
if norm <= 0:
msg = "fidelity reference ket must have nonzero norm."
raise ValueError(msg)
ket = array / norm
return cast(Array, np.outer(ket, ket.conj()))
if array.ndim == 2 and array.shape == (dim, dim):
rho = np.asarray(0.5 * (array + array.conj().T), dtype=np.complex128)
trace = np.trace(rho)
if abs(trace) <= 1e-15:
msg = "fidelity reference density matrix must have nonzero trace."
raise ValueError(msg)
rho = cast(Array, rho / trace)
purity = float(np.real(np.trace(rho @ rho)))
if not np.isclose(purity, 1.0, atol=1e-8):
msg = (
"mcsolve can aggregate fidelity state observables only for pure "
"reference states; mixed-state fidelity is nonlinear."
)
raise NotImplementedError(msg)
return rho
msg = (
"fidelity reference must be a ket or density matrix with dimension "
"matching H."
)
raise ValueError(msg)
def _validate_state_observable_name(name: object) -> None:
if not isinstance(name, str):
msg = "state observable names must be strings."
raise TypeError(msg)
if not name:
msg = "state observable names must not be empty."
raise ValueError(msg)
if "/" in name:
msg = "state observable names must not contain '/'."
raise ValueError(msg)
def _validate_mcsolve_inputs(
H: Operator,
psi0: Array,
times: FloatArray,
c_ops: Sequence[Any],
e_ops: Sequence[Any],
n_traj: int,
max_step: float,
n_jobs: int,
checkpoint_every: int,
) -> None:
if H.shape[0] != H.shape[1]:
msg = "H must be square."
raise ValueError(msg)
if psi0.shape != (H.shape[0],):
msg = "psi0 must be a ket vector with length matching H."
raise ValueError(msg)
if times.ndim != 1 or len(times) == 0:
msg = "tlist must be a non-empty one-dimensional sequence."
raise ValueError(msg)
if np.any(np.diff(times) < 0):
msg = "tlist must be sorted in ascending order."
raise ValueError(msg)
if n_traj <= 0:
msg = "n_traj must be positive."
raise ValueError(msg)
if max_step <= 0:
msg = "max_step must be positive."
raise ValueError(msg)
if n_jobs != -1 and n_jobs <= 0:
msg = "n_jobs must be -1 or positive."
raise ValueError(msg)
if checkpoint_every <= 0:
msg = "checkpoint_every must be positive."
raise ValueError(msg)
for op in [*c_ops, *e_ops]:
if op.shape != H.shape:
msg = "collapse and expectation operators must have the same shape as H."
raise ValueError(msg)
def _field(raw: object, name: str) -> object:
try:
return getattr(raw, name)
except AttributeError:
return raw[name] # type: ignore[index]
def _density_states(raw_states: Array, n_times: int) -> list[Array]:
if raw_states.ndim != 3 or raw_states.shape[2] != n_times:
msg = "backend did not return density states with shape (d, d, n_times)."
raise RuntimeError(msg)
return [raw_states[:, :, idx].copy() for idx in range(raw_states.shape[2])]
def _ket_states(raw_states: Array, n_times: int) -> list[Array]:
if raw_states.ndim != 2 or raw_states.shape[1] != n_times:
msg = "backend did not return ket states with shape (d, n_times)."
raise RuntimeError(msg)
return [raw_states[:, idx].copy() for idx in range(raw_states.shape[1])]
_SOLVER_STAT_NAMES = (
"nsteps",
"nfev",
"wall_time",
"retcode",
"n_traj",
"max_step",
"n_jobs_requested",
"n_workers",
"threaded",
"checkpoint_file",
"checkpoint_every",
"checkpoint_completed",
"checkpoint_start_completed",
"checkpoint_previous_target_n_traj",
"progress",
"resumed",
"method",
"requested_method",
"krylov_dim",
"compute_entropy",
"time_dependent",
)
def _to_python_dict(raw: object) -> dict[str, object]:
if isinstance(raw, Mapping):
return dict(raw)
try:
present = set(dir(raw))
except Exception:
present = set(_SOLVER_STAT_NAMES)
values: dict[str, object] = {}
for name in _SOLVER_STAT_NAMES:
if name not in present:
continue
values[name] = getattr(raw, name)
return values