Metadata-Version: 2.4
Name: aquakin
Version: 0.1.0
Summary: Reactive scalar transport kinetics for aqueous environmental systems
Author-email: Christopher DeGroot <christopher.degroot@uwo.ca>
License-Expression: MIT
Project-URL: Homepage, https://github.com/DeGrootResearchGroup/aquakin
Project-URL: Documentation, https://aquakin.readthedocs.io
Project-URL: Repository, https://github.com/DeGrootResearchGroup/aquakin
Project-URL: Issues, https://github.com/DeGrootResearchGroup/aquakin/issues
Project-URL: Changelog, https://github.com/DeGrootResearchGroup/aquakin/blob/main/CHANGELOG.md
Keywords: chemistry,kinetics,ozone,bromate,jax,diffrax
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Scientific/Engineering :: Chemistry
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: jax>=0.10
Requires-Dist: jaxlib>=0.10
Requires-Dist: diffrax>=0.7
Requires-Dist: optimistix>=0.1
Requires-Dist: pydantic>=2.0
Requires-Dist: pyyaml>=6.0
Requires-Dist: numpy>=1.24
Requires-Dist: scipy>=1.10
Provides-Extra: dataframe
Requires-Dist: pandas>=2.0; extra == "dataframe"
Provides-Extra: plot
Requires-Dist: matplotlib>=3.5; extra == "plot"
Provides-Extra: lint
Requires-Dist: ruff<0.16,>=0.15; extra == "lint"
Provides-Extra: docs
Requires-Dist: sphinx>=8.0; extra == "docs"
Requires-Dist: myst-parser>=4.0; extra == "docs"
Requires-Dist: linkify-it-py>=2.0; extra == "docs"
Requires-Dist: pydata-sphinx-theme>=0.16; extra == "docs"
Requires-Dist: sphinx-copybutton>=0.5; extra == "docs"
Provides-Extra: test
Requires-Dist: pytest>=7; extra == "test"
Requires-Dist: pytest-xdist>=3; extra == "test"
Requires-Dist: pytest-timeout>=2; extra == "test"
Requires-Dist: pytest-split>=0.8; extra == "test"
Requires-Dist: pytest-cov>=5; extra == "test"
Requires-Dist: aquakin[dataframe,plot]; extra == "test"
Dynamic: license-file

# aquakin

[![CI](https://github.com/DeGrootResearchGroup/aquakin/actions/workflows/ci.yml/badge.svg)](https://github.com/DeGrootResearchGroup/aquakin/actions/workflows/ci.yml)
[![codecov](https://codecov.io/gh/DeGrootResearchGroup/aquakin/branch/main/graph/badge.svg)](https://codecov.io/gh/DeGrootResearchGroup/aquakin)
[![Docs](https://readthedocs.org/projects/aquakin/badge/?version=latest)](https://aquakin.readthedocs.io/en/latest/)

`aquakin` is a Python library for modelling reactive scalar transport in
aqueous environmental systems. Reaction models are declared at runtime in
YAML and compiled to JAX-native, automatic-differentiable rate functions
integrated with [Diffrax](https://github.com/patrick-kidger/diffrax).

Shipped models span chemistry (ozonation/bromate after Acero & von Gunten,
2001; UV/H₂O₂) and biology (the ASM activated-sludge family, including a
two-step nitrification/denitrification variant with explicit nitrite, a
two-pathway AOB nitrous-oxide (N₂O) model, an anammox / deammonification
variant, and a comammox complete-nitrifier variant; ADM1 anaerobic
digestion in its BSM2 form, with gas headspace; the WATS sewer-process models
`wats_sewer_extended` and the paper-faithful `wats_sewer_khalil_paper`, the
latter with structural variants for model-structure studies) as well as
chemistry coupled to acid-base speciation (a charge-balance state-derived pH,
and SI-driven mineral precipitation/dissolution — struvite + calcite, and
iron/aluminium chemical-phosphorus removal — after Kazadi Mbamba et al. 2015 and
Flores-Alsina et al. 2016. A very insoluble mineral's stiff kinetics defeat every
sensitivity method, so two opt-in differentiable variants are provided: an
**algebraic equilibrium** mode that solves `IAP = Ksp` directly
(`model.precipitation_equilibrium(...)`) and a **bounded-driver** kinetic form
for differentiable dynamics). The model
YAML files live under `aquakin/models/`; see `CLAUDE.md` for the full list.
Future models include UV/TiO₂ and chlorine decay.

## Features

- Reaction models declared in YAML — no recompilation required.
- Full automatic differentiation through `solve()` via JAX.
- JAX-native stiff ODE integration via Diffrax (`Kvaerno5` by default).
- Safe rate expression evaluation via a custom AST (no `eval()`).
- Decoupled transport / reaction operator splitting at all scales (0D, 1D, 3D).
- Reactors for batch (0D), plug flow (1D), Lagrangian particle tracks, and a
  layered biofilm (`BiofilmReactor`: 1-D diffusion-reaction over biofilm depth,
  for penetration-controlled processes).
- Plant-wide flowsheets (`aquakin.plant`): the IWA benchmark plants BSM1 and
  BSM2 — reactors, clarifiers, mixers/splitters and an ADM1 digester integrated
  under one monolithic solve, with run-to-steady-state, a fast differentiable
  algebraic steady-state solver (`plant.steady_state`, pseudo-transient
  continuation — ~10× faster than integrating to settle), dynamic influents,
  EQI/OCI performance metrics, and GHG (N₂O / CO₂e) + monetised-cost reporting
  with standardized scenario-comparison KPI tables. Includes an `IFASUnit` / 
  `MBBRUnit` (a CSTR bulk coupled to a depth-resolved attached biofilm) for 
  modelling MBBR/IFAS intensification retrofits.
- Located events / discontinuities (`Event` + `solve(events=...)`): time events
  and state root-crossings with exact state resets / mode switches (on/off pumps,
  SBR phases, dosing on/off, level limits) — time-scheduled events keep
  `jax.grad` finite.
- Full automatic differentiation everywhere, including cap-free forward
  sensitivity and reverse-mode gradients through stiff plant solves (see
  [Advanced: differentiation & sensitivity](#advanced-differentiation--sensitivity)).

## Installation

```bash
pip install -e ".[test]"
```

> **⚠️ `import aquakin` enables JAX 64-bit (x64) mode process-wide.** The stiff
> implicit ODE solves need double precision, so at import aquakin runs
> `jax.config.update("jax_enable_x64", True)` — which is **global JAX state**.
> Any other JAX code in the same process will use float64 afterward (more memory,
> different numerics). This is required, not optional. If you are co-running JAX
> code that needs float32, run aquakin in a separate process. aquakin emits a
> one-time warning if it overrides an explicit float32 preference (JAX already
> imported, or `JAX_ENABLE_X64` set off), so the side effect is never silent.

## Quickstart

```python
import jax.numpy as jnp
import aquakin

model = aquakin.load_model("asm1")   # Activated Sludge Model No. 1 (IAWQ)

# 0-D (a single well-mixed tank): start from the model's declared condition
# defaults (ASM1 runs at a temperature T) and change only what differs.
conditions = model.default_conditions()                   # YAML defaults (T = 20 C)
# conditions = model.default_conditions().with_(T=283.15)   # ...or run it at 10 C
# (OperatingConditions(T=293.15) is the scalar 0-D shorthand; use SpatialConditions
#  for a spatially varying PFR/CFD case.)

reactor = aquakin.BatchReactor(model, conditions)

# Build the initial state by name -- no .at[species_index[...]].set() chains.
# A simple aerobic batch: activated-sludge biomass + substrate + ammonia.
# (A dict, since some ASM species names aren't valid kwargs; rest = YAML defaults.)
C0 = model.concentrations({
    "SS": 60.0, "SNH": 25.0, "XB_H": 500.0, "XB_A": 80.0, "SO": 2.0})

# For a FEED composition use base="zero" (or model.influent): unlisted species
# are absent, not silently left at their YAML reference value.
feed = model.concentrations({"SS": 60.0, "SNH": 25.0}, base="zero")
influent = model.influent({"SS": 60.0, "SNH": 25.0}, Q=18446.0)   # InfluentSeries

# Or characterize an influent from lab measurements (total COD, TKN, ammonia,
# alkalinity, ...): the SUMO-style fractionation splits them into the ASM1 states.
influent = aquakin.characterize_influent(model, flow=24000.0, total_cod=420.0,
                                         tkn=34.4, ammonia=24.0, alkalinity=330.0)
# A lab/SCADA CSV with arbitrary headers maps + fractionates per row, no renaming:
#   aquakin.read_influent_csv("plant_log.csv", model,
#       column_map={"t": "day", "Q": "flow_m3d", "total_cod": "COD",
#                   "tkn": "TKN", "ammonia": "NH4-N", "alkalinity": "Alk"})

# There is NO global time unit: t_span / t_eval are in whatever unit the
# model's rate constants use, and it differs by model -- ozone/UV are in
# SECONDS (M-1 s-1), the biological models (ASM/ADM/WATS) in DAYS (1/d). Check
# it before choosing a span:
model.time_unit                    # "d" for asm1 (the ozone/UV models are in "s")

# ...or pass time_unit= to work in a unit of your choice: the input times are
# converted to the model's native unit for the solve and solution.t comes back
# in the unit you asked for (solution.time_unit reports it). Works the same on
# BatchReactor / BiofilmReactor / Plant.solve. e.g. an ASM run in hours:
#   sol = reactor.solve(C0, t_span=(0.0, 48.0), t_eval=..., time_unit="h")

# params is optional and defaults to model.default_parameters().
t_eval = jnp.linspace(0.0, 1.0, 121)        # one day, in days (asm1's native unit)
solution = reactor.solve(C0, t_span=(0.0, 1.0), t_eval=t_eval)

print("[SNH] at t=1 d:", float(solution.C_named("SNH")[-1]))   # final effluent ammonia

# Reporting last-point values without the per-species [-1] slice:
solution.final_named(["SS", "SNH", "SNO"])  # {name: float} at the final time (None = all)
solution.final                          # == final_named(): every species' last value
solution.C_named_many(["SNH", "SNO"])   # several full trajectories -> {name: array}

# Species units and descriptions are carried from the YAML to results, so you
# never have to re-derive units by string-matching names.
model.units_of("SNH")              # e.g. "g_N/m³"
model.description_of("SNH")
solution.units_named("SNH")          # same, for axis/column labels
model.summary()                    # tabulates every species with its units

# Dimensional ("unit") consistency check of the rate expressions. Currency-aware:
# g_COD/m3 and g_N/m3 are different dimensions, so it catches a dropped
# concentration factor, a wrong rate-constant exponent, or a Monod term mixing
# two currencies -- bugs a plain SI dimension check misses. Opt-in and advisory
# (never raises; unknown/unparseable units are skipped).
for w in model.check_units():      # -> list of (reaction, location, detail)
    print(w)

# Conservation (mass / electron balance) check. The companion to check_units:
# each species declares its content of the conserved quantities in the YAML
# (`composition: {COD: 1.0}` for an organic, `{COD: -1.0}` for oxygen, `{COD:
# -2.86, N: 1.0}` for nitrate-N, ...), and check_conservation dots that table
# against the stoichiometry -- so a wrong electron-acceptor demand breaks COD and
# a wrong product split breaks an elemental (S/N/P/Fe) balance. Opt-in and
# advisory (never raises). The ASM/ADM families fall back to a shipped table.
model.composition()                # -> {species: {quantity: content}}
for r, q, residual in model.check_conservation(quantities=["COD"]):
    print(r, q, residual)            # reactions whose COD content does not balance
# (For ASM1 this lists only `anoxic_growth_heterotrophs`: denitrification's
# electrons leave as N2 gas, which ASM1 does not track -- a known, intentional
# exception. The WATS sewer models declare composition in their YAML and close
# COD/S/N/Fe exactly.)

# Better than checking after the fact: write a conservation-determined coefficient
# as `auto` and let it be SOLVED from the declared balances at load -- so it can
# never be typed wrong. With composition declared on each species:
#   reactions:
#     - name: growth
#       conserved_for: [COD]                 # or a model-level `conserved_for:`
#       rate: "mu * [SS] * [XBH]"
#       stoichiometry: {SS: "-1/Y_H", XBH: 1.0, SO: auto}   # O2 demand solved from COD
# When a neighbour is a parameter expression (here -1/Y_H), the solved `auto`
# coefficient is itself yield-dependent and stays differentiable -- calibrating
# Y_H flows through it and the reaction conserves at *every* parameter value.

# Export results to a table instead of float()-casting one species at a time.
# Requires the optional `pandas` extra: pip install aquakin[dataframe]
df = solution.to_dataframe()         # time-indexed, one column per species
df.attrs["units"]                    # {species: unit} (units kept off the labels)
solution.to_csv("run.csv")           # units embedded in the CSV header

# Plot a species (or several) without matplotlib boilerplate -- the x-axis is
# labelled with the model's time unit, the y-axis with the species' units.
# Returns a matplotlib Axes. Requires the optional `plot` extra: aquakin[plot]
ax = solution.plot("SNH")            # one line; y-axis "SNH [g_N/m³]"
solution.plot(["SNH", "SNO"])        # several, legended; pass ax= to overlay
```

Discontinuous operations -- on/off pumps, SBR phases, dosing, level limits --
are handled with **located events** (`solve(events=...)`, on reactors and the
plant). A time event fires at a known time (AD-safe), a state event when a
`cond_fn` crosses zero; each can reset the state or terminate the solve:

```python
# An anoxic (denitrification) batch: dose external carbon partway through to
# drive denitrification, and stop once the nitrate has been removed.
i_ss = model.species_index["SS"]
i_sno = model.species_index["SNO"]
anoxic = model.concentrations(
    {"SS": 70.0, "SNO": 12.0, "SNH": 20.0, "XB_H": 150.0, "SO": 0.0})
events = [
    aquakin.Event(at_times=[0.1],                           # dose carbon at t = 0.1 d
                  apply=lambda t, C, p: C.at[i_ss].add(60.0)),
    aquakin.Event(cond_fn=lambda t, C, p: C[i_sno] - 0.5,   # stop when nitrate is gone
                  direction=-1, terminal=True, name="denitrified"),
]
sol = reactor.solve(anoxic, t_span=(0.0, 0.5),
                    t_eval=jnp.linspace(0.0, 0.5, 101), events=events)
sol.events_log                       # [(0.1, 'event0'), (~0.13, 'denitrified')] -- the audit trail
```

## Plant-wide simulation

Beyond a single reactor, `aquakin.plant` assembles full treatment-plant
flowsheets. The IWA benchmark plants ship ready to build: load BSM1, drive it to
steady state, and read the effluent — no autodiff or solver tuning required.

```python
import jax.numpy as jnp
import aquakin
from aquakin.plant.bsm import build_bsm1, load_bsm1_influent, evaluate_bsm1

model = aquakin.load_model("asm1")
plant = build_bsm1(model)            # 5 reactors + secondary clarifier + recycles

# A constant average-load influent. add_influent wires it to the plant's
# canonical front -- no "unit.port" string to hard-code.
plant.add_influent("feed", model.influent(
    {"SI": 30.0, "SS": 69.5, "XI": 51.2, "XS": 202.32, "XB_H": 28.17,
     "SNH": 31.56, "SND": 6.95, "XND": 10.59, "SALK": 7.0}, Q=18446.0))

# Integrate until the plant settles (a self-terminating steady-state event --
# no horizon to guess). Sensible solver defaults; nothing to tune.
ss = plant.run_to_steady_state()
print("converged:", ss.converged, "after", round(ss.time), "days")

# ...or snap straight to steady state algebraically (pseudo-transient
# continuation: ~10x faster, robust on stiff topologies, and differentiable --
# jax.grad of a loss on ss.state flows to the plant parameters for design sweeps).
ss = plant.steady_state()
print("converged:", ss.converged, "in", int(ss.iterations), "iterations")

# Reconstruct the clarified effluent and read its quality.
eff = plant.stream(ss.solution, plant.effluent_endpoint)
print("effluent SNH:", round(float(eff.C_named("SNH")[-1]), 2), "g N / m³")   # ~0.5
```

For a dynamic run, drive a fresh plant with a diurnal dry-weather influent,
warm-started from the steady state, and score the headline performance indices:

```python
dyn = build_bsm1(model)
dyn.add_influent("feed", load_bsm1_influent("dry", model))   # a 14-day diurnal load
sol = dyn.solve(t_span=(0.0, 14.0), t_eval=jnp.linspace(0.0, 14.0, 15), y0=ss.state)

ev = evaluate_bsm1(dyn, sol)           # Effluent Quality / Operational Cost indices
print(f"EQI = {ev.eqi:.0f} kg/d   OCI = {ev.oci:.0f}")
```

`build_bsm2(...)` assembles the full BSM2 sludge train (primary clarifier,
thickener, ADM1 digester with the ASM1↔ADM1 interfaces, dewatering, reject
recycle) the same way; see `examples/` and `CLAUDE.md` for the BSM2 steady state,
dynamic/seasonal runs, DO control, and the SRT/HRT/F:M design helpers.

Reactor temperature is a selectable `TemperatureModel`: the default
`AlgebraicTemperature` takes each reactor's temperature to be its instantaneous
flow-weighted inlet temperature, while `HeatBalanceTemperature` gives every
finite-volume unit a dynamic temperature *state* governed by the first-order heat
balance `V dT/dt = Q_in (T_in − T)` (the heated digester stays fixed), so the
reactor temperature lags and damps the influent. Select it with
`plant.set_temperature_model(aquakin.HeatBalanceTemperature())` or
`build_bsm2(temperature_model=aquakin.HeatBalanceTemperature())`.

Aeration energy can be scored from **blower/diffuser physics** instead of the
Copp-2002 correlation. An `AerationSystem` (diffuser submergence, SOTE, fouling,
blower efficiency) turns the `kLa` a solve produced into the required air flow
(via the standard oxygen transfer rate `kLa·C_s·V` and the diffuser SOTE) and the
adiabatic blower power against the submergence head — keeping the `kLa` kinetic
interface unchanged. Pass it to size a tank
(`aquakin.design_summary(kla, volume, system)`) or to the evaluators, where it
replaces the aeration-energy term and reports the air flow:

```python
syst = aquakin.AerationSystem(depth=5.0, sote=0.20)      # 5 m diffusers, 20 % SOTE
ev   = evaluate_bsm2(plant, sol, params, aeration_system=syst)
print(ev.aeration_energy, "kWh/d   air:", ev.air_flow, "m3/d")   # mechanistic AE
```

**Disinfection** unit ops reduce a pathogen indicator at the end of the train:
a `UVUnit` (dose = intensity × exposure × UVT-factor → log-linear inactivation)
and a `ChlorineContactUnit` (a chlorine residual that decays first-order; the CT
credit `residual × T10` → log-removal, with `T10` from a baffling factor or a
residence-time distribution). Both pass the process stream through and reduce the
indicator-organism density carried on the stream (the `"org"` entry of a stream's
`scalars` side-channel map, the disinfection analogue of the temperature scalar),
so the reconstructed effluent reports it:

```python
p.add_unit(aquakin.ChlorineContactUnit("cl", net, volume=500.0, dose=5.0,
                                       ct_per_log=8.0, decay_rate=2.0,
                                       inlet_density=1e6))   # CFU/100 mL in
sol = p.solve(t_span=(0.0, 2.0), t_eval=t)
print(p.stream(sol, "cl.out").org[-1])                      # effluent indicator
```

### GHG, cost and scenario reporting

On top of the EQI / OCI evaluation, `aquakin` reports a **carbon footprint**
(CO₂e/d) and a monetised **operating cost** (currency/d), and tabulates
scenarios side by side:

```python
ev  = evaluate_bsm2(plant, sol, params)          # EQI / OCI + physical flows
n2o = aquakin.direct_n2o_emission(plant, sol)    # stripped N₂O (0 unless the AS
                                                 # model resolves an SN2O state)

fp = aquakin.carbon_footprint(                   # kg CO₂e/d, with breakdown
    ev.total_energy(), grid_factor=0.4, n2o_emission=n2o,
    methane_production=ev.methane_production, ch4_fugitive_fraction=0.015)

oc = aquakin.operating_cost(                     # currency/d OPEX (+ optional CAPEX)
    energy_kwh_per_d=ev.total_energy(), carbon_kg_cod_per_d=ev.carbon_mass,
    sludge_kg_tss_per_d=ev.sludge_production,
    methane_kg_per_d=ev.methane_production,
    factors=aquakin.CostFactors(energy_price=0.12), co2e_per_d=fp.total_co2e)

print(fp)            # labeled CO₂e breakdown
print(oc)            # labeled cost breakdown
print(aquakin.kpi_comparison({"baseline": ev, "low-DO": ev_b}).table())
```

`carbon_footprint` weights direct N₂O (GWP ~273), grid-energy CO₂e and fugitive
biogas methane (GWP ~27), crediting recovered biogas energy; `operating_cost`
prices energy / carbon / sludge / biogas (and an optional CAPEX + carbon charge);
`kpi_comparison` puts any report objects (`BSM2Evaluation`, `CarbonFootprint`,
`OperatingCost`) into one standardized KPI table. The scenario-orchestration
primitives `monte_carlo`, `compare_scenarios` and `optimize_design` propagate
input uncertainty and size designs to a permit at minimum cost. See
`examples/bsm2_ghg_cost_report.py`.

## Advanced: differentiation & sensitivity

Every solve — reactor or whole plant — is differentiable. The machinery below is
for parameter estimation and sensitivity analysis; a plain forward simulation
(above) never needs it.

> **Heads-up — silent non-finite reverse gradients.** A reverse-mode gradient
> (`jax.grad` / `jax.jacrev`) taken *directly* through a stiff model's `solve`
> (ASM / ADM / WATS) returns silent `NaN`/`Inf` when the reactor's `dtmax` is
> uncapped — no exception, so the garbage gradient flows into your optimizer and
> the fit never converges. `aquakin.calibrate` and `aquakin.sensitivity` guard
> this for you; if you roll your own loss + optimizer, either cap `dtmax` (via
> `integrator=aquakin.IntegratorConfig(dtmax=...)`), use forward mode (`jax.jacfwd`
> through a reactor built with
> `diff=aquakin.DifferentiationConfig(mode="forward", method="through_solve")`),
> or wrap your gradient in `reactor.check_gradient_finite(jax.grad(loss)(p))`
> (equivalently the free `aquakin.check_finite_gradient`) to get an actionable
> error instead of silent `NaN`.

### Forward sensitivity (cap-free stiff gradients)

Differentiating *through* a stiff reaction-model solve with ordinary AD goes
non-finite above an integrator-step threshold, and the usual workaround — a
global `dtmax` cap — forces tiny steps over the whole solve. `solve_sensitivity`
avoids both: it integrates the sensitivity `S = dC/dθ` *alongside* the state and
lets the adaptive step controller bound the sensitivity error too, so the step
tightens only where the sensitivity is stiff and the result is exact with no cap.

```python
import jax.numpy as jnp
import aquakin

model = aquakin.load_model("uv_h2o2")
conditions = model.default_conditions(1)
reactor = aquakin.BatchReactor(model, conditions)

C0 = model.default_concentrations()
params = model.default_parameters()
t_eval = jnp.linspace(0.0, 5.0, 6)

sol, S = reactor.solve_sensitivity(
    C0, t_span=(0.0, 5.0), t_eval=t_eval, params=params,
    sens_params=["H2O2_photolysis.k_photo", "OH_target.k_OH_target"],
)
# sol : the usual solution; S : dC/dθ, shape (n_t, n_species, n_sens_params)

# A richer wrapper with by-name accessors:
res = aquakin.forward_sensitivity(
    reactor, C0, params=params, t_span=(0.0, 5.0), t_eval=t_eval,
    sens_params=["H2O2_photolysis.k_photo"],
)
res.dC_dparam("target", "H2O2_photolysis.k_photo")   # (n_t,)
```

`solve_sensitivity` is available on `BatchReactor`, `PlugFlowReactor` and
`BiofilmReactor`. For more than one parameter it defaults to a CVODES-style
*simultaneous corrector* (`shared_factor=True`): the augmented Jacobian is
block-lower-triangular with one shared diagonal block, so that block is
factorised once per step and reused across the sensitivity columns instead of
factorising the full augmented system. This is several times faster than the
dense augmented solve on large stiff systems (e.g. the layered biofilm) and
gives bit-identical results.

The **plant** has it too: `Plant.solve_sensitivity(params, wrt, t_span=, t_eval=,
y0=)` runs the same augmented `[y; S]` solve on the plant's enhanced solver
config (Kvaerno3 + decoupled Newton + cached recycle map + the simultaneous
corrector), so a forward-mode sensitivity of a stiff dynamic plant stays **finite
over long horizons** where `jacfwd` through `plant.solve` goes non-finite. It is
exact for kinetic parameters and backs the forward mode of
`plant.dynamic_sensitivity`.

### Cap-free reverse-mode gradients (stable adjoint)

`solve_sensitivity` scales with the parameter count, so for a scalar-loss
gradient over many parameters — the calibration case — reverse mode is wanted,
and that is the mode the `dtmax` cap exists for. The cap-free alternative there is
a hand-written discrete adjoint: the forward is an ordinary robust adaptive ESDIRK
(Kvaerno5) solve and the reverse is a per-step transposed solve over the saved
trajectory, finite at any step size with no cap. **`Plant.solve` uses it
automatically:** the differentiation default is
`diff=DifferentiationConfig(mode="reverse", method="stable")`, which keeps a plain
forward solve on the fast cached path but routes a solve under `jax.grad` to the
cap-free stable adjoint — so a stiff plant gradient is finite by default with
nothing to tune:

```python
sol = plant.solve(t_span=(0.0, T), t_eval=t_eval, params=params, y0=y0)
g = jax.grad(loss)(params)   # finite through the stiff, coupled BSM2 plant — no dtmax
```

This is what lets a reverse-mode gradient flow through the whole monolithic BSM2
solve — across the ASM↔ADM interface and the recycle loops — where differentiating
*through* the stiff solve is non-finite. It is exact through a transient solve:
`plant.solve` carries the integration time in the state, so the explicit time
dependence of a time-varying influent is captured exactly in the gradient.

### Calibrating a plant (`plant.calibrate`)

The cap-free gradient above powers a MAP calibration of plant parameters against
a measured output stream — the plant analogue of `aquakin.calibrate` (a reactor
fit). It reuses the same machinery (parameter transforms, priors, multistart, an
optional Laplace posterior) behind a forward-model seam, so only the forward
solve differs. Fit by-name plant parameters (`"<model>.<param>"`, see
`plant.parameter_names()`) so a target stream's channels match observations:

```python
result = plant.calibrate(
    observations,               # (n_t, n_channels) measured effluent data
    t_obs,                      # observation times
    ["asm1.muH", "asm1.bH"],    # plant parameters to fit
    target="effluent",          # a registered stream (or "unit.port")
    observed_channels=["SNH", "SNO"],
    y0=bsm2_warm_start(plant),  # warm start (recommended for a stiff plant)
)
result.params_named            # {"asm1.muH": ..., "asm1.bH": ...}
```

The reverse gradient runs through `plant.solve`'s stable adjoint, so it is finite
for a stiff plant with no `dtmax` to tune. Fit against **several streams at once**
by passing `observables=` instead of a single `target`:

```python
from aquakin.plant import PlantObservable

result = plant.calibrate(
    observations, t_obs, ["asm1.muH", "adm1.k_m_ac"],
    observables=[
        PlantObservable("effluent", ["SNH", "SNO"]),   # water-line N
        PlantObservable("wastage", ["XS"]),            # sludge line
    ],
)
```

The observation columns run in observable order (channels within a stream first).

Fit **assembled-state initial conditions** alongside the parameters with
`free_ic` — each entry names a `(unit, species)` slot of the plant's initial
state `y0` (a concentration unit: a CSTR or the digester), fit in log space:

```python
from aquakin import FreeICConfig

result = plant.calibrate(
    observations, t_obs, ["asm1.muH"],
    target="effluent", observed_channels=["SNH"],
    free_ic=FreeICConfig(["reactor1.S_S", "reactor1.X_BH"]),   # fit these initial pools too
    y0=bsm1_warm_start(plant),
)
result.ic_named[0]     # {"reactor1.S_S": ..., "reactor1.X_BH": ...}
result.C0_fitted[0]    # the fitted full initial state
```

Run a **joint multi-batch fit** — several plant runs from different initial
states sharing the parameters — by passing list-valued `observations` / `t_obs` /
`y0` (one entry per batch); the batches' data terms are summed:

```python
result = plant.calibrate(
    [obs_run1, obs_run2],          # per-batch effluent data
    [t_obs1, t_obs2],
    ["asm1.muH"],
    target="effluent", observed_channels=["SNH"],
    y0=[y0_run1, y0_run2],         # each batch its own initial state
)
```

This fits kinetic parameters (optionally with initial conditions, or across
several batches) against one or more streams. `free_ic` and multi-batch are not
yet combinable in one call.

### Choosing the integrator (`integrator=IntegratorConfig(...)`)

The integrator / step machinery is configured with one value object,
`aquakin.IntegratorConfig`, passed as `integrator=`:

```python
IntegratorConfig(order=3, factormax=3.0, colored_jacobian="auto",
                 dtmax=None, max_steps=100_000, solver=None)
```

`order` selects the ESDIRK method (`3` → `Kvaerno3`, the fast default; `5` →
`Kvaerno5`), `factormax` caps the PID step-growth, `dtmax` caps the step (set it
only for a reverse gradient *through* the solve), `max_steps` is the step budget,
`solver` is an explicit override, and `colored_jacobian` selects sparse
colored-AD Jacobian materialisation. `rtol` / `atol` stay separate arguments of
`plant.solve` (the accuracy contract, not the machinery).

A long dynamic plant run — the multi-hundred-day dynamic BSM2 simulation — is
**stiffness-bound**: the step count barely depends on the tolerance, and the
per-step cost is dominated by the implicit Jacobian factorisation of the whole
167-state plant. Every implicit solve uses a **decoupled root finder** (the
per-stage Newton tolerance loosened 10× from the step tolerance, ~15–20% cheaper
per step at preserved accuracy) automatically. The default
`IntegratorConfig(order=3, factormax=3.0)` is the fast stack — `Kvaerno3` (4
stages, less linear algebra per step) with a capped step growth — which together
run **~40% faster** than a plain `Kvaerno5`/no-cap solve:

```python
sol = plant.solve((0.0, 609.0), t_eval=t_eval, params=params, y0=y0)  # fast default

# the older robust higher-order stack:
sol = plant.solve((0.0, 609.0), t_eval=t_eval, params=params, y0=y0,
                  integrator=aquakin.IntegratorConfig(order=5, factormax=None))
```

`colored_jacobian` (default `"auto"`) forms the per-step implicit Jacobian by
**sparse column compression** instead of densely:

```python
sol = plant.solve((0.0, 609.0), t_eval=t_eval, params=params, y0=y0,
                  integrator=aquakin.IntegratorConfig(colored_jacobian=True))

The plant Jacobian is sparse (dense per-unit kinetic blocks + sparse inter-unit
coupling), so it is built in a handful of colored Jacobian-vector products (~45
for BSM2) rather than one per state (167) — a dominant per-step linear-algebra
cost. The reconstructed matrix equals the dense Jacobian, so the trajectory and
gradient are unchanged to integration tolerance; only the cost of forming it
drops, and it stacks with `order`/`factormax`. It applies to **both** the
forward solve (where `J` is built once per step) and the reverse stable discrete
adjoint (where `J` is rebuilt many times per step, so it dominates). It is built
and guarded against the dense Jacobian once per plant, falling back to the dense
solver if the guard fails.

The three settings:

- **`"auto"`** (default) governs the reverse stable-adjoint **backward** only: on
  a concrete solve it *measures* whether the colored `df/dy` build is actually
  cheaper than dense (it can be slower on a small plant) and enables it only when
  it pays — so a large plant (BSM2) gets the reverse-gradient speedup while a
  small one (BSM1) stays dense. The measured decision is reported by
  `plant.colored_jacobian_decision()`. `"auto"` leaves the forward solve dense.
- **`True`** forces coloring on both the forward solve and the reverse backward
  (skipping the measurement). Most worthwhile on a large stiff plant (BSM2); on
  the small BSM1 the materialisation is not the bottleneck.
- **`False`** disables it entirely.

### A non-AD fast lane (`forward_fast=True`)

If a solve never needs gradients (`jax.grad` / `calibrate` / `sensitivity` of the
result), `forward_fast=True` runs a lean integrator that skips the diffrax
adjoint / optimistix / lineax machinery entirely — a plain `lax.while_loop`
adaptive ESDIRK with a simplified Newton and the colored Jacobian:

```python
sol = plant.solve(t_span=(0.0, 609.0), t_eval=t_eval, params=params, y0=y0,
                  forward_fast=True)   # ~3x compile, ~1.3-1.9x run (dynamic BSM2)
```

That machinery exists to make the *whole solve* differentiable, and tracing it
dominates compile time — so dropping it gives **~3× faster compile** (a big deal
for the multi-minute full-BSM2 compile, which file-caching can't help because the
cost is Python-level tracing) **and ~1.3–1.9× faster run** on the dynamic BSM2
(the run gain narrows over a long run with dense `t_eval`, the compile win is the
robust benefit), at the same accuracy. The
per-step Jacobian is *still* colored forward-mode AD (the exact same matrix), so
the trajectory matches a valid adaptive solution to the same `rtol`; only
end-to-end differentiability is given up. It is opt-in and forward-only: it needs
concrete `params`/`y0` and raises a clear error under `jax.grad`/`jax.jit` or with
`events=`, and falls back to the diffrax path if the colored-Jacobian guard fails.

For reactor-level fits, the adjoint plumbing is hidden too: `aquakin.calibrate`
and `aquakin.sensitivity` take `diff=DifferentiationConfig(mode=...)` and build
the right adjoint internally (no `diffrax` import), and
`DifferentiationConfig(check_finite=True)` (the
default) raises a friendly error with the remedy instead of returning silent
`NaN` gradients on a stiff model.

## Testing

```bash
pytest -m "not validation"   # unit + integration (fast)
pytest -m validation          # scientific validation against published data
pytest                        # everything
```

## License

MIT.
