Back to AI News
arxiv

KPPax converts KPP chemistry mechanisms into differentiable JAX code

KPPax is a kinetic preprocessor that turns mechanisms written in the Kinetic PreProcessor (KPP) language into differentiable, just-in-time compiled, and vectorisable chemistry code in JAX. It evaluates rate laws with sparse operations in log space and integrates stiff systems using diffrax implicit solvers and forward- and reverse-mode adjoints, producing exact sensitivities without hand-derived adjoint code. On the POLLU benchmark, a Fisher information matrix from an automatically differentiated Jacobian predicted which of 25 rate coefficients a flow-tube-like design can identify, and noisy synthetic estimation matched the predicted uncertainties of the nine identifiable coefficients to within about 30%. A neural-network emission-inventory correction, trained offline on per-scenario 4D-Var retrievals and fine-tuned online through the chemical model, reduced biogenic correction error by a factor of 2.4–5 and recovered low-temperature emission information lost by the individual retrievals.

AI Brief

It gives atmospheric-chemistry researchers a concrete way to run KPP mechanisms inside differentiable JAX workflows for parameter estimation and hybrid physics–machine-learning models.

arXiv:2610.37655v1 Announce Type: new 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.
Be the first to rate this article.

Comments & Questions (0)

Captcha
Please be respectful — let's keep the conversation friendly.

No comments yet

Be the first to comment!