Back to AI Research

AI Research

KPPax brings differentiable atmospheric chemistry mechanisms into JAX

Key Takeaways

  • A kinetic preprocessor generates chemistry code for sensitivity analysis and hybrid physics–machine-learning models.
  • [KPPax](https://arxiv.org/abs/2610.37655) converts atmospheric chemistry mechanisms written in the Kinetic PreProcessor language into differentiable JAX code.
  • The authors aim to make existing chemical mechanisms usable in machine-learning and parameter-estimation workflows without manually deriving separate adjoint code.
  • ## A different output for existing mechanisms
  • Traditional KPP workflows generate efficient Fortran or C.

KPPax converts atmospheric chemistry mechanisms written in the Kinetic PreProcessor language into differentiable JAX code. The authors aim to make existing chemical mechanisms usable in machine-learning and parameter-estimation workflows without manually deriving separate adjoint code.

A different output for existing mechanisms

Traditional KPP workflows generate efficient Fortran or C. KPPax instead produces code that can be just-in-time compiled and vectorised in JAX. It evaluates rate laws using sparse operations in log space to retain differentiability.
The system uses the diffrax library’s implicit solvers and forward- and reverse-mode adjoints to integrate stiff chemical systems. This connects the chemistry calculation with automatic differentiation, allowing sensitivities to be computed within the same programming framework used for learning.

Test which parameters can be identified

The first demonstrated application uses the POLLU benchmark mechanism. An automatically differentiated Jacobian supports a Fisher information calculation to predict which of 25 rate coefficients a flow-tube-like setup can identify.
Noisy synthetic estimation then checks those predictions. The authors report agreement within about 30% for the predicted uncertainties of nine identifiable coefficients, while some fast reactions constrain ratios rather than individual values. These are controlled identification experiments, not field validation of every atmospheric mechanism.

Learn a correction through the chemical model

A second application trains a neural correction to an emission inventory using offline retrievals, then fine-tunes it through the chemical model. The paper reports a 2.4–5-fold reduction in biogenic-correction error in that setup.
For researchers, the useful contribution is a route from an existing mechanism description to sensitivity-aware estimation and hybrid modelling. Applying it to a new mechanism still requires checking numerical behaviour, parameter identifiability and the information available in the intended observations. Automatic differentiation does not make an otherwise unidentifiable parameter identifiable.

Comments