"""CFD-coupled batch chemistry reactor.
:class:`CFDReactor` is the Python entry point for runtime coupling with a
CFD solver (e.g. OpenFOAM Option C in CLAUDE.md). It vectorises the
single-cell stiff-chemistry sub-problem over every CFD cell using
``jax.vmap``, then returns post-reaction concentrations as a NumPy array.
The intended usage from a C++ ``fvOptions``-style plugin is::
reactor = aquakin.CFDReactor(model)
# ... per timestep ...
C_new = reactor.step(
C, # (n_cells, n_species) float64 NumPy
conditions, # {name: (n_cells,)} float64 NumPy
dt, # scalar float (seconds)
params, # (n_params,) NumPy or None
)
The NumPy boundary keeps the pybind11 binding straightforward — the C++
side hands over contiguous ``double[]`` buffers and receives the same.
Columns of ``C`` follow ``model.species`` order; the C++ side is
responsible for assembling that array from its OpenFOAM volScalarFields in
the right order. The reactor's own
:attr:`CFDReactor.species_field_order` attribute exposes this contract.
"""
from __future__ import annotations
from collections.abc import Mapping
import diffrax
import jax
import jax.numpy as jnp
import numpy as np
from aquakin.core.model import CompiledModel
from aquakin.integrate._common import (
DifferentiationConfig,
IntegratorConfig,
init_solver_settings,
resolve_state_atol,
solve_chemistry,
)
[docs]
class CFDReactor:
"""
Vectorised batch reactor for CFD operator-splitting.
Each call to :meth:`step` advances the chemistry from the supplied
cell-state for a transport sub-step ``dt``. Inside, every cell runs an
independent stiff Diffrax integration via ``jax.vmap``.
Parameters
----------
model : CompiledModel
Compiled reaction model.
rtol : float, optional
Relative tolerance for the per-cell 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:`BatchReactor` for the per-species rationale).
integrator : IntegratorConfig, optional
Integrator / step-size configuration (ESDIRK ``order``, ``factormax``,
``dtmax``, ``max_steps``, an explicit ``solver``) for the per-cell solve.
See :class:`BatchReactor`.
diff : DifferentiationConfig, optional
Autodiff configuration (``mode``, ``method``) for the per-cell solve.
Note that the only public method, :meth:`step`, returns NumPy (it
converts at the pybind11 boundary and runs a host-side finiteness check),
so gradients do not flow through ``step`` and this argument is inert for
that path. It is retained for advanced use of the internal jitted
per-cell step in a differentiable context, and for symmetry with the
other reactors.
on_nan : {"raise", "ignore"}, optional
Policy when any cell produces a non-finite concentration (NaN or
Inf) after the chemistry step.
- ``"raise"`` (default) raises ``RuntimeError`` with the offending
cell indices; the C++ caller is then expected to retry with a
smaller transport timestep or otherwise recover.
- ``"ignore"`` returns the array as-is, non-finite cells included,
with **no signal** to the caller. This is the fast path (no
per-step finiteness scan), but a non-finite cell crossing the
pybind11 boundary is then undetectable unless the caller checks.
When using ``"ignore"`` you should pass ``return_mask=True`` to
:meth:`step` (or check finiteness yourself) so corrupted cells
can be detected and recovered.
Attributes
----------
model : CompiledModel
species_field_order : list[str]
Convenience: the order in which species columns of ``C`` must be
supplied. Equal to ``model.species``.
"""
[docs]
def __init__(
self,
model: CompiledModel,
*,
rtol: float = 1e-6,
atol=None,
integrator: IntegratorConfig = IntegratorConfig(),
diff: DifferentiationConfig = DifferentiationConfig(),
on_nan: str = "raise",
) -> None:
if on_nan not in ("raise", "ignore"):
raise ValueError(f"on_nan must be 'raise' or 'ignore', got {on_nan!r}")
init_solver_settings(self, model, rtol=rtol, integrator=integrator, diff=diff)
self.atol = resolve_state_atol(model, atol)
self.on_nan = on_nan
# Cache jit-compiled vmapped step keyed on n_cells.
self._jit_cache: dict[int, callable] = {}
@property
def species_field_order(self) -> list[str]:
"""Order in which species columns of ``C`` must be supplied."""
return list(self.model.species)
@property
def condition_field_names(self) -> list[str]:
"""Names of condition fields expected in ``conditions`` dict."""
return list(self.model.conditions_required)
[docs]
def step(
self,
C: np.ndarray,
conditions: Mapping[str, np.ndarray],
dt: float,
params: np.ndarray | None = None,
*,
return_mask: bool = False,
) -> np.ndarray:
"""
Advance chemistry by ``dt`` for every cell.
Parameters
----------
C : np.ndarray
Cell concentrations, shape ``(n_cells, n_species)``. Columns
follow ``self.species_field_order``.
conditions : mapping str -> np.ndarray
Per-cell condition arrays, each shape ``(n_cells,)``. Must
include every name in ``self.condition_field_names``.
dt : float
Transport sub-step length over which to integrate chemistry.
Must be positive.
params : np.ndarray, optional
Flat parameter vector, shape ``(n_params,)``. Defaults to
``model.default_parameters()``.
return_mask : bool, optional
If ``True``, also return a per-cell boolean mask of finite
results, so the caller can detect non-finite cells even under
``on_nan="ignore"`` (where ``step`` otherwise gives no signal).
Default ``False`` keeps the plain-array return.
Returns
-------
np.ndarray or tuple[np.ndarray, np.ndarray]
Post-reaction concentrations, same shape as input ``C``. If
``return_mask=True``, a ``(C_new, finite_mask)`` tuple where
``finite_mask`` has shape ``(n_cells,)`` and is ``True`` for
cells whose every species value is finite.
"""
C_np = np.ascontiguousarray(np.asarray(C, dtype=np.float64))
if C_np.ndim != 2:
raise ValueError(f"C must be 2-D (n_cells, n_species); got shape {C_np.shape}")
n_cells, n_species_in = C_np.shape
if n_species_in != self.model.n_species:
raise ValueError(
f"C has {n_species_in} species columns but model has "
f"{self.model.n_species} species ({self.model.species})."
)
if n_cells < 1:
raise ValueError(f"C must have at least 1 row; got {n_cells}")
missing = set(self.model.conditions_required) - set(conditions)
if missing:
raise ValueError(
f"conditions is missing required field(s): {sorted(missing)}. "
f"Provided: {sorted(conditions)}"
)
cond_jax: dict[str, jnp.ndarray] = {}
for name in self.model.conditions_required:
arr = np.asarray(conditions[name], dtype=np.float64)
if arr.shape != (n_cells,):
raise ValueError(
f"conditions[{name!r}] has shape {arr.shape}, expected ({n_cells},)."
)
cond_jax[name] = jnp.asarray(arr)
dt_f = float(dt)
if not (dt_f > 0):
raise ValueError(f"dt must be positive; got {dt_f}")
if params is None:
params_jax = self.model.default_parameters()
else:
params_np = np.asarray(params, dtype=np.float64)
if params_np.shape != (self.model.n_params,):
raise ValueError(
f"params has shape {params_np.shape}, expected ({self.model.n_params},)."
)
params_jax = jnp.asarray(params_np)
inner = self._jit_cache.get(n_cells)
if inner is None:
inner = self._build_step()
self._jit_cache[n_cells] = inner
C_new = inner(jnp.asarray(C_np), cond_jax, jnp.asarray(dt_f), params_jax)
C_new_np = np.asarray(C_new)
# The finiteness scan is only needed to raise or to build the mask;
# under on_nan="ignore" without return_mask we skip it (fast path).
if self.on_nan == "raise" or return_mask:
finite_mask = np.all(np.isfinite(C_new_np), axis=1)
if self.on_nan == "raise":
bad_rows = np.where(~finite_mask)[0]
if bad_rows.size:
raise RuntimeError(
f"CFDReactor.step produced non-finite concentrations in "
f"{bad_rows.size} cell(s); offending indices (first 10): "
f"{bad_rows[:10].tolist()}. Consider reducing dt."
)
if return_mask:
return C_new_np, finite_mask
return C_new_np
def _build_step(self):
"""Construct the jit-compiled vmapped per-cell step."""
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
def _per_cell(C_cell, cond_cell, dt, params):
# cond_cell has scalar values (vmap stripped the cells axis);
# wrap each in a length-1 array so ConditionNode can index with
# loc_idx=0.
cond_arrays = {name: v[None] for name, v in cond_cell.items()}
sol = solve_chemistry(
model,
C_cell,
params,
cond_fn=lambda t: cond_arrays,
saveat=diffrax.SaveAt(t1=True),
t0=0.0,
t1=dt,
rtol=rtol,
atol=atol,
adjoint=adjoint,
dtmax=dtmax,
max_steps=max_steps,
order=order,
factormax=factormax,
solver=solver,
)
return sol.ys[-1]
vmapped = jax.vmap(_per_cell, in_axes=(0, 0, None, None))
return jax.jit(vmapped)