Coverage for jaxquantum/core/solvers.py: 93%
182 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-27 22:28 +0000
1import logging
2import warnings
3from collections.abc import Callable
4from typing import Any, Literal
6import diffrax
7import jax.numpy as jnp
8import tqdm
9from flax import struct
10from jax import Array
11from jax.experimental import sparse as _sparse
13from jaxquantum.core.dims import Qdims
14from jaxquantum.core.operators import identity_like, multi_mode_basis_set
15from jaxquantum.core.qarray import DenseImpl, Qarray, Qtypes, dag_data
16from jaxquantum.utils.utils import robust_isscalar
19def _is_dense_array(x) -> bool:
20 """True for a dense JAX array (not BCOO, not SparseDIA data)."""
21 return not isinstance(x, _sparse.BCOO) and not getattr(x, "_is_sparse_dia", False)
24def _default_stepsize_controller() -> diffrax.AbstractStepSizeController:
25 return diffrax.PIDController(rtol=1e-7, atol=1e-9)
28@struct.dataclass
29class SolverOptions:
30 """Options forwarded to :func:`diffrax.diffeqsolve`.
32 Attributes:
33 solver: Native Diffrax solver; strings are deprecated.
34 stepsize_controller: Native Diffrax controller; strings are deprecated.
35 stepsize_controller_kwargs: Deprecated controller constructor arguments.
36 saveat: Custom save policy; overrides JAXQuantum's save-time handling.
37 dt0: Initial step, ``"tlist"`` for the first interval, or ``None`` for
38 Diffrax's automatic choice.
39 adjoint: Differentiation strategy. ``None`` uses Diffrax's default.
40 event: Native Diffrax termination event.
41 max_steps: Maximum solver steps.
42 throw: Whether unsuccessful solves raise an exception.
43 progress_meter: ``None``, ``"default"``, or a native progress meter.
44 Booleans are deprecated.
45 solver_state: Solver state used to continue a previous solve.
46 controller_state: Controller state used to continue a previous solve.
47 made_jump: Previous jump state used when continuing a solve.
49 ``saveat=None`` saves at ``saveat_tlist`` (or ``tlist`` when omitted).
50 ``adjoint=None`` and ``progress_meter=None`` preserve Diffrax's defaults.
51 Native Diffrax objects pass through unchanged. Legacy values still work and
52 issue a ``FutureWarning``.
53 """
55 progress_meter: bool | Literal["default"] | diffrax.AbstractProgressMeter | None = (
56 struct.field(pytree_node=False, default="default")
57 )
58 solver: diffrax.AbstractSolver | str = struct.field(
59 pytree_node=False, default_factory=diffrax.Tsit5
60 )
61 max_steps: int | None = struct.field(pytree_node=False, default=100_000)
62 stepsize_controller: diffrax.AbstractStepSizeController | str = struct.field(
63 pytree_node=False, default_factory=_default_stepsize_controller
64 )
65 stepsize_controller_kwargs: dict[str, Any] | None = struct.field(
66 pytree_node=False, default=None
67 )
68 saveat: diffrax.SaveAt | None = struct.field(pytree_node=False, default=None)
69 dt0: float | Array | None | Literal["tlist"] = struct.field(
70 pytree_node=False, default="tlist"
71 )
72 adjoint: diffrax.AbstractAdjoint | None = struct.field(
73 pytree_node=False, default=None
74 )
75 event: diffrax.Event | None = struct.field(pytree_node=False, default=None)
76 throw: bool = struct.field(pytree_node=False, default=True)
77 solver_state: Any = None
78 controller_state: Any = None
79 made_jump: bool | Array | None = None
81 @classmethod
82 def create(
83 cls,
84 progress_meter: bool = True,
85 solver: str = "Tsit5",
86 max_steps: int = 100_000,
87 stepsize_controller: str = "PIDController",
88 stepsize_controller_kwargs: dict[str, Any] | None = None,
89 ) -> "SolverOptions":
90 """Create options with the deprecated string-based interface."""
91 warnings.warn(
92 "SolverOptions.create() is deprecated; use SolverOptions with native "
93 "objects, such as solver=diffrax.Tsit5() and "
94 "progress_meter='default'.",
95 FutureWarning,
96 stacklevel=2,
97 )
98 return cls(
99 solver=_diffrax_object(solver, diffrax.AbstractSolver),
100 stepsize_controller=_diffrax_object(
101 stepsize_controller,
102 diffrax.AbstractStepSizeController,
103 _legacy_controller_kwargs(
104 stepsize_controller, stepsize_controller_kwargs
105 ),
106 ),
107 max_steps=max_steps,
108 progress_meter="default" if progress_meter else None,
109 )
112class CustomProgressMeter(diffrax.TqdmProgressMeter):
113 """JAXQuantum's default Diffrax progress bar."""
115 @staticmethod
116 def _init_bar() -> tqdm.tqdm:
117 bar_format = (
118 "{desc}: {percentage:3.0f}% |{bar}| "
119 "[{elapsed}<{remaining}, {rate_fmt}{postfix}]"
120 )
121 return tqdm.tqdm(
122 total=100, bar_format=bar_format, unit="%", colour="MAGENTA", ascii="░▒█"
123 )
126def _resolve_saveat(
127 options: SolverOptions,
128 tlist: Array,
129 saveat_tlist: Array | None,
130) -> diffrax.SaveAt:
131 if options.saveat is not None:
132 if saveat_tlist is not None:
133 raise ValueError(
134 "Pass save times through saveat_tlist or SolverOptions.saveat, not both."
135 )
136 return options.saveat
138 times = tlist if saveat_tlist is None else jnp.atleast_1d(saveat_tlist)
139 return diffrax.SaveAt(t1=True) if len(times) == 0 else diffrax.SaveAt(ts=times)
142def _resolve_progress_meter(
143 progress_meter: (bool | Literal["default"] | diffrax.AbstractProgressMeter | None),
144) -> diffrax.AbstractProgressMeter | None:
145 if progress_meter is None:
146 return None
147 if isinstance(progress_meter, bool):
148 return CustomProgressMeter() if progress_meter else None
149 if isinstance(progress_meter, str):
150 if progress_meter == "default":
151 return CustomProgressMeter()
152 raise ValueError(
153 "progress_meter must be None, 'default', or a Diffrax progress meter."
154 )
155 if not isinstance(progress_meter, diffrax.AbstractProgressMeter):
156 raise TypeError(
157 "progress_meter must be a Diffrax AbstractProgressMeter instance."
158 )
159 return progress_meter
162def _diffrax_object(
163 name: str,
164 expected_type: type,
165 kwargs: dict[str, Any] | None = None,
166):
167 try:
168 value = getattr(diffrax, name)(**(kwargs or {}))
169 except AttributeError as error:
170 raise ValueError(f"Unknown Diffrax type: {name!r}.") from error
171 if not isinstance(value, expected_type):
172 raise TypeError(f"diffrax.{name} is not a {expected_type.__name__}.")
173 return value
176def _legacy_controller_kwargs(
177 name: str, kwargs: dict[str, Any] | None
178) -> dict[str, Any]:
179 if kwargs is not None:
180 return kwargs
181 return {"rtol": 1e-7, "atol": 1e-9} if name == "PIDController" else {}
184def _uses_legacy_options(options: SolverOptions) -> bool:
185 return (
186 isinstance(options.solver, str)
187 or isinstance(options.stepsize_controller, str)
188 or options.stepsize_controller_kwargs is not None
189 or isinstance(options.progress_meter, bool)
190 )
193def _resolve_solver(solver: diffrax.AbstractSolver | str) -> diffrax.AbstractSolver:
194 if isinstance(solver, str):
195 return _diffrax_object(solver, diffrax.AbstractSolver)
196 if not isinstance(solver, diffrax.AbstractSolver):
197 raise TypeError("solver must be a Diffrax AbstractSolver instance.")
198 return solver
201def _resolve_stepsize_controller(
202 options: SolverOptions,
203) -> diffrax.AbstractStepSizeController:
204 controller = options.stepsize_controller
205 if isinstance(controller, str):
206 return _diffrax_object(
207 controller,
208 diffrax.AbstractStepSizeController,
209 _legacy_controller_kwargs(controller, options.stepsize_controller_kwargs),
210 )
211 if options.stepsize_controller_kwargs is not None:
212 raise ValueError(
213 "Pass controller arguments when constructing stepsize_controller."
214 )
215 if not isinstance(controller, diffrax.AbstractStepSizeController):
216 raise TypeError(
217 "stepsize_controller must be a Diffrax AbstractStepSizeController instance."
218 )
219 return controller
222def _resolve_dt0(options: SolverOptions, tlist: Array) -> float | Array | None:
223 if isinstance(options.dt0, str):
224 if options.dt0 != "tlist":
225 raise ValueError("dt0 must be a number, None, or 'tlist'.")
226 return tlist[1] - tlist[0]
227 return options.dt0
230def solve(
231 f: Callable,
232 y0: Array,
233 tlist: Array,
234 saveat_tlist: Array | None = None,
235 args: Any = None,
236 solver_options: SolverOptions | None = None,
237) -> diffrax.Solution:
238 """Solve an ODE using native Diffrax configuration from ``SolverOptions``."""
239 options = SolverOptions() if solver_options is None else solver_options
240 if _uses_legacy_options(options):
241 warnings.warn(
242 "String and boolean SolverOptions values are deprecated; use native "
243 "objects such as solver=diffrax.Tsit5(), "
244 "stepsize_controller=diffrax.PIDController(...), and "
245 "progress_meter='default' or None.",
246 FutureWarning,
247 stacklevel=2,
248 )
249 kwargs = {
250 "saveat": _resolve_saveat(options, tlist, saveat_tlist),
251 "stepsize_controller": _resolve_stepsize_controller(options),
252 "args": args,
253 "max_steps": options.max_steps,
254 "throw": options.throw,
255 }
256 optional = {
257 "adjoint": options.adjoint,
258 "event": options.event,
259 "progress_meter": _resolve_progress_meter(options.progress_meter),
260 "solver_state": options.solver_state,
261 "controller_state": options.controller_state,
262 "made_jump": options.made_jump,
263 }
264 kwargs.update(
265 (name, value) for name, value in optional.items() if value is not None
266 )
268 with warnings.catch_warnings():
269 warnings.filterwarnings(
270 "ignore",
271 message="Complex dtype support in Diffrax",
272 category=UserWarning,
273 )
274 return diffrax.diffeqsolve(
275 diffrax.ODETerm(f),
276 _resolve_solver(options.solver),
277 t0=tlist[0],
278 t1=tlist[-1],
279 dt0=_resolve_dt0(options, tlist),
280 y0=y0,
281 **kwargs,
282 )
285def mesolve(
286 H: Qarray | Callable[[float], Qarray],
287 rho0: Qarray,
288 tlist: Array,
289 saveat_tlist: Array | None = None,
290 c_ops: Qarray | None = None,
291 solver_options: SolverOptions | None = None,
292) -> Qarray:
293 """Solve a Lindblad master equation and return the saved states.
295 Args:
296 H: Static Hamiltonian or callable ``H(t)``.
297 rho0: Initial ket or density matrix.
298 tlist: Integration interval; also the default save times.
299 saveat_tlist: Save times. An empty array saves only the final state.
300 c_ops: Collapse operators.
301 solver_options: Native Diffrax configuration.
303 Returns:
304 Saved density matrices as a batched ``Qarray``.
306 See Also:
307 :func:`mesolve_result` returns the complete Diffrax solution.
308 """
309 solution = mesolve_result(
310 H,
311 rho0,
312 tlist,
313 saveat_tlist=saveat_tlist,
314 c_ops=c_ops,
315 solver_options=solver_options,
316 )
317 qdims = Qdims((rho0.space_dims, rho0.space_dims))
318 return Qarray._from_impl(DenseImpl._make(solution.ys), qdims)
321def mesolve_result(
322 H: Qarray | Callable[[float], Qarray],
323 rho0: Qarray,
324 tlist: Array,
325 saveat_tlist: Array | None = None,
326 c_ops: Qarray | None = None,
327 solver_options: SolverOptions | None = None,
328) -> diffrax.Solution:
329 """Solve a Lindblad master equation and return its Diffrax solution.
331 Use this form for solver statistics, events, dense interpolation, custom
332 ``SaveAt`` functions, or continuation state.
333 """
334 collapse_ops = c_ops if c_ops is not None else Qarray.from_list([])
336 if len(collapse_ops) == 0 and rho0.qtype != Qtypes.oper:
337 logging.warning( # noqa: LOG015
338 "Consider sesolve(): no collapse operators were provided and the "
339 "initial state is not a density matrix."
340 )
342 rho_data = rho0.to_dm().to_dense()
344 if robust_isscalar(H):
345 H = H * identity_like(rho_data)
347 if isinstance(H, Qarray):
348 H_data = lambda t: H.data
349 else:
350 H_data = lambda t: H(t).data
352 return _mesolve_result_data(
353 H_data,
354 rho_data.data,
355 tlist,
356 saveat_tlist,
357 collapse_ops.data,
358 solver_options=solver_options,
359 )
362def _mesolve_result_data(
363 H: Callable[[float], Array],
364 rho0: Array,
365 tlist: Array,
366 saveat_tlist: Array | None,
367 c_ops: Array | None = None,
368 solver_options: SolverOptions | None = None,
369) -> diffrax.Solution:
370 """Array-level master-equation implementation."""
372 c_ops = c_ops if c_ops is not None else jnp.array([])
374 # Shape inference: when c_ops contains batched operators (e.g. shape
375 # (1, B, N, N)), the initial state ρ0 must be broadcast to (B, N, N) so
376 # that the ODE RHS produces consistently shaped output.
377 #
378 # The output batch shape is the broadcast of:
379 # c_ops[0] batch dims → c_ops.shape[1:-2] (outer batch index stripped)
380 # H batch dims → H(tlist[0]).shape[:-2]
381 # ρ0 batch dims → ρ0.shape[:-2]
382 # This is a pure shape calculation — no array values are materialised.
383 H0_shape = H(tlist[0]).shape
384 if len(c_ops) == 0:
385 batch_shape = jnp.broadcast_shapes(H0_shape[:-2], rho0.shape[:-2])
386 else:
387 # c_ops.shape[1:-2]: strip the outermost (c_op index) dim and the two
388 # matrix dims to get the batch dims that will be broadcast into ρ.
389 batch_shape = jnp.broadcast_shapes(
390 c_ops.shape[1:-2], H0_shape[:-2], rho0.shape[:-2]
391 )
392 rho = jnp.broadcast_to(rho0, batch_shape + rho0.shape[-2:])
394 # Precompute the adjoint once, outside the ODE hot-loop.
395 # dag_data dispatches to the correct impl (dense or sparse) automatically,
396 # so c_ops_dag is BCOO when c_ops is sparse and a dense array otherwise.
397 c_ops_dag = dag_data(c_ops) if len(c_ops) != 0 else c_ops
399 def f(
400 t: float,
401 rho: Array,
402 args,
403 ):
404 c_ops_val, c_ops_dag_val = args
405 H_val = H(t) # type: ignore
407 rho_dot = -1j * (H_val @ rho - rho @ H_val)
409 if len(c_ops_val) == 0:
410 return rho_dot
412 # Compute the Lindblad dissipator D[L](ρ) = L ρ L† - ½(L†L ρ + ρ L†L)
413 # using only (sparse L) @ (dense rho) operations to support BCOO
414 # collapse operators natively — no dense @ sparse required:
415 #
416 # L ρ L† = dag( L @ dag(L @ ρ) ) avoids the dense @ L† step
417 # L†L ρ = L† @ (L @ ρ) BCOO @ dense → dense ✓
418 # ρ L†L = dag(L†L ρ) dag of dense ✓ (ρ Hermitian)
419 Lrho = c_ops_val @ rho
420 LrhoLdag = dag_data(c_ops_val @ dag_data(Lrho))
421 LdagLrho = c_ops_dag_val @ Lrho
422 rhoLdagL = dag_data(LdagLrho)
424 rho_dot_delta = 0.5 * (2 * LrhoLdag - LdagLrho - rhoLdagL)
426 rho_dot_delta = jnp.sum(rho_dot_delta, axis=0)
428 rho_dot += rho_dot_delta
430 return rho_dot
432 return solve(
433 f,
434 rho,
435 tlist,
436 saveat_tlist,
437 (c_ops, c_ops_dag),
438 solver_options=solver_options,
439 )
442def sesolve(
443 H: Qarray | Callable[[float], Qarray],
444 rho0: Qarray,
445 tlist: Array,
446 saveat_tlist: Array | None = None,
447 solver_options: SolverOptions | None = None,
448) -> Qarray:
449 """Solve a Schrödinger equation and return the saved states.
451 Args:
452 H: Static Hamiltonian or callable ``H(t)``.
453 rho0: Initial ket.
454 tlist: Integration interval; also the default save times.
455 saveat_tlist: Save times. An empty array saves only the final state.
456 solver_options: Native Diffrax configuration.
458 Returns:
459 Saved kets as a batched ``Qarray``.
461 See Also:
462 :func:`sesolve_result` returns the complete Diffrax solution.
463 """
464 solution = sesolve_result(
465 H,
466 rho0,
467 tlist,
468 saveat_tlist=saveat_tlist,
469 solver_options=solver_options,
470 )
471 return Qarray._from_impl(DenseImpl._make(solution.ys), rho0.qdims)
474def sesolve_result(
475 H: Qarray | Callable[[float], Qarray],
476 rho0: Qarray,
477 tlist: Array,
478 saveat_tlist: Array | None = None,
479 solver_options: SolverOptions | None = None,
480) -> diffrax.Solution:
481 """Solve a Schrödinger equation and return its Diffrax solution.
483 Use this form for solver statistics, events, dense interpolation, custom
484 ``SaveAt`` functions, or continuation state.
485 """
486 if rho0.qtype == Qtypes.oper:
487 raise ValueError("Use mesolve() for an initial density matrix.")
489 state = rho0.to_ket().to_dense()
491 if robust_isscalar(H):
492 H = H * identity_like(state)
494 if isinstance(H, Qarray):
495 H_data = lambda t: H.data
496 else:
497 H_data = lambda t: H(t).data
499 return _sesolve_result_data(
500 H_data,
501 state.data,
502 tlist,
503 saveat_tlist,
504 solver_options=solver_options,
505 )
508def _sesolve_result_data(
509 H: Callable[[float], Array],
510 rho0: Array,
511 tlist: Array,
512 saveat_tlist: Array | None,
513 solver_options: SolverOptions | None = None,
514) -> diffrax.Solution:
515 """Array-level Schrödinger-equation implementation."""
517 def f(t: float, ψₜ: Array, _):
518 H_val = H(t) # type: ignore
520 # State vectors live on a single trailing axis (..., N). For a dense
521 # Hamiltonian contract that axis directly via einsum (batch-safe, and no
522 # (N,1) is ever materialised — keeps the scan carry 1-D). For a sparse
523 # Hamiltonian use a transient column local to this RHS (sparse is not the
524 # TPU-padding path).
525 if _is_dense_array(H_val):
526 ψₜ_dot = -1j * jnp.einsum("...ij,...j->...i", H_val, ψₜ)
527 else:
528 ψₜ_dot = -1j * (H_val @ ψₜ[..., None])[..., 0]
530 return ψₜ_dot
532 batch_shape = jnp.broadcast_shapes(H(tlist[0]).shape[:-2], rho0.shape[:-1])
533 state = jnp.broadcast_to(rho0, batch_shape + rho0.shape[-1:])
535 return solve(
536 f,
537 state,
538 tlist,
539 saveat_tlist,
540 solver_options=solver_options,
541 )
544# propagators
547def propagator(
548 H: Qarray | Callable[[float], Qarray],
549 ts: float | Array,
550 saveat_tlist: Array | None = None,
551 solver_options: SolverOptions | None = None,
552) -> Qarray:
553 """Generate a propagator for a Hamiltonian.
555 Args:
556 H (Qarray or callable):
557 A Qarray static Hamiltonian OR
558 a function that takes a time argument and returns a Hamiltonian.
559 ts (float or Array):
560 A single time point or
561 an Array of time points.
562 saveat_tlist: Times at which to save the propagator.
563 solver_options: Native Diffrax configuration for time-dependent input.
565 Returns:
566 The propagator at each saved time.
567 """
568 ts_is_scalar = robust_isscalar(ts)
569 H_is_qarray = isinstance(H, Qarray)
571 if H_is_qarray:
572 return (-1j * H * ts).expm()
573 else:
574 if ts_is_scalar:
575 H_first = H(0.0)
576 if ts == 0:
577 return identity_like(H_first)
578 ts = jnp.array([0.0, ts])
579 else:
580 H_first = H(ts[0])
582 basis_states = multi_mode_basis_set(H_first.space_dims)
583 results = sesolve(
584 H,
585 basis_states,
586 ts,
587 saveat_tlist=saveat_tlist,
588 solver_options=solver_options,
589 )
590 # results.data is (T, M, M): T times, M evolved basis kets (batch), M
591 # ket components. Transpose the last two axes so each time slice is a
592 # propagator whose columns are the evolved basis states. No squeeze:
593 # kets no longer carry a trailing singleton.
594 propagators_data = results.data.mT
595 return Qarray.create(propagators_data, dims=H_first.space_dims)