PaperlayerView exact revision on arXiv

Source-linked AI summary

KPPax: preprocess differentiable kinetic ODE for machine learning in atmospheric chemistry

Wenqing Peng, Michael Boy, Zhi-Song Liu

arXiv:2610.37655v1cs.AI

TL;DR

Atmospheric chemical mechanisms distributed in KPP are not directly suited to differentiable machine-learning frameworks, limiting their use in gradient-based inverse and hybrid modelling. KPPax generates differentiable JAX chemistry code from KPP mechanisms and applies its sensitivities to coefficient identifiability and online emission correction. In the POLLU benchmark, nine of 25 rate coefficients were identifiable, and online training recovered emission information that individual retrievals lose.

  • Problem

    KPP mechanisms generate efficient Fortran or C code but do not fit directly into differentiable machine-learning frameworks used for gradient-based inverse and hybrid modelling.

    Source: §1 · pp. 1–3

  • Method

    KPPax generates differentiable JAX chemistry code from KPP mechanisms and uses differentiable ODE solves to obtain sensitivities for parameter estimation and hybrid modelling.

    Source: §1 · pp. 2–3, §6 · p. 15

  • Results

    Nine of the 25 POLLU rate coefficients were identifiable, and estimates from noisy synthetic observations scattered as predicted to within about 30%.

    Source: §6 · p. 15

  • Takeaways & Limitations

    Online fine-tuning through the chemical model recovers emission information that individual retrievals cannot resolve.

    Source: §4.4 · p. 11, §5 · p. 15

    For mechanisms with thousands of species, dense Jacobian factorisation is expected to dominate computational cost as it grows cubically with species count.

    Source: §5 · p. 13

Abstract

from arXiv · show

Atmospheric chemical mechanisms are commonly distributed in the language of the Kinetic PreProcessor (KPP), which generates efficient Fortran or C code but does not fit into the differentiable programming frameworks in which machine-learning models are developed. We present KPPax, a kinetic preprocessor that generates differentiable, just-in-time compiled, and vectorisable chemistry code in JAX from mechanisms written in the KPP language. KPPax evaluates the rate law with sparse operations in log space, which keeps it differentiable, and integrates the stiff system with the implicit solvers and the forward- and reverse-mode adjoints of the \texttt{diffrax} library, so that exact sensitivities are obtained without hand-derived adjoint code. We demonstrate two applications with the POLLU benchmark mechanism. In the first, the Fisher information matrix computed from an automatically differentiated Jacobian predicts, before any fit, which of the 25 rate coefficients a flow-tube-like design can identify, and noisy synthetic estimation confirms the predicted uncertainties of the nine identifiable rate coefficients to within about 30\% while revealing clusters of fast reactions of which only ratios are constrained. In the second, a neural-network correction of an emission inventory is trained offline on per-scenario 4D-Var retrievals and then fine-tuned online through the chemical model, which reduces the error of the biogenic correction by a factor of 2.4--5 and recovers information on low-temperature emissions that the individual retrievals lose. KPPax makes KPP mechanisms directly usable for parameter estimation and hybrid physics--machine-learning modelling.

1 Introduction

Atmospheric chemistry models combine many coupled reactions, making stiff ODE systems and their Jacobians tedious and error-prone to hand-code. KPPax brings KPP mechanisms into differentiable JAX workflows, addressing a gap for parameter estimation and hybrid modelling.

1 Introduction

  • KPP automates stiff chemistry code and sensitivity generation, but its Fortran/C output inherits constraints of those environments and CPU hardware.Generated KPP code includes sparse Jacobians and stiff integrators.Source: §1 · pp. 1–2
  • Existing chemistry tools target varied software and hardware, while process-based models remain costly to amend or tune against observations.The cited tools include Python/SciPy, C++, CUDA, and CPU/GPU implementations.Source: §1 · p. 2
  • Neural networks offer fast atmospheric-process emulation but can lack interpretability, anomaly prediction, and feasibility when observations are sparse.Their gradient-based training commonly uses Julia or Python frameworks such as PyTorch and JAX.Source: §1 · p. 2
  • Hybrid models combine neural flexibility with process-model interpretability, but require differentiable process models; no KPP-language generator was known to the authors.The authors identify this missing tool for mechanisms distributed in KPP, including atmospheric mechanisms such as the MCM.Source: §1 · p. 2
  • Rate-coefficient estimation can be ill-posed because observations constrain only some coefficient combinations, while online emission training requires gradients through chemistry.The cited motivation contrasts offline retrieval-based training, which discards information individual retrievals cannot resolve, with online training through the chemical model.Source: §1 · p. 2
  • KPPax generates differentiable, just-in-time compiled, vectorisable JAX chemistry code from KPP mechanisms, using sparse log-space rate evaluation and diffrax adjoints.The approach targets inverse and hybrid modelling without hand-derived adjoint code.Source: §1 · pp. 2–3

2 Implementation of differentiable kinetic pre-processing

KPPax represents chemical kinetics with sparse arrays and evaluates rates in a differentiable log-space formulation. It uses JAX-traceable ODE solves and configurable adjoints to obtain exact discrete sensitivities without deriving solver-specific adjoint code by hand.

2.1 The chemistry kinetic problem

  • For n species and m reactions, the chemical system is an initial-value problem whose species concentrations evolve through reaction rates and stoichiometric production and consumption.The concentration vector begins at an initial state, and rate coefficients may vary with environmental conditions.Source: §2.1 · p. 3

2.2 Reaction parsing

  • KPPax parses mechanism sections with pykpp, stores stoichiometry in COO sparse arrays, and converts time-dependent Fortran rate expressions to JAX functions.Users supply initial concentrations and construct the customized simulation; KPPax generates chemistry code alone.Source: §2.2 · p. 4
  • Sparse gather and segment-reduction operations evaluate rates without materializing dense stoichiometric matrices, while log-space segment sums replace a non-differentiable scatter-multiplication.The log-space formulation restores differentiation with respect to concentrations and rate coefficients, but requires strictly positive concentrations.Source: §2.2 · p. 4
  • KPPax integrates stiff chemistry with implicit Kvaerno solvers, adaptive error control, and a positive concentration floor for the log-space rate law.The solver uses 64-bit double precision; the floor prevents evaluations at zero or negative concentrations.Source: §2.2 · pp. 4–5

2.3 Automatic differentiation

  • Automatic differentiation through JAX-traceable operations supplies sensitivities without hand-derived tangent-linear or adjoint equations, avoiding solver-specific adjoint derivations.KPP's discretize-then-optimize adjoints require analytic derivations for individual integrators.Source: §2.3 · p. 5
  • Diffrax differentiates the numerical solve step by step, and recursive checkpointing returns exact discrete gradients with memory growing sub-linearly in the number of solver steps.Checkpointing recomputes states during the backward pass, trading additional computation for reduced memory.Source: §2.3 · p. 5

3 Application I: identifiability and parameter estimation

The section uses KPPax sensitivities to assess which POLLU rate coefficients a flow-tube-like design can identify before fitting, then tests those predictions with noisy synthetic data. It finds nine individually identifiable coefficients and fast-reaction clusters constrained only through rate ratios.

3 Application I: identifiability and parameter estimation

  • Rate fitting is often non-unique because observations constrain only combinations, motivating identifiability analysis before parameter estimation.Automatic differentiation through the stiff solver provides exact sensitivities at a cost comparable to a few forward integrations.Source: §3 · pp. 5–6

3.1 Method

  • Before fitting, the Fisher information matrix uses local sensitivities to identify constrained parameter combinations, while eigenvector loadings reveal confounded rates.The analysis relies on the Jacobian at a single parameter point.Source: §3.1 · p. 6, §3 · p. 6

3.2 Experimental design

  • The POLLU flow-tube-like design measures 20 species once after 60 minutes across eight initial conditions; continuous observation would make all 25 rates structurally identifiable.The single-time design’s non-identifiability is therefore practical, arising from the observation design.Source: §3.2 · pp. 6–7

3.3 Results

  • Nine individual rate coefficients are identifiable, and their RMS estimation errors agree with Cramér–Rao predictions to within about 30%.Figure 1 compares the estimated coefficients with the predicted uncertainty intervals.Source: Figure 1 · p. 7, §3.3 · p. 8
  • Weakly constrained coefficients group into three fast-reaction clusters where observed concentrations constrain rate ratios, not absolute rates.The clusters are the NO–NO2–O3 cycle, the NO3–NO2–N2O5 equilibrium, and O3–O(1D)–O(3P) interconversion.Source: §3.3 · p. 8
  • The FIM identifies unconstrained combinations but can underestimate drift along curved cost valleys; one cluster’s RMS deviation was 6.3, seven times its Cramér–Rao prediction.Fits continued along flat valleys, while estimates within clusters moved together and kept their ratios constant.Source: §3.3 · p. 8

4 Application II: hybrid neural emission inversion

The experiment compares offline learning from per-scenario 4D-Var retrievals with online fine-tuning of a neural emission correction through the chemical ODE. Online fine-tuning most improves the biogenic correction, including at low temperatures, where individual retrievals lose information.

4.1 Experimental setup

  • The synthetic experiment uses 150 training, 30 validation, and 30 interpolation-test scenarios, with additional extrapolation scenarios and hourly observations over a one-day window.It compares all 20 species with five routinely measured species, each with and without observation noise.Source: §4.1 · p. 9

4.2 Neural emission correction

  • Each sector’s correction is represented by a separate two-hidden-layer MLP with 32 softplus units per layer, bounded output, and zero-initialized final layer.Together the networks have 2434 parameters and initially reproduce the inventory correction of zero.Source: §4.2 · p. 10

4.3 Offline two-step training

  • Offline training retrieves correction knots for each scenario by 4D-Var, then regresses the network on those retrievals using weighted mean squared error.Weighting improves anthropogenic regression in every setting and biogenic regression when all species are observed; the methods perform similarly for station-only biogenic observations.Source: §4.3 · pp. 10–11
  • The offline retrieval’s selected correlation length is 192 h with station observations, indicating those observations cannot resolve within-day correction variation.With all species observed, selected correlation lengths are 48 h or 96 h.Source: §4.3 · p. 11

4.4 Online fine-tuning through the ODE

  • Online training fits shared network parameters directly to observations across scenarios, allowing information from multiple scenarios to combine when individual retrievals cannot constrain a correction.The residual Jacobian factors into knot-value sensitivities and the network Jacobian; Gauss–Newton fine-tuning is initialized from the offline solution.Source: §4.4 · p. 11

4.5 Results

  • A factor of 2.4–5 reduction in biogenic correction RMSE inside training ranges followed online fine-tuning; with all species observed, cold-band error fell from 0.26–0.29 to 0.04–0.06.The anthropogenic extrapolation error was unchanged in one setting: noise-free observations of all species.Source: §4.5 · p. 12
  • For noisy station observations, online fine-tuning left most correction error below 0.25 within training ranges, with remaining bias concentrated below about 10 °C.The offline network missed the low-temperature increase and the true response optimum; the anthropogenic error was small in both cases.Source: §4.5 · p. 12
  • Errors in extrapolation bands are about 2–4 times larger than within training ranges for both approaches; noise-free all-species fine-tuning also produced a larger biogenic error than the noisy case.The authors suggest a proximity term toward offline parameters or early stopping for noise-free data.Source: §4.5 · p. 13

5 Discussion

The discussion links KPPax's online emission fine-tuning to recovery of information lost in per-scenario retrievals, while detailing differentiation choices and limits from clipping, local identifiability, and computational cost.

5 Discussion

  • Forward mode suits few parameters and full Jacobians, reverse mode suits scalar losses over many parameters, and online Gauss–Newton combines forward sensitivities at knot values with network Jacobians.KPPax makes the differentiation mode configurable through diffrax, without hand-derived adjoints.Source: §5 · p. 13
  • Dense Jacobian factorization is negligible for POLLU's 20 species but grows cubically and will dominate for mechanisms with thousands of species.JIT compilation adds a first-call cost that is amortized over repeated integrations; sparse or matrix-free implicit solvers are the key extension for large mechanisms.Source: §5 · p. 13
  • Clipping concentrations to a positive floor preserves log-space rate evaluation but sets gradients to zero for clipped species, including species starting at zero or fully consumed.Applications differentiating with respect to near-zero concentrations should account for this effect.Source: §5 · p. 13
  • The single-point FIM identifies unconstrained rate-coefficient combinations but underestimates drift along curved cost-function valleys; profile likelihoods or MCMC can quantify that drift.Because the FIM is quick to compute, it can also guide choices of observed species, times, and precursor concentrations before an experiment.Source: §5 · p. 13
  • The emission experiment assumes synthetic truth, matched chemistry, a transport-free box, and known initial states; initial-state errors could be aliased into early-morning emissions.Gauss–Newton forms a Hessian over all 2,434 network parameters, feasible here but requiring matrix-free or approximate second-order methods for larger networks.Source: §5 · p. 13
  • Online training through the chemical model recovers observational information lost in per-scenario retrievals, as Figure 3 contrasts offline and online correction errors.The figure shows the most realistic setting with noisy station observations.Source: §5 · pp. 14–15, Figure 3 · p. 14

6 Conclusions

KPPax generates differentiable JAX chemistry from KPP mechanisms and provides exact sensitivities for stiff simulations. POLLU demonstrations show applications to rate identifiability and emission correction, while large mechanisms and real observations remain future work.

6 Conclusions

  • KPPax converts KPP mechanisms into sparse, log-space JAX rate laws and uses configurable diffrax adjoints for exact stiff-chemistry sensitivities without hand-derived code.The generated code can be vectorised over scenarios and run on GPUs.Source: §6 · p. 15
  • On POLLU, the FIM identified 9 of 25 rate coefficients with estimates matching predicted uncertainties within about 30%, while online fine-tuning reduced biogenic-correction error by a factor of 2.4–5.Online training recovered low-temperature emission information lost in individual retrievals and took about two minutes on one GPU.Source: §6 · p. 15
  • KPPax brings KPP mechanisms into differentiable modelling for combining process-based chemistry and machine learning; large mechanisms and real chamber or flow-tube observations remain future work.Large mechanisms require sparse linear solvers for implicit integration.Source: §6 · p. 15

Related papers

Loading 2610.37655v1…