Atmospheric chemistry: KPPax makes the KPP mechanisms differentiable and divides the error of correction of biogenic emissions up to five

The chemical mechanisms of the atmosphere circulate in a format designed for Fortran and C. KPPax translate them into driftable JAX code, connecting them directly to neuron networks.

Air quality and climate models describe their chemical reactions in the language of the Kinetic PreProcessor (KPP), which produces Fortran or C code that is powerful but incompatible with machine learning tools, which need to calculate gradients. Wenqing Peng, Michael Boy and Zhi-Song Liu present KPPax, a preprocessor that generates differentiable, on-the-fly compiled and vectorizable JAX code from the same files.

The velocity laws are evaluated in logarithmic space with sparse operations, and the system, which is very stiff, is integrated by the implicit solvers of the diffrax library, which provides exact sensitivities without writing an adjoint by hand.

Two demonstrations on the POLLU mechanism

First, the automatically calculated Fisher information matrix predicts, before any adjustment, which of the 25 reaction coefficients a flow tube experiment will actually be able to estimate: nine are, with uncertainties confirmed to within about 30%. Second, a neural network responsible for correcting an emissions inventory, pre-trained offline and then refined through the chemical model, reduces the error on biogenic emissions by a factor of 2.4 to 5 and retrieves information on the low-temperature emissions that individual inversions were losing.