KPPax: Differentiable Kinetic ODE Preprocessing in JAX for Atmospheric Chemistry Machine Learning

Forum topic · 小凯 · 2026-10-28

Summary

KPPax is a kinetic preprocessor that converts atmospheric chemical mechanisms written in the KPP (Kinetic PreProcessor) language into differentiable, JIT-compiled, vectorisable JAX code, enabling integration with machine-learning frameworks. It evaluates rate laws using sparse operations in log space to preserve differentiability and integrates stiff chemical systems with implicit solvers and forward/reverse-mode adjoints from the diffrax library, providing exact sensitivities without hand-derived adjoint code. Using the POLLU benchmark mechanism, the authors demonstrate two applications: first, a Fisher information matrix computed via automatic differentiation predicts, before any fitting, which of 25 rate coefficients a flow-tube-like experiment can identify; noisy synthetic estimation confirms predicted uncertainties of the nine identifiable coefficients to within roughly 30%. Second, a neural-network emission inventory correction model is trained offline on per-scenario 4D-Var retrievals and fine-tuned online through the chemical model, reducing biogenic correction error by a factor of 2.4-5 and recovering information on low-temperature emissions that individual retrievals lose. KPPax makes KPP mechanisms directly usable for parameter estimation and hybrid physics-machine-learning modeling.

Paper Overview

Field: AI4Science / Atmospheric Chemistry Authors: Wenqing Peng, Michael Boy, Zhi-Song Liu Published: 2026-10-27 arXiv: 2610.37655

Abstract

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 diffrax library, so that exact sensitivities are obtained without hand-derived adjoint code.

Key Applications (POLLU Benchmark)

1. Fisher Information Analysis

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. 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.

2. Neural-Network Emission Correction

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. This 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.

Conclusion

KPPax makes KPP mechanisms directly usable for parameter estimation and hybrid physics-machine-learning modelling.