"""Reactor that integrates kinetics along a single time-varying-condition track.
This is the runtime side of the offline OpenFOAM coupling (Option A in
CLAUDE.md). Each Lagrangian particle, as it traverses the CFD domain,
samples the condition fields (pH, T, scavenging, fluence_rate, ...) along
its path. ``ParticleTrackReactor`` integrates the chemistry along that path
by linearly interpolating each condition field over time inside the ODE
term.
"""
from __future__ import annotations
from collections.abc import Mapping
from dataclasses import dataclass, field
import diffrax
import jax
import jax.numpy as jnp
from aquakin.core.model import CompiledModel
from aquakin.integrate._common import (
DifferentiationConfig,
ExportableSolutionMixin,
GradientCheckMixin,
IntegratorConfig,
PlottableSolutionMixin,
_HasNamedSpecies,
_interp_fields_to_scalar,
cached_jitted_solver,
friendly_solve_errors,
init_solver_settings,
reactor_settings_key,
resolve_state_atol,
solve_chemistry,
validate_C0_params,
)
[docs]
@dataclass
class Track:
"""
A single particle track: condition fields sampled at successive times.
Attributes
----------
t : jnp.ndarray
Ascending sample times, shape ``(n_points,)``.
fields : dict[str, jnp.ndarray]
Mapping from condition field name to a 1-D array of length
``n_points`` giving that field along the track.
"""
t: jnp.ndarray
fields: dict[str, jnp.ndarray] = field(default_factory=dict)
def __post_init__(self) -> None:
t = jnp.asarray(self.t)
if t.ndim != 1:
raise ValueError(f"Track.t must be 1-D, got shape {t.shape}")
if t.shape[0] < 2:
raise ValueError("Track must have at least 2 sample points")
if not bool(jnp.all(jnp.diff(t) > 0)):
raise ValueError("Track.t must be strictly ascending")
n = int(t.shape[0])
normalised: dict[str, jnp.ndarray] = {}
for name, value in self.fields.items():
arr = jnp.asarray(value)
if arr.shape != (n,):
raise ValueError(f"Track field '{name}' has shape {arr.shape}, expected ({n},)")
normalised[name] = arr
self.t = t
self.fields = normalised
@property
def n_points(self) -> int:
return int(self.t.shape[0])
[docs]
@dataclass
class TrackSolution(_HasNamedSpecies, PlottableSolutionMixin, ExportableSolutionMixin):
"""Solution returned by :meth:`ParticleTrackReactor.solve`.
``C`` (shape ``(n_t, n_species)``) is the raw integrated state. If the
model sets ``clip_negative_states``, entries may be **small transient
negatives** -- the ``max(C, 0)`` clamp is applied only when evaluating the
reaction rates, not to the saved state. A normal numerical transient, not an
error; clip with ``jnp.maximum(sol.C, 0.0)`` for display if needed.
"""
t: jnp.ndarray
C: jnp.ndarray
model: CompiledModel
[docs]
class ParticleTrackReactor(GradientCheckMixin):
"""
Integrate chemistry along a single particle track.
Parameters
----------
model : CompiledModel
track : Track
Time series of condition values along the particle path. Must supply
every field declared in ``model.conditions_required``.
n_save : int, optional
Number of times at which to record the solution between ``t[0]`` and
``t[-1]``. Defaults to the number of track sample points.
rtol : float, optional
Relative tolerance for the ODE solver.
atol : float or jnp.ndarray, optional
Absolute tolerance, scalar or shape ``(n_species,)``. Defaults to
``None`` -> a per-component noise floor scaled off the model reference
concentrations (see :class:`~aquakin.BatchReactor`).
integrator : IntegratorConfig, optional
Integrator / step-size configuration (ESDIRK ``order``, ``factormax``,
``dtmax``, ``max_steps``, an explicit ``solver``). See
:class:`~aquakin.BatchReactor`. Set ``dtmax`` for reverse-mode
differentiation of a stiff model.
diff : DifferentiationConfig, optional
Autodiff configuration (``mode``, ``method``). See
:class:`~aquakin.BatchReactor`.
"""
[docs]
def __init__(
self,
model: CompiledModel,
track: Track,
*,
n_save: int | None = None,
rtol: float = 1e-6,
atol=None,
integrator: IntegratorConfig = IntegratorConfig(),
diff: DifferentiationConfig = DifferentiationConfig(),
) -> None:
missing = sorted(set(model.conditions_required) - set(track.fields))
if missing:
raise ValueError(
f"Track is missing required condition fields: {missing}. "
f"Provided: {sorted(track.fields)}"
)
init_solver_settings(self, model, rtol=rtol, integrator=integrator, diff=diff)
self.track = track
self.n_save = int(n_save) if n_save is not None else track.n_points
if self.n_save < 2:
raise ValueError(f"n_save must be >= 2, got {self.n_save}")
self.atol = resolve_state_atol(model, atol)
[docs]
def solve(self, C0: jnp.ndarray, *, params: jnp.ndarray | None = None) -> TrackSolution:
"""
Integrate the model along the track.
Parameters
----------
C0 : jnp.ndarray
Inlet concentration vector, shape ``(n_species,)``.
params : jnp.ndarray, optional
Flat parameter vector. Defaults to ``model.default_parameters()``.
Returns
-------
TrackSolution
"""
C0 = jnp.asarray(C0)
params = self.model.default_parameters() if params is None else jnp.asarray(params)
validate_C0_params(self.model, C0, params)
t_grid = jnp.asarray(self.track.t)
t0 = float(t_grid[0])
t1 = float(t_grid[-1])
t_save = jnp.linspace(t0, t1, self.n_save)
# Shared across reactor instances: the track-specific arrays (sample
# times and condition fields) are runtime arguments rather than baked
# into the closure, so an ensemble of same-shape tracks for one model
# reuses a single compiled solver (see cached_jitted_solver /
# integrate_ensemble). JAX's per-shape cache covers tracks that differ in
# length, so the key need only carry the model + settings.
settings = reactor_settings_key(self)
cache_key = None if settings is None else ("particle", id(self.model), settings)
jitted = cached_jitted_solver(cache_key, self._build_jitted_solve, self.model, self.adjoint)
with friendly_solve_errors(self.max_steps, what="particle-track solve"):
ts, ys = jitted(C0, params, t_grid, self.track.fields, t_save)
return TrackSolution(t=ts, C=ys, model=self.model)
def _build_jitted_solve(self):
"""Build a jit-compiled inner solver that takes the track as arguments.
``t_grid`` / ``fields`` / ``t_save`` are passed at call time (not closed
over), so one compiled program serves every track of a given shape -- the
cross-instance reuse the per-instance cache could not give. The
integration bounds are read from the track inside the trace.
"""
model = self.model
rtol = self.rtol
atol = self.atol
adjoint = self.adjoint
dtmax = self.dtmax
max_steps = self.max_steps
order = self.order
factormax = self.factormax
solver = self.solver
@jax.jit
def _solve(C0, params, t_grid, fields, t_save):
sol = solve_chemistry(
model,
C0,
params,
cond_fn=lambda t: _interp_fields_to_scalar(t, t_grid, fields),
saveat=diffrax.SaveAt(ts=t_save),
t0=t_grid[0],
t1=t_grid[-1],
rtol=rtol,
atol=atol,
adjoint=adjoint,
dtmax=dtmax,
max_steps=max_steps,
order=order,
factormax=factormax,
solver=solver,
)
return sol.ts, sol.ys
return _solve
[docs]
def integrate_ensemble(
model: CompiledModel,
tracks: Mapping[int, Track],
C0_fn,
params: jnp.ndarray,
*,
rtol: float = 1e-6,
atol=None,
n_save: int | None = None,
integrator: IntegratorConfig = IntegratorConfig(),
diff: DifferentiationConfig = DifferentiationConfig(),
) -> dict[int, TrackSolution]:
"""
Integrate the model along an ensemble of particle tracks.
Parameters
----------
model : CompiledModel
tracks : mapping int -> Track
``particle_id -> Track``.
C0_fn : callable
Maps ``particle_id`` to its inlet concentration vector. Often
``lambda pid: model.default_concentrations()``.
params : jnp.ndarray
Flat parameter vector shared across all particles.
rtol, atol, n_save, integrator, diff : passed through to each
:class:`ParticleTrackReactor`. ``integrator`` / ``diff`` let an ensemble
of stiff tracks be differentiated the same way a single track can.
Returns
-------
dict[int, TrackSolution]
One solution per particle, keyed by id.
"""
results: dict[int, TrackSolution] = {}
for pid, track in tracks.items():
reactor = ParticleTrackReactor(
model,
track,
n_save=n_save,
rtol=rtol,
atol=atol,
integrator=integrator,
diff=diff,
)
results[pid] = reactor.solve(C0_fn(pid), params=params)
return results