aquakin

aquakin models reactive scalar transport in aqueous environmental systems. Reaction models are declared at runtime in YAML, parsed into an AST, and compiled to JAX-native, automatically-differentiable rate functions integrated by Diffrax.

It ships a library of ready-to-use models — ozonation and UV/H₂O₂ chemistry, the ASM activated-sludge family, ADM1 anaerobic digestion, WATS sewer processes, and mineral precipitation — and the tools to simulate them in batch, plug-flow, biofilm, and full plant-wide flowsheets, with automatic differentiation throughout for sensitivity analysis and parameter calibration.

New here? Start with Getting started.

Guide

• Getting started
  • Installation
  • Your first solve
  • Setting the initial state by name
  • Time units
  • Reading results
  • Inspecting a model
  • Where to go next
• Reactors
  • Operating conditions
  • Batch reactor (0-D)
  • Plug-flow reactor (1-D)
  • Biofilm reactor (1-D over depth)
  • Particle-track and CFD reactors
  • Solutions
  • Events: pumps, dosing, phase switches
• Plant-wide simulation
  • The benchmark plants
  • Steady state, faster
  • Dynamic runs and performance indices
  • Reporting: GHG, cost, and scenarios
  • Building a custom flowsheet
• Sensitivity & calibration
  • Differentiating a solve
  • Local sensitivity
  • Forward sensitivity (cap-free stiff gradients)
  • Global sensitivity (DGSM)
  • Fitting parameters to data
  • Calibrating a plant
  • Uncertainty propagation and design
  • Choosing the integrator

Authoring models

• YAML model format
  • Shared parameters vs reaction-local parameters
  • Named rate expressions
  • Rate expression grammar
  • Stoichiometry coefficients
  • Built-in functions
  • Notes
• Adding a built-in model
  • Authoring tips

Reference

• Model catalog
  • Oxidation chemistry
  • Activated sludge (ASM family)
  • Anaerobic digestion
  • Sewer processes (WATS)
  • Mineral precipitation
• Public API
  • Loading and inspecting a model
  • Conditions
  • Reactors and solutions
  • Plant-wide simulation
  • Influent characterisation
  • Sensitivity analysis
  • Calibration and identifiability
  • Uncertainty propagation and design
  • Model checks: units and conservation
  • Benchmark evaluation and reporting
  • Design and sizing
• API reference
  • Classes
  • Functions

Installation

pip install aquakin

Note

import aquakin enables JAX 64-bit (x64) mode process-wide — the stiff implicit ODE solves require double precision. This is global JAX state: other JAX code in the same process will use float64 afterward. aquakin emits a one-time warning if it overrides an explicit float32 preference, so the side effect is never silent.

Quickstart

import jax.numpy as jnp
import aquakin

model = aquakin.load_model("ozone_bromate")
conditions = model.default_conditions().with_(pH=7.5, T=293.15)   # 0-D batch case
reactor = aquakin.BatchReactor(model, conditions)
sol = reactor.solve(
    model.default_concentrations(),
    params=model.default_parameters(),
    t_span=(0.0, 600.0),
    t_eval=jnp.linspace(0.0, 600.0, 121),
)
print("[BrO3-] at 10 min:", float(sol.C_named("BrO3-")[-1]))

See Getting started for a step-by-step walkthrough.

Architecture

aquakin uses a two-layer data model:

1. Schema layer (load time) — Pydantic models validate the YAML and produce a clean spec object. Pydantic never appears on the hot path.

2. Runtime layer — a CompiledModel dataclass, built once from the spec, holds the stoichiometry matrix, the per-reaction compiled rate callables, and the parameter index map. This is what the integrators operate on.

Each rate callable has the signature rates(C, params, condition_arrays, loc_idx) returning a (n_reactions,) vector, and the reaction right-hand side is stoich.T @ rates(...). Rate constants are always passed in via params, never baked in — which is what makes the whole solve differentiable for sensitivity analysis and calibration.