"""Shared internals for the integrator submodules.
Not part of the public API. Reactors depend on this; this module depends only
on JAX and Diffrax.
"""
from __future__ import annotations
import contextlib
import copy
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from typing import TYPE_CHECKING, Protocol, runtime_checkable
import diffrax
import equinox as eqx
import jax
import jax.numpy as jnp
import numpy as np
from aquakin.core.hints import did_you_mean
from aquakin.core.model import CompiledModel
if TYPE_CHECKING: # pragma: no cover
from aquakin.core.conditions import SpatialConditions
# --- Public solver / autodiff configuration objects --------------------------
[docs]
@dataclass(frozen=True)
class IntegratorConfig:
"""The integrator / step-size configuration for a solve.
A single frozen value object centralising the per-step solver settings that
used to be scattered as loose keyword arguments on ``Plant.solve`` and the
reactor constructors. Passed as ``integrator=IntegratorConfig(...)``; the
default flips the public solve stack to the fast configuration (Kvaerno3 +
a capped step-growth factor + an auto colored Jacobian).
``rtol`` / ``atol`` are deliberately **not** here -- they stay separate
arguments of ``Plant.solve`` / the reactor constructors (a tolerance is the
accuracy contract, not the integrator machinery).
Parameters
----------
order : int
ESDIRK order: ``3`` -> ``Kvaerno3`` (the new fast default, fewer stages),
``5`` -> ``Kvaerno5`` (the robust higher-order method). Ignored when
``solver`` is given.
factormax : float, optional
Cap on the PID controller's per-step growth factor (diffrax default 10).
A smaller cap (default 3) damps the overshoot-then-reject oscillation on a
stiff forced system. ``None`` keeps the diffrax default.
colored_jacobian : {"auto", True, False}
Sparse (colored-AD) materialisation of the per-step implicit Jacobian.
See ``Plant.solve`` for the full semantics. Plant-only -- reactors ignore
it.
dtmax : float, optional
Maximum integrator step size (``None`` = uncapped). Set it for a
reverse-mode gradient *through* a stiff solve (the ``through_solve``
differentiation method); the cap-free ``stable`` method does not need it.
max_steps : int
Maximum number of internal solver steps.
solver : diffrax.AbstractSolver, optional
An explicit solver object that overrides ``order`` and is honoured
verbatim (the one escape hatch). ``None`` builds the canonical solver of
``order``.
"""
order: int = 3
factormax: float | None = 3.0
colored_jacobian: bool | str = "auto"
dtmax: float | None = None
max_steps: int = 100_000
solver: diffrax.AbstractSolver | None = None
[docs]
@dataclass(frozen=True)
class DifferentiationConfig:
"""How a gradient / sensitivity flows through a solve.
A single frozen value object replacing the old AD selectors (``adjoint=``,
``gradient=``, ``ad_mode=``, ``check_finite=``, ``stable_adjoint_max_steps=``).
Passed as ``diff=DifferentiationConfig(...)``.
Parameters
----------
mode : {"reverse", "forward"}
Autodiff direction. ``"reverse"`` (default) is the calibration-gradient
direction (``jax.grad`` / a discrete adjoint); ``"forward"`` is the
sensitivity-screen direction (``jax.jacfwd`` / a variational solve).
method : {"stable", "through_solve"}
How the adjoint / sensitivity is *formed* for the chosen ``mode``.
- ``mode="reverse", method="stable"`` (the default) -- a plant uses the
cap-free hand-written discrete adjoint (the old
``gradient="stable_adjoint"`` / ``"auto"`` routing: a concrete forward
solve still takes the fast cached path, a differentiated solve takes the
discrete adjoint). A **reactor** has no stable reverse adjoint, so it
always uses ``RecursiveCheckpointAdjoint`` regardless of ``method``.
- ``mode="reverse", method="through_solve"`` -- differentiate *through*
the diffrax solve (``RecursiveCheckpointAdjoint``; the old
``gradient="jax_adjoint"``). Needs a ``dtmax`` cap for a stiff model.
- ``mode="forward", method="stable"`` -- the augmented ``[y; S]``
variational solve (cap-free, exact). Routed through
``solve_sensitivity`` / ``augmented_forward_sensitivity``.
- ``mode="forward", method="through_solve"`` -- ``jax.jvp`` / ``jacfwd``
via a forward-capable adjoint (``DirectAdjoint`` / ``ForwardMode``).
check_finite : bool
Raise a friendly error if a freshly computed gradient/Jacobian is
non-finite, instead of returning silent ``NaN`` values.
adjoint_max_steps : int
Buffer length for the discrete-adjoint backward scan (the old
``stable_adjoint_max_steps``); set it to a tight upper bound on the
forward step count.
adjoint_low_memory : bool
Low-memory mode for the ``method="stable"`` discrete adjoint. ``False``
(default) saves each forward step's dense-output stage increments, so the
backward reconstructs the stages cheaply -- but the saved buffer is
``~n_stages`` × the trajectory. ``True`` stores only the step states and
**recomputes** each step's stages in the backward pass (a per-step Newton
scan), so the dense buffer is never allocated -- trading compute for
memory when that buffer is the binding constraint on a long, large-state
plant solve. The recompute matches the saved-stage gradient (it is the
same well-conditioned per-stage solve), so it is a pure
memory/compute trade. Reverse ``method="stable"`` only -- ignored for a
``through_solve`` or forward solve.
"""
mode: str = "reverse"
method: str = "stable"
check_finite: bool = True
adjoint_max_steps: int = 100_000
adjoint_low_memory: bool = False
# -- Decoding the (mode, method) vocabulary -----------------------------
#
# The public (mode, method) meaning is owned here rather than re-decoded in
# each entry point (calibrate / sensitivity / dgsm / Plant.solve), so a change
# to the config's semantics -- or a new (mode, method) pairing -- lands in one
# place instead of being mirrored across every consumer.
[docs]
def validated(self) -> DifferentiationConfig:
"""Validate the ``mode`` / ``method`` vocabulary; return ``self``.
Raises ``ValueError`` for an out-of-range ``mode`` or ``method``. This is
the vocabulary check only -- whether a *particular* consumer serves a
given valid pairing (e.g. ``Plant.solve`` rejecting the ``forward`` +
``stable`` augmented solve) is that consumer's own guard.
"""
if self.mode not in ("reverse", "forward"):
raise ValueError(f"diff.mode must be 'reverse' or 'forward'; got {self.mode!r}.")
if self.method not in ("stable", "through_solve"):
raise ValueError(
f"diff.method must be 'stable' or 'through_solve'; got {self.method!r}."
)
return self
[docs]
def gradient_backend(self) -> str:
"""The reverse-adjoint backend this config selects.
``"stable_adjoint"`` for the cap-free hand-written discrete adjoint
(``method="stable"``), or ``"jax_adjoint"`` for differentiating *through*
the diffrax solve (``method="through_solve"``). This is the canonical
``method`` -> backend meaning the calibration entry points consume; the
plant *solve* routes ``method="stable"`` through its own ``"auto"``
dispatch instead (see :meth:`Plant._resolve_diff_config`).
"""
return "stable_adjoint" if self.method == "stable" else "jax_adjoint"
[docs]
def reactor_adjoint(self):
"""The diffrax adjoint object a reactor / plant solve should carry.
Reactors and the plant have no hand-written ("stable") reverse adjoint on
the reactor object -- reverse mode always uses the default
``RecursiveCheckpointAdjoint`` (signalled by ``None``). ``method`` matters
only for *forward* mode: ``through_solve`` -> a forward-capable
``DirectAdjoint`` (:func:`forward_adjoint`); ``stable`` -> the augmented
``solve_sensitivity`` backend, invoked explicitly (not via this adjoint),
so ``None`` is still returned.
"""
if self.mode == "forward" and self.method == "through_solve":
return forward_adjoint()
return None
# --- AD-mode helpers (hide the diffrax adjoint plumbing) ---------------------
[docs]
def forward_adjoint() -> diffrax.AbstractAdjoint:
"""Return the diffrax adjoint that supports forward-mode autodiff.
A thin, dependency-free alias for ``diffrax.DirectAdjoint()`` so a user
script that needs forward-mode AD through a reactor solve (e.g. the reactor
inside a ``dgsm(ad_mode="forward")`` ``fn``) can write
``adjoint=aquakin.forward_adjoint()`` without importing ``diffrax`` or
knowing that the default ``RecursiveCheckpointAdjoint`` registers a
``custom_vjp`` that rejects forward mode.
"""
return diffrax.DirectAdjoint()
def with_adjoint(reactor, adjoint):
"""Return a shallow copy of ``reactor`` with its adjoint strategy replaced.
Reactors are stateless after construction and read ``self.adjoint`` at solve
time, so a shallow copy with a swapped ``adjoint`` is a valid forward-/
reverse-capable variant of the same reactor. Used by ``calibrate`` /
``sensitivity`` to build the right adjoint internally from an ``ad_mode``
string, so ``diffrax`` never appears in user code.
"""
clone = copy.copy(reactor)
clone.adjoint = adjoint
return clone
[docs]
def check_finite_gradient(value, *, what: str, remedy: str) -> None:
"""Raise a friendly ``RuntimeError`` if ``value`` is non-finite.
The silent-NaN footgun of differentiating a stiff solve: the gradient comes
back ``NaN``/``Inf`` and nothing says why. Call this on a freshly computed
gradient/Jacobian to convert that into an actionable error.
Parameters
----------
value : array-like
The gradient or Jacobian to check.
what : str
Short noun for the message (e.g. ``"calibration gradient"``).
remedy : str
The concrete fix to suggest.
"""
if not bool(np.isfinite(np.asarray(value)).all()):
raise RuntimeError(
f"The {what} is non-finite (NaN/Inf). This is almost always the "
f"reverse-mode adjoint of a stiff solve overflowing, not a bug in "
f"your model. {remedy}"
)
class GradientCheckMixin:
"""Give every reactor a finiteness check for a hand-rolled reverse gradient.
A reverse-mode gradient (``jax.grad``/``jax.jacrev``) taken *directly*
through a stiff model's :meth:`solve` -- a user's own loss + optimizer,
outside :func:`aquakin.calibrate` / :func:`aquakin.sensitivity`, which guard
this internally -- can return silent ``NaN``/``Inf`` when the integrator step
is uncapped (the backward accumulation overflows; see the ``dtmax`` note on
the reactor). Nothing raises, so the garbage gradient flows into the
optimizer and the run "works" but never converges.
This mixin adds :meth:`check_gradient_finite`, a one-call guard the DIY user
wraps their gradient in. The remedy it suggests is tailored to whether the
reactor already caps ``dtmax``.
"""
def check_gradient_finite(self, grad_value, *, what: str = "gradient"):
"""Raise a friendly error if a reverse-mode gradient is non-finite.
Wrap a freshly computed ``jax.grad`` result through this reactor's
``solve`` to turn a silent non-finite gradient into an actionable error::
g = reactor.check_gradient_finite(jax.grad(loss)(params))
Parameters
----------
grad_value : array-like
The gradient/Jacobian to check (returned unchanged when finite, so
the call composes inline).
what : str, optional
Short noun for the message (default ``"gradient"``).
Returns
-------
The ``grad_value`` argument, unchanged.
Raises
------
RuntimeError
If ``grad_value`` contains any non-finite entry.
"""
if getattr(self, "dtmax", None) is None:
remedy = (
"A reverse-mode gradient through a stiff solve overflows in the "
"backward pass when the integrator step is uncapped. Build the "
"reactor with a dtmax cap (a small multiple of the fastest "
"reaction timescale), or differentiate in forward mode "
"(jax.jacfwd with adjoint=aquakin.forward_adjoint()). "
"aquakin.calibrate (gradient='stable_adjoint', cap-free) and "
"aquakin.sensitivity handle this for you."
)
else:
remedy = (
f"The reactor already caps dtmax={self.dtmax}; if the gradient is "
"still non-finite, tighten dtmax further, or check the model, "
"data and parameter ranges."
)
check_finite_gradient(grad_value, what=what, remedy=remedy)
return grad_value
# --- Tabular export helpers (optional pandas) --------------------------------
def require_pandas():
"""Import and return pandas, with a helpful message if it is missing.
pandas is an optional dependency, used only by the ``to_dataframe()`` /
``to_csv()`` result exporters.
"""
try:
import pandas as pd
except ImportError as e: # pragma: no cover - exercised only without pandas
raise ImportError(
"to_dataframe() / to_csv() require pandas, an optional dependency. "
"Install it with `pip install pandas` or `pip install "
"aquakin[dataframe]`."
) from e
return pd
def require_matplotlib():
"""Import and return ``matplotlib.pyplot``, with a helpful message if missing.
matplotlib is an optional dependency, used only by the ``plot()`` result
helpers.
"""
try:
import matplotlib.pyplot as plt
except ImportError as e: # pragma: no cover - exercised only without mpl
raise ImportError(
"plot() requires matplotlib, an optional dependency. Install it with "
"`pip install matplotlib` or `pip install aquakin[plot]`."
) from e
return plt
def build_dataframe(
index,
columns,
*,
index_name=None,
units=None,
units_in_columns=False,
extra=None,
):
"""Assemble a pandas ``DataFrame`` from an index and named value columns.
Shared by every result exporter so the column/units conventions stay
consistent.
Parameters
----------
index : array-like or pandas.Index
The row index. A plain array is wrapped in a ``pd.Index`` named
``index_name``; a pre-built ``Index``/``MultiIndex`` is used as-is.
columns : list of (str, array)
``(species_name, 1-D values)`` pairs, in display order.
index_name : str, optional
Name for the index when ``index`` is a plain array.
units : dict, optional
``{species_name: unit_string}``; stored in ``df.attrs["units"]`` and,
when ``units_in_columns`` is set, appended to the column labels.
units_in_columns : bool, optional
Append ``" [unit]"`` to each column label.
extra : list of (str, array), optional
Additional non-species columns to place before the value columns (e.g.
a flow ``Q`` or a ``depth`` column). Never relabelled with units.
Returns
-------
pandas.DataFrame
"""
pd = require_pandas()
units = units or {}
if isinstance(index, (pd.Index, pd.MultiIndex)):
idx = index
else:
idx = pd.Index(np.asarray(index), name=index_name)
data = {}
for name, arr in extra or []:
data[name] = np.asarray(arr)
for name, arr in columns:
unit = units.get(name, "")
label = f"{name} [{unit}]" if (units_in_columns and unit) else name
data[label] = np.asarray(arr)
df = pd.DataFrame(data, index=idx)
df.attrs["units"] = dict(units)
return df
# --- Cross-instance compiled-solver cache ------------------------------------
#
# Each ``reactor.solve(...)`` jit-compiles an inner ``_solve`` closure that
# captures the model and the solver settings. Compilation (trace + lower +
# XLA) dominates the cost of a solve -- the run itself is comparatively free --
# so rebuilding that closure for every reactor instance means every fresh
# reactor pays a full from-scratch compile, even for an identical model and
# settings. That is the dominant cost of the test suite (build a reactor, solve
# once) and of any code that constructs many short-lived reactors.
#
# This module-level cache shares the compiled solver across *all* reactor
# instances keyed by everything the compiled computation depends on: the
# model identity, the solver settings, and the call signature. A repeat solve
# of the same (model, settings, signature) then reuses the compiled graph and
# runs in milliseconds. The cache holds a reference to the model (and any
# settings objects, e.g. a custom adjoint) so their ``id()`` cannot be reused by
# a later object while the entry is live -- keying on ``id`` is therefore safe.
#
# Built-in models are themselves cached by name (``load_model``), so the
# same model object is reused across calls and the ``id()`` key is stable
# across the whole process. The key NEVER omits anything that changes the
# compiled result, so a hit always returns a solver compiled for the exact same
# computation (no false hits); argument shapes/dtypes (C0, params, conditions,
# t_eval) are handled by JAX's own per-function cache and need not be keyed here.
_SOLVER_CACHE: dict = {}
def atol_cache_key(atol):
"""A hashable, value-based key for an (array or scalar) absolute tolerance.
``atol`` is often a per-component array derived from the model, so a fresh
array object each reactor; key on its *values* so identical tolerances share
a cache entry.
"""
a = jnp.asarray(atol)
return (tuple(a.shape), tuple(float(x) for x in np.asarray(a).reshape(-1)))
def _assemble_settings_key(rtol, atol_key, adjoint, dtmax, max_steps, order, factormax, solver):
"""Assemble the solver-settings cache-key tuple from an already-materialised
``atol_key`` (from :func:`atol_cache_key`).
The **single** field-ordering of the settings key, shared by
:func:`settings_cache_key` (the plant path, raw ``atol``) and
:func:`reactor_settings_key` (the reactor path, memoised ``atol`` key). A new
compile-affecting setting is added here **once**, so the two entry points
cannot drift out of lockstep -- a drift would be a silent wrong-result cache
hit. ``adjoint`` is keyed by identity, ``solver`` by class name; ``None`` for
``order`` / ``factormax`` / ``solver`` keeps the historic key shape.
"""
return (
float(rtol),
atol_key,
None if adjoint is None else id(adjoint),
None if dtmax is None else float(dtmax),
int(max_steps),
None if order is None else int(order),
None if factormax is None else float(factormax),
None if solver is None else type(solver).__name__,
)
def settings_cache_key(
rtol, atol, adjoint, dtmax, max_steps, *, order=None, factormax=None, solver=None
):
"""Hashable key for the solver settings that affect the compiled solve.
``adjoint`` is keyed by identity (``None`` for the default, which is shared);
a custom adjoint object keys by ``id`` -- safe (never a false hit), and it
shares whenever the same object is reused. ``order`` / ``factormax`` /
``solver`` (the integrator-config fields that change the compiled reactor
solve) are keyed in when supplied -- ``solver`` by class name, matching the
plant path. ``None`` for any of them keeps the historic key shape.
"""
return _assemble_settings_key(
rtol, atol_cache_key(atol), adjoint, dtmax, max_steps, order, factormax, solver
)
def concrete_settings_key(
rtol, atol, adjoint, dtmax, max_steps, *, order=None, factormax=None, solver=None
):
"""Return :func:`settings_cache_key`, or ``None`` if a value is traced.
The key materialises ``atol`` to Python floats, which is impossible when
``solve`` runs *under tracing* (e.g. a calibration loss differentiating
through the solve). In that case the shared cache gives no benefit anyway --
the solve is being traced into an outer computation that JAX compiles as a
whole -- so return ``None`` to signal "build without caching".
"""
try:
return settings_cache_key(
rtol, atol, adjoint, dtmax, max_steps, order=order, factormax=factormax, solver=solver
)
except jax.errors.TracerArrayConversionError:
return None
def reactor_settings_key(reactor):
"""Return the reactor's concrete solver-settings cache key.
Equivalent to ``concrete_settings_key(reactor.rtol, reactor.atol, ...)`` but
memoises the expensive part. The ``atol`` value-based key materialises the
per-component tolerance array to Python floats -- a device->host sync that
otherwise runs *before* the cache lookup on every solve, including hits, on
exactly the many-short-solve loop the cache exists for. ``atol`` is fixed at
construction, so its key is computed once and stored on the reactor.
The cheap parts (``rtol`` / ``adjoint`` id / ``dtmax`` / ``max_steps``) are
re-read every call rather than cached: :func:`with_adjoint` produces a
*shallow copy* with a swapped ``adjoint`` (the forward-/reverse-mode variant
``calibrate`` / ``sensitivity`` build), and that copy inherits the cached
``_atol_key`` (correctly -- same ``atol`` array) while its adjoint differs,
so the adjoint must not be baked into the cached portion.
Returns ``None`` under tracing (``atol`` is a tracer); not cached in that
case, so a later concrete solve still computes and stores a real key.
"""
atol_key = reactor.__dict__.get("_atol_key")
if atol_key is None:
try:
atol_key = atol_cache_key(reactor.atol)
except jax.errors.TracerArrayConversionError:
return None
reactor._atol_key = atol_key
return _assemble_settings_key(
reactor.rtol,
atol_key,
reactor.adjoint,
reactor.dtmax,
reactor.max_steps,
getattr(reactor, "order", None),
getattr(reactor, "factormax", None),
getattr(reactor, "solver", None),
)
def cached_jitted_solver(key, build, *keep_alive):
"""Return the cached compiled solver for ``key``, building it once.
``build`` is a zero-arg factory returning the jitted ``_solve``; it is called
only on a cache miss. ``keep_alive`` objects (the model, a custom adjoint)
are retained by the cache so their ``id()`` stays valid for the lifetime of
the entry. A ``key`` of ``None`` (an un-cacheable, traced call) bypasses the
cache and just builds.
"""
if key is None:
return build()
entry = _SOLVER_CACHE.get(key)
if entry is None:
fn = build()
_SOLVER_CACHE[key] = (fn, *keep_alive)
return fn
return entry[0]
def validate_t_eval(t_eval_arr: jnp.ndarray, t0: float, t1: float) -> None:
"""Validate output times before handing them to ``SaveAt(ts=...)``.
Diffrax silently returns NaN for save times outside ``[t0, t1]`` or a
non-ascending sequence, so check here for a clear error instead. Value
checks run only for concrete (non-traced) ``t_eval``; a traced array
(e.g. differentiating with respect to the save times) skips them.
Parameters
----------
t_eval_arr : jnp.ndarray
Candidate save times.
t0, t1 : float
Integration interval bounds.
Raises
------
ValueError
If ``t_eval`` is not 1-D, lies outside ``[t0, t1]``, or is not
strictly ascending.
"""
if t_eval_arr.ndim != 1:
raise ValueError(f"t_eval must be 1-D; got shape {tuple(t_eval_arr.shape)}.")
if isinstance(t_eval_arr, jax.core.Tracer):
return
t_eval_np = np.asarray(t_eval_arr)
if t_eval_np.size == 0:
return
lo, hi = float(t_eval_np.min()), float(t_eval_np.max())
if lo < t0 or hi > t1:
raise ValueError(f"t_eval must lie within t_span [{t0}, {t1}]; got values in [{lo}, {hi}].")
if t_eval_np.size > 1 and not np.all(np.diff(t_eval_np) > 0):
raise ValueError("t_eval must be strictly ascending.")
def require_increasing_t_span(t_span) -> tuple[float, float]:
"""Coerce a ``(t_start, t_end)`` pair to floats and require ``t_end > t_start``.
The single home for the interval check every temporal ``solve`` /
``solve_sensitivity`` runs before integrating (a zero- or negative-width
span makes the solver take no step / step backwards).
Parameters
----------
t_span : tuple of float
``(t_start, t_end)`` integration interval.
Returns
-------
tuple of float
``(t0, t1)`` as Python floats.
Raises
------
ValueError
If ``t_end`` does not strictly exceed ``t_start``.
"""
t0, t1 = float(t_span[0]), float(t_span[1])
if not (t1 > t0):
raise ValueError(f"t_span end must exceed start; got ({t0}, {t1}).")
return t0, t1
def prepare_sensitivity(model, C0, params, sens_params, shared_factor):
"""Shared preamble for the flat-state reactors' ``solve_sensitivity``.
``BatchReactor`` and ``PlugFlowReactor`` open their forward-sensitivity solve
with the same four steps: coerce ``C0`` / ``params`` to arrays and validate
their shapes against the model, resolve the ``sens_params`` names/indices to
a flat index array, and auto-select the ``shared_factor`` linear-solve
strategy when the caller left it ``None`` (the CVODES simultaneous corrector
pays off only for more than one sensitivity parameter). (``BiofilmReactor``
coerces its *layered* ``(n_comp, n_species)`` state differently and so keeps
its own preamble.)
Parameters
----------
model : CompiledModel
The reactor's model, for the shape validation and index resolution.
C0, params : jnp.ndarray
Initial state and full parameter vector (coerced here).
sens_params : list of str or int
Namespaced parameter names, or integer indices into ``params``.
shared_factor : bool or None
The caller's ``shared_factor``; ``None`` auto-selects.
Returns
-------
tuple
``(C0, params, free_idx, shared_factor)`` -- the coerced arrays, the
resolved ``(n_sens,)`` index array, and the resolved boolean.
"""
from aquakin.integrate.forward_sensitivity import resolve_sens_indices
C0 = jnp.asarray(C0)
params = jnp.asarray(params)
validate_C0_params(model, C0, params)
free_idx = resolve_sens_indices(model, sens_params)
if shared_factor is None:
shared_factor = free_idx.shape[0] > 1
return C0, params, free_idx, shared_factor
class _HasNamedSpecies:
"""Mixin: the pure species-by-name accessor over a ``.C`` array and ``.model``.
Solution dataclasses inherit from this to share the species-by-name accessor
(``C_named`` / ``C_named_many`` / ``final_named`` / ``final`` / ``units_named``
/ ``time_unit``) plus the independent-axis description (``_table_index`` /
``_independent_axis_label``) that the presentation mixins build on. The two
presentation concerns are **separate mixins** -- :class:`PlottableSolutionMixin`
(matplotlib) and :class:`ExportableSolutionMixin` (pandas / CSV) -- composed
alongside this one, so a change to the plotting API or the CSV schema (and
their optional dependencies) never touches the core numeric accessor every
reactor depends on.
"""
C: jnp.ndarray # set by the dataclass subclass
model: CompiledModel # set by the dataclass subclass
def C_named(self, species: str) -> jnp.ndarray:
"""Return the trajectory of a single species by name."""
if species not in self.model.species_index:
suffix = did_you_mean(species, self.model.species)
raise KeyError(f"Unknown species '{species}'. Available: {self.model.species}.{suffix}")
return self.C[:, self.model.species_index[species]]
def C_named_many(self, species) -> dict[str, jnp.ndarray]:
"""Trajectories of several species by name, as ``{name: array}``.
The multi-species companion to :meth:`C_named` -- read a handful of
species in one call (``sol.C_named_many(["SNH", "SNO", "SO"])``) instead
of a separate slice each. Each value has the same shape ``C_named``
returns. An unknown name raises the same hinted ``KeyError``.
"""
return {sp: self.C_named(sp) for sp in species}
def final_named(self, species=None) -> dict[str, float]:
"""Values at the **last** recorded point, as ``{name: float}``.
The reporting shortcut for a steady-state / end-of-run value: instead of
``float(sol.C_named("SNH")[-1])`` per species, ``sol.final_named(["SNH",
"SNO"])`` returns them in one dict. With ``species=None`` (default) every
model species is returned. Values are plain Python floats (this is a
post-processing read on an already-solved solution -- use
``C_named(sp)[-1]`` if you need a differentiable last value).
"""
names = list(self.model.species) if species is None else list(species)
return {sp: float(self.C_named(sp)[-1]) for sp in names}
@property
def final(self) -> dict[str, float]:
"""Every species' value at the last recorded point (``{name: float}``).
The no-argument attribute form of :meth:`final_named` -- ``sol.final``."""
return self.final_named()
def units_named(self, species: str) -> str:
"""Return the declared units of a species (for axis/column labels).
Convenience for plotting and tabulating results without re-deriving
units by string-matching species names. Equivalent to
``self.model.units_of(species)``.
"""
return self.model.units_of(species)
@property
def time_unit(self) -> str | None:
"""The time unit of ``self.t`` (``"s"``, ``"d"``, ... or ``None``).
Defaults to the model's native unit (:attr:`CompiledModel.time_unit`),
surfaced here so a plot or table can label the time axis unambiguously.
When the solve was called with an explicit ``time_unit=`` the result times
are reported in *that* unit, and this returns it."""
override = getattr(self, "_requested_time_unit", None)
return override if override is not None else self.model.time_unit
def _table_index(self) -> tuple[str, jnp.ndarray]:
"""Return ``(name, array)`` for the dataframe index. Time by default;
space-indexed solutions (PFR) override this."""
return "t", self.t
def _independent_axis_label(self) -> str:
"""Label for the plot's independent axis. Time (with the model's
time unit) by default; a space-indexed PFR overrides this."""
unit = self.time_unit
return f"time [{unit}]" if unit else "time"
class PlottableSolutionMixin:
"""Mixin: matplotlib rendering of a solution's species trajectories.
Composed onto a solution dataclass **alongside** :class:`_HasNamedSpecies`,
whose ``C_named`` / ``units_named`` accessors and ``_table_index`` /
``_independent_axis_label`` it reads. Kept separate from the pure accessor so
a change to the plotting API does not touch the numeric ``C_named`` every
reactor depends on, and so the optional matplotlib dependency stays isolated.
"""
def plot(self, species=None, *, ax=None, **kwargs):
"""Plot one or more species against the independent axis (time, or axial
position for a PFR).
A thin wrapper over matplotlib so "plot SNH over time" needs no manual
``C_named`` / unit / axis-label boilerplate. The x-axis is labelled with
the model's time unit (or position for a PFR), and a single-species
plot labels the y-axis with that species' units.
Parameters
----------
species : str or iterable of str, optional
Species to plot. A single name plots one line; an iterable plots
several with a legend; ``None`` (default) plots every species.
ax : matplotlib.axes.Axes, optional
Axes to draw on. If ``None``, a new figure/axes is created.
**kwargs
Forwarded to ``ax.plot`` (e.g. ``lw``, ``ls``, ``color``).
Returns
-------
matplotlib.axes.Axes
The axes drawn on (for further customisation / saving).
Raises
------
ImportError
If matplotlib is not installed (optional dependency; install with
``pip install aquakin[plot]``).
KeyError
If a species name is unknown (with a "did you mean?" hint).
"""
plt = require_matplotlib()
names = (
[species]
if isinstance(species, str)
else list(self.model.species)
if species is None
else list(species)
)
if ax is None:
_, ax = plt.subplots()
_, x = self._table_index()
x = np.asarray(x)
for sp in names:
ax.plot(x, np.asarray(self.C_named(sp)), label=sp, **kwargs)
ax.set_xlabel(self._independent_axis_label())
if len(names) == 1:
ax.set_ylabel(f"{names[0]} [{self.units_named(names[0])}]")
else:
ax.set_ylabel("concentration")
ax.legend()
return ax
class ExportableSolutionMixin:
"""Mixin: pandas / CSV export of a solution's species trajectories.
Composed onto a solution dataclass **alongside** :class:`_HasNamedSpecies`,
whose ``C`` / ``model`` / ``_table_index`` it reads. Kept separate from the
pure accessor so a change to the CSV schema or the optional pandas dependency
does not touch the core ``C_named`` accessor.
"""
def to_dataframe(self, *, units_in_columns: bool = False):
"""Return the solution as a pandas ``DataFrame``.
One row per recorded point, one column per species (in model
ordering), indexed by the independent axis (time ``t`` for batch /
track / biofilm solutions, axial position ``x`` for a PFR).
Parameters
----------
units_in_columns : bool, optional
If ``True``, append ``" [unit]"`` to each species column label
(e.g. ``"SNH [g_N/m³]"``). If ``False`` (default), columns are bare
species names and the per-species units are stored in
``df.attrs["units"]`` instead, which keeps columns selectable by
species name.
Returns
-------
pandas.DataFrame
Raises
------
ImportError
If pandas is not installed (it is an optional dependency; install
with ``pip install aquakin[dataframe]``).
"""
model = self.model
columns = [(sp, self.C[:, j]) for j, sp in enumerate(model.species)]
units = {sp: model.units_of(sp) for sp in model.species}
name, index = self._table_index()
return build_dataframe(
index,
columns,
index_name=name,
units=units,
units_in_columns=units_in_columns,
)
def to_csv(self, path_or_buf=None, *, units_in_columns: bool = True, **kwargs):
"""Write the solution to CSV (delegates to :meth:`to_dataframe`).
Parameters
----------
path_or_buf : str or path or file-like, optional
Destination passed to ``DataFrame.to_csv``. If ``None``, the CSV is
returned as a string.
units_in_columns : bool, optional
Defaults to ``True`` here (unlike :meth:`to_dataframe`): a CSV
cannot carry ``df.attrs``, so the units are embedded in the column
headers by default so the written file is self-describing.
**kwargs
Forwarded to ``pandas.DataFrame.to_csv``.
"""
return self.to_dataframe(units_in_columns=units_in_columns).to_csv(path_or_buf, **kwargs)
@runtime_checkable
class Reactor(Protocol):
"""Structural type for the solve-based reactors that ``sensitivity`` /
``fit`` / ``calibrate`` / ``profile_likelihood`` consume.
Declares the contract those consumers actually rely on: the compiled
``model``, the five solver settings every reactor exposes (set uniformly
by :func:`init_solver_settings` + :func:`resolve_state_atol`), and a
``solve(C0, params=None, ...)`` whose extra arguments are reactor-specific
(a batch/biofilm reactor takes ``t_span`` / ``t_eval``; a PFR fixes its grid
at construction; a particle reactor takes neither). ``CFDReactor`` is
deliberately **not** a ``Reactor`` -- it exposes ``step()``, not ``solve()``.
A reactor that also carries spatially-varying ``conditions`` (batch / PFR /
biofilm) satisfies the narrower :class:`ConditionedReactor`; a particle
reactor carries a ``track`` instead and does not.
"""
model: CompiledModel
rtol: float
atol: float | jnp.ndarray
adjoint: diffrax.AbstractAdjoint | None
dtmax: float | None
max_steps: int
order: int
factormax: float | None
solver: diffrax.AbstractSolver | None
def solve(self, C0, params=None, *args, **kwargs): # pragma: no cover
...
@runtime_checkable
class ConditionedReactor(Reactor, Protocol):
"""A :class:`Reactor` that exposes spatially-varying ``conditions``.
The batch / PFR / biofilm reactors carry a :class:`SpatialConditions`, so a
consumer that differentiates through condition fields (e.g.
:func:`aquakin.sensitivity`) requires this narrower type; a particle reactor,
which carries a ``track`` instead, is a :class:`Reactor` but not a
``ConditionedReactor``.
"""
conditions: SpatialConditions
def _coerce_atol(atol, n_species: int):
"""Validate and normalise an ``atol`` argument.
Returns either the original scalar or a ``(n_species,)`` JAX array.
Raises ``ValueError`` if an array of the wrong shape is supplied.
"""
arr = jnp.asarray(atol)
if arr.ndim == 0:
# A concrete scalar (the reactor-construction path) is returned as a
# Python float, but a traced value is returned as the 0-d array: calling
# ``float()`` on a tracer raises a concretization error, which would
# otherwise prevent jitting a solve whose ``atol`` flows in under tracing.
if isinstance(arr, jax.core.Tracer):
return arr
return float(arr)
if arr.shape != (n_species,):
raise ValueError(f"atol array must have shape ({n_species},), got {arr.shape}")
return arr
def is_forward_mode_ad_error(exc: BaseException) -> bool:
"""True if ``exc`` is JAX's rejection of forward-mode autodiff through a
``custom_vjp``.
The default ``RecursiveCheckpointAdjoint`` registers a ``custom_vjp``, so a
``jax.jvp`` / ``jax.jacfwd`` through the solve fails with *"can't apply
forward-mode autodiff (jvp) to a custom_vjp function"*. This is the one
specific failure the forward-mode sensitivity / DGSM paths convert into
``forward_adjoint()`` guidance; matching it narrowly keeps unrelated errors
(bad shapes, OOM, a bug in the user's ``fn``) from being masked by that hint.
"""
msg = str(exc).lower()
return "forward-mode autodiff" in msg and "custom_vjp" in msg
@contextlib.contextmanager
def friendly_solve_errors(max_steps, *, what: str = "solve"):
"""Re-raise two opaque solve-time failures as domain-level errors.
Both failures surface from JAX/Diffrax/Equinox internals with messages a
process engineer cannot map to a fix. Wrap the *execution* of a solve in
this context manager -- the call to the jitted solve / ``diffeqsolve``, where
each failure actually surfaces -- to re-raise it as a plain ``RuntimeError``
naming the remedy, with the noisy traceback suppressed (``from None``). Any
other exception propagates unchanged.
1. **Step-budget exhaustion.** An adaptive solve that exhausts ``max_steps``
raises a verbose ``EquinoxRuntimeError`` ("The maximum number of solver
steps was reached") wrapped in JAX/Equinox debugging chatter
(``EQX_ON_ERROR``, ``kidger.site``, ...). Re-raised naming the
warm-start / loosen-rtol / raise-max_steps remedies.
2. **Forward-mode AD through the reverse-only adjoint.** The default
``RecursiveCheckpointAdjoint`` registers a ``custom_vjp``, so a
``jax.jacfwd`` / ``jax.jvp`` through the solve fails with JAX's
"can't apply forward-mode autodiff (jvp) to a custom_vjp function" --
which never mentions the cure. Re-raised naming
``aquakin.forward_adjoint()``.
Parameters
----------
max_steps : int
The step budget that was hit (quoted back in the step-budget message).
what : str
A short label for the failing solve (e.g. ``"plant solve"``), used in
the message.
"""
try:
yield
except Exception as exc:
msg = str(exc).lower()
if "maximum number of solver steps" in msg:
raise RuntimeError(
f"The {what} hit its integrator step budget (max_steps={max_steps}) "
"before completing. This is almost always a stiff transient, not a "
"bug. Try, in order: (1) warm-start from a settled state -- for a "
"plant, pass y0 from plant.run_to_steady_state (or a previous run); "
"(2) loosen rtol (the default atol already auto-scales to the state "
"magnitudes); or (3) raise max_steps. If none help, the model may be "
"genuinely unstable at these parameters/inputs."
) from None
if is_forward_mode_ad_error(exc):
raise RuntimeError(
f"Forward-mode autodiff (jax.jacfwd / jax.jvp) cannot flow through "
f"the {what}: the default adjoint (RecursiveCheckpointAdjoint) is "
"reverse-mode only -- it registers a custom_vjp that rejects forward "
"mode. Build the reactor with adjoint=aquakin.forward_adjoint() to "
"use a forward-mode-capable solve, or take a reverse-mode gradient "
"(jax.grad / jax.jacrev) instead. (aquakin.sensitivity and "
"aquakin.dgsm accept ad_mode='forward' and set this adjoint for you.)"
) from None
raise
def default_atol(
scale_like, reference=None, *, atol_factor: float = 1e-6, floor_frac: float = 1e-6
):
"""Per-component absolute tolerance scaled off the states' operating magnitudes.
The error test every adaptive solver uses weights each component by
``atol_i + rtol*|y_i|``; ``atol_i`` is the **noise floor** below which
component ``i`` is treated as negligible. When components span very different
scales (e.g. ADM1 from ~1e-13 to ~17, or an OH radical at ~1e-12 beside a
bulk reactant at ~1e-4) a single scalar floor is wrong for most of them, so
the floor is set **per component** -- the SUNDIALS "vector atol" guidance and
Hairer & Wanner's rule of ``atol_i`` proportional to the typical magnitude of
component ``i``.
Returns ``atol_i = atol_factor * max(|scale_i|, |reference_i|, floor_frac*char)``,
where ``char = max_j(typical_j)`` is the system's bulk magnitude (so a
component whose typical value is ~0, e.g. a product not present initially,
gets a small floor tied to the system scale rather than ``atol_i = 0``, which
the solver literature explicitly warns against).
Parameters
----------
scale_like : array
A representative state vector -- typically the initial condition ``C0`` /
``y0`` (the operating magnitudes).
reference : array, optional
A second magnitude source merged in via elementwise max -- typically the
model's ``default_concentrations`` (the YAML reference values), so a
component that starts at 0 but has a nonzero reference is still floored
sensibly.
atol_factor : float
Fraction of each component's typical magnitude used as its noise floor.
floor_frac : float
Floor for near-zero components, as a fraction of the bulk scale ``char``.
Returns
-------
jnp.ndarray
Per-component absolute tolerance, same shape as ``scale_like``.
"""
typ = jnp.abs(jnp.asarray(scale_like, dtype=float))
if reference is not None:
typ = jnp.maximum(typ, jnp.abs(jnp.asarray(reference, dtype=float)))
char = jnp.max(typ)
# When every magnitude is zero there is no scale to floor against, so the
# near-zero floor floor_frac*char would itself be 0 -- returning atol_i = 0,
# which the solver literature warns against (the invariant this floor exists
# to uphold). Fall back to unit scale in that case. Identity for any input
# with a nonzero magnitude (the common path), since char > 0 there.
char = jnp.where(char > 0.0, char, 1.0)
typ = jnp.maximum(typ, floor_frac * char)
# An absolute tolerance is a solver noise floor, never a differentiated
# quantity, so detach it from autodiff. This matters when ``scale_like`` is the
# initial state ``y0`` and the gradient is taken *with respect to* ``y0``: a
# traced atol would otherwise be baked into the integrator's step controller
# and -- inside the discrete-adjoint (``gradient="stable_adjoint"``) custom-VJP
# forward, which re-runs ``diffrax.diffeqsolve`` -- escape that inner solve as
# a leaked tracer (issue #420). ``stop_gradient`` collapses the ``y0``-derived
# tracer to its concrete primal; it is the identity for a concrete
# ``scale_like`` (the common path), so every existing solve is unchanged.
return jax.lax.stop_gradient(atol_factor * typ)
# Seconds per time unit, for converting a user-supplied ``time_unit`` to a
# model's native (rate-constant) time unit. The keys match
# ``CompiledModel.time_unit`` / ``utils.units._TIME_TOKENS``.
_TIME_UNIT_SECONDS = {"s": 1.0, "min": 60.0, "h": 3600.0, "d": 86400.0}
def native_time_factor(model_time_unit, requested_unit) -> float:
"""Factor ``f`` such that ``t_native = f * t_requested``.
Converts a time expressed in ``requested_unit`` into the model's native
(rate-constant) time unit, so a caller can pass ``t_span`` / ``t_eval`` in a
convenient unit (e.g. hours) and have the solve run with the model's
unchanged rate constants. Returns ``1.0`` when ``requested_unit is None``
(no conversion -- the default, native-unit path).
Parameters
----------
model_time_unit : str or None
The model's native time unit (``CompiledModel.time_unit``).
requested_unit : str or None
The unit the caller's times are in. ``None`` means "native".
Raises
------
ValueError
If ``requested_unit`` is unknown, or it is given but the model's own
time unit could not be inferred (``None``) and so there is no native
unit to convert to.
"""
if requested_unit is None:
return 1.0
if requested_unit not in _TIME_UNIT_SECONDS:
raise ValueError(
f"Unknown time_unit {requested_unit!r}; expected one of {sorted(_TIME_UNIT_SECONDS)}."
)
if model_time_unit is None:
raise ValueError(
f"Cannot convert times to time_unit={requested_unit!r}: this "
"model's own time unit could not be inferred from its "
"rate-constant units (model.time_unit is None), so there is no "
"native unit to convert to. Pass t_span / t_eval already in the "
"model's rate-constant time unit and omit time_unit."
)
# model_time_unit comes from CompiledModel.time_unit, which only ever
# returns a known token, so it is guaranteed to be in the table.
return _TIME_UNIT_SECONDS[requested_unit] / _TIME_UNIT_SECONDS[model_time_unit]
def to_native_time(model_time_unit, requested_unit, t_span, t_eval):
"""Scale ``(t_span, t_eval)`` from ``requested_unit`` into native time.
Returns ``(t_span_native, t_eval_native, factor)`` where ``factor`` is the
:func:`native_time_factor` (``t_native = factor * t_requested``); divide a
native-unit result time by it to report back in ``requested_unit``. A
``factor`` of ``1.0`` (the ``requested_unit is None`` default, or a unit
equal to the native one) leaves the inputs untouched.
"""
factor = native_time_factor(model_time_unit, requested_unit)
if factor == 1.0:
return t_span, t_eval, factor
if t_span is not None:
t_span = (t_span[0] * factor, t_span[1] * factor)
if t_eval is not None:
t_eval = jnp.asarray(t_eval) * factor
return t_span, t_eval, factor
def resolve_state_atol(model, atol):
"""Resolve the ``atol`` for a reactor whose state is one ``(n_species,)``
concentration vector.
``atol=None`` -> the per-component :func:`default_atol` noise floor scaled off
the model's reference concentrations (so a g/m³ ASM model and a mol/L
ozone model each get sensible tolerances without hand-tuning, instead of a
fixed scalar that is ~9 orders too tight for g/m³ states). An explicit scalar
or ``(n_species,)`` array is validated and returned verbatim. Shared by every
reactor with a single-concentration-vector state (Batch / PFR / Particle /
CFD); the layered :class:`~aquakin.BiofilmReactor`, whose state spans several
compartments, uses :func:`resolve_layered_atol` (the same per-component floor,
tiled across compartments) instead.
"""
return (
default_atol(model.default_concentrations())
if atol is None
else _coerce_atol(atol, model.n_species)
)
def resolve_layered_atol(model, atol, n_compartments):
"""Resolve the ``atol`` for a reactor whose state stacks ``n_compartments``
copies of the model's ``(n_species,)`` concentration vector (the layered
:class:`~aquakin.BiofilmReactor`: a bulk plus per-depth-layer states).
``atol=None`` -> the per-component :func:`default_atol` noise floor tiled
across every compartment, so each (layer, species) slot gets the same
per-species floor a single-vector reactor would use -- not a fixed scalar
that is ~9 orders too tight for g/m3 ASM/ADM states. A scalar broadcasts to
the whole state; a ``(n_species,)`` array is tiled across the compartments; a
``(n_compartments, n_species)`` or flat ``(n_compartments*n_species,)`` array
is accepted verbatim. Returns the ``(n_compartments, n_species)`` array
matching the reactor's 2-D compartment state (the sensitivity path flattens
it for the augmented solve).
"""
n = model.n_species
nc = int(n_compartments)
if atol is None:
per_species = default_atol(model.default_concentrations()) # (n,)
return jnp.tile(per_species, (nc, 1))
arr = jnp.asarray(atol, dtype=float)
if arr.ndim == 0:
return jnp.full((nc, n), float(arr))
if arr.shape == (n,):
return jnp.tile(arr, (nc, 1))
if arr.shape == (nc, n):
return arr
if arr.shape == (nc * n,):
return arr.reshape(nc, n)
raise ValueError(
f"atol must be a scalar, a ({n},) per-species array, or a "
f"({nc}, {n}) / ({nc * n},) full-state array; got shape {tuple(arr.shape)}."
)
def init_solver_settings(reactor, model, *, rtol, integrator, diff):
"""Store the solver settings every reactor shares, on ``reactor``.
Resolves the public :class:`IntegratorConfig` + :class:`DifferentiationConfig`
into the primitive attributes the cache + ``_run_diffeqsolve`` path read:
``model``, ``rtol``, ``adjoint`` (the diffrax adjoint object, resolved from
the differentiation config), ``dtmax``, ``max_steps``, ``order``,
``factormax`` and ``solver``. ``atol`` is **not** set here because its
resolution depends on the reactor's state shape (see
:func:`resolve_state_atol` for the single-vector case and
:func:`resolve_layered_atol` for the layered biofilm); the caller sets
``reactor.atol`` itself.
"""
reactor.model = model
reactor.rtol = float(rtol)
reactor.integrator = integrator
reactor.diff = diff
reactor.adjoint = diff.reactor_adjoint()
reactor.dtmax = integrator.dtmax
reactor.max_steps = int(integrator.max_steps)
reactor.order = int(integrator.order)
reactor.factormax = integrator.factormax
reactor.solver = integrator.solver
def validate_C0_params(model, C0, params):
"""Raise ``ValueError`` if ``C0`` / ``params`` do not match the model.
The shared shape check every single-vector-state reactor runs at the top of
``solve`` (``C0`` is ``(n_species,)``, ``params`` is ``(n_params,)``).
"""
if C0.shape != (model.n_species,):
raise ValueError(f"C0 has shape {C0.shape}, expected ({model.n_species},)")
if params.shape != (model.n_params,):
raise ValueError(f"params has shape {params.shape}, expected ({model.n_params},)")
# The per-stage Newton (root-find) tolerance is decoupled from the step tolerance
# by this factor: diffrax's stock Kvaerno root finder copies the step tolerances,
# driving each stage solve to full-step accuracy (more Newton iterations than the
# embedded error estimate needs). The step controller still enforces the solution
# accuracy at (rtol, atol); loosening only the per-stage Newton ends each stage
# solve sooner (~15-20% faster on the stiff BSM2 plant, final-state agreement
# ~1e-6).
NEWTON_TOL_FACTOR = 10.0
# The fixed set of ESDIRK orders any implicit solve may use, keyed by order. This
# is the ONLY place aquakin constructs a diffrax ESDIRK solver *object*: every
# solve path gets its solver from build_implicit_solver, which builds from this
# table, so the solver mode cannot drift between the reactor and plant paths
# (Kvaerno5 is the robust default; Kvaerno3 does less linear algebra per step).
# Enforced by tests/unit/test_solver_config_single_source.py.
_CANONICAL_SOLVERS = {3: diffrax.Kvaerno3, 5: diffrax.Kvaerno5}
def build_implicit_solver(
rtol,
atol,
*,
order=5,
solver=None,
colored_root_finder=None,
force_root_finder=False,
linear_solver=None,
):
"""The single source of truth for the implicit ODE solver, shared by EVERY
aquakin solve so the forward and discrete-adjoint paths cannot silently drift.
Both :func:`_run_diffeqsolve` (the forward path) and the discrete-adjoint
forward pass (:func:`~aquakin.integrate.discrete_adjoint.esdirk_adjoint_solve`)
build their solver here, so a per-step optimization -- the decoupled Newton
tolerance, a colored root finder, a different ESDIRK order -- lands in one
place and automatically reaches both modes.
Parameters
----------
rtol, atol : float
Step-size tolerances; the per-stage Newton tolerance is these loosened by
:data:`NEWTON_TOL_FACTOR` unless a ``colored_root_finder`` is given.
order : int
ESDIRK order to build when ``solver`` is ``None`` -- a key into
:data:`_CANONICAL_SOLVERS` (``5`` -> ``Kvaerno5``, the robust default;
``3`` -> ``Kvaerno3``, less linear algebra per step). This is how an
internal caller selects the solver order without constructing the diffrax
object itself (so construction stays funneled here).
solver : diffrax.AbstractSolver, optional
A user-supplied ESDIRK solver -- the one escape hatch (e.g.
``Plant.solve(solver=...)``). Overrides ``order``. ``None`` builds the
canonical solver of ``order``.
colored_root_finder : diffrax root finder, optional
A sparsity-colored root finder (``ColoredVeryChord``) that materializes the
per-step implicit Jacobian by one JVP per color instead of a dense
``jacfwd``. When given it is used verbatim (its own tolerances stand).
force_root_finder : bool
When ``True`` (the discrete-adjoint forward, which pairs the solver with a
step-clipping controller that needs explicit root-finder tolerances), the
decoupled root finder is injected even into a *supplied* ``solver``. When
``False`` (the forward path) a supplied ``solver`` is honoured verbatim --
its own root finder is left untouched -- and only the *default* solver gets
the decoupled root finder.
linear_solver : lineax.AbstractLinearSolver, optional
A custom linear solver for the decoupled root finder's per-stage solve.
The forward-sensitivity augmented ``[y; S]`` solve passes the block-arrow
:class:`~aquakin.integrate._simultaneous_corrector.SimultaneousCorrector`
here so the decoupled Newton uses the factor-once-reuse linear algebra.
Ignored when ``colored_root_finder`` is given (that root finder carries
its own solver).
"""
if colored_root_finder is not None:
rf = colored_root_finder
else:
rf_kw = {} if linear_solver is None else {"linear_solver": linear_solver}
rf = diffrax.VeryChord(
rtol=NEWTON_TOL_FACTOR * rtol, atol=NEWTON_TOL_FACTOR * jnp.asarray(atol), **rf_kw
)
if solver is None:
if order not in _CANONICAL_SOLVERS:
raise ValueError(f"order must be one of {sorted(_CANONICAL_SOLVERS)}; got {order!r}.")
return _CANONICAL_SOLVERS[order](root_finder=rf)
if force_root_finder or colored_root_finder is not None:
return eqx.tree_at(lambda s: s.root_finder, solver, rf)
return solver # forward path: a user-supplied solver is honoured verbatim
def build_step_controller(rtol, atol, *, factormax=None, dtmax=None):
"""The single source of truth for the PID step-size controller core, shared by
the forward solve (which uses it directly) and the discrete adjoint (which
wraps it in a ``ClipStepSizeController``) -- so ``factormax`` / the tolerances
cannot drift between the two paths."""
kw = {} if factormax is None else {"factormax": factormax}
return diffrax.PIDController(rtol=rtol, atol=atol, dtmax=dtmax, **kw)
def _run_diffeqsolve(
rhs: Callable,
*,
t0: float,
t1: float,
y0: jnp.ndarray,
args,
saveat: diffrax.SaveAt,
rtol: float,
atol,
adjoint: diffrax.AbstractAdjoint | None = None,
max_steps: int = 100_000,
dtmax: float | None = None,
event: diffrax.Event | None = None,
progress_meter: diffrax.AbstractProgressMeter | None = None,
solver: diffrax.AbstractSolver | None = None,
factormax: float | None = None,
order: int = 5,
):
"""Wrapper around the canonical Kvaerno + PIDController + adjoint setup.
All reactors call this with their own ``rhs``. Adjusting the default
solver, controller, or adjoint here changes behaviour for every reactor.
``solver`` overrides the default ``Kvaerno5`` integrator with any diffrax
solver (``None`` keeps ``Kvaerno5``). A lower-order ESDIRK such as
``Kvaerno3`` does less implicit linear algebra per step (fewer stages), which
can be faster on a large stiff system whose per-step cost is dominated by the
Jacobian factorisation, at the cost of taking more (cheaper) steps. The
*default* ``Kvaerno5`` is built with a **decoupled root finder** -- its
per-stage Newton tolerance is 10x looser than the step tolerance (see the
inline note), which makes each step cheaper at preserved accuracy; a
user-supplied ``solver`` is used as-is.
``factormax`` caps the per-step growth factor of the ``PIDController`` (its
own default is 10). A smaller cap (e.g. 3) damps the overshoot-then-reject
oscillation that drives the step-rejection rate up on a stiff, forced system.
``None`` keeps the diffrax default.
``dtmax`` caps the integrator step size. It is ``None`` (uncapped) by
default, which is fastest for plain forward solves. For *reverse-mode*
differentiation of a stiff model it must be set. An L-stable solver may
take steps far larger than the fastest reaction timescale and simply damp
the unresolved fast modes in the primal (which stays accurate). The two AD
modes then diverge: **forward mode** (``jax.jvp`` / ``jax.jacfwd``) stays
finite at any step, losing only accuracy when the fast modes are
unresolved; **reverse mode** (``jax.grad``, the discrete adjoint) returns
**non-finite** values above a step-size threshold, an overflow in the
backward accumulation governed by the per-step stiffness ``gamma*dt*||J||``
(not by operator conditioning). Capping ``dtmax`` to a small multiple of
the fastest reaction timescale bounds that product; the resulting reverse
gradient is finite and matches both forward mode and finite differences.
This is reverse-mode-specific and independent of the adjoint flavour. See
the "Differentiating stiff models" discussion in CLAUDE.md.
"""
term = diffrax.ODETerm(rhs)
# Build the solver + controller from the shared single-source-of-truth helpers
# (the decoupled-Newton default solver and the PID controller) so this forward
# path and the discrete-adjoint forward pass cannot drift in their per-step
# configuration. A user-supplied ``solver`` is honoured verbatim here.
solver = build_implicit_solver(rtol, atol, order=order, solver=solver)
controller = build_step_controller(rtol, atol, factormax=factormax, dtmax=dtmax)
return diffrax.diffeqsolve(
term,
solver,
t0=t0,
t1=t1,
dt0=None,
y0=y0,
args=args,
saveat=saveat,
stepsize_controller=controller,
adjoint=adjoint if adjoint is not None else diffrax.RecursiveCheckpointAdjoint(),
max_steps=max_steps,
event=event,
progress_meter=(
progress_meter if progress_meter is not None else diffrax.NoProgressMeter()
),
)
def make_chemistry_rhs(
model: CompiledModel,
params: jnp.ndarray,
*,
cond_fn: Callable[[jnp.ndarray], Mapping[str, jnp.ndarray]],
rate_scale=None,
) -> Callable:
"""Build the canonical reaction right-hand side ``rhs(t, C, args)``.
Hoists the (parameter-dependent) stoichiometry once, then returns
``dC/dt = rate_scale * dCdt(C, params, cond_fn(t))``. This is the **single
source** for the reaction RHS: both the plain solve (:func:`solve_chemistry`)
and the event-driven segmented solve build their RHS through here, so the two
paths cannot drift (a change to how the RHS is assembled reaches both). The
threaded ``args`` carries the live ``params`` for the rate constants while the
stoichiometry is the hoisted constant.
Parameters
----------
model : CompiledModel
params : jnp.ndarray
Parameter vector; the stoichiometry is hoisted from it once.
cond_fn : callable
``cond_fn(t)`` returning the condition arrays at independent value ``t``
(a constant dict for batch/CFD; an interpolation for PFR/particle).
rate_scale : optional
``None`` (identity) or a scalar multiplying the rate (e.g. ``1/velocity``
for the steady-state PFR).
"""
stoich = model.compute_stoich(params)
if rate_scale is None:
def rhs(t, C, args):
return model.dCdt(C, args, cond_fn(t), 0, stoich=stoich)
else:
def rhs(t, C, args):
return model.dCdt(C, args, cond_fn(t), 0, stoich=stoich) * rate_scale
return rhs
def solve_chemistry(
model: CompiledModel,
C0: jnp.ndarray,
params: jnp.ndarray,
*,
cond_fn: Callable[[jnp.ndarray], Mapping[str, jnp.ndarray]],
saveat: diffrax.SaveAt,
t0,
t1,
rtol: float,
atol,
adjoint: diffrax.AbstractAdjoint | None = None,
dtmax: float | None = None,
max_steps: int = 100_000,
rate_scale=None,
order: int = 5,
factormax: float | None = None,
solver: diffrax.AbstractSolver | None = None,
):
"""The canonical chemistry sub-solve shared by every reactor.
Hoists the (parameter-dependent) stoichiometry out of the per-step RHS ---
so dynamic coefficients are evaluated once per solve, not per step --- builds
the right-hand side ``dC/dt = rate_scale * dCdt(C, params, cond_fn(t))`` and
runs the Kvaerno5 + ``PIDController`` solve via :func:`_run_diffeqsolve`.
The reactors differ only in three traced-time choices, passed in here:
- ``cond_fn(t)`` returns the condition arrays at independent-variable value
``t``. A batch / CFD cell passes a constant dict; a PFR or particle track
passes an interpolation of its spatially / temporally varying fields.
- ``rate_scale`` (``None`` = identity) multiplies the rate, e.g. ``1/velocity``
for the steady-state PFR whose independent variable is axial position.
- ``saveat`` / ``t0`` / ``t1`` select the output points and the span.
Returns the diffrax ``Solution``; callers read ``sol.ts`` / ``sol.ys`` (or
``sol.ys[-1]`` for a single-endpoint step).
"""
rhs = make_chemistry_rhs(model, params, cond_fn=cond_fn, rate_scale=rate_scale)
return _run_diffeqsolve(
rhs,
t0=t0,
t1=t1,
y0=C0,
args=params,
saveat=saveat,
rtol=rtol,
atol=atol,
adjoint=adjoint,
dtmax=dtmax,
max_steps=max_steps,
order=order,
factormax=factormax,
solver=solver,
)
def _interp_fields_to_scalar(
t: jnp.ndarray,
t_grid: jnp.ndarray,
fields: Mapping[str, jnp.ndarray],
) -> dict[str, jnp.ndarray]:
"""Linearly interpolate every field to a single-location array at ``t``.
Returns a fresh dict whose values are shape-``(1,)`` arrays. Reactors
pair this with ``loc_idx=0`` to satisfy the canonical rate-callable
signature.
"""
return {name: jnp.asarray([jnp.interp(t, t_grid, arr)]) for name, arr in fields.items()}