"""Backend protocol and backend-agnostic solver-result containers.
quchip keeps the *physics description* separate from the *solver conversion*.
The engine (``quchip.engine``) emits structured IR (``EngineResult``,
``SolveProblem`` and ``SolveBatch``). This IR carries ordinary-GHz frequencies
and backend-free operator payloads (``CanonicalOperator``), which each backend
translates into the native form best suited to its solver. This module sets
only the *contract* (see :class:`Backend`), not the storage.
Unit convention
---------------
All frequencies that cross the backend boundary are **ordinary** (not angular)
GHz, with time in ns. The single ``2π`` conversion occurs at the engine
boundary (``assembly.py``), and backends never rescale.
Aliases
-------
``Operator`` and ``State`` are both ``typing.Any``. Concrete backends bind
them to their native type: ``qutip.Qobj`` for :class:`QuTiPBackend`,
``dynamiqs.QArray`` for :class:`DynamiqsBackend`.
References
----------
* Johansson, Nation, Nori — *QuTiP 2*, Comput. Phys. Commun. 183, 1760 (2012)
* Guilmin et al. — *dynamiqs: an open-source Python library for GPU-accelerated
and differentiable simulation of quantum systems* (2024)
* Bradbury et al. — *JAX: composable transformations of Python+NumPy programs*
(2018)
* Lindblad — Commun. Math. Phys. 48, 119 (1976)
* Breuer & Petruccione — *The Theory of Open Quantum Systems* (OUP, 2002)
"""
from __future__ import annotations
import math
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Any, Callable, Iterator, Sequence
import numpy as np
from quchip.backend._dims import normalize_dims_from_list
from quchip.backend.containers import (
_DEFAULT_SOLVE_OPTIONS,
DeferredBatch,
EagerBatch,
EigensystemData,
PreparedBatch,
PreparedHamiltonian,
PreparedStationary,
LinearResponseSolverResult,
CoherenceResponseSolverResult,
SolverResult,
SteadyStateSolverResult,
VmappedBatch,
)
if TYPE_CHECKING:
from quchip.engine.ir import (
CanonicalOperator,
EngineResult,
SolveBatch,
SolveProblem,
)
Operator = Any
State = Any
def _is_unit(coefficient: Any) -> bool:
"""True for a Python ``1``, which leaves an operator unscaled in a linear combination."""
return type(coefficient) in (int, float) and coefficient == 1
[docs]
class Backend(ABC):
r"""Abstract contract implemented by every quchip operator backend.
A backend wraps a quantum-dynamics library (QuTiP, dynamiqs, ...) and
exposes three orthogonal surfaces:
* **Operator algebra**: ``destroy``/``create``/``number``/``identity``,
``tensor``, ``embed``, ``dag``, ``matmul``, and eigensystem helpers.
Concrete operators support native Python arithmetic (``+ - * @``), so
quchip never adds ``scale``/``add`` wrappers.
* **State algebra**: ``basis``/``coherent``, ``overlap``, ``ptrace``,
``expect``, and ket ↔ density-matrix conversion.
* **IR lowering & solve dispatch**: ``prepare_hamiltonian`` /
``prepare_batch`` convert engine IR into native RHS form, and ``sesolve``
/ ``mesolve`` run the solver. ``solve_problem`` /
``parallel_solve_problems`` / ``solve_batch`` drive the full pipeline
from engine to result.
The NumPy-based defaults are correct for all backends. Subclasses override
them only where they can do better (e.g. QuTiP reuses ``Qobj.eigenstates``
and dynamiqs reuses ``dq.expect``).
Solver Hamiltonians use angular frequencies, while public device parameters use
ordinary GHz. For operator and circuit-QED conventions, see Blais et al.,
`Rev. Mod. Phys. 93, 025005 (2021) <https://doi.org/10.1103/RevModPhys.93.025005>`_.
"""
# ------------------------------------------------------------------
# Array / scalar surface
# ------------------------------------------------------------------
@property
@abstractmethod
def array_module(self) -> Any:
"""Return the array module used by backend-aware numeric code.
Returns ``numpy`` for CPU-only backends and ``jax.numpy`` for
JAX-traceable backends. Engine code that must stay differentiable
(signal evaluation, frame assembly) routes array operations through
this module, so the dynamiqs backend stays JAX-traceable end to end.
"""
...
[docs]
@abstractmethod
def to_array(self, op: Operator) -> Any:
r"""Return a dense array (native to :attr:`array_module`) for *op*.
Parameters
----------
op : Operator
Backend-native operator.
"""
...
[docs]
def overlap(self, a: State, b: State) -> complex:
r"""Return the scalar inner product ⟨a|b⟩ for two kets.
Parameters
----------
a, b : State
Kets in the same Hilbert space, where the first is conjugated.
"""
a_arr = np.asarray(self.to_array(a), dtype=complex)
b_arr = np.asarray(self.to_array(b), dtype=complex)
return complex(np.conj(a_arr).T @ b_arr)
[docs]
def norm(self, state_or_op: State | Operator) -> Any:
r"""Return the Frobenius / ℓ²-norm of a state or operator.
A concrete ``float`` for CPU-only backends. A JAX-traceable backend can
return a native 0-d array instead, so a traced amplitude stays
differentiable. A caller that needs a concrete value passes the result
through :func:`~quchip.utils.jax_utils.maybe_concrete_scalar`.
Parameters
----------
state_or_op : State or Operator
Ket or matrix whose Euclidean/Frobenius norm is required.
"""
arr = np.asarray(self.to_array(state_or_op), dtype=complex)
return float(np.linalg.norm(arr))
[docs]
def is_native_state(self, state: Any) -> bool:
r"""Return whether *state* is already represented by this backend.
Parameters
----------
state : State
Ket or density matrix to inspect or convert.
"""
return False
[docs]
def trace(self, op: Operator) -> complex:
r"""Return the scalar trace ``Tr(op)``.
Parameters
----------
op : Operator
Backend-native operator.
"""
arr = np.asarray(self.to_array(op), dtype=complex)
return complex(np.trace(arr))
# ------------------------------------------------------------------
# Single-mode operator factories (Fock-basis defaults)
# ------------------------------------------------------------------
def _single_mode(self, data: np.ndarray) -> Operator:
"""Wrap a dense ``n × n`` Fock-basis matrix as a native operator.
Concrete backends route through :meth:`from_canonical_operator` so
they can choose the optimal native layout (dense, DIA, CSR).
"""
from quchip.engine.ir import CanonicalOperator
n = data.shape[0]
return self.from_canonical_operator(
CanonicalOperator.from_dense(data, dims=(n,), basis="fock", subsystem_labels=("0",))
)
[docs]
def eager_operators(self) -> Any:
"""Return a context manager for operator construction during JAX compile-time evaluation.
The default implementation leaves backend state unchanged.
"""
from contextlib import nullcontext
return nullcontext()
[docs]
def destroy(self, n: int) -> Operator:
r"""Build the annihilation operator :math:`\hat a` for an *n*-level Fock space.
Parameters
----------
n : int
Positive Hilbert-space dimension. Number states go from 0 to n - 1.
"""
data = np.zeros((n, n), dtype=complex)
for k in range(1, n):
data[k - 1, k] = np.sqrt(k)
return self._single_mode(data)
[docs]
def create(self, n: int) -> Operator:
r"""Build the creation operator :math:`\hat a^\dagger` for an *n*-level Fock space.
Parameters
----------
n : int
Positive Hilbert-space dimension. Number states go from 0 to n - 1.
"""
data = np.zeros((n, n), dtype=complex)
for k in range(1, n):
data[k, k - 1] = np.sqrt(k)
return self._single_mode(data)
[docs]
def number(self, n: int) -> Operator:
r"""Build the number operator :math:`\hat n = \hat a^\dagger \hat a`.
Parameters
----------
n : int
Positive Hilbert-space dimension. Number states go from 0 to n - 1.
"""
return self._single_mode(np.diag(np.arange(n, dtype=complex)))
[docs]
def identity(self, n: int) -> Operator:
r"""Build the identity operator for an *n*-level space.
Parameters
----------
n : int
Positive Hilbert-space dimension. Number states go from 0 to n - 1.
"""
return self._single_mode(np.eye(n, dtype=complex))
[docs]
def from_array(self, data: Any, dims: list[list[int]] | None = None) -> Operator:
r"""Construct a native operator from a dense matrix.
*dims* accepts quchip's row/col layout (``[[row_dims], [col_dims]]``)
or a flat list. ``None`` means a single subsystem of size
``data.shape[0]``.
Parameters
----------
data : array_like
Dense matrix in the requested tensor-product basis.
dims : list or None, default None
Subsystem dimensions, flat or [[row_dims], [column_dims]]. None treats the row space as one subsystem.
"""
from quchip.engine.ir import CanonicalOperator
arr = np.asarray(data, dtype=complex)
dim_tuple = normalize_dims_from_list(dims, arr.shape[0])
labels = tuple(str(i) for i in range(len(dim_tuple)))
return self.from_canonical_operator(
CanonicalOperator.from_dense(arr, dims=dim_tuple, basis="fock", subsystem_labels=labels)
)
[docs]
def diag(self, values: Any, dims: list[list[int]] | None = None) -> Operator:
r"""Construct a backend-native diagonal operator from main-diagonal *values*.
A backend with sparse layouts (for example, dynamiqs sparse-DIA) should
return its sparse representation so that element-wise composition with
other diagonal operators stays sparse. The default builds a dense
``np.diag`` and then calls :meth:`from_array`.
Parameters
----------
values : array_like
Main-diagonal entries, flattened in input order.
dims : list or None, default None
Subsystem dimensions as in from_array. None uses one subsystem.
"""
from quchip.engine.ir import CanonicalOperator
arr = np.asarray(values).reshape(-1)
dim_tuple = normalize_dims_from_list(dims, arr.shape[0])
labels = tuple(str(i) for i in range(len(dim_tuple)))
return self.from_canonical_operator(
CanonicalOperator.from_dia(
arr.astype(complex)[None, :],
np.array([0], dtype=int),
shape=(arr.shape[0], arr.shape[0]),
dims=dim_tuple,
basis="fock",
subsystem_labels=labels,
)
)
[docs]
def compact_operator(self, rows: Any, columns: Any, values: Any, size: int) -> Operator:
r"""Construct an operator on ``size`` compact excitation states from its entries.
Repeated entries add. Each backend chooses the storage. The default
stores a dense matrix, which keeps traced values traced, and raises
:class:`MemoryError` before it allocates one that cannot fit.
Parameters
----------
rows, columns : array_like of int
Row and column index of each entry.
values : array_like
Value of each entry.
size : int
Number of compact states.
"""
from quchip.utils.jax_utils import concrete_array_module
from quchip.utils.memory import dense_bytes, require_memory
require_memory(dense_bytes(size, size, copies=2), task=f"A dense compact operator (dimension {size})",
remedy="Reduce the model, the initial excitation support, or an excitation cutoff's max_total.")
xp = concrete_array_module(values)
result = xp.zeros((size, size), dtype=complex)
if xp is np:
np.add.at(result, (rows, columns), values)
else:
result = result.at[rows, columns].add(values)
return self.from_array(result, dims=[[size], [size]])
[docs]
def compact_operators_fit(self, size: int, count: int) -> bool:
r"""Return whether ``count`` compact operators on ``size`` states fit in the available memory.
The default stores each as a dense matrix, as :meth:`compact_operator` does.
Parameters
----------
size : int
Number of compact states.
count : int
Number of operators that exist at the same time.
"""
from quchip.utils.memory import dense_bytes, fits_memory
return fits_memory(dense_bytes(size, size, copies=count))
[docs]
@abstractmethod
def to_canonical_operator(self, op: Operator) -> "CanonicalOperator":
r"""Serialize a native operator into the backend-agnostic canonical IR.
Parameters
----------
op : Operator
Backend-native operator.
"""
...
[docs]
@abstractmethod
def from_canonical_operator(self, canonical: "CanonicalOperator") -> Operator:
r"""Reconstruct a native operator from the canonical IR payload.
Parameters
----------
canonical : CanonicalOperator
Backend-neutral matrix payload with its tensor dimensions and basis metadata.
"""
...
# ------------------------------------------------------------------
# Operator algebra
# ------------------------------------------------------------------
[docs]
def coerce_operator(self, op: Operator) -> Operator:
r"""Coerce an array-like local operator into backend-native form.
Mirrors :meth:`coerce_state` for operators. Device-side accessors for
physical coupling (for example,
:meth:`~quchip.devices.protocols.ChargeCoupled.charge_coupling_operator`)
return dense, trace-safe arrays, not backend-native operators.
Composition entry points (:meth:`tensor`, :meth:`dag`) send their
operands through this hook, so both forms compose freely. A backend whose
native operators are already arrays overrides this method with a
trace-safe passthrough.
Parameters
----------
op : Operator
Backend-native operator.
"""
return self.from_array(self.to_array(op))
[docs]
def dag(self, op: Operator) -> Operator:
r"""Return the Hermitian conjugate ``op†``.
Parameters
----------
op : Operator
Backend-native operator.
"""
arr = np.asarray(self.to_array(op), dtype=complex)
return self.from_array(np.conj(arr).T)
[docs]
def matmul(self, a: Operator, b: Operator) -> Operator:
r"""Return the matrix product ``a @ b``.
Parameters
----------
a, b : Operator
Left and right matrix factors with compatible dimensions.
"""
return a @ b
[docs]
def linear_combination(self, terms: Sequence[tuple[Any, Operator]]) -> Operator:
r"""Return ``Σᵢ cᵢ·Aᵢ``, adding the terms in order.
A backend can accumulate concrete terms in its native layout, whereas
this default uses operator arithmetic. A coefficient of exactly ``1``
adds its operator unscaled.
Parameters
----------
terms : sequence of (scalar, Operator)
At least one coefficient and operator pair on a common space.
"""
total = None
for coefficient, op in terms:
term = op if _is_unit(coefficient) else coefficient * op
total = term if total is None else total + term
if total is None:
raise ValueError("linear_combination needs at least one term.")
return total
[docs]
def eigenenergies(self, op: Operator) -> Any:
r"""Return the ascending eigenvalues of a Hermitian operator.
Parameters
----------
op : Operator
Hermitian operator to diagonalize.
"""
dense = np.asarray(self.to_array(op), dtype=complex)
return np.linalg.eigvalsh(dense)
[docs]
def eigenstates(self, op: Operator) -> tuple[Any, Any]:
r"""Return ``(eigenvalues, eigenstates)`` of a Hermitian operator, ascending.
Parameters
----------
op : Operator
Hermitian operator to diagonalize.
"""
data = self.eigensystem_data(op)
return data.eigenvalues, data.eigenstates
[docs]
def eigensystem_data(self, op: Operator) -> EigensystemData:
r"""Return ascending eigenvalues, stacked eigenvector matrix, and lazy eigenstates.
Parameters
----------
op : Operator
Hermitian operator to diagonalize.
"""
dense = np.asarray(self.to_array(op), dtype=complex)
evals, evecs = np.linalg.eigh(dense)
def build_states() -> list[Any]:
return [self.from_array(evecs[:, idx:idx + 1]) for idx in range(evecs.shape[1])]
return EigensystemData(
eigenvalues=evals,
eigenvector_matrix=evecs,
_states_builder=build_states,
)
[docs]
def expect(self, op: Operator, state: State) -> complex:
r"""Return the expectation value ⟨state|op|state⟩ for a ket or a density matrix.
Parameters
----------
op : Operator
Observable with the same tensor dimensions as state.
state : State
Ket or density matrix, not normalized by the method.
"""
op_arr = np.asarray(self.to_array(op), dtype=complex)
state_arr = np.asarray(self.to_array(state), dtype=complex)
if state_arr.ndim == 2 and state_arr.shape[1] == 1:
return complex(np.conj(state_arr).T @ op_arr @ state_arr)
return complex(np.einsum("ij,ji->", op_arr, state_arr))
[docs]
def ptrace(self, state: State, keep: int | list[int], dims: list[int]) -> State:
"""Reduce a composite state onto subsystem(s) *keep* with a partial trace.
Parameters
----------
state
Composite ket or density matrix over the subsystems in *dims*.
keep
Subsystem index (or indices) to keep. All other subsystems are traced out.
dims
Subsystem dimensions of the full composite space, in order.
"""
if isinstance(keep, int):
keep = [keep]
state_arr = np.asarray(self.to_array(state), dtype=complex)
n = len(dims)
if state_arr.ndim == 2 and state_arr.shape[1] == 1:
state_arr = state_arr @ np.conj(state_arr).T
rho = state_arr.reshape(list(dims) + list(dims))
for idx in reversed(sorted(set(range(n)) - set(keep))):
rho = np.trace(rho, axis1=idx, axis2=idx + n)
n -= 1
kept_dim = int(np.prod([dims[k] for k in keep]))
return self.from_array(rho.reshape(kept_dim, kept_dim))
[docs]
def permute_state(self, state: State, dims: Sequence[int], order: Sequence[int]) -> State:
r"""Reorder the subsystems of a composite state.
``dims`` are the subsystem levels in *state*'s current tensor order.
``order`` follows the ``numpy.transpose`` convention, where
``order[i]`` is the subsystem index in the *current* order that becomes
position ``i`` after the permutation, so
``order == list(range(len(dims)))`` is a no-op.
Default, correct for all backends: densify, reshape to one axis per
subsystem, ``transpose``, and reshape back. A ket reshapes/transposes
once. A density matrix applies the same subsystem permutation to the
row index group and to the column index group. Concrete backends can
override this method with a native reorder (e.g. QuTiP's
``Qobj.permute``), which stays sparse-friendly and keeps the dims
metadata.
Every array operation goes through :attr:`array_module` (never bare
``numpy``), so a traced state stays traced end to end under the
dynamiqs backend.
Parameters
----------
state : State
Ket or density matrix in the current tensor order.
dims : sequence of int
Current subsystem dimensions.
order : sequence of int
Permutation of subsystem indices, listing their positions in the new order.
"""
# `dims`/`order` are always concrete Python ints (subsystem structure,
# never a physics parameter), so their product is computed with
# `math.prod`, not with `array_module.prod`. If a shape value goes
# through the array module, it becomes a tracer under `jit`, and the
# next `reshape` fails.
xp = self.array_module
arr = xp.asarray(self.to_array(state), dtype=complex)
dims = list(dims)
is_ket = arr.ndim == 2 and arr.shape[1] == 1
if is_ket:
ket_axes = arr.reshape(dims)
permuted_ket_axes = xp.transpose(ket_axes, order)
return self.from_array(permuted_ket_axes.reshape(-1, 1))
n_subsystems = len(dims)
row_then_col_axes = arr.reshape(dims + dims)
# The row axes take `order` directly; the column axes need the same
# permutation, shifted by `n_subsystems` since they occupy the second
# half of the reshaped axes.
row_and_col_order = list(order) + [n_subsystems + i for i in order]
permuted = xp.transpose(row_then_col_axes, row_and_col_order)
total_dim = math.prod(dims)
return self.from_array(permuted.reshape(total_dim, total_dim))
# ------------------------------------------------------------------
# Batched-over-time extraction
# ------------------------------------------------------------------
#
# A trajectory is ``T`` saved states. The user-facing extractors
# (``population``, ``populations``, ``overlap``,
# ``amplitude_array``) need one scalar/vector per save point. Doing that
# with a Python loop over ``T`` single-state ``expect``/``ptrace`` calls
# makes per-point dispatch the dominant cost on long ``tlist``s and
# clean-``O(N)`` sweeps, rather than the matrix work itself. The methods
# below collapse that loop into a single batched op over the leading
# time axis.
#
# The native stacked form is the backend's own (a stacked ``QArray`` for
# dynamiqs, a ``np.stack`` for QuTiP) so neither layout is forced on the
# other. dynamiqs implementations stay ``jnp`` — storing/operating on one
# stacked array is strictly more traceable than a Python list of
# per-time tracers.
[docs]
def stack_states(self, states: Any) -> Any:
r"""Stack the saved states of a trajectory into a single ``(T, …)`` array.
The default densifies each state and applies ``np.stack`` along a new
leading time axis. A backend whose solver already returns a stacked
native state (dynamiqs) overrides this method with a no-op
pass-through.
Parameters
----------
states : sequence of State
Saved states in time order, with matching shapes and tensor dimensions.
"""
return np.stack([np.asarray(self.to_array(s), dtype=complex) for s in states])
@staticmethod
def _is_ket_stack(stacked: Any) -> bool:
"""Return whether a stacked trajectory is kets ``(T, n, 1)`` vs DMs ``(T, n, n)``."""
return stacked.ndim == 3 and stacked.shape[2] == 1
[docs]
def expect_over_time(self, op: Operator, stacked_states: Any) -> Any:
r"""Return ⟨op⟩(t) for every save point, as one ``(T,)`` array in :attr:`array_module`.
Accepts a ket stack ``(T, n, 1)`` (returns ⟨ψ|op|ψ⟩) or a density-matrix
stack ``(T, n, n)`` (returns ``Tr(op·ρ)``). One matrix product replaces
the per-point ``expect`` loop. Every intermediate stays in
:attr:`array_module`, so the JAX backend keeps the trace differentiable.
Parameters
----------
op : Operator
Observable in the tensor-product basis of the states.
stacked_states : array_like
Backend-stacked kets or density matrices from stack_states.
"""
xp = self.array_module
op_arr = xp.asarray(self.to_array(op), dtype=complex)
rho = xp.asarray(self.to_array(stacked_states), dtype=complex)
if self._is_ket_stack(rho):
psi = rho[..., 0]
return xp.einsum("ti,ti->t", xp.conj(psi), psi @ op_arr.T)
return xp.einsum("tij,ji->t", rho, op_arr)
[docs]
def overlap_over_time(self, target: State, stacked_states: Any) -> Any:
r"""Return ⟨target|ψ(t)⟩ for every save point, as one ``(T,)`` array.
For a ket stack, the result is the **phase-sensitive complex
amplitude** ⟨target|ψ(t)⟩ (never ``|·|²``), so phase-dependent
gradients flow through it. A density-matrix stack has no single phase,
so the method returns the target-state population ⟨target|ρ(t)|target⟩.
The calculation stays in :attr:`array_module` for JAX traceability.
Parameters
----------
target : State
Reference ket in the same Hilbert space.
stacked_states : array_like
Backend-stacked saved states from stack_states.
"""
xp = self.array_module
tgt = xp.asarray(self.to_array(target), dtype=complex).reshape(-1)
arr = xp.asarray(self.to_array(stacked_states), dtype=complex)
if self._is_ket_stack(arr):
return xp.einsum("i,ti->t", xp.conj(tgt), arr[..., 0])
return xp.einsum("i,tij,j->t", xp.conj(tgt), arr, tgt)
[docs]
def populations_over_time(self, stacked_states: Any) -> Any:
r"""Return full-chip diagonal populations ``(T, ∏dims)`` (real) for every save point.
For a ket stack, the result is ``|ψ(t)|²`` along the level axis (the
density matrix is never built). For a DM stack, it is the real
diagonal. The calculation stays in :attr:`array_module` for JAX
traceability.
Parameters
----------
stacked_states : array_like
Backend-stacked saved states. The trailing dimensions encode ket or matrix entries.
"""
xp = self.array_module
arr = xp.asarray(self.to_array(stacked_states), dtype=complex)
if self._is_ket_stack(arr):
return xp.abs(arr[..., 0]) ** 2
return xp.real(xp.diagonal(arr, axis1=1, axis2=2))
[docs]
def ptrace_over_time(self, stacked_states: Any, keep: int | list[int], dims: list[int]) -> Any:
r"""Reduce onto subsystem(s) *keep* at every save point, and return a ``(T, k, k)`` DM stack.
Accepts a ket or DM trajectory and uses no per-point Python loop.
Parameters
----------
stacked_states : array_like
Backend-stacked saved states from stack_states.
keep : int or list of int
Subsystem positions to keep.
dims : list of int
Full tensor-product dimensions before tracing.
"""
if isinstance(keep, int):
keep = [keep]
rho = np.asarray(stacked_states, dtype=complex)
n = len(dims)
if self._is_ket_stack(rho):
psi = rho[..., 0]
rho = np.einsum("ti,tj->tij", psi, np.conj(psi))
t_count = rho.shape[0]
rho = rho.reshape([t_count] + list(dims) + list(dims))
offset = 1 # leading time axis
live = n
for idx in reversed(sorted(set(range(n)) - set(keep))):
rho = np.trace(rho, axis1=offset + idx, axis2=offset + idx + live)
live -= 1
kept_dim = int(np.prod([dims[k] for k in keep]))
return self.array_module.asarray(rho.reshape(t_count, kept_dim, kept_dim))
[docs]
def embed(self, op: Operator, device_index: int, dims: Sequence[int]) -> Operator:
"""Embed a single-device operator at *device_index* into the full tensor space.
Validation is shared, and the identity factors and tensor product go
through the backend's :meth:`identity` / :meth:`tensor`, so each
backend keeps its native-optimal layout.
Parameters
----------
op
Single-device operator whose dimension matches ``dims[device_index]``.
device_index
Position of *op* in the tensor product.
dims
Subsystem dimensions of the full space, in order.
"""
n_devices = len(dims)
if device_index < 0 or device_index >= n_devices:
raise ValueError(f"device_index {device_index} out of range for {n_devices} devices")
if op.shape[0] != dims[device_index]:
raise ValueError(
f"Operator dimension {op.shape[0]} does not match dims[{device_index}] = {dims[device_index]}"
)
factors = [op if i == device_index else self.identity(d) for i, d in enumerate(dims)]
return self.tensor(*factors)
[docs]
@abstractmethod
def embed_two_body(
self,
op_ab: Operator,
index_a: int,
index_b: int,
dims: Sequence[int],
) -> Operator:
r"""Embed a two-body operator on devices *index_a* ⊗ *index_b* into the full space.
Handles the subsystem SWAP when ``index_a > index_b`` and the
identity-padding for non-adjacent devices. Backends with sparse layouts
do not materialize the dense full matrix.
Parameters
----------
op_ab : Operator
Two-subsystem operator, in index_a then index_b tensor order.
index_a, index_b : int
Distinct positions of the target subsystems in dims.
dims : sequence of int
Full tensor-product dimensions in chip order.
"""
...
# ------------------------------------------------------------------
# State factories
# ------------------------------------------------------------------
[docs]
def basis(self, n: int, k: int) -> State:
r"""Build the Fock basis ket :math:`|k\rangle` in an *n*-level space.
Parameters
----------
n : int
Hilbert-space dimension.
k : int
Number-state index, satisfying 0 <= k < n.
"""
vec = np.zeros((n, 1), dtype=complex)
vec[k, 0] = 1.0
return self.from_array(vec)
[docs]
def tensor_states(self, *states: State) -> State:
r"""Return the tensor product of states (defaults to :meth:`tensor`).
Parameters
----------
*states : State
Factors in tensor-product order.
"""
return self.tensor(*states)
[docs]
def coherent(self, n: int, alpha: complex) -> State:
r"""Build the coherent state :math:`|\alpha\rangle` truncated to *n* Fock levels.
Built from the analytic series
:math:`|\alpha\rangle = e^{-|\alpha|^2/2} \sum_k \alpha^k/\sqrt{\Gamma(k+1)}\,|k\rangle`.
Parameters
----------
n : int
Fock-space truncation.
alpha : complex
Dimensionless coherent-state amplitude. Its squared magnitude sets the untruncated mean occupation.
"""
import math
coeffs = np.array(
[alpha**k / np.sqrt(math.factorial(k)) for k in range(n)],
dtype=complex,
)
coeffs *= np.exp(-0.5 * abs(alpha) ** 2)
return self.from_array(coeffs.reshape(n, 1))
[docs]
def state_to_dm(self, state: State) -> State:
r"""Return a density matrix, or pass *state* through if *state* is already one.
Parameters
----------
state : State
Ket or density matrix to inspect or convert.
"""
if not self.is_ket(state):
return state
arr = np.asarray(self.to_array(state), dtype=complex).reshape(-1, 1)
return self.from_array(arr @ np.conj(arr).T)
[docs]
def is_ket(self, state: State) -> bool:
r"""Return whether *state* is a ket (flat or column vector) rather than a density matrix.
Parameters
----------
state : State
Ket or density matrix to inspect or convert.
"""
arr = np.asarray(self.to_array(state), dtype=complex)
return arr.ndim == 1 or (arr.ndim == 2 and arr.shape[1] == 1)
[docs]
def as_density_matrix(self, state: State) -> State:
r"""Promote a ket to its density matrix, and pass a density matrix through unchanged.
Parameters
----------
state : State
Ket or density matrix to inspect or convert.
"""
return self.state_to_dm(state) if self.is_ket(state) else state
[docs]
@abstractmethod
def tensor(self, *operators: Operator) -> Operator:
r"""Return the tensor product of operators, and keep the subsystem dims metadata.
Parameters
----------
*operators : Operator
Factors in tensor-product order.
"""
...
# ------------------------------------------------------------------
# Solver dispatch
# ------------------------------------------------------------------
[docs]
@abstractmethod
def sesolve(
self,
H: Any,
psi0: State,
tlist: Any,
e_ops: list[Operator] | None = None,
options: dict[str, Any] | None = None,
) -> SolverResult:
"""Solve the Schrödinger equation :math:`i\\hbar \\partial_t |\\psi\\rangle = H |\\psi\\rangle`.
Parameters
----------
H
Native Hamiltonian (``Qobj``/``QobjEvo`` or dynamiqs
``TimeQArray``).
psi0
Initial ket.
tlist
Save times in ns.
e_ops
Optional observables. The solver accumulates ⟨ψ|Oᵢ|ψ⟩ along the trajectory.
options
Backend-specific options dict. For the supported keys, see each
backend's ``resolve_solver_options``.
"""
...
[docs]
@abstractmethod
def mesolve(
self,
H: Any,
rho0: State,
tlist: Any,
c_ops: list[Operator] | None = None,
e_ops: list[Operator] | None = None,
options: dict[str, Any] | None = None,
) -> SolverResult:
r"""Solve the Lindblad master equation (Lindblad 1976; Breuer & Petruccione 2002).
Integrates
:math:`\dot\rho = -i[H,\rho] + \sum_k \mathcal D[L_k]\rho`
with
:math:`\mathcal D[L]\rho = L\rho L^\dagger - \tfrac12\{L^\dagger L, \rho\}`.
Parameters
----------
H : object
Native Hamiltonian, already scaled to angular frequency in rad/ns. Do not multiply by 2*pi again.
rho0 : State
Initial density matrix, or a ket that the concrete solver accepts.
tlist : array_like, shape (nt,)
Save times in ns, strictly increasing.
c_ops : list of Operator or None, default None
Collapse operators including square-root rates in 1/sqrt(ns). None gives no dissipators.
e_ops : list of Operator or None, default None
Observables in output order. None requests no expectation traces.
options : dict or None, default None
Native integration options. See the concrete backend's resolve_solver_options method.
Returns
-------
SolverResult
Saved states, expectation traces and solver diagnostics.
"""
...
[docs]
def batched_sesolve(
self,
problems: list[dict[str, Any]],
*,
n_jobs: int = -1,
progress: bool = True,
) -> list[SolverResult]:
r"""Solve multiple sesolve problems, sequentially by default.
Backends that can vectorize (dynamiqs ``vmap``) or parallelize (QuTiP
with loky) override this method. Each problem dict carries the same
kwargs as :meth:`sesolve`.
Parameters
----------
problems : list of dict
One keyword-argument mapping per sesolve call.
n_jobs : int, default -1
Worker count, where -1 requests all CPUs. The base sequential implementation ignores it.
progress : bool, default True
Request a progress display if the backend supports it.
Returns
-------
list of SolverResult
Results in input order.
"""
return [self.sesolve(**p) for p in problems]
[docs]
def batched_mesolve(
self,
problems: list[dict[str, Any]],
*,
n_jobs: int = -1,
progress: bool = True,
) -> list[SolverResult]:
r"""Solve multiple mesolve problems, sequentially by default.
Parameters
----------
problems : list of dict
One keyword-argument mapping per mesolve call.
n_jobs : int, default -1
Worker count, where -1 requests all CPUs. The base sequential implementation ignores it.
progress : bool, default True
Request a progress display if the backend supports it.
Returns
-------
list of SolverResult
Results in input order.
"""
return [self.mesolve(**p) for p in problems]
# ------------------------------------------------------------------
# IR lowering
# ------------------------------------------------------------------
[docs]
def prepare_hamiltonian(
self,
description: "EngineResult",
tlist: Any,
) -> "PreparedHamiltonian":
r"""Convert a :class:`EngineResult` into a native solver RHS.
The engine supplies frequencies already converted by ``2π`` (angular
rad/ns), so backends must not rescale. Concrete backends override this
method to select their native time-dependence representation.
Parameters
----------
description : EngineResult
Captured engine Hamiltonian and channels in solver units.
tlist : array_like
Requested save times in ns, used by native time-dependent lowering.
Returns
-------
PreparedHamiltonian
Native right-hand side and numerical integration hints.
"""
raise NotImplementedError(f"{type(self).__name__} must implement prepare_hamiltonian()")
[docs]
def resolve_solver_options(
self,
options: dict[str, Any],
*,
metadata: dict[str, Any],
tlist: Any,
) -> dict[str, Any]:
r"""Merge user options with backend-side defaults and metadata heuristics.
The default is a shallow copy. Concrete backends use *metadata* (for
example, ``spectral_bound_ghz``) to fill in integrator step budgets the
user did not specify.
Parameters
----------
options : dict
User integration settings, copied before missing defaults are filled.
metadata : dict
Lowering hints, including ordinary-GHz spectral and carrier bounds when available.
tlist : array_like
Save-time grid in ns, used to estimate integration budgets.
Returns
-------
dict
Resolved options. The input mapping is unchanged.
"""
return dict(options)
def _resolve_automatic_solver_options(
self,
options: dict[str, Any],
*,
user_options: dict[str, Any],
engine_result: "EngineResult",
solver_name: str,
tlist: Any,
) -> dict[str, Any]:
"""Select a backend-native fast path when the assembled problem is eligible.
The protocol default preserves the resolved options. Concrete backends
may use the engine result's static structure and fixed Hilbert-space
dimensions to select a native constant-generator method. Physics values
must not be concretized here.
"""
_ = user_options, engine_result, solver_name, tlist
return dict(options)
[docs]
def coerce_state(self, state: State, dims: tuple[int, ...] | None = None) -> State:
r"""Convert a foreign-native *state* into this backend's native form.
The solve boundary calls this method so that a per-call ``backend=``
override (see :meth:`QuantumSequence.simulate`) accepts initial states
built under another backend. For example, a QuTiP ``Qobj`` from
``chip.state(...)`` given to a dynamiqs gradient solve. Default: pass
through unchanged. ``dims`` are the chip's subsystem levels, for
backends whose state type carries tensor structure.
Parameters
----------
state : State or array_like
Ket or density matrix to convert to this backend.
dims : tuple of int or None, default None
Target tensor dimensions. None lets the backend infer the dimensions.
"""
_ = dims
return state
[docs]
def solve_problem(self, problem: "SolveProblem") -> SolverResult:
r"""Lower and solve a :class:`SolveProblem` (the single-element entry point).
Delegates to :meth:`SolveProblem.solver_name`. It uses ``sesolve`` only
for a ket with no collapse operators and ``mesolve`` otherwise, unless
``problem.solver`` forces a choice.
Parameters
----------
problem : SolveProblem
Complete captured Hamiltonian, initial state, time grid, observables and solver settings.
Returns
-------
SolverResult
Native solver payload for the requested calculation.
"""
if problem.stochastic:
raise NotImplementedError(f"{type(self).__name__} does not implement native solver {problem.solver!r}.")
prepared = self.prepare_hamiltonian(problem.engine_result, problem.tlist)
tlist_arr, c_ops, solver, opts, e_ops_arg = self._resolve_solve_config(
problem, prepared
)
psi0 = self.coerce_state(problem.initial_state, dims=problem.engine_result.dims)
if solver == "sesolve":
return self.sesolve(prepared.rhs, psi0, tlist_arr,
e_ops=e_ops_arg, options=opts)
return self.mesolve(prepared.rhs, psi0, tlist_arr,
c_ops=c_ops, e_ops=e_ops_arg, options=opts)
[docs]
def trajectory_view(
self, native: Any, problem: "SolveProblem", *, index: int | None, states: bool, deterministic: bool,
) -> tuple[Any, Any, Any]:
"""Read one view of a native stochastic result of this backend.
Parameters
----------
native : object
Native stochastic result that :meth:`solve_problem` returned.
problem : SolveProblem
Captured request, with its solver, state storage and native options.
index : int or None
Kept trajectory, or ``None`` for the native ensemble average.
states : bool
Whether to read the saved states.
deterministic : bool
Whether to read the deterministic no-click run instead.
Returns
-------
tuple
Expectation values, state history and final state. Each entry is
``None`` when the native result does not keep it.
"""
raise NotImplementedError(f"{type(self).__name__} does not implement native solver {problem.solver!r}.")
[docs]
def trajectory_records(self, native: Any, problem: "SolveProblem") -> tuple[Any, str]:
"""Read the diffusive measurement currents of a native stochastic result.
Parameters
----------
native : object
Native stochastic result that :meth:`solve_problem` returned.
problem : SolveProblem
Captured request, with its solver and native options.
Returns
-------
tuple
Currents of shape ``(runs, channels, intervals)``, or
``(runs, channels, 2, intervals)`` for heterodyne detection, and
the sampling of the signal in each interval: ``"start"``,
``"middle"``, ``"end"`` or ``"mean"``.
"""
raise NotImplementedError(f"{type(self).__name__} does not implement native solver {problem.solver!r}.")
[docs]
def trajectory_clicks(self, native: Any, problem: "SolveProblem") -> tuple[list[Any], Any, Any, Any]:
"""Read the jump events of a native stochastic result.
Parameters
----------
native : object
Native stochastic result that :meth:`solve_problem` returned.
problem : SolveProblem
Captured request, with its solver and native options.
Returns
-------
tuple
For each run, its click times and the native jump-operator index of
each click; the sampling weight of each run; the weight of the
no-click run that the native sampling solves apart; and whether
each run filled its native click buffer.
"""
raise NotImplementedError(f"{type(self).__name__} does not implement native solver {problem.solver!r}.")
[docs]
def trajectory_run_expect(self, native: Any, problem: "SolveProblem") -> list[Any]:
"""Read the expectation values of every kept run of a native stochastic result.
Parameters
----------
native : object
Native stochastic result that :meth:`solve_problem` returned.
problem : SolveProblem
Captured request, with its native expectation operators.
Returns
-------
list
One array of shape ``(runs, times)`` for each native expectation operator.
"""
raise NotImplementedError(f"{type(self).__name__} does not implement native solver {problem.solver!r}.")
def _stationary_liouvillian(self, engine_result: "EngineResult") -> Any:
"""Lower a static engine description to the backend's native generator."""
raise NotImplementedError(f"{type(self).__name__} must implement stationary lowering")
[docs]
def prepare_stationary(
self, engine_result: "EngineResult", *, prepared: PreparedStationary | None = None,
) -> PreparedStationary:
r"""Prepare once, or reuse the same backend and captured operating point.
Parameters
----------
engine_result : EngineResult
Captured static Hamiltonian and collapse channels.
prepared : PreparedStationary or None, default None
Matching prepared generator. ``None`` builds it.
Returns
-------
PreparedStationary
Native stationary generator tied to this operating point.
"""
if prepared is None:
return PreparedStationary(self, engine_result, self._stationary_liouvillian(engine_result))
if prepared.backend is not self or prepared.engine_result is not engine_result:
raise ValueError("Stationary preparation belongs to a different backend or operating point.")
return prepared
[docs]
def stationary_bytes(self, engine_result: "EngineResult") -> int | None:
"""Return the estimated peak bytes of a stationary solve of a static model, or ``None``.
Parameters
----------
engine_result : EngineResult
Captured static Hamiltonian and collapse channels.
"""
return None
[docs]
def steadystate(
self, problem: Any, *, prepared: PreparedStationary | None = None, guess: State | None = None,
) -> SteadyStateSolverResult:
r"""Solve one static Lindblad problem in the backend's native representation.
Parameters
----------
problem : SteadyStateProblem
Captured static model, observables, and stationary solver options.
prepared : PreparedStationary or None, default None
Matching prepared generator. ``None`` builds it.
guess : State or None, default None
Native stationary state of a related generator, such as the same
model in a neighbouring frame. If this state is stationary here to
round-off, a backend can return it instead of solving. The
diagnostics are then evaluated on this problem's generator.
Returns
-------
SteadyStateSolverResult
Stationary density matrix and convergence diagnostics.
"""
raise NotImplementedError(f"{type(self).__name__} must implement steadystate()")
[docs]
def linear_response(self, problem: Any) -> LinearResponseSolverResult:
r"""Solve one passive-linear input-output problem in the backend's native arrays.
Parameters
----------
problem : LinearResponseProblem
Captured passive-linear model, selected ports, probe frequencies, and incident amplitudes.
Returns
-------
LinearResponseSolverResult
Scattering response, internal amplitudes and numerical diagnostics.
"""
raise NotImplementedError(f"{type(self).__name__} must implement linear_response()")
[docs]
def coherence_response(self, problem: Any) -> CoherenceResponseSolverResult:
"""Solve a frequency response on adjacent excitation-sector coherences.
Parameters
----------
problem : CoherenceResponseProblem
Sector Hamiltonians, channels and frozen-reference probe sources.
"""
raise NotImplementedError(f"{type(self).__name__} must implement coherence_response()")
[docs]
def sector_generator(self, sizes: Sequence[int], hamiltonian: Any, lowering: Sequence[Any],
conserving: Sequence[Any]) -> Any:
r"""Return the truncated Lindblad generator of a conserving model on the totals ``N <= K``.
The generator solves stationary states, resolvents and regression
propagations matrix-free on excitation sectors. This default solves on
host NumPy and SciPy arrays for every backend, so it needs concrete
values.
Parameters
----------
sizes : sequence of int
Dimension of each sector ``N = 0..K``.
hamiltonian : scipy.sparse matrix
Static Hamiltonian in rad/ns, in the lab frame and relative to the
vacuum energy, with states ordered by ``N``.
lowering, conserving : sequence of scipy.sparse matrix
Channel couplings in ``1/sqrt(ns)`` that lower ``N`` by one or
conserve it.
Returns
-------
SectorGenerator
Generator with ``stationary()``, ``resolvent()`` and ``propagate()``.
"""
from quchip.backend._sector_stationary import SectorGenerator
return SectorGenerator(sizes, hamiltonian, lowering, conserving)
[docs]
def sector_solve_bytes(self, sizes: Sequence[int], *, smallest: bool = False) -> int:
"""Return the estimated peak bytes of a stationary solve on excitation sectors.
Parameters
----------
sizes : sequence of int
Dimension of each sector ``N = 0..K``.
smallest : bool, default=False
Estimate at the shortest Krylov restart, which is the least memory
that a solve can run with, instead of the default restart.
"""
from quchip.backend._sector_stationary import _MINIMUM_RESTART, _RESTART, sector_solve_bytes
return sector_solve_bytes(sizes, _MINIMUM_RESTART if smallest else _RESTART)
[docs]
def sector_unique(self, sizes: Sequence[int], lowering: Sequence[Any]) -> bool:
"""Return whether the lowering channels prove one stationary state on the sectors ``N <= K``.
The proof holds whatever the Hamiltonian, a probe and the conserving
channels. Without a probe, that state is the vacuum.
Parameters
----------
sizes : sequence of int
Dimension of each sector ``N = 0..K``.
lowering : sequence of scipy.sparse matrix
Channel couplings in ``1/sqrt(ns)`` that lower ``N`` by one.
"""
from quchip.backend._sector_stationary import has_dark_state
return not has_dark_state(sizes, lowering)
[docs]
def stationary_resolvent(
self,
engine_result: "EngineResult",
sources: tuple[tuple[str, "CanonicalOperator"], ...],
observables: tuple[tuple[str, "CanonicalOperator"], ...],
frequencies: Any,
*,
prepared: PreparedStationary | None = None,
) -> dict[tuple[str, str], Any]:
r"""Evaluate trace-zero stationary resolvents in native backend arrays.
For each ordinary-GHz offset ``f``, form and factor the trace-constrained
shifted Liouvillian once, then solve ``(L + i 2π f) X = source`` with
``Tr(X) = 0`` for every named source column. Return ``Tr(observable X)``
keyed by ``(source, observable)``. The engine supplies canonical operators.
The backend owns Liouvillian construction, factorization, and solves.
Parameters
----------
engine_result : EngineResult
Captured static operating point.
sources : tuple of (str, CanonicalOperator)
Named source columns for the constrained linear systems.
observables : tuple of (str, CanonicalOperator)
Named operators evaluated against each solution.
frequencies : array_like
Offset frequencies in ordinary GHz.
prepared : PreparedStationary or None, default None
Matching prepared generator. ``None`` builds it.
Returns
-------
dict
Arrays keyed by (source_name, observable_name), indexed by frequency.
"""
raise NotImplementedError(f"{type(self).__name__} must implement stationary_resolvent()")
[docs]
def stationary_propagate(
self,
engine_result: "EngineResult",
initial: "CanonicalOperator",
observables: tuple[tuple[str, "CanonicalOperator"], ...],
times: Any,
*,
prepared: PreparedStationary | None = None,
) -> dict[str, Any]:
r"""Propagate an operator under a static Liouvillian and evaluate observables.
``initial`` does not have to be a normalized density matrix. This
supports quantum-regression queries, and the backend keeps ownership of
the propagation algorithm and the numerical Liouvillian.
Parameters
----------
engine_result : EngineResult
Captured static operating point.
initial : CanonicalOperator
Initial operator. Regression queries need neither normalization nor positivity.
observables : tuple of (str, CanonicalOperator)
Named operators evaluated after propagation.
times : array_like
Propagation delays in ns.
prepared : PreparedStationary or None, default None
Matching prepared generator. ``None`` builds it.
Returns
-------
dict
Observable names mapped to arrays indexed by delay.
"""
raise NotImplementedError(f"{type(self).__name__} must implement stationary_propagate()")
[docs]
def parallel_solve_problems(
self,
problems: list["SolveProblem"],
*,
progress: bool = True,
max_workers: int | None = None,
) -> list[SolverResult] | None:
"""Dispatch structurally heterogeneous problems in parallel.
Parameters
----------
problems
Problems that cannot be merged into one structural batch.
progress
Display solver progress if the backend supports it.
max_workers
Worker processes, for backends that solve in parallel. ``None``
lets the backend choose, and ``1`` solves in this process.
Returns
-------
list[SolverResult] | None
Backend results in input order, or ``None`` to keep the engine's
structural-group dispatch.
"""
_ = problems, progress, max_workers
return None
[docs]
def prepare_batch(self, batch: "SolveBatch") -> "PreparedBatch":
r"""Lower each explicit batch problem into a prepared batch.
The return type declares the batching strategy:
:class:`~quchip.backend.containers.EagerBatch` (one RHS per element),
:class:`~quchip.backend.containers.VmappedBatch` (one prebuilt native
RHS), or :class:`~quchip.backend.containers.DeferredBatch`
(backend-private payload for an overridden :meth:`solve_batch`).
Default: lowers each element independently with
:meth:`prepare_hamiltonian` into an :class:`EagerBatch`.
Parameters
----------
batch : SolveBatch
Batch of captured problems with compatible save grids.
Returns
-------
PreparedBatch
Eager, native-vectorized, or deferred representation selected by this backend.
"""
rhs_list: list[Any] = []
shared_metadata: dict[str, Any] = {}
for idx in range(batch.batch_size):
prepared = self.prepare_hamiltonian(batch.element(idx).engine_result, batch.tlist)
rhs_list.append(prepared.rhs)
if not shared_metadata:
shared_metadata = dict(prepared.metadata)
return EagerBatch(
rhs_list=rhs_list,
batch_size=batch.batch_size,
metadata=shared_metadata,
tlist=batch.tlist,
)
[docs]
def solve_batch(
self, batch: "SolveBatch", *, progress: bool = True, max_workers: int | None = None,
) -> list[SolverResult]:
r"""Lower and solve a :class:`SolveBatch`.
Default: prepares the batch, then dispatches each element's RHS through
:meth:`batched_sesolve` / :meth:`batched_mesolve`. Native backends
override this method to avoid the per-element unpack.
Parameters
----------
batch : SolveBatch
Captured problems and batch parameter values.
progress : bool, default True
Request a batch progress display.
max_workers : int or None, default None
Worker processes, for backends that solve in parallel. ``None``
lets the backend choose, and ``1`` solves in this process.
Returns
-------
list of SolverResult
Results in batch order.
"""
if batch.batch_size == 0:
return []
prepared = self.prepare_batch(batch)
if isinstance(prepared, DeferredBatch):
raise RuntimeError(
f"{type(self).__name__}.prepare_batch produced a DeferredBatch; "
"the default solve_batch cannot dispatch it. A backend that "
"defers RHS construction must override solve_batch to consume "
"its own deferred payload."
)
dict_problems: list[dict[str, Any]] = []
solver_names: list[str] = []
for idx in range(batch.batch_size):
problem = batch.element(idx)
rhs = prepared.rhs if isinstance(prepared, VmappedBatch) else prepared.rhs_list[idx]
tlist_arr = self.array_module.asarray(problem.tlist, dtype=float)
c_ops = self._collapse_operators(problem.engine_result)
solver_name = problem.solver_name(self)
solver_names.append(solver_name)
opts = self._resolve_problem_options(
problem,
metadata=problem.engine_result.metadata,
tlist=tlist_arr,
solver_name=solver_name,
)
dict_problems.append(
self._element_solver_kwargs(
solver_name,
rhs,
self.coerce_state(problem.initial_state, dims=problem.engine_result.dims),
tlist_arr,
e_ops=problem.e_ops if isinstance(problem.e_ops, list) else None,
c_ops=c_ops,
options=dict(opts),
)
)
if len(set(solver_names)) != 1:
raise ValueError("Every SolveBatch point must resolve to the same solver.")
solver_name = solver_names[0]
runner = self.batched_mesolve if solver_name == "mesolve" else self.batched_sesolve
return runner(dict_problems, n_jobs=-1 if max_workers is None else max_workers, progress=progress)
# ------------------------------------------------------------------
# Internal helpers
# ------------------------------------------------------------------
def _make_op_cache(self) -> Callable[[Any], Operator]:
"""Return an id-keyed memoizing wrapper over :meth:`from_canonical_operator`.
A fresh cache per call: it is shared across the static and dynamic
conversions of one batch preparation, never across calls (so JAX
tracers cannot leak between traces).
"""
op_cache: dict[int, Operator] = {}
def cached(canonical: Any) -> Operator:
key = id(canonical)
op = op_cache.get(key)
if op is None:
op = self.from_canonical_operator(canonical)
op_cache[key] = op
return op
return cached
@staticmethod
def _sum_terms(terms: Any, to_native: Callable[[Any], Operator]) -> Operator | None:
"""Sum coefficient-scaled terms into one native operator (``None`` if empty)."""
rhs: Operator | None = None
for term in terms:
scaled = term.coefficient * to_native(term.operator)
rhs = scaled if rhs is None else rhs + scaled
return rhs
@staticmethod
def _scalar_dynamic_terms(description: "EngineResult") -> Iterator[tuple[Any, Any]]:
"""Yield ``(operator, signal)`` for each ``ScalarModulation`` dynamic term.
Centralizes the time-dependence filter both backends apply when
lowering dynamic terms. ``DynamicTerm.time_dependence`` is typed as
``ScalarModulation``; the ``isinstance`` guard prevents a non-scalar
modulation from contributing a malformed RHS.
"""
from quchip.engine.ir import ScalarModulation
for term in description.dynamic_terms:
td = term.time_dependence
if isinstance(td, ScalarModulation):
yield term.operator, td.signal
def _merge_options(
self,
user_options: dict[str, Any],
*,
metadata: dict[str, Any],
tlist: Any,
) -> dict[str, Any]:
"""Apply :data:`_DEFAULT_SOLVE_OPTIONS` then backend heuristics.
The single option-merge boundary: defaults are applied here exactly
once, and the backend's :meth:`resolve_solver_options` finalizes the
dict (metadata-derived step budgets, portability key-stripping).
"""
merged = dict(_DEFAULT_SOLVE_OPTIONS)
merged.update(user_options)
return self.resolve_solver_options(merged, metadata=metadata, tlist=tlist)
def _resolve_problem_options(
self,
problem: Any,
*,
metadata: dict[str, Any],
tlist: Any,
solver_name: str,
) -> dict[str, Any]:
"""Merge options and apply one backend-owned automatic-method decision."""
options = {
**problem.options,
"store_states": problem.states == "all",
"store_final_state": problem.states == "final",
}
resolved = self._merge_options(options, metadata=metadata, tlist=tlist)
return self._resolve_automatic_solver_options(
resolved,
user_options=problem.options,
engine_result=problem.engine_result,
solver_name=solver_name,
tlist=tlist,
)
def _resolve_solve_config(
self,
problem: Any,
prepared: Any,
) -> tuple[Any, list[Any], str, dict[str, Any], list[Any] | None]:
"""Resolve one problem's backend-independent solve configuration.
Coerces the save grid (via :attr:`array_module`), assembles collapse
operators, selects the solver (``mesolve`` for a density matrix or when
collapse operators are present, unless ``problem.solver`` forces a
choice), and merges options
through the single boundary. Returns
``(tlist_arr, c_ops, solver_name, opts, e_ops_arg)``; each backend
contributes only its RHS-sourcing + native-solve dispatch tail.
``prepared`` supplies advisory metadata; physics and user options come
from the frozen problem.
"""
engine_result = problem.engine_result
tlist_arr = self.array_module.asarray(problem.tlist, dtype=float)
c_ops = self._collapse_operators(engine_result)
solver_name = problem.solver_name(self)
opts = self._resolve_problem_options(
problem,
metadata=prepared.metadata,
tlist=tlist_arr,
solver_name=solver_name,
)
e_ops_arg = problem.e_ops if isinstance(problem.e_ops, list) else None
return tlist_arr, c_ops, solver_name, opts, e_ops_arg
def _collapse_operators(self, engine_result: "EngineResult") -> list[Operator]:
"""Materialize canonical collapse terms for this backend."""
return [
self.array_module.sqrt(term.rate) * self.from_canonical_operator(term.operator)
for term in engine_result.collapse_terms
]
@staticmethod
def _element_solver_kwargs(
solver_name: str,
rhs: Any,
state: State,
tlist: Any,
*,
e_ops: list[Operator] | None,
c_ops: list[Operator],
options: dict[str, Any],
) -> dict[str, Any]:
"""Assemble one batch element's solver kwargs (pure dict construction).
Selects the state key (``rho0`` for ``mesolve``, ``psi0`` for
``sesolve``) and attaches ``c_ops`` only on the open-system path. No
backend native type is introduced here — *rhs* / *state* / *options*
are passed through unchanged.
"""
kwargs: dict[str, Any] = {
"H": rhs,
"tlist": tlist,
"e_ops": e_ops,
"options": options,
}
if solver_name == "mesolve":
kwargs["rho0"] = state
if c_ops:
kwargs["c_ops"] = list(c_ops)
else:
kwargs["psi0"] = state
return kwargs
