aquakin.BatchReactor#

class aquakin.BatchReactor(model, conditions, *, 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))[source]#

Bases: GradientCheckMixin

Stateless 0-D (batch) reactor.

Parameters:

  • model (CompiledModel) – Compiled reaction model.

  • conditions (SpatialConditions) – Condition fields. For a batch reactor only the single location loc_idx=0 is used.

  • rtol (float, optional) – Relative tolerance for the ODE solver.

  • atol (float or jnp.ndarray, optional) – Absolute tolerance for the ODE solver. If an array, it must have shape (n_species,) and gives a per-species tolerance. Useful when some species (e.g. radical intermediates) sit several orders of magnitude below the bulk concentrations.

  • integrator (IntegratorConfig, optional) – Integrator / step-size configuration (ESDIRK order, factormax, dtmax, max_steps, an explicit solver). The default flips the reactor to the fast stack (Kvaerno3 + capped step growth). dtmax is always in the model’s native time unit (the unit of its rate constants – seconds for ozone/UV, days for ASM/ADM/WATS), independent of any time_unit= passed to solve().

  • diff (DifferentiationConfig, optional) – Autodiff configuration. mode="reverse" (default) differentiates the solve with RecursiveCheckpointAdjoint (reactors have no stable reverse adjoint – that is plant-only – so method is ignored for reverse mode). mode="forward", method="through_solve" builds a forward-capable DirectAdjoint so jax.jvp / jax.jacfwd flow through the solve; mode="forward", method="stable" is the augmented solve_sensitivity path (cap-free).

Notes

Differentiating a stiff solve. A reverse-mode gradient (jax.grad / jax.jacrev) taken directly through solve on a stiff model (ASM / ADM / WATS) returns silent NaN / Inf when dtmax is uncapped – the backward accumulation overflows (see dtmax above). No exception is raised, so the non-finite gradient flows into an optimizer as garbage and the fit never converges with no indication why. aquakin.calibrate() and aquakin.sensitivity() guard this (check_finite=True), but a hand-rolled loss + optimizer through solve is exposed. The remedies, in order of convenience:

  • build the reactor with a dtmax cap (the simplest fix);

  • differentiate in forward mode (jax.jacfwd with adjoint=aquakin.forward_adjoint()), finite at any step;

  • use aquakin.calibrate() (gradient="stable_adjoint", cap-free) or aquakin.sensitivity(), which handle it internally;

  • guard your own gradient with check_gradient_finite() (g = reactor.check_gradient_finite(jax.grad(loss)(params))), which raises an actionable error instead of returning silent NaN.

__init__(model, conditions, *, 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))[source]#
Parameters:
Return type:

None

Methods

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

check_gradient_finite(grad_value, *[, what])

Raise a friendly error if a reverse-mode gradient is non-finite.

solve(C0, t_span[, t_eval, params, ...])

Integrate the reaction model over a time span.

solve_sensitivity(C0, t_span[, t_eval, ...])

Solve and return the forward sensitivity dC/dtheta alongside C.
solve(C0, t_span, t_eval=None, *, params=None, conditions=None, time_unit=None, events=None)[source]#

Integrate the reaction model over a time span.

Parameters:

  • C0 (jnp.ndarray) – Initial concentration vector, shape (n_species,).

  • t_span (tuple of float) – (t_start, t_end) integration interval, in the model’s time unit unless time_unit is given. The required second positional argument (solve(C0, (0.0, 600.0))).

  • t_eval (jnp.ndarray, optional) – Time points at which to record solution. If None the solver returns endpoints only.

  • params (jnp.ndarray, optional, keyword-only) – Rate constant vector, shape (n_params,). Defaults to model.default_parameters() – pass it (as a keyword) only to override rate constants (e.g. a what-if run; see model.parameter_values). Keyword-only so a positional t_span tuple can never land in it.

  • conditions (SpatialConditions, optional) – Override the conditions stored on the reactor for this call. Used by aquakin.sensitivity() to differentiate through condition fields without mutating the reactor.

  • time_unit (str, optional) – The time unit t_span / t_eval are expressed in ("s", "min", "h", "d"). aquakin has no global time unit – the native unit is set by the model’s rate constants (model.time_unit: seconds for ozone/UV, days for ASM/ADM/WATS). Pass time_unit to work in a different unit: the input times are converted into the native unit for the solve (rate constants unchanged) and the returned solution.t is reported back in time_unit (with solution.time_unit set to it). Default None uses the model’s native unit. Raises if the model’s own time unit is undeclared (model.time_unit is None), since there is then no native unit to convert to.

  • events (sequence of Event, optional) – Located discontinuities (on/off switches, SBR phases, level limits) applied during the solve. Each Event fires at a known time (at_times=) or when a state cond_fn crosses zero, and may reset the state (apply=) or terminate the solve (terminal=). The solve is split into segments at the firings; the returned solution.events_log records them. Time-only events keep jax.grad finite (static segment boundaries); a state event makes the solve a forward simulation (the firing count is data-dependent). See aquakin.solve_with_events().

Return type:

BatchSolution

solve_sensitivity(C0, t_span, t_eval=None, *, params, sens_params, conditions=None, sens_rtol=None, sens_atol=None, param_scale=None, shared_factor=None)[source]#

Solve and return the forward sensitivity dC/dtheta alongside C.

Integrates the augmented [C; S] system (state plus sensitivity) with adaptive control over both, so the sensitivity is exact and finite without the dtmax cap that ordinary AD through a stiff solve needs (see aquakin.integrate.forward_sensitivity).

Parameters:

  • C0 (Array) – As for solve().

  • t_span (tuple[float, float]) – As for solve().

  • t_eval (Array | None) – As for solve().

  • conditions (SpatialConditions | None) – As for solve().

  • params (jnp.ndarray) – Full parameter vector; keyword-only, matching solve().

  • sens_params (list of str or int) – Namespaced parameter names (or integer indices into params) to differentiate with respect to.

  • sens_rtol (float | None) – Sensitivity error-control tolerances. Defaults follow CVODES: rtol_S = rtol and atol_S = atol / |theta_k|. See augmented_forward_sensitivity().

  • sens_atol – Sensitivity error-control tolerances. Defaults follow CVODES: rtol_S = rtol and atol_S = atol / |theta_k|. See augmented_forward_sensitivity().

  • param_scale – Sensitivity error-control tolerances. Defaults follow CVODES: rtol_S = rtol and atol_S = atol / |theta_k|. See augmented_forward_sensitivity().

  • shared_factor (bool, optional) – Use the CVODES simultaneous-corrector linear solve (factorise the shared diagonal block once, forward-substitute across the sensitivity columns). None (default) auto-selects: True for more than one sensitivity parameter (where it is markedly cheaper than the dense augmented solve), False for a single parameter.

Returns:

  • sol (BatchSolution) – The usual state trajectory (uncapped, exact).

  • S (jnp.ndarray) – Sensitivity dC/dtheta at the saved times, shape (n_t, n_species, n_sens_params).

Return type:

tuple[BatchSolution, Array]