aquakin.CFDReactor#

class aquakin.CFDReactor(model, *, rtol=1e-06, atol=None, integrator=IntegratorConfig(order=3, factormax=3.0, colored_jacobian='auto', dtmax=None, max_steps=100000, solver=None), diff=DifferentiationConfig(mode='reverse', method='stable', check_finite=True, adjoint_max_steps=100000, adjoint_low_memory=False), on_nan='raise')[source]#

Bases: object

Vectorised batch reactor for CFD operator-splitting.

Each call to 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 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 BatchReactor.

  • diff (DifferentiationConfig, optional) – Autodiff configuration (mode, method) for the per-cell solve. Note that the only public method, 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 step() (or check finiteness yourself) so corrupted cells can be detected and recovered.

Variables:

  • model (CompiledModel)

  • species_field_order (list[str]) – Convenience: the order in which species columns of C must be supplied. Equal to model.species.

__init__(model, *, rtol=1e-06, atol=None, integrator=IntegratorConfig(order=3, factormax=3.0, colored_jacobian='auto', dtmax=None, max_steps=100000, solver=None), diff=DifferentiationConfig(mode='reverse', method='stable', check_finite=True, adjoint_max_steps=100000, adjoint_low_memory=False), on_nan='raise')[source]#
Parameters:
Return type:

None

Methods

__init__(model, *[, rtol, atol, integrator, ...])

step(C, conditions, dt[, params, return_mask])

Advance chemistry by dt for every cell.

Attributes

condition_field_names

Names of condition fields expected in conditions dict.

species_field_order

Order in which species columns of C must be supplied.
property species_field_order: list[str]#

Order in which species columns of C must be supplied.

property condition_field_names: list[str]#

Names of condition fields expected in conditions dict.

step(C, conditions, dt, params=None, *, return_mask=False)[source]#

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:

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.

Return type:

np.ndarray or tuple[np.ndarray, np.ndarray]