KPPax enables differentiable chemistry models for air pollution studies
KPPax: preprocess differentiable kinetic ODE for machine learning in atmospheric chemistry
Artificial Intelligence
Summary
Atmospheric chemistry models are important for understanding air pollution, but their code is often hard to use in machine learning tools. The authors present KPPax, a new tool that changes chemistry models into a format that works smoothly with advanced machine learning software called JAX. This allows exact calculations of how changing chemicals affects outcomes without extra manual coding. They show that KPPax helps identify which chemical reactions can be measured and improves corrections to air pollution data using machine learning.
What this means in practice
- •For environmental model developers: Create atmospheric chemistry models that integrate directly with differentiable machine learning frameworks for improved parameter estimation and sensitivity analysis.
- •For air quality monitoring teams: Use hybrid physics and machine learning approaches to correct emission inventories online, improving accuracy of pollution source estimates.
Authors
Wenqing Peng, Michael Boy, Zhi-Song Liu
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 \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.