aquakin.sensitivity#

aquakin.sensitivity(reactor, C0, params=None, output_fn=None, *, t_span=None, t_eval=None, solve_kwargs=None, diff=DifferentiationConfig(mode='reverse', method='stable', check_finite=True, adjoint_max_steps=100000, adjoint_low_memory=False))[source]#

Compute gradients of a scalar output with respect to parameters and condition fields, via autodiff through reactor.solve.

Parameters:

  • reactor (BatchReactor or PlugFlowReactor) – Any reactor exposing .solve(C0, t_span, ..., params=...) and a .conditions attribute.

  • C0 (jnp.ndarray) – Initial concentration vector.

  • params (jnp.ndarray, optional) – Parameter vector at which to evaluate sensitivity. Defaults to reactor.model.default_parameters().

  • output_fn (callable) – Maps a solution object to a scalar JAX value, e.g. lambda sol: sol.C_named("BrO3-")[-1].

  • t_span (optional) – Integration window / save times, passed straight to reactor.solve (the common batch case). Equivalent to putting them in solve_kwargs; provide whichever reads better.

  • t_eval (optional) – Integration window / save times, passed straight to reactor.solve (the common batch case). Equivalent to putting them in solve_kwargs; provide whichever reads better.

  • solve_kwargs (dict, optional) – Any further keyword arguments forwarded to reactor.solve – including time_unit= if t_span / t_eval are in a non-native unit (solve converts them, so the sensitivities stay consistent).

  • diff (DifferentiationConfig, optional) – Autodiff configuration. mode="reverse" (default) uses jax.grad. mode="forward" uses jax.jacfwd and rebuilds the reactor internally with a forward-capable adjoint, so a stiff reactor whose reverse adjoint is non-finite can be differentiated without a dtmax cap and without the caller touching diffrax. check_finite (default True) raises a friendly RuntimeError if the computed sensitivities are non-finite, instead of returning silent NaN values.

Return type:

SensitivityResult