aquakin.CompiledModel#

class aquakin.CompiledModel(name, description, references, species, parameters, conditions_required, stoich_matrix, reaction_names, rate_callables, rate_asts, param_index, species_index, _default_concentrations, _default_parameters, _condition_defaults=<factory>, species_units=<factory>, species_descriptions=<factory>, species_composition=<factory>, parameter_units=<factory>, condition_units=<factory>, parameter_bounds=<factory>, parameter_transforms=<factory>, parameter_priors=<factory>, stoich_dynamic=<factory>, _stoich_dynamic_rows=None, _stoich_dynamic_cols=None, derived_condition_fn=None, derived_fields=<factory>, precipitation_equilibrium_fn=None, positivity_threshold=None, clip_negative_states=False, temperature_corrections=<factory>)[source]#

Bases: object

Runtime representation of a reaction model.

Variables:

name (str) – Model identifier (e.g. "ozone_bromate").
description (str) – Free-text description of the model.
references (list[str]) – Literature citations associated with the model.
species (list[str]) – Ordered species names. Index in this list is the index used in C.
parameters (list[str]) – Ordered, namespaced parameter names (e.g. "O3_Br_direct.k1"). Index in this list is the index used in params.
conditions_required (list[str]) – Names of condition fields the model needs at runtime.
stoich_matrix (jnp.ndarray) – Shape (n_reactions, n_species). Stoichiometric coefficient of species j in reaction i.
reaction_names (list[str]) – Ordered reaction names corresponding to rows of stoich_matrix.
rate_callables (list[Callable]) – Per-reaction compiled rate functions with the canonical signature (C, params, condition_arrays, loc_idx) -> scalar.
rate_asts (list[ASTNode]) – Per-reaction parsed AST roots. Retained for inspection (to_latex) and not used in the runtime hot path.
param_index (dict[str, int]) – Map from namespaced parameter name to its position in params.
species_index (dict[str, int]) – Map from species name to its position in C.
_default_concentrations (jnp.ndarray) – Default initial concentrations, shape (n_species,).
_default_parameters (jnp.ndarray) – Default parameter values, shape (n_params,).
parameter_bounds (dict[str, tuple[float, float]]) – (low, high) bounds per namespaced parameter name. Parameters without declared bounds are absent from the mapping (no None sentinel values).

Parameters:

name (str)
description (str)
references (list[str])
species (list[str])
parameters (list[str])
conditions_required (list[str])
stoich_matrix (jnp.ndarray)
reaction_names (list[str])
rate_callables (list[RateCallable])
rate_asts (list[ASTNode])
param_index (dict[str, int])
species_index (dict[str, int])
_default_concentrations (jnp.ndarray)
_default_parameters (jnp.ndarray)
_condition_defaults (dict[str, float])
species_units (dict[str, str])
species_descriptions (dict[str, str])
species_composition (dict[str, dict[str, float]])
parameter_units (dict[str, str])
condition_units (dict[str, str])
parameter_bounds (dict[str, tuple[float, float]])
parameter_transforms (dict[str, str])
parameter_priors (dict[str, tuple[float, float]])
stoich_dynamic (list[tuple[int, int, Callable]])
_stoich_dynamic_rows (jnp.ndarray | None)
_stoich_dynamic_cols (jnp.ndarray | None)
derived_condition_fn (Callable | None)
derived_fields (list[str])
precipitation_equilibrium_fn (Callable | None)
positivity_threshold (float | None)
clip_negative_states (bool)
temperature_corrections (list)

__init__(name, description, references, species, parameters, conditions_required, stoich_matrix, reaction_names, rate_callables, rate_asts, param_index, species_index, _default_concentrations, _default_parameters, _condition_defaults=<factory>, species_units=<factory>, species_descriptions=<factory>, species_composition=<factory>, parameter_units=<factory>, condition_units=<factory>, parameter_bounds=<factory>, parameter_transforms=<factory>, parameter_priors=<factory>, stoich_dynamic=<factory>, _stoich_dynamic_rows=None, _stoich_dynamic_cols=None, derived_condition_fn=None, derived_fields=<factory>, precipitation_equilibrium_fn=None, positivity_threshold=None, clip_negative_states=False, temperature_corrections=<factory>)#

Parameters:

name (str)
description (str)
references (list[str])
species (list[str])
parameters (list[str])
conditions_required (list[str])
stoich_matrix (jnp.ndarray)
reaction_names (list[str])
rate_callables (list[RateCallable])
rate_asts (list[ASTNode])
param_index (dict[str, int])
species_index (dict[str, int])
_default_concentrations (jnp.ndarray)
_default_parameters (jnp.ndarray)
_condition_defaults (dict[str, float])
species_units (dict[str, str])
species_descriptions (dict[str, str])
species_composition (dict[str, dict[str, float]])
parameter_units (dict[str, str])
condition_units (dict[str, str])
parameter_bounds (dict[str, tuple[float, float]])
parameter_transforms (dict[str, str])
parameter_priors (dict[str, tuple[float, float]])
stoich_dynamic (list[tuple[int, int, Callable]])
_stoich_dynamic_rows (jnp.ndarray | None)
_stoich_dynamic_cols (jnp.ndarray | None)
derived_condition_fn (Callable | None)
derived_fields (list[str])
precipitation_equilibrium_fn (Callable | None)
positivity_threshold (float | None)
clip_negative_states (bool)
temperature_corrections (list)

Return type:

None

Methods

`__init__`(name, description, references, ...)
`atol`([overrides, default])	Per-species absolute-tolerance vector for a reactor's `atol=`.
`check_conservation`(*[, tol, params, ...])	Conservation violations `(reaction, quantity, residual)` above `tol`.
`check_nitrogen`(*[, tol, params, ...])	Nitrogen-balance violations `(reaction, residual)` above `tol`.
`check_units`(*[, check_root])	Check the rate expressions for dimensional ("unit") consistency.
`composition`(*[, params, electron_acceptor_cod])	The per-species conserved-quantity content table for this model.
`compute_stoich`(params)	Evaluate the stoichiometry matrix at the given parameter vector.
`concentrations`([overrides, base])	Initial-concentration vector with named species set.
`condition_defaults`()	The model's declared default condition values as a plain scalar dict.
`dCdt`(C, params, condition_arrays, loc_idx, *)	Return the fully post-processed chemistry RHS for one location.
`default_concentrations`()	Return a copy of the default initial-concentration vector.
`default_conditions`([n_locations])	Build a `SpatialConditions` from the model's declared defaults.
`default_parameters`()	Return a copy of the default parameter vector.
`description_of`(species)	Return the human-readable description of a species.
`influent`([overrides, base, T])	Build a constant-in-time influent stream from a feed composition.
`parameter_values`([overrides])	Parameter vector: defaults with named (namespaced) parameters set.
`precipitation_equilibrium`([C, conditions, ...])	Project a composition onto its mineral precipitation equilibrium.
`rates`(C, params, condition_arrays, loc_idx)	Evaluate all reaction rates at the given state.
`summary`()	Return a human-readable table summarising the model.
`to_latex`()	Return a mapping `reaction_name -> LaTeX rate expression`.
`units_of`(species)	Return the declared concentration units of a species.

Attributes

`clip_negative_states`
`derived_condition_fn`
`n_params`
`n_reactions`
`n_species`
`positivity_threshold`
`precipitation_equilibrium_fn`
`time_unit`	Integration time unit, inferred from the rate-constant units.
`name`
`description`
`references`
`species`
`parameters`
`conditions_required`
`stoich_matrix`
`reaction_names`
`rate_callables`
`rate_asts`
`param_index`
`species_index`
`species_units`
`species_descriptions`
`species_composition`
`parameter_units`
`condition_units`
`parameter_bounds`
`parameter_transforms`
`parameter_priors`
`stoich_dynamic`
`derived_fields`
`temperature_corrections`

default_concentrations()[source]#

Return a copy of the default initial-concentration vector.

This is the model’s reference state – every species at its YAML default_concentration – not a blank slate. For most models that means many species are nonzero (e.g. a full biomass inventory). It is a sensible starting initial condition, but it is the wrong base for a feed composition: an influent species you do not list should be absent, not sitting at its reference value. Build feeds with influent() or concentrations(..., base="zero") instead – see concentrations() for the trap in full.

Return type:: Array

default_parameters()[source]#

Return a copy of the default parameter vector.

Return type:: Array

concentrations(overrides=None, /, *, base='defaults', **kwargs)[source]#

Initial-concentration vector with named species set.

A by-name builder that avoids manual default_concentrations().at[species_index[name]].set(value) chains.

Warning

Building a feed? Use influent() or base="zero". With the default base="defaults" every species you do not list keeps its YAML reference value, so concentrations({"SS": 100.0}) silently carries a full biomass/inert inventory (XB_H, XS, XI, …) into the result. That is correct for an initial condition (the reactor starts from the reference state with a few species adjusted) but wrong for an influent (an unlisted species should be absent). For a feed, pass base="zero" so the vector contains only what you list, or use influent(), which defaults to the zero base.

Parameters:

overrides (dict[str, float], optional) – Species name -> concentration. Positional-only. Use the dict for names that are not valid Python identifiers ("Br-", "BrO3-").
base ({"defaults", "zero"}, optional) – Starting point for unlisted species. "defaults" (the default) keeps each unspecified species at its YAML reference value; "zero" starts every species at 0, so the result contains only what was passed – the correct base for building a feed composition (where an unspecified species means “absent”, not “at its reference value”). A species literally named base must be passed via the overrides dict.
**kwargs (float) – Convenience overrides for identifier-safe species names (O3=1e-4).

Returns:

Concentration vector of shape (n_species,).

Return type:

jnp.ndarray

Examples

>>> model.concentrations({"O3": 1e-4, "Br-": 1e-5})
>>> model.concentrations(SS=50.0)
>>> model.concentrations({"SS": 50.0, "SNH": 25.0}, base="zero")

influent(overrides=None, /, *, Q, base='zero', T=None, **kwargs)[source]#

Build a constant-in-time influent stream from a feed composition.

Convenience for the common “constant feed of known composition” case: a one-call, zero-based InfluentSeries so an unspecified species is absent from the feed rather than sitting at its YAML reference value. The returned series is constant in time, so it can be passed straight to plant.add_influent(...).

Parameters:

overrides (dict[str, float], optional) – Species name -> feed concentration. Positional-only.
Q (float) – Volumetric flow rate of the feed (constant), required.
base ({"zero", "defaults"}, optional) – Composition base, defaulting to "zero" (see concentrations()).
T (float, optional) – Constant feed temperature (Kelvin). None (default) leaves the influent temperature-agnostic.
**kwargs (float) – Convenience overrides for identifier-safe species names.

Returns:

A constant-in-time influent.

Return type:

InfluentSeries

Examples

>>> net.influent({"SS": 60.0, "SNH": 25.0}, Q=18446.0)
>>> net.influent(SS=400.0, Q=2.0)            # carbon dose

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

Parameter vector: defaults with named (namespaced) parameters set.

The parameter analogue of concentrations(). Names are the namespaced keys ("O3_Br_direct.k1"), so the dict form is the usual one; kwargs works for the rare bare model-level parameter.

Examples

>>> model.parameter_values({"O3_Br_direct.k1": 175.0})

Return type:: Array

atol(overrides=None, /, default=None, **kwargs)[source]#

Per-species absolute-tolerance vector for a reactor’s atol=.

default everywhere, with named species overridden – the by-name replacement for jnp.full((n_species,), d).at[species_index[s]].set(v) when a trace species needs a tighter tolerance.

default=None (the default) starts from the per-component noise floor default_atol() scaled off the model’s reference concentrations – the same floor a reactor built with atol=None uses – rather than a fixed scalar, which is ~9 orders too tight for g/m3 ASM/ADM states. Pass an explicit scalar default= for a uniform floor.

Examples

>>> reactor = BatchReactor(net, conds, atol=net.atol({"OH": 1e-20}))
>>> reactor = BatchReactor(net, conds, atol=net.atol({"OH": 1e-20}, default=1e-12))

Return type:: Array

default_conditions(n_locations=1)[source]#

Build a SpatialConditions from the model’s declared defaults.

Convenience for the common case of “use the defaults as written in the YAML”. Each required condition is broadcast to n_locations.

Parameters:: n_locations (int)

condition_defaults()[source]#

The model’s declared default condition values as a plain scalar dict.

The scalar-dict counterpart to default_conditions() (which returns a location-broadcast SpatialConditions): {field: default} for each condition the YAML declares a default for. The public accessor for the raw defaults – e.g. a plant builder seeding a reactor unit’s scalar conditions= from the model’s declared operating point. A copy, so the caller cannot mutate the model’s stored defaults.

Return type:: dict[str, float]

precipitation_equilibrium(C=None, conditions=None, *, loc_idx=0)[source]#

Project a composition onto its mineral precipitation equilibrium.

For a precipitation: model with mode: equilibrium minerals, solve the coupled algebraic equilibrium – every precipitated mineral on its solubility (IAP = Ksp), every absent mineral undersaturated, mass balanced across the shared ions – and return the equilibrium-projected state: each equilibrium solid set to its equilibrium amount and the dissolved ions rebalanced. This is the differentiable, non-stiff alternative to integrating an ultra-insoluble mineral’s kinetics (whose ~1e13 rate Jacobian defeats every sensitivity method): the solve is well conditioned and jax.grad flows through it via the implicit-function-theorem sensitivity, so it composes with sensitivity() / calibrate() w.r.t. the dose (the composition) and the operating conditions (pH, T).

Parameters:

C (jnp.ndarray, optional) – Composition to project, shape (n_species,). Defaults to the model’s default_concentrations().
conditions (SpatialConditions, optional) – Conditions supplying pH / T. Defaults to default_conditions().
loc_idx (int, optional) – Spatial location index into conditions (default 0).

Returns:

The equilibrium-projected composition, shape (n_species,). Read a residual dissolved ion or a solid amount with species_index()-based indexing or a BatchSolution-style accessor.

Return type:

jnp.ndarray

Raises:

ValueError – If the model declares no mode: equilibrium minerals.

rates(C, params, condition_arrays, loc_idx)[source]#

Evaluate all reaction rates at the given state.

Parameters:

C (jnp.ndarray) – Concentration vector, shape (n_species,).
params (jnp.ndarray) – Flat parameter vector, shape (n_params,).
condition_arrays (dict[str, jnp.ndarray]) – Mapping field_name -> (n_locations,) array.
loc_idx (int or jnp.ndarray) – Spatial location index.

Returns:

Reaction rate vector, shape (n_reactions,).

Return type:

jnp.ndarray

compute_stoich(params)[source]#

Evaluate the stoichiometry matrix at the given parameter vector.

For models whose stoichiometry is purely numeric this returns the cached stoich_matrix unchanged. For models with parameter-dependent coefficients (the dynamic entries listed in stoich_dynamic), this scatters the per-call values onto the static base.

Reactors typically call this once per solve() and hoist the result as a closure constant for the duration of the integration — see the rhs builders in aquakin.integrate.batch etc.

Parameters:: params (Array)
Return type:: Array

dCdt(C, params, condition_arrays, loc_idx, *, stoich=None, rate_scale=None)[source]#

Return the fully post-processed chemistry RHS for one location.

This is the single canonical reaction term — stoich.T @ rates(...) with clip_negative_states applied to the rate inputs (inside rates()) and the positivity_limiter applied to the net term. Reactors that add their own transport build their RHS as dCdt(...) + transport rather than re-deriving this, so the clip and limiter are applied identically everywhere.

stoich may be precomputed via compute_stoich() and passed in to avoid re-evaluating parameter-dependent coefficients on every ODE step. If omitted, it is computed from params here.

rate_scale is an optional per-reaction multiplier applied to the rate vector before stoich.T @ r — for a reaction-level throttle such as the biofilm density cap, where uptake and production must scale together (a post-multiply on the net term would not conserve mass).

Parameters:

C (Array)
params (Array)
condition_arrays (dict[str, Array])
stoich (Array | None)
rate_scale (Array | None)

Return type:

Array

units_of(species)[source]#

Return the declared concentration units of a species.

Parameters:: species (str) – Species name.
Returns:: The units string from the YAML species: block (e.g. "g_COD/m3", "g_N/m3", "mol/L").
Return type:: str
Raises:: KeyError – If species is not a declared species.

description_of(species)[source]#

Return the human-readable description of a species.

Parameters:: species (str) – Species name.
Returns:: The description string from the YAML species: block ("" if none was declared).
Return type:: str
Raises:: KeyError – If species is not a declared species.

property time_unit: str | None#

Integration time unit, inferred from the rate-constant units.

aquakin has no global time unit: t_span and t_eval are interpreted in whatever time unit the model’s rate constants are written in, and that differs by model. The chemistry models (ozone, UV/H₂O₂) use seconds (rate constants in M-1 s-1); the biological models (ASM1/2d/3, ADM1, WATS) use days (1/d). So reactor.solve(C0, t_span=(0, 600)) integrates 600 seconds for ozone_bromate but 600 days for asm1 — same code, no warning.

This property recovers that unit by parsing the declared parameter units: strings and reading the inverse-time token the rate constants share (the s in M-1 s-1, the d in 1/d), so a caller can check it before choosing a t_span.

Returns:: The shared inverse-time token ("s", "d", "h" or "min"), or None when it cannot be determined unambiguously — no parameter declares a time unit, or different rate constants disagree.
Return type:: str or None

summary()[source]#

Return a human-readable table summarising the model.

Return type:: str

to_latex()[source]#

Return a mapping reaction_name -> LaTeX rate expression.

Return type:: dict[str, str]

check_units(*, check_root=True)[source]#

Check the rate expressions for dimensional (“unit”) consistency.

A currency-aware dimensional analysis of every rate: expression: it catches a dropped concentration factor, a wrong rate-constant exponent, or a Monod term that compares two different “currencies” (g_COD/m3 vs g_N/m3), which a plain SI dimension check waves through because both are mass/volume. Units are taken from the species, parameter, and condition units: declarations.

The check is advisory: a blank or unparseable unit is treated as unknown and skipped, so an empty result is “no inconsistency among the declared, parseable units”, not a proof of correctness. Stoichiometry (a conservation question) is out of scope – use aquakin.check_conservation() for that.

Parameters:: check_root (bool, default True) – Also assert each rate resolves to currency / volume / time (e.g. g_COD/m3/d or mol/L/s). Set False to run only the local operand- and Monod-consistency rules.
Returns:: One entry per finding, as (reaction, location, detail) named tuples (empty when nothing is flagged).
Return type:: list of aquakin.utils.units.UnitWarning

Examples

>>> net = aquakin.load_model("asm1")
>>> for w in net.check_units():
...     print(w)

composition(*, params=None, electron_acceptor_cod=True)[source]#

The per-species conserved-quantity content table for this model.

Returns {species: {quantity: content}} – the content of each conserved quantity (COD / N / P / S / Fe / …) per unit of the species’ own measure. This is the table a conservation check dots against the stoichiometry (check_conservation()) and that a results-level balance dots against concentrations.

Resolution order:

the model’s own declared species[].composition metadata, if any;
otherwise the shipped role-based table (aquakin.composition_table()) for the ASM / ADM families, which reads the model’s composition parameters (so a calibrated N / P fraction flows through);
otherwise an empty table (no metadata available).

Parameters:

params (array-like, optional) – Parameter vector passed to the shipped role-based fallback so the table tracks a calibrated / run composition. Ignored when the model declares its own (literal) composition: metadata.
electron_acceptor_cod (bool, default True) – Passed to the shipped role-based fallback (False selects the lab-COD convention; see aquakin.composition_table()). Ignored for declared metadata.

Return type:

dict[str, dict[str, float]]

check_conservation(*, tol=0.01, params=None, quantities=None, composition=None, electron_acceptor_cod=True)[source]#

Conservation violations (reaction, quantity, residual) above tol.

Dots each reaction’s stoichiometry against the per-species composition() table; a quantity whose stoichiometry-weighted content does not sum to zero (beyond tol) is reported. Catches the conservation-determined-coefficient errors that are otherwise easy to miss – a wrong electron-acceptor (O2 / NO3) demand breaks the COD balance, a wrong product split breaks an elemental (S / N / P / Fe) balance.

Advisory and opt-in, like check_units(): it never runs at load and never raises on a violation – it returns the list for you to inspect. Restrict to specific quantities (e.g. ["COD"]) if desired, or pass an explicit composition to override the model’s table.

Raises ValueError only if no composition table is available (the model declares none and there is no shipped fallback).

Parameters:

tol (float)
electron_acceptor_cod (bool)

Return type:

list

check_nitrogen(*, tol=0.01, params=None, composition=None, nitrate='S_NO', n_key='N')[source]#

Nitrogen-balance violations (reaction, residual) above tol.

The nitrogen analogue of check_conservation(), accounting for the nitrate reduced to (untracked) N2 gas: a reaction conserves nitrogen when its tracked-species N content plus the gassed-off nitrate is zero. Exact for both nitrification (no nitrate consumed) and denitrification. Uses the model’s composition() table unless composition is passed.

Parameters:

tol (float)
nitrate (str)
n_key (str)

Return type:

list