aquakin.Plant#

class aquakin.Plant(name, *, recycle_passes=3, recycle_tol=1e-08, recycle_max_passes=100)[source]#

Bases: object

A plant flowsheet of Unit components.

Build a plant by:

Constructing each unit (CSTRUnit, MixerUnit, etc.).
Adding them via add_unit() in the order you want their RHS evaluated (downstream-first if there’s a recycle).
Adding influent sources via add_influent().
Connecting units via connect().
Calling solve().

Parameters:

name (str)
recycle_passes (int, optional) – Number of Gauss-Seidel mop-up passes per RHS evaluation (default 3). Both the recycle flows AND concentrations are first resolved exactly and gain-independently – flows by _resolve_flows(), concentrations by _resolve_recycle_concentrations() (an affine probe + linear solve) – so for any linear topology (every shipped plant) the recycle back-edges are seeded at their exact fixed point and these passes do no work. They only refine the residual of a genuinely non-affine in-cycle unit (an ASM↔ADM translator inside a pure-stateless loop, which the shipped units cannot form), so the default 3 is ample and rarely matters. It is a fixed count (not iterate-to-tolerance) because the RHS is jitted and differentiated. The first non-traced solve() runs a one-time convergence diagnostic on the mop-up residual and warns only if even the exact pre-solve plus these passes has not converged (a non-affine loop); raise recycle_passes until it clears, or set recycle_tol.
recycle_tol (float, optional) – Relative tolerance for the adaptive recycle resolution, on by default (1e-8). The recycle back-edge streams (flow, concentration, temperature) are iterated to this tolerance – correct for any topology, not just the low-gain BSM reject loop. The alternative fixed recycle_passes mop-up converges in log(tol)/log(rho) passes where rho is the nonlinear flow<->concentration coupling’s spectral radius (~0.0066 for BSM, so 3 passes is ample), but rho is topology-dependent and not bounded below 1: a recycle-heavy plant with a strong concentration-dependent in-loop flow can leave the fixed count silently under-converged. The adaptive solve (warm-started from the exact affine seed, an adaptive jax.lax.while_loop() wrapped in jax.lax.custom_root()) iterates until the actual residual clears, so it converges for any rho < 1, stops early on a low-gain plant, and – via the implicit-function-theorem tangent – gives a gradient that is exact and O(1) in the pass count. 1e-8 is well below the typical solver rtol and a strict improvement on the old fixed-3-pass default (~1e-6 for BSM) at ~neutral cost (~3 iterations from the affine seed). Set recycle_tol=None to fall back to the fixed recycle_passes path. See _adaptive_recycle_refine().
recycle_max_passes (int, optional) – Cap on the adaptive recycle_tol iteration (default 100), the worst-case guard for a near-unit-gain loop.

__init__(name, *, recycle_passes=3, recycle_tol=1e-08, recycle_max_passes=100)[source]#

Parameters:

name (str)
recycle_passes (int)
recycle_tol (float | None)
recycle_max_passes (int)

Return type:

None

Methods

`__init__`(name, *[, recycle_passes, ...])
`activated_sludge_reactors`(*[, require_volume])	The activated-sludge reactor units (CSTR / MBR), in plant order.
`add_influent`(name, series[, to, translator])	Register an external time-varying influent stream, optionally wiring it.
`add_unit`(unit)	Register a unit.
`calibrate`(observations, t_obs, free_params, *)	MAP-calibrate a plant's parameters against an output stream.
`check`(*[, raise_on_error])	Validate the wiring before solving: find unfed ports / dangling outputs.
`colored_jacobian_decision`()	Report the auto colored-backward-Jacobian decision: `("colored" \| "dense", n_states)`, or `None` before the state layout can be built.
`connect`(source, dest, *[, translator, ...])	Wire a stream from one unit's output to another unit's input.
`default_parameters`()	Concatenated default parameters: kinetic models then flow setpoints.
`derivative`(state[, params, t])	Evaluate the assembled flowsheet RHS once: `dstate/dt` at `state`.
`digester_gas`(solution[, params])	The anaerobic digester's biogas trajectory.
`dynamic_dgsm`(ranges, *, output_fn, t_span[, ...])	Derivative-based global sensitivity (DGSM) of a transient plant output.
`dynamic_sensitivity`([params, t_eval, wrt, ...])	Sensitivity of a transient (time-window) output to the parameters.
`effluent_stream`(solution[, params])	The plant's final effluent stream (a shortcut for the most-read one).
`initial_state`([overrides])	Concatenated initial state from each unit's `initial_state()`.
`list_ports`([unit, role])	The `"unit.port"` endpoint strings, for discovering stream args.
`list_species`(unit)	The species names of a concentration-vector unit's model.
`list_streams`()	The registered semantic stream names `{name: "unit.port"}`.
`list_units`()	The plant's unit names, in the order added.
`mass_balance`(solution, *[, components, ...])	Results-level mass-balance closure of this plant over a solved window.
`outputs_at`(t, state_full[, params])	Reconstruct every unit's output streams at one `(t, state)`.
`parameter_index`(name)	Flat index of `name` (`"<model>.<param>"`) in the parameter vector.
`parameter_names`()	The plant's calibratable parameter names, `"<model>.<param>"`.
`parameter_values`([overrides])	Plant parameter vector: the defaults with named entries overridden.
`register_stream`(name, endpoint)	Register a semantic name for an output `"unit.port"` endpoint.
`run_to_steady_state`([params, y0, max_time, ...])	Integrate forward until the plant settles to steady state.
`set_temperature`(celsius, *[, units])	Set the operating temperature of the reactors, in degrees Celsius.
`set_temperature_model`(model)	Select how reactor temperature is handled (see `aquakin.plant.temperature`).
`signals_at`(t, state_full[, params])	Reconstruct the control-signal bus at one `(t, state)`.
`sludge_age`(solution[, params])	Achieved SRT / HRT / F:M of this activated-sludge plant.
`solve`(t_span[, t_eval, params, rtol, atol, ...])	Integrate the plant over `t_span`.
`solve_sensitivity`(params, wrt, *[, ...])	Stable forward (variational) sensitivity `dy/dtheta` of the plant.
`states_by_unit`(state_full)	Split a flat plant vector into a `{unit_name: sub-vector}` map.
`steady_state`([params, y0, influent_time, ...])	Solve the plant steady state algebraically by pseudo-transient continuation.
`steady_state_dgsm`(ranges, *, output_fn[, ...])	Derivative-based global sensitivity (DGSM) of the plant steady state.
`steady_state_sensitivity`([params, y0, ...])	Exact steady-state output sensitivities via the implicit function theorem.
`stream`(solution, endpoint[, params])	Reconstruct a named output stream's trajectory from a solution.

Attributes

time_unit

The plant's native integration time unit, or None.

add_unit(unit)[source]#

Units may be added in any order: the plant computes the RHS evaluation order itself by topologically sorting the connection graph at solve time (recycles are the graph back-edges, detected automatically – you do not have to add a downstream unit before its upstream consumer). The add order is used only as a deterministic tie-break for the sort and to order the per-model parameter blocks.

Parameters:: unit (Unit)
Return type:: None

add_influent(name, series, to=None, *, translator=None)[source]#

Parameters:

name (str) – Identifier for this influent (used in stream keys and error messages).
series (InfluentSeries) – The time-varying influent data.
to (str, optional) – Destination endpoint to feed this influent into, as "unit.port" (or bare "unit" to use the unit’s sole input port). When given, the connection is made here, so no separate connect() call is needed. The destination unit must already be added. When omitted it defaults to the plant’s influent_endpoint (set by the builders, so plant.add_influent("feed", series) wires to the canonical front without hard-coding a port); if that is also unset the influent is registered but inert – influents are not valid connect() sources, so wiring happens here.
translator (StateTranslator, optional) – Translator for the influent -> unit connection. Defaults to identity when the influent and destination share a model.

Return type:

None

set_temperature(celsius, *, units=None)[source]#

Set the operating temperature of the reactors, in degrees Celsius.

One knob for “run the plant at this temperature”: converts to Kelvin and writes the static T condition of every temperature-bearing reactor, so a re-solve runs the kinetics – including the Arrhenius temperature_corrections – at that temperature. (For a model whose corrections are referenced to this temperature, the correction is unity and the run reproduces the calibrated operating point; a colder/warmer setting drives the kinetics away from it.)

By default it targets the activated-sludge reactors (the units exposing CSTRUnit.set_temperature() with a T condition) and leaves a fixed-temperature unit like the heated anaerobic digester untouched – the digester is a separate ADM1 unit without that method. Pass units to target a specific set by name.

Parameters:

celsius (float) – Operating temperature in °C (converted to Kelvin internally; the ASM/ADM condition fields are Kelvin).
units (iterable of str, optional) – Unit names to set. Defaults to every reactor that supports it.

Returns:

self (so calls can be chained after build_*).

Return type:

Plant

Raises:

UnknownUnitError – If a name in units is not a unit of this plant.
WiringError – If a named unit does not support a temperature (e.g. a separator).

set_temperature_model(model)[source]#

Select how reactor temperature is handled (see aquakin.plant.temperature).

Pass HeatBalanceTemperature to give each finite-volume liquid unit a dynamic temperature state (a first-order heat balance), or the default AlgebraicTemperature for the instantaneous flow-weighted behaviour. Changes the flat state-vector length (an appended temperature block), so it clears the compiled-solve cache; rebuild y0 (e.g. via bsm2_warm_start) after calling it. Returns self for chaining.

Parameters:: model (TemperatureModel)
Return type:: Plant

connect(source, dest, *, translator=None, initial_value=None)[source]#

Wire a stream from one unit’s output to another unit’s input.

Both endpoints are "unit.port" strings, read as source -> dest. The port may be omitted (bare "unit") when the unit has exactly one port for that role – a single output (source) or single input (dest) – so only multi-port units like mixers and splitters need an explicit port. To feed an external influent, use add_influent() with to=... rather than connect().

Parameters:

source (str) – Source endpoint, "unit.port" or bare "unit" for the unit’s sole output port.
dest (str) – Destination endpoint, "unit.port" or bare "unit" for the unit’s sole input port.
translator (StateTranslator, optional) – State translator to apply as the stream crosses the connection. Defaults to IdentityTranslator when source and destination share a model; required when they differ.
initial_value (Stream, optional) – Seed for the stream’s first RHS pass on a recycle edge. Recycles are detected automatically as the back-edges of the connection graph (you do not mark them, and the units may be added in any order); a recycle with no initial_value is auto-seeded with a zero-flow stream of the source model. Pass it only to override that with a non-zero warm start (e.g. a temperature-carrying seed to ignite temperature propagation around the loop on the first pass). Ignored for feed-forward edges. Both endpoints must already be added.

Return type:

None

Examples

>>> plant.connect("tank1", "tank2")                 # sole ports inferred
>>> plant.connect("tank5_split.to_clarifier", "clarifier")
>>> plant.connect("split.recycle", "mix.recycle")   # recycle: auto-detected

check(*, raise_on_error=False)[source]#

Validate the wiring before solving: find unfed ports / dangling outputs.

Every unit input port must be fed by exactly one stream (an influent or another unit’s output); an unfed input port has no source, so the RHS sweep cannot resolve it and the solve fails. Output ports consumed by no connection are reported too, but as information – a terminal stream that leaves the plant (final effluent, wasted sludge, disposal cake, biogas) is legitimately unconsumed.

Parameters:: raise_on_error (bool, optional) – If True, raise ValueError when any input port is unfed (the actionable error) instead of only reporting it. Default False.
Returns:: unfed_ports (errors), dangling_outputs (info), recycles (the detected back-edges), with .ok and .summary().
Return type:: PlantCheck

property time_unit: str | None#

The plant’s native integration time unit, or None.

t_span / t_eval passed to solve() are in this unit (unless a time_unit= override is given). It is the unit shared by every kinetic model’s rate constants (CompiledModel.time_unit) – "d" for the BSM-family plants (ASM1 + ADM1 are both in days). Returns None if the plant has no kinetic model, or its models disagree on (or do not declare) a time unit – in which case a time_unit= conversion in solve() cannot be applied.

default_parameters()[source]#

Concatenated default parameters: kinetic models then flow setpoints.

Return type:: Array

parameter_names()[source]#

The plant’s calibratable parameter names, "<model>.<param>".

Each kinetic model contributes its (namespaced) parameter names prefixed by the model name – e.g. "asm1.muH" or "adm1.k_hyd_ch" – which are the keys accepted by parameter_values() / parameter_index().

Return type:: list[str]

parameter_index(name)[source]#

Flat index of name ("<model>.<param>") in the parameter vector.

The companion to parameter_values() for code that needs the position rather than a new vector – e.g. jax.grad with respect to one parameter, which can’t go through parameter_values() (that materialises concrete values) – without hand-computing the model block offset. Raises KeyError with a close-match hint for an unknown name.

Examples

>>> gidx = plant.parameter_index("adm1.k_m_ac")
>>> grad = jax.grad(lambda th: f(params.at[gidx].set(th)))(theta0)

Parameters:: name (str)
Return type:: int

parameter_values(overrides=None, /)[source]#

Plant parameter vector: the defaults with named entries overridden.

The plant analogue of CompiledModel.parameter_values. Keys are "<model>.<param>" – the model name plus the model’s own (namespaced) parameter name – so you can bump one rate in a multi-model plant (e.g. BSM2’s ASM1 water line + ADM1 digester) without hunting the block offset and index by hand.

Parameters:: overrides (dict[str, float], optional) – Map of "<model>.<param>" to a new value. Unknown names raise a KeyError with a close-match hint. None returns the defaults.
Returns:: The flat parameter vector to pass to solve().
Return type:: jnp.ndarray

Examples

>>> params = plant.parameter_values({"asm1.muH": 4.0, "adm1.k_hyd": 10.0})
>>> plant.solve(t_span=(0.0, 200.0), params=params)

See also

parameter_names: the valid keys.

list_units()[source]#

The plant’s unit names, in the order added.

Each name keys units, is the unit part of the "unit.port" endpoints stream() / connect() accept, and is a key of states_by_unit() / PlantSolution.unit_state(). See list_ports() for the ports and list_species() for a kinetic unit’s species. Available before solving (it is plant structure).

Return type:: list[str]

list_ports(unit=None, *, role='output')[source]#

The "unit.port" endpoint strings, for discovering stream args.

role="output" (default) returns every unit-output endpoint – the strings stream() reconstructs and connect() reads from as a source. role="input" returns the input endpoints (connect() destinations). Pass unit to restrict to one unit. No reading the builder source to find a port string.

Parameters:

unit (str | None)
role (str)

Return type:

list[str]

activated_sludge_reactors(*, require_volume=True)[source]#

The activated-sludge reactor units (CSTR / MBR), in plant order.

Identified by the CSTR-only aeration attribute – the digester and the other volumed units lack it. require_volume (the default) additionally requires a volume field: warm-starts and sizing need the volume, whereas the mixing-energy term wants every mechanically mixed reactor regardless. The single source of truth behind the warm-start / design / evaluation reactor heuristics.

Parameters:: require_volume (bool)
Return type:: list[str]

list_species(unit)[source]#

The species names of a concentration-vector unit’s model.

These are the valid names for PlantSolution.C_named() and PlantSolution.to_dataframe() on that unit. Raises KeyError for an unknown unit (with a hint), or for a unit whose state is not a concentration vector (a mixer/splitter/clarifier) – use stream() to read those as flow streams instead.

Parameters:: unit (str)
Return type:: list[str]

initial_state(overrides=None)[source]#

Concatenated initial state from each unit’s initial_state().

Parameters:

overrides (dict[str, array-like], optional) –

Map of unit name to an initial-state vector that replaces that unit’s own initial_state(). Each vector must have length equal to the unit’s state_size. Use this to warm-start a plant – e.g. seed the activated-sludge reactors with a healthy biomass before settling a slow digester – without reaching into the internal state layout:

warm = asm1.concentrations(XB_H=2244.0, ...)   # (n_species,)
y0 = plant.initial_state(overrides={t: warm for t in tanks})

Returns:

The flat initial-state vector, shape (total_state_size,).

Return type:

jnp.ndarray

Raises:

KeyError – If an override names a unit not in the plant.
ValueError – If an override vector length does not match the unit’s state size.

states_by_unit(state_full)[source]#

Split a flat plant vector into a {unit_name: sub-vector} map.

The inverse of initial_state() with overrides: that assembles a flat vector from per-unit pieces, this splits one back apart. Works on any flat plant vector in the plant’s state layout – an initial state, a solution snapshot (PlantSolution.final_state), or a derivative from derivative():

dig = plant.states_by_unit(sol.final_state)["digester"]
rate = plant.states_by_unit(plant.derivative(y0))["tank5"]

Parameters:: state_full (jnp.ndarray) – A flat plant vector, shape (total_state_size,).
Returns:: {unit_name: sub-vector}, each of shape (unit.state_size,) in the unit-addition order.
Return type:: dict

outputs_at(t, state_full, params=None)[source]#

Reconstruct every unit’s output streams at one (t, state).

The plant integrates unit states, not the inter-unit streams, so a stream such as the clarifier effluent is recomputed from the state on demand. This resolves the full output sweep (including recycle edges) and returns the {(unit, port): Stream} map. See stream() for a named-stream trajectory over a whole solution.

Parameters:

t (float) – Time (plant units), used to interpolate the influents.
state_full (jnp.ndarray) – Flat plant state at t (e.g. one row of PlantSolution.state).
params (jnp.ndarray, optional) – Plant parameters (defaults to default_parameters()).

Returns:

{(unit_name, port_name): Stream} for every unit output.

Return type:

dict

stream(solution, endpoint, params=None)[source]#

Reconstruct a named output stream’s trajectory from a solution.

The plant integrates unit states, not the inter-unit streams, so a stream such as the secondary-clarifier effluent is not stored – it is recomputed (flow + concentration over time) from the saved states on demand. The whole output sweep is reconstructed once and cached on the solution (see _cached_streams()), so a sequence of stream calls for different ports reuses one reconstruction.

Parameters:

solution (PlantSolution) – A solution returned by solve() (carries t and state).
endpoint (str) – A registered semantic name ("effluent", "ras", "internal_recycle", "primary_sludge", "reject" … – see list_streams()) or a literal "unit.port" (e.g. "clarifier.overflow"; the port may be omitted for a single-output unit). The semantic name is resolved to its port first, so an engineer asks for "effluent" without knowing the internal port.
params (jnp.ndarray, optional) – The plant parameters used for the run (defaults to default_parameters()). Output reconstruction is parameter-independent for the shipped units, so the default is usually fine.

Returns:

t, Q, C (shape (n_t, n_species)) and model, with a C_named(species) accessor.

Return type:

StreamSeries

register_stream(name, endpoint)[source]#

After this, plant.stream(sol, name) reads that port – so an engineer asks for "effluent" / "ras" / "digester_gas" instead of the internal port. Builders (build_bsm1 / build_bsm2) register the plant’s named streams; call this to add your own. Returns self for chaining.

Parameters:

name (str)
endpoint (str)

Return type:

Plant

list_streams()[source]#

The registered semantic stream names {name: "unit.port"}.

The engineering shortcuts stream() accepts ("effluent", "ras", …) and the internal port each resolves to – the discoverable companion to list_ports(), so you never read the builder source to find a port string. Empty on a plant whose builder registered none.

Return type:: dict[str, str]

effluent_stream(solution, params=None)[source]#

The plant’s final effluent stream (a shortcut for the most-read one).

Reads the builder-recorded effluent_endpoint – the right port even when an option moved it (a BSM2 influent bypass relocates the effluent to effluent_mix.out). Equivalent to plant.stream(sol, plant.effluent_endpoint).

Parameters:

solution (PlantSolution)
params (Array | None)

Return type:

StreamSeries

digester_gas(solution, params=None)[source]#

The anaerobic digester’s biogas trajectory.

A DigesterGas with the biogas flow Q (m³/d), the CH₄/CO₂/H₂ partial pressures and the CH₄ mass flow ch4 (kg/d), plus .methane_production() (time-averaged kg CH₄/d). Unlike stream(), the biogas is derived from the ADM1 headspace state (not a material port). Raises if the plant has no ADM1 digester.

Parameters:

solution (PlantSolution)
params (Array | None)

mass_balance(solution, *, components=('COD', 'N', 'P'), influent_ports=None, effluent_ports=None, params=None)[source]#

Results-level mass-balance closure of this plant over a solved window.

Accounts, per component (COD / N / P), the component that flowed in (influents), out (terminal material streams), left as gas (O₂ transferred by aeration, digester biogas COD, denitrification N₂ + its oxidised COD) and accumulated (inventory change across every unit) – and reports the closure imbalance = in − out − gas − accumulation, which is ~0 for a trustworthy result. The gas terms are computed independently of the in/out/accumulation bookkeeping (from the aeration term, the digester headspace and an activated-sludge reaction integral), so the imbalance is a genuine check rather than a definition.

Everything is reported on one canonical gram basis (g COD / g N / g P), so inventories and fluxes sum across models of different units (the ASM water line in g/m³, the ADM digester in kg/m³ and kmol/m³) via the shipped aquakin.composition_table().

Parameters:

solution (PlantSolution) – A solved trajectory (the closure is over solution.t[0]..t[-1]). The gas integrals are exact at steady state and otherwise accurate to the t_eval sampling.
components (tuple of str, optional) – Which components to balance (default ("COD", "N", "P")); a model that does not carry one contributes nothing to it.
influent_ports (list of str, optional) – Override the boundary streams. By default the influents are every registered influent and the effluents are the plant’s terminal (dangling) output ports (final effluent, wasted sludge, disposal cake) – the biogas is handled as a gas term, not a material port.
effluent_ports (list of str, optional) – Override the boundary streams. By default the influents are every registered influent and the effluents are the plant’s terminal (dangling) output ports (final effluent, wasted sludge, disposal cake) – the biogas is handled as a gas term, not a material port.
params (jnp.ndarray, optional) – Plant parameters used for the run (defaults to default_parameters()).

Returns:

Per-component ComponentBalance (index by name, e.g. mb["N"]), with .closed(rtol) and .summary().

Return type:

MassBalance

sludge_age(solution, params=None, **kwargs)[source]#

Achieved SRT / HRT / F:M of this activated-sludge plant.

Convenience wrapper for aquakin.plant.design.sludge_metrics() – the design loop’s closing half: the solids retention time (sludge age) is an emergent property of the wastage flow, so this reports what the solved model actually achieved instead of requiring it be guessed.

Parameters:

solution (PlantSolution) – A solution from solve() (ideally near steady state).
params (jnp.ndarray, optional) – Plant parameters used for the run.
**kwargs – Forwarded to sludge_metrics() (reactor_units, influent_name, effluent_port, waste_port, substrate).

Returns:

SRT, HRT, F:M and the intermediate inventories / loads.

Return type:

SludgeMetrics

derivative(state, params=None, *, t=0.0)[source]#

Evaluate the assembled flowsheet RHS once: dstate/dt at state.

A public single evaluation of the same right-hand side solve() integrates (recycles resolved by the fixed-pass sweep). Useful for inspecting the dynamics – sign, magnitude, finiteness – at a state without running a full solve. Split the result with states_by_unit():

d = plant.derivative(y0, params)
snh_rate = plant.states_by_unit(d)["tank5"][net.species_index["SNH"]]

Parameters:

state (jnp.ndarray) – Flat plant state, shape (total_state_size,).
params (jnp.ndarray, optional) – Plant parameters (defaults to default_parameters()).
t (float, optional) – Time at which to evaluate, for time-varying influents (default 0).

Returns:

dstate/dt, shape (total_state_size,) – the same layout as state.

Return type:

jnp.ndarray

signals_at(t, state_full, params=None)[source]#

Reconstruct the control-signal bus at one (t, state).

The plant integrates unit states, not the control signals, so a signal such as a DO controller’s manipulated kLa is recomputed from the state on demand – the signal analogue of outputs_at(). Returns {} for an open-loop plant (no controllers).

Parameters:

t (float) – Time (plant units), used to interpolate the influents.
state_full (jnp.ndarray) – Flat plant state at t (e.g. one row of PlantSolution.state).
params (jnp.ndarray, optional) – Plant parameters (defaults to default_parameters()).

Returns:

{signal_name: scalar} for every published control signal.

Return type:

dict

colored_jacobian_decision()[source]#

Report the auto colored-backward-Jacobian decision: ("colored" | "dense", n_states), or None before the state layout can be built.

Deterministic from the plant size (_COLORED_BACKWARD_MIN_STATES): the colored backward is auto-enabled iff the plant has at least that many states. (With colored_jacobian=True/False the choice is forced and this size gate does not apply.)

solve(t_span, t_eval=None, params=None, *, rtol=1e-06, atol=None, y0=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), time_unit=None, event=None, events=None, progress_meter=None, forward_fast=False)[source]#

Integrate the plant over t_span.

Parameters:

t_span ((float, float)) – (t_start, t_end) in the plant’s time unit (plant.time_unit, typically days for BSM-family) unless time_unit is given.
t_eval (jnp.ndarray, optional) – Save times. If None only the endpoint is saved.
params (jnp.ndarray, optional) – Flat plant parameter vector. Defaults to default_parameters().
rtol (float or array, optional) – Solver tolerances. atol=None auto-scales a per-component noise floor off the operating magnitudes.
atol (float or array, optional) – Solver tolerances. atol=None auto-scales a per-component noise floor off the operating magnitudes.
y0 (jnp.ndarray, optional) – Warm-start initial state (e.g. a previously-computed steady state), shape (total_state_size,). Defaults to the per-unit initial states.
integrator (IntegratorConfig, optional) – Integrator / step configuration. order selects Kvaerno3 (default, fast) or Kvaerno5; factormax caps the PID step-growth (default 3); dtmax caps the step (set it only for a reverse gradient through the solve, diff.method='through_solve'); max_steps is the step budget; solver is an explicit override (honoured verbatim); colored_jacobian ({“auto”, True, False}) selects sparse colored-AD materialisation of the per-step Jacobian. "auto" (default) governs the mode='reverse', method='stable' backward decision – it measures whether the colored df/dy build pays and enables it only then (a large plant like BSM2 yes, a small one like BSM1 no), and leaves the forward solve dense; True forces coloring on both paths; False disables it. Reported by colored_jacobian_decision().
diff (DifferentiationConfig, optional) – How a gradient / sensitivity flows through the solve. mode='reverse', method='stable' (the default) routes a plain concrete forward solve to the fast cached path and a reverse-mode differentiated solve to the cap-free hand-written discrete adjoint (esdirk_adjoint_solve()) – finite for a stiff plant with no dtmax to tune. mode='reverse', method='through_solve' differentiates through the diffrax solve (RecursiveCheckpointAdjoint; needs a dtmax cap for a stiff plant). mode='forward', method='through_solve' builds a forward-capable adjoint so jax.jvp / jacfwd flow through the solve. mode='forward', method='stable' is the augmented variational solve – not supported here; call solve_sensitivity() / dynamic_sensitivity(). check_finite raises on a non-finite gradient; adjoint_max_steps bounds the discrete-adjoint backward-scan buffer (it must exceed the forward step count); adjoint_low_memory recomputes the backward stages instead of saving the ~n_stages``x dense buffer (memory for compute, gradient unchanged). Passing ``event= pins the forward jax_adjoint path.
time_unit (str, optional) – The unit t_span / t_eval are in ("s"/"min"/"h"/ "d"). Default None uses the plant’s native time unit. When given, the input times are converted to the native unit for the solve and solution.t is reported back in time_unit.
event (diffrax.Event, optional) – The low-level single terminating event (used internally by run_to_steady_state()); pins the forward jax_adjoint path.
events (sequence of Event, optional) – Located plant-wide discontinuities (on/off pumps, SBR phase switches, dosing on/off, tank-level limits). The monolithic solve is split into segments at the firings and solution.events_log records them. Time-only events keep jax.grad finite; a state event makes the run a forward simulation. Mutually exclusive with event= and with diff=DifferentiationConfig(method='stable') reverse differentiation.
progress_meter (diffrax.AbstractProgressMeter, optional) – A diffrax progress meter for a long forward solve.
forward_fast (bool, optional) – Use the lean non-AD forward integrator (no diffrax adjoint machinery): a much faster compile + run for a one-off forward solve, but the result is NOT differentiable and it needs concrete params/y0. It runs its own hand-rolled ESDIRK and honours only rtol / atol / integrator.max_steps; a non-default integrator solver / order / factormax / dtmax / colored_jacobian is rejected (rather than silently dropped).

Return type:

PlantSolution

run_to_steady_state(params=None, y0=None, *, max_time=1000.0, ss_rtol=0.001, ss_atol=0.001, rtol=1e-06, atol=None, atol_factor=1e-06, max_steps=500000)[source]#

Integrate forward until the plant settles to steady state.

A single continuous adaptive solve that terminates itself the instant the dynamics die out – diffrax’s steady-state event halts when the vector field is approximately zero (||dstate/dt|| <= ss_atol + ss_rtol*||state||), the standard “march in time until the residual dies” criterion. There is no fixed horizon to guess and no chunked re-integration: max_time is only a safety cap, reached only if the plant has not settled (then converged is False).

The absolute tolerance for the underlying solve defaults to a per-component noise floor scaled off the operating magnitudes (atol_i = atol_factor * max(|y0_i|, |default_i|), with a small global floor) – the SUNDIALS/Hairer “vector atol” recommendation for states spanning many orders of magnitude – so the engineer does not hand-tune tolerances for a g/m³ plant.

Parameters:

params (jnp.ndarray, optional) – Plant parameters (defaults to default_parameters()).
y0 (jnp.ndarray, optional) – Initial state / warm start (defaults to initial_state()). A healthy warm start reaches steady state far faster than a cold start.
max_time (float) – Safety cap on integration time (plant units, typically days).
ss_rtol (float) – Steady-state event tolerances on ||dstate/dt|| (the convergence criterion). ss_rtol=1e-3 ~ “no state changes faster than ~0.1% per unit time”.
ss_atol (float) – Steady-state event tolerances on ||dstate/dt|| (the convergence criterion). ss_rtol=1e-3 ~ “no state changes faster than ~0.1% per unit time”.
rtol (float) – Relative tolerance of the underlying integrator.
atol (float or jnp.ndarray, optional) – Absolute tolerance of the integrator. None (default) auto-scales per component (see above); pass a value to override.
atol_factor (float) – Scale factor for the auto per-component atol.
max_steps (int) – Maximum integrator steps.

Returns:

state (the operating point), converged, time, and the underlying solution.

Return type:

SteadyStateResult

Examples

>>> ss = plant.run_to_steady_state(params, y0=warm)
>>> ss.converged
True
>>> snh = plant.states_by_unit(ss.state)["tank5"][asm1.species_index["SNH"]]

steady_state(params=None, y0=None, *, influent_time=0.0, dt0=0.01, dt_max=10000000000.0, growth_cap=10.0, max_iter=400, tol=1e-06, scale_floor=None, nonneg=True, design=None, colored_jacobian=False, fallback=True, fallback_kwargs=None, continuation_from=None, continuation_kwargs=None, arclength=True, arclength_kwargs=None)[source]#

Solve the plant steady state algebraically by pseudo-transient continuation.

Finds the root of the plant right-hand side F(y) = dy/dt = 0 directly, rather than integrating forward until the dynamics die out (run_to_steady_state()). Pseudo-transient continuation takes damped-Newton steps (V/dt - J) dy = F with the exact AD Jacobian J and a per-state pseudo-time V/dt that ramps from a stable time-stepping move (far from the root) to a full Newton step (near it), so it is robust on stiff plants (long-SRT digesters) where a plain Newton root-find stalls. Typically reaches steady state ~10x faster than the forward solve and to a tighter residual; see aquakin.plant.steady.

Differentiable. The returned state carries the implicit-function-theorem gradient with respect to params (the iteration itself is gradient-blocked), so jax.grad of a loss on the steady state flows to the plant parameters – the intended path for design sweeps. The diagnostic fields (converged, iterations, residual) and the forward fallback are only available in eager use; under an outer jit/grad trace they are traced values and the fallback is skipped.

Parameters:

params (jnp.ndarray, optional) – Plant parameters (defaults to default_parameters()).
y0 (jnp.ndarray, optional) – Warm start (defaults to initial_state()). A healthy warm start (e.g. bsm2_warm_start) converges in fewer iterations.
influent_time (float) – Time at which to read the (constant) influent for the steady residual. The steady state is only well defined for a constant load; for a time-varying influent this samples it at influent_time.
dt0 (float) – Pseudo-transient controls: initial / maximum pseudo-timestep and the per-step SER growth cap (see aquakin.plant.steady.ptc_forward()).
dt_max (float) – Pseudo-transient controls: initial / maximum pseudo-timestep and the per-step SER growth cap (see aquakin.plant.steady.ptc_forward()).
growth_cap (float) – Pseudo-transient controls: initial / maximum pseudo-timestep and the per-step SER growth cap (see aquakin.plant.steady.ptc_forward()).
max_iter (int) – PTC iteration cap.
tol (float) – Convergence tolerance on the scaled residual.
scale_floor (float or array, optional) – Floor on |y| in the per-state pseudo-time / residual scaling (so near-zero states do not distort the damping or the convergence criterion). Default None builds a per-state floor max(|y0|, 1e-6) – each state scaled by its own warm-start magnitude, which roughly halves the iteration count on a stiff multi-model plant (the flat scalar floor over-damps small-magnitude states). Pass a scalar or per-state array to override.
nonneg (bool) – Clamp the state to >= 0 each step (concentrations are non-negative).
design (dict, optional) –
Differentiable design-variable overrides folded into the quantity the implicit-function-theorem gradient is taken w.r.t., so a design sweep can jax.grad / jacobian the steady state through them. Currently supports the influent load: design={"influent": {port: {"Q": ..., "C": ..., "T": ...}}} (plain arrays; T optional), which replaces the recorded influent at influent_time. Example – sensitivity of effluent ammonia to the influent ammonia load:
```
jax.grad(lambda c: plant.steady_state(
    p, y0, design={"influent": {"feed": {"Q": Q, "C": c}}}
).state[eff_idx])(C_influent)
```
colored_jacobian (bool) – Materialize the PTC iteration Jacobian dF/dy by column-compressed colored AD (one Jacobian-vector product per color rather than per state; BSM2 46 colors vs 167) instead of dense jax.jacfwd. The operating point is unchanged (the colored matrix equals the dense one on its sparsity-pattern support: bit-identical on a single-model plant, identical to PTC tolerance on a multi-model plant where the recycle solve differs by round-off); only the per-iteration Jacobian cost drops. Built and guarded once concretely (falls back to dense with a warning on a start-state mismatch, or under a jit/grad trace where the probe cannot run). Benefit is regime-specific (measured BSM2): the Jacobian build is ~2.4x cheaper and the whole solve ~1.9x faster run-only under jit, but an un-jitted one-shot call is compile/trace-bound, so the one-time pattern build makes it ~0.8x (slower). Worth enabling only when the steady solve is run repeatedly under jit (differentiable design sweeps, where compile amortizes) or for a much larger plant; not for a single steady state. Default False.
fallback (bool) – If PTC does not converge within max_iter (eager use only), fall back to run_to_steady_state() and return that result (with method="ptc->forward").
fallback_kwargs (dict, optional) – Extra keyword arguments for the forward fallback.
continuation_from (tuple | None)
continuation_kwargs (dict | None)
arclength (bool)
arclength_kwargs (dict | None)

Returns:

state (the operating point, IFT-differentiable), converged, method="ptc", iterations and residual.

Return type:

SteadyStateResult

Examples

>>> ss = plant.steady_state(params, y0=warm)
>>> ss.converged, ss.iterations
(True, 75)
>>> g = jax.grad(lambda p: plant.steady_state(p, y0=warm).state[idx])(params)

steady_state_sensitivity(params=None, y0=None, *, state=None, output_fn=None, wrt=None, operating=None, mode='auto', elasticity=False, return_jacobian=False, **steady_kwargs)#

Exact steady-state output sensitivities via the implicit function theorem.

Solves the plant to steady state once, then returns d(output)/d(params) directly from the right-hand-side Jacobians – no time integration and no differentiation through the solve. A single steady-state solve and a single dF/dy factorisation are reused for every output and parameter, so this is far cheaper than jacfwd/jacrev through steady_state() (which re-solves per call).

The two AD directions are the two ways to read the same sensitivity: forward mode (one solve per parameter, all outputs follow) is efficient when the outputs outnumber the parameters; reverse mode (one transposed solve plus a vector–Jacobian product per output, all parameters follow) is efficient when the parameters outnumber the outputs.

Parameters:

params (jnp.ndarray, optional) – Plant parameters (defaults to default_parameters()).
y0 (jnp.ndarray, optional) – Warm start for the steady-state solve.
state (jnp.ndarray, optional) – A pre-solved steady state to evaluate the sensitivity at, skipping the internal solve. Use this when the operating point is already known (e.g. a warm-start steady state) and to read the sensitivity in several ways without re-solving each time. Must satisfy F(state, params) = 0.
output_fn (callable, optional) – Maps the flat plant state (total_state_size,) to a length-m vector of scalar outputs. Defaults to the identity (the full state, so the result is the (n_states, n_params) sensitivity dy*/dtheta).
wrt (sequence of int or str, optional) – The parameters to differentiate with respect to – flat indices or "<model>.<param>" names (resolved by parameter_index()). Defaults to all parameters. Restricting to a subset of k parameters makes forward mode cost k solves (one per chosen parameter) rather than n_params; reverse mode returns all parameters from the per-output solve regardless and simply selects this subset.
mode ({"auto", "forward", "reverse"}) – AD direction. "auto" picks "forward" when the number of selected parameters <= m and "reverse" otherwise. Both give the same exact sensitivity.
elasticity (bool) – If True return the dimensionless elasticity (dg/dtheta)(theta/g) instead of the raw derivative.
return_jacobian (bool) – If True also return the steady-state Jacobian dF/dy (so a caller can assess its conditioning without recomputing it, as steady_state_dgsm() does to flag near-singular operating points).
**steady_kwargs – Forwarded to steady_state() (e.g. max_iter).
operating (Sequence | None)

Returns:

(m, k) sensitivity of the m outputs to the k selected parameters (k == n_params when wrt is None).

Return type:

jnp.ndarray

Notes

Exact when the steady Jacobian dF/dy is full rank (true for the shipped models at their operating point; see solve_steady_state()).

steady_state_dgsm(ranges, *, output_fn, output_names=None, wrt=None, y0=None, n_samples=256, seed=0, mode='auto', cond_factor=None, input_dist='uniform', input_transforms=None, continuation=True, progress=None, **steady_kwargs)#

Derivative-based global sensitivity (DGSM) of the plant steady state.

Samples the screened parameters over their ranges (scrambled-Sobol QMC), solves the steady state at each sample, and reads each output’s sensitivity to those parameters through the implicit-function-theorem helper (steady_state_sensitivity()) – reusing one dF/dy factorisation per sample. That makes it markedly cheaper than a generic dgsm() over steady_state (whose jacfwd/jacrev recompute the steady-state structure for every input tangent / output). Aggregates to the Sobol total-index upper bound per (output, parameter); the retained per-sample data lets SteadyStateDGSMResult.convergence() report the sample-size study without re-solving.

Parameters:

ranges (array-like, shape (k, 2)) – [a_j, b_j] sampling range for each screened parameter, aligned with wrt. For input_dist="uniform" (default) inputs are sampled uniformly within. For input_dist="normal" the range is read as the +/-2 sigma band of the input’s prior in its calibration-transform space, i.e. [a_j, b_j] = [t^{-1}(m_j - 2 sigma_j), t^{-1}(m_j + 2 sigma_j)], so m_j = t(nominal_j) and sigma_j = (t(b_j) - t(a_j))/4 with t the parameter’s transform.
output_fn (callable) – Maps the flat plant state to a length-m vector of scalar outputs.
output_names (sequence of str, optional) – Names for the m outputs (default "output0" …).
wrt (sequence of int or str, optional) – The screened parameters – flat indices or "<model>.<param>" names (default: all parameters; usually pass an explicit subset).
y0 (jnp.ndarray, optional) – Warm start for each steady-state solve.
n_samples (int, optional) – Sobol points (rounded to a power of two). Increase until SteadyStateDGSMResult.convergence() flattens.
seed (int, optional) – Scrambled-Sobol seed (fixing it makes the screen reproducible).
mode ({"auto", "forward", "reverse"}, optional) – AD direction for the per-sample sensitivity (passed to steady_state_sensitivity()). "auto" is reverse for the usual many-parameter / few-output screen.
cond_factor (float, optional) – Drop any sample whose steady-state Jacobian dF/dy is more than cond_factor times as ill-conditioned as the sample median – a near-singular (near-bifurcation) operating point, where the sensitivity is not well-defined and would give the DGSM a heavy tail. None (default) keeps every sample (then the bounds match dgsm() exactly). A value around 1e2 removes the marginal-stability outliers of a stiff plant; the dropped count is reported in n_valid.
input_dist ({"uniform", "normal"}, optional) – Input sampling distribution. "uniform" (default) draws uniformly over ranges and bounds the Sobol total index with the uniform Poincare constant (b_j-a_j)^2/pi^2. "normal" draws each input from its Gaussian prior in the calibration-transform space (a log-normal for a positive rate, logit-normal for a fraction), reads the sensitivity in that space (chain-ruling dg/dz = (dg/dtheta) (dtheta/dz)), and uses the Gaussian Poincare constant sigma_j^2 – the established generalization of the DGSM bound to non-uniform inputs (Sobol & Kucherenko 2010, Sec. 8; Lamboni et al. 2013, Thm 3.1). It samples the prior we believe rather than a box, so an improbable corner is weighted by its prior density. Requires input_transforms.
input_transforms (sequence of str, optional) – Required for input_dist="normal": the calibration transform of each screened parameter ("positive_log" / "logit" / "none"), aligned with wrt – the space in which its prior is Gaussian. Ignored for "uniform".
progress (int, optional) – If set, emit a logging.INFO progress record every progress samples (on the aquakin.plant.sensitivity logger). It is silent unless the caller enables INFO logging, e.g. logging.basicConfig(level=logging.INFO).
**steady_kwargs – Forwarded to steady_state() (e.g. max_iter).
continuation (bool)

Returns:

The (m, k) Sobol total-index bounds, standard errors, and the retained per-sample data.

Return type:

SteadyStateDGSMResult

solve_sensitivity(params, wrt, *, operating=None, t_span, t_eval=None, y0=None, rtol=1e-06, atol=None, sens_rtol=None, factormax=None, dtmax=None, max_steps=1000000)#

Stable forward (variational) sensitivity dy/dtheta of the plant.

Integrates the augmented system z = [y; S] (the state plus its sensitivity S = dy/dtheta) with the step controller’s error norm bounding S directly, so the sensitivity stays finite over long horizons where jacfwd through the stiff plant solve goes non-finite. This is the plant counterpart of the reactors’ solve_sensitivity – the forward-mode sensitivity tool the dynamic plant lacked.

The augmented [y; S] solve uses the SAME enhanced solver config the plant’s forward / discrete-adjoint solves use – a decoupled Newton root finder, the factormax step-growth cap, and the cached recycle / flow maps – with a lean Kvaerno3 base solver and the block-arrow SimultaneousCorrector for the per-stage linear algebra (factor the shared diagonal block D = I - gamma.dt.J once, reuse it across every sensitivity column).

Parameters:

params (jnp.ndarray) – Plant parameter vector.
wrt (sequence of str or int) – The sensitivity parameters, as "<model>.<param>" names or flat indices. The cached recycle map drops the dM/dtheta term, so this is exact for kinetic parameters; a flow-setpoint sensitivity would need the per-call map (not wired here).
operating (sequence of dict, optional) – Operating-parameter sensitivities computed alongside wrt, each a differentiable multiplicative scale on an influent (nominal 1.0, so the sensitivity is d output / d(scale) at the recorded influent). Each spec is {"kind": "influent_flow", "port": p} (scale that influent’s flow) or {"kind": "influent_concentration", "port": p, "species": s} (scale that species’ load). They append columns to S after the wrt columns, in order. Exact through the variational solve – the cached recycle/flow maps are influent-independent, so no term is dropped (unlike a flow setpoint).
t_span (tuple of float) – (t0, t1) integration interval (plant time units).
t_eval (jnp.ndarray, optional) – Times to record. None records the endpoint only.
y0 (jnp.ndarray, optional) – Initial state (defaults to initial_state()); warm-start a stiff plant.
rtol (float or array, optional) – State step tolerances (atol defaults to the per-component plant floor).
atol (float or array, optional) – State step tolerances (atol defaults to the per-component plant floor).
sens_rtol (float, optional) – Relative tolerance on S (defaults to rtol).
factormax (float, optional) – Step-growth cap and maximum step, passed to the augmented solve.
dtmax (float, optional) – Step-growth cap and maximum step, passed to the augmented solve.
max_steps (int) – Step budget for the augmented solve.

Returns:

ts (jnp.ndarray) – Save times, shape (n_t,).
ys (jnp.ndarray) – State trajectory, shape (n_t, ndof) (split with states_by_unit()).
S (jnp.ndarray) – Sensitivity dy/dtheta, shape (n_t, ndof, k) for the k wrt parameters.

Notes

The block-arrow SimultaneousCorrector exploits the augmented system’s structure (each sensitivity column couples only to y); it is specific to this [y; S] solve and does not apply to the single-state forward / adjoint solves, whose per-step lever is the colored Jacobian.

dynamic_sensitivity(params=None, *, output_fn, t_span, t_eval=None, wrt=None, operating=None, mode='reverse', y0=None, elasticity=False, **solve_kwargs)#

Sensitivity of a transient (time-window) output to the parameters.

The sensitivity of a transient output – the dynamic counterpart of steady_state_sensitivity(), for outputs that depend on the trajectory rather than the operating point (an effluent time series, a window average, a peak). Unlike the steady-state case there is no implicit-function-theorem shortcut, so the cost is one stiff solve per direction; this wrapper’s value is that it uses the stable method for each AD direction (the footgun it removes – a naive differentiation of solve() is non-finite on a stiff plant). Reverse mode differentiates the solve through the cap-free stable discrete adjoint. Forward mode integrates the augmented [y; S] variational system (solve_sensitivity()), whose step controller bounds S so it stays finite over long horizons where forward-mode jacfwd through the solve goes non-finite, then chains the full-state sensitivity through output_fn. Forward is also the memory-light direction over a long horizon: solve_sensitivity carries the parameter tangents in lockstep with the state, so its memory is independent of the integration length, whereas reverse mode stores the trajectory to replay it. Neither is wrapped in an enclosing jit, so the primal runs with concrete parameters (which some plant setup requires); the solve’s compiled-solve cache reuses the integrator compile across calls.

Parameters:

params (jnp.ndarray, optional) – Plant parameters (defaults to default_parameters()).
output_fn (callable) – Maps the PlantSolution to a length-m vector of scalar outputs (e.g. lambda sol: sol.C_named("tank5", "SNO") for the effluent-nitrate trajectory, or a window average of it).
t_span (tuple) – Integration interval and the times at which the trajectory is saved (passed straight to solve()).
t_eval (Array | None) – Integration interval and the times at which the trajectory is saved (passed straight to solve()).
wrt (sequence of int or str, optional) – Parameters to differentiate (flat indices or "<model>.<param>" names; default all).
mode ({"reverse", "forward", "auto"}) – AD direction. "reverse" (default; "auto" resolves to it) suits many parameters / few outputs and uses the stable discrete adjoint; "forward" suits few parameters / many outputs and uses the augmented variational solve (solve_sensitivity()) – finite and memory-light over long horizons. In forward mode **solve_kwargs are forwarded to solve_sensitivity() (max_steps, dtmax, factormax, rtol, atol, sens_rtol), not to solve().
y0 (jnp.ndarray, optional) – Initial plant state (e.g. a warm-start steady state).
elasticity (bool) – If True return the dimensionless (dg/dtheta)(theta/g).
**solve_kwargs – Forwarded to solve() (e.g. max_steps).
operating (Sequence | None)

Returns:

(m, k) sensitivity of the m outputs to the k parameters.

Return type:

jnp.ndarray

dynamic_dgsm(ranges, *, output_fn, t_span, t_eval=None, wrt=None, mode='reverse', y0=None, n_samples=256, seed=0, output_names=None, progress=None, **solve_kwargs)#

Derivative-based global sensitivity (DGSM) of a transient plant output.

The dynamic counterpart of steady_state_dgsm(): scrambled-Sobol QMC over the screened parameters, each sample’s sensitivity from dynamic_sensitivity() (differentiated through the dynamic solve, with the adjoint selected by mode), aggregated to the Sobol total-index upper bound. The solve’s compiled-solve cache amortizes the integrator compile across samples, but each sample is a full differentiated solve, so a dynamic screen is far heavier per sample than the steady-state one; use a modest n_samples / parameter count.

Parameters mirror steady_state_dgsm() (ranges aligned with wrt, output_fn mapping the PlantSolution to m outputs) plus t_span / t_eval for the window. Returns a DynamicDGSMResult (sobol_total_bound / std_error shape (m, k), .ranked(output), .convergence()).

Parameters:

output_fn (Callable)
t_span (tuple)
t_eval (Array | None)
wrt (Sequence | None)
mode (str)
y0 (Array | None)
n_samples (int)
seed (int)
output_names (Sequence | None)
progress (int | None)

Return type:

DynamicDGSMResult

calibrate(observations, t_obs, free_params, *, target='effluent', observed_channels=None, observables=None, t_span=None, y0=None, params=None, transforms=None, free_ic=None, time_unit=None, loss='mse', sigma=None, priors=None, use_priors=True, optimizer=OptimizerConfig(method='lbfgsb', n_starts=1, jitter=0.5, jitter_schedule=None, seed=0, max_iter=500, tol=1e-06, param_halfwidth=None), laplace=False, check_finite=True, integrator=None, diff=DifferentiationConfig(mode='reverse', method='stable', check_finite=True, adjoint_max_steps=100000, adjoint_low_memory=False))#

MAP-calibrate a plant’s parameters against an output stream.

Bound onto Plant as plant.calibrate(...). The plant analogue of aquakin.calibrate(): it fits plant parameters (by "<model>.<param>" name – see Plant.parameter_names()) so a target stream’s channels match observations. The forward solve is the cap-free stable adjoint by default, so a reverse-mode gradient through a stiff plant is finite with no dtmax to tune.

Parameters:

plant (Plant) – The plant to calibrate. It must already have its influent(s) added.
observations (array-like or list of array-like) – Observed values, shape (n_t,) for a single channel or (n_t, n_channels). Pass a list of such arrays for a joint multi-batch fit: each entry is one run of the plant, the batches share the parameter vector and prior and their data terms are summed. t_obs and y0 must then be matching lists (one per batch).
t_obs (array-like or list of array-like) – Observation times, shape (n_t,), in the plant’s time unit (or time_unit if given). The solve integrates over t_span and reports at t_obs. A list in multi-batch mode (the batches may differ in n_t).
free_params (list of str) – Plant parameter names to calibrate ("<model>.<param>"). Others fixed.
target (str, optional) – The single output stream to compare against – a registered stream name ("effluent", "ras", …; see Plant.list_streams()) or a "unit.port" / unit name. Default "effluent". Ignored when observables is given.
observed_channels (list of str, optional) – Species of the target stream’s model that observations columns correspond to. None observes every stream species. Ignored when observables is given.
observables (list, optional) – Fit against several streams at once. Each entry is a PlantObservable (stream + channels), a {"stream": ..., "channels": ...} dict, a (stream, channels) pair, or a bare stream name. The observations columns then run in observable order, channels within a stream first – e.g. observables=[PlantObservable("effluent", ["SNH", "SNO"]), PlantObservable("wastage", ["XS"])] expects 3 columns. Overrides target / observed_channels.
t_span (tuple or list of tuple, optional) – (t0, t1) integration window. Defaults to (t_obs[0], t_obs[-1]). A list of windows (one per batch) in multi-batch mode; a single window broadcasts to every batch.
y0 (jnp.ndarray or list of jnp.ndarray, optional) – Warm-start plant state (e.g. bsm2_warm_start(plant) or a saved steady state). Strongly recommended for a stiff plant. In multi-batch mode this is required and must be a list of initial states, one per batch (the batches differ in their initial state).
params (jnp.ndarray, optional) – Starting parameter vector. Defaults to Plant.default_parameters().
transforms (dict, optional) – Per-parameter transform override ("positive_log" / "logit" / "none"). Unspecified free params fall back to the parameter’s model-declared transform.
free_ic (FreeICConfig, optional) – Assembled-state slots to fit alongside the parameters. species are "unit.species" strings (or (unit, species) pairs) naming an initial concentration of a concentration unit (a CSTR, the digester); they are fit in log space, box-bounded by FreeICConfig.bounds, with an optional log-space prior via prior_log_std pulling each toward its starting value (read from y0 or the plant’s default initial state). The fitted state is returned as result.C0_fitted[0] and the pools as result.ic_named[0]. Default None (no free ICs).
time_unit (str, optional) – Unit t_obs / t_span are expressed in; passed to plant.solve.
loss (str) – As in aquakin.calibrate() (the shared machinery): optimizer is an OptimizerConfig and laplace a bool or LaplaceConfig. laplace defaults to False here (a plant Hessian is expensive). OptimizerConfig.param_halfwidth is not used by plant fits.
sigma (jnp.ndarray | None) – As in aquakin.calibrate() (the shared machinery): optimizer is an OptimizerConfig and laplace a bool or LaplaceConfig. laplace defaults to False here (a plant Hessian is expensive). OptimizerConfig.param_halfwidth is not used by plant fits.
priors (dict | None) – As in aquakin.calibrate() (the shared machinery): optimizer is an OptimizerConfig and laplace a bool or LaplaceConfig. laplace defaults to False here (a plant Hessian is expensive). OptimizerConfig.param_halfwidth is not used by plant fits.
use_priors (bool) – As in aquakin.calibrate() (the shared machinery): optimizer is an OptimizerConfig and laplace a bool or LaplaceConfig. laplace defaults to False here (a plant Hessian is expensive). OptimizerConfig.param_halfwidth is not used by plant fits.
optimizer (OptimizerConfig) – As in aquakin.calibrate() (the shared machinery): optimizer is an OptimizerConfig and laplace a bool or LaplaceConfig. laplace defaults to False here (a plant Hessian is expensive). OptimizerConfig.param_halfwidth is not used by plant fits.
laplace (bool | LaplaceConfig) – As in aquakin.calibrate() (the shared machinery): optimizer is an OptimizerConfig and laplace a bool or LaplaceConfig. laplace defaults to False here (a plant Hessian is expensive). OptimizerConfig.param_halfwidth is not used by plant fits.
check_finite (bool) – As in aquakin.calibrate() (the shared machinery): optimizer is an OptimizerConfig and laplace a bool or LaplaceConfig. laplace defaults to False here (a plant Hessian is expensive). OptimizerConfig.param_halfwidth is not used by plant fits.
integrator (IntegratorConfig, optional) – Plant integrator configuration passed to plant.solve.
diff (DifferentiationConfig, optional) – How the gradient flows through plant.solve. Default mode='reverse', method='stable' – the cap-free discrete adjoint.

Returns:

Same result type as aquakin.calibrate(). predictive_band (which takes a reactor) does not apply to a plant fit.

Return type:

CalibrationResult

Notes

Fits kinetic parameters (and, optionally, assembled-state initial conditions via free_ic) against one or more output streams, over one run of the plant or several joined in a multi-batch fit (pass list-valued observations / t_obs / y0). free_ic and multi-batch are not yet combinable in one call.