How to train an ML chemistry-rate surrogate

aeolus1d ships a framework for learning the surface-chemistry step instead of solving it with a stiff integrator every macro step. This recipe walks the full offline loop — generate data → train → inspect → tune → use — for the two-stage linearized surrogate, the most promising of the framework's heads.

Status: experimental, opt-in

A trained surrogate can now be dropped into the quasi-steady driver behind the experimental chem_solver = "linop" switch (see Step 5) — but it is not a validated default. The wiring keeps a per-cell relaxation-fraction gate that falls the pipe back to the banded LSODA solve whenever the linear model would take too aggressive a step, so the result is bounded-error by construction. Before it is trusted as a default it still has to beat that LSODA baseline on cost and hold its error budget on a full drive-cycle run — see Performance history for why the earlier ISAT approach was a measured no-go against the same baseline. Treat the numbers you get here as a feasibility signal.

The idea in one line

Freeze the gas temperature over the step and model each cell's chemistry as relaxation toward an equilibrium anchor:

\[\frac{dY}{dt} = -K\,(Y - Y^\*) \quad\Rightarrow\quad Y(t+\Delta t) = Y^\* + (Y - Y^\*)\,e^{-K\,\Delta t}\]

Two small MLPs predict the pieces from the cell state x = [T, ρ, Y(12), θ_Ce] (the 12 canonical Brinkmeier-TWC species — CO, CO2, C3H6, C3H8, CH4, H2, H2O, N2, NO, NO2, O2, NH3 — plus one surface coverage):

• an equilibrium net x → Y* (the state the cell relaxes toward), and
• a rate net x → log k (per-component; exp() guarantees k > 0, so the step is unconditionally stable and stays between Y and Y* — positivity-preserving, no overshoot, valid at any Δt).

Runtime inference is pure numpy (no torch/sklearn needed); only the training needs torch.

Prerequisites

# from the repo root, in your venv
pip install "aeolus1d[ml]"    # pulls torch + scikit-learn
python -c "import torch; print(torch.__version__, torch.cuda.is_available())"

A GPU is optional — training is small (seconds to a couple of minutes). All commands below run from the repo root.

The files

Stage	File	What it does
Data	ml_chemistry/generate_linop_data.py	samples the TWC envelope, integrates the real Brinkmeier kinetics, emits Y* + k targets
Train	ml_chemistry/train_linop.py	fits the two nets (torch), exports one .npz, scores end-to-end
Runtime	src/aeolus1d/closures/chem_linop.py	ChemLinOp — pure-numpy inference
Feasibility	ml_chemistry/linop_feasibility.py	ideal-k upper-bound probe (no training)

Datasets and models are written under ml_chemistry/data/ and ml_chemistry/models/ — both are git-ignored (regenerable).

Step 1 — generate training data

python -m ml_chemistry.generate_linop_data \
    --n 6000 --dt 7e-4 --t-ref 0.1 \
    --out ml_chemistry/data/linop.npz

Knobs:

• --n — number of sampled states. More = better coverage of the hard light-off band (start at 6000; push to 15000+ when tuning the tail).
• --dt — the step the rate is fit at. 7e-4 is the AVL TWC macro Δt; the model extrapolates to other Δt at runtime, so this is mostly the fitting reference.
• --t-ref — the anchor time for Y*. 0.1 s is the recommended reachable/slow-manifold anchor (for most cells it is already the equilibrium, but it stays well-conditioned for slow below-light-off cells). The anchor is clamped up to at least 3·dt internally, so a too-small t_ref (e.g. 0.01) collapses onto the probe window, lands mid-transient, and degrades the fit — keep it well above dt.

The run prints the fraction of state components with a clean monotone single-exponential relaxation (the rest are masked out of the rate loss and fall back to LSODA at runtime):

wrote ml_chemistry/data/linop.npz  (n=6000, dt=0.0007, t_ref=0.1; k well-defined 66% of components)

Generate once, train many — Step 1 is the slow part (one stiff solve per sample). Re-run Step 2 repeatedly against the cached .npz.

Step 2 — train + read the score

python -m ml_chemistry.train_linop \
    --data ml_chemistry/data/linop.npz \
    --hidden 128,128,128 --epochs 400 \
    --out ml_chemistry/models/linop.npz

The decisive output is the end-to-end validation on a held-out split: it applies the NN-predicted relaxation step and compares to the true LSODA Δt map — split by light-off temperature, with the residual-fallback rate as the headline tuning target:

[train_linop] END-TO-END NN step error vs LSODA, dt=7.00e-04 (n_val=1200) [median/p95/max]:
    all cells       : 2.81e-03 / 3.72e-02 / 8.46e-02
    below light-off : 5.60e-05 / 1.45e-04 / 2.85e-04
    above light-off : 4.60e-03 / 3.92e-02 / 8.46e-02
  >> TUNE TARGET: residual-fallback to hold <=2% = 16.9% of cells (above light-off 19.6%); <=5% = 1.7%  [lower is better]

How to read it:

• below light-off is easy — the net nails it (max ~0.03 %).
• above light-off is the hard regime — sharp Y*, fast/variable k. That tail is what drives the fallback rate.
• fallback rate is the cost number: the fraction of cells that must fall back to the exact LSODA solve to keep error ≤ 2 %. Lower is the goal. (The baseline run above needs ~17 %.)

Step 3 — tune (this is the loop to iterate)

Re-run Step 2 with different knobs against the same dataset and watch the above-light-off p95/max and the fallback rate drop:

# more capacity + longer training (the first thing to try)
python -m ml_chemistry.train_linop --hidden 256,256,256 --epochs 800

# learning-rate / batch sweep
python -m ml_chemistry.train_linop --hidden 256,256,256 --epochs 800 --lr 5e-4 --batch 256

# force CPU / pick a device
python -m ml_chemistry.train_linop --device cpu

Knob	Default	Tune for
--hidden	128,128,128	capacity for the sharp Y* at light-off — try 256,256,256
--epochs	400	underfit tail — try 800–1500
--lr	1e-3	stability of the fit — drop to 5e-4/3e-4 if it plateaus high
--batch	512	smaller (128–256) often sharpens the light-off transition
--seed	0	check run-to-run variance

If capacity/epochs stop helping, the lever moves upstream:

• more data — re-run Step 1 with --n 15000 (better light-off coverage). This is usually the highest-leverage change for the tail.
• anchor — try --t-ref 0.05 vs 0.1 vs 0.2 in Step 1.
• the genuinely hard part is the rate net at light-off (k spans ~100–10⁴× and Y* is near step-like there). Stratifying the training by temperature band, or a dedicated light-off sub-net, is the documented next step if a single global MLP plateaus.

Goal: drive the above-light-off p95 under 2 % so the fallback rate collapses toward the below-light-off value (≈ 0 %).

Step 4 — use the trained model

Runtime inference is pure numpy:

import numpy as np
from aeolus1d.closures.chem_linop import ChemLinOp

m = ChemLinOp.load("ml_chemistry/models/linop.npz")

# a batch of cells: T (n,), rho (n,), Y (n, 12), theta (n,)
Y_new, theta_new = m.apply_step(T, rho, Y, theta, dt=7e-4)
# Y_new >= 0, theta_new in [0, 1], by construction.

# the pieces, if you want them:
Ystar = m.equilibrium(T, rho, Y, theta)   # (n, 13) anchor [Y*(12), theta*]
k     = m.rate_k(T, rho, Y, theta)        # (n, 13) > 0 relaxation rate

apply_step accepts any dt — the relaxation form is exact for the linear model at any step size, which is the property that lets a single evaluation replace many stiff sub-steps. The QS-driver wiring (Step 5) pairs it with a relaxation-fraction fallback to LSODA for the high-error light-off cells (the fallback rate from Step 2 tells you how often that fires).

Step 5 (optional) — wire it into the QS driver

The trained .npz can be handed to the quasi-steady driver behind the experimental chem_solver = "linop" switch. Point chem_linop_model at the model file in your case's [case] header:

[case]
solver = "quasi_steady"
chem_solver = "linop"
chem_linop_model = "ml_chemistry/models/linop.npz"
chem_linop_tol = 0.4        # max per-step relaxation fraction before fallback

Behaviour:

• The surrogate only fires on pipes matching its Brinkmeier-TWC contract (canonical 12-species seam, one surface coverage theta_Ce); any other reacting pipe stays on the banded LSODA solve.
• chem_linop_tol (default 0.4) is the validity gate: if any cell's per-step relaxation 1 - exp(-k·dt) exceeds it — the fast light-off transient the single-rate linear model can't capture — the whole pipe falls back to LSODA for that macro step. Lower = stricter (more fallback, safer).
• An empty chem_linop_model silently disables the surrogate; a model that fails to load warns (RuntimeWarning) and reverts to the banded LSODA solve. Either way the switch can never do worse than the bounded-error fallback.

This path is opt-in and unvalidated as a default — treat a run with it on as a measurement, not a production result, and compare the tailpipe / theta_Ce / q_rxn trajectory against a pure-LSODA baseline.

Quick feasibility check (no training)

To see the upper bound — how accurate the form is with an ideal k estimated from the true trajectory, independent of NN quality:

python -m ml_chemistry.linop_feasibility --n 600 --dt 7e-4 --yref 0.1

The gap between this (ideal k) and Step 2 (NN-predicted k) is exactly the learnability gap you are tuning to close.