KPPax: preprocess differentiable kinetic ODE for machine learning in atmospheric chemistry

KPPax: preprocess differentiable kinetic ODE for machine learning in atmospheric chemistry Wenqing Peng∗1,3 , Michael Boy1,2,3 , and Zhi-Song Liu2,3

arXiv:2610.37655v1 [cs.AI] 27 Oct 2026

1

Institute for Atmospheric and Earth Systems Research, University of Helsinki, Helsinki, Finland 2 LUT University, Department of Computational Engineering, Lahti, Finland 3 Atmospheric Modelling Centre Lahti (AMC-Lahti), Lahti, Finland

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

1

Introduction

Chemicals and aerosols in the ambient air affect human health, weather, and climate (Seinfeld and Pandis, 2016). Their atmospheric concentrations vary strongly in space and time, governed by emission, chemical transformation, aerosol formation, transport by meteorological conditions, and eventual removal by deposition (Brasseur and Jacob, 2017). Comprehensive atmospheric chemical mechanisms, such as the near-explicit Master Chemical Mechanism (Jenkin et al., 1997; Saunders et al., 2003), contain thousands of species coupled through more than ten thousand reactions, and hand-coding the resulting stiff ODE system and its Jacobian is tedious and error-prone. The Kinetic PreProcessor (KPP) automates this process (Damian et al., 2002). Given a chemical mechanism specified in a dedicated description syntax, the KPP language, KPP parses the reaction list with a Flex/Yacc scanner–parser and generates Fortran or C code that evaluates the time derivative of the concentrations and its Jacobian. The generated code exploits the sparsity of the Jacobian through species reordering and sparse linear algebra routines, which keeps the implicit time stepping required by stiff chemistry efficient, and it is paired with a suite of stiff numerical integrators, including Rosenbrock, BDF, and Runge–Kutta methods. KPP also generates direct-decoupled and adjoint sensitivity code, which efficiently computes derivatives of the model solution with respect to rate coefficients, initial concentrations, and other parameters. These sensitivities support applications such as parameter estimation, uncertainty analysis, and variational data assimilation (Sandu et al., 2003). This tight coupling between mechanism parsing, sparse Jacobian generation, and forward and adjoint sensitivity code generation has made KPP a standard tool in atmospheric chemistry modelling, but it also means that the resulting code ∗ Corresponding author: wenqing.peng@helsinki.fi

1

inherits the constraints of the Fortran and C environments and of the CPU hardware for which it was designed. Other tools to preprocess and solve chemical kinetics have been developed for other software and hardware environments. pykpp (Henderson) is a preprocessor written in Python, which parses mechanisms written in the KPP language and solves the kinetic ODE with SciPy solvers. PyBox generates 0-dimensional box models that include both atmospheric chemistry and aerosol processes, in which the chemistry is solved by the Assimulo ODE solver (Topping et al., 2018), and PyCHAM reuses the mechanism parsing of PyBox but solves the ODE with SciPy (O’Meara et al., 2021). The Multi-Scale Infrastructure for Chemistry and Aerosols (MUSICA) is implemented in C++ and provides an option to run on accelerators such as GPUs (Pfister et al., 2020). For the EMAC model, a source-to-source tool translates the KPP-generated Fortran code into CUDA kernels to solve the chemical kinetics on GPUs (Alvanos and Christoudias, 2017). TChem-atm couples gas-phase and aerosol chemistry and runs on CPUs and GPUs through the Kokkos performance-portability library, constructing the Jacobian automatically for a range of stiff ODE solvers (Díaz-Ibarra et al., 2026). These developments target the efficient forward simulation of process-based atmospheric models, which are built on explicit physical and chemical processes and are therefore interpretable, since the driving forces in the model can be perturbed to test their influence on the output. However, process-based models are limited by the processes that they do not resolve, and amending the model structure or tuning its parameters against observations remains laborious and computationally expensive. Data-driven approaches, for example neural networks trained on historical atmospheric observations, offer flexible and fast emulation of atmospheric processes with comparable accuracy. ChemNNE is a neural-network emulator of the atmospheric chemical ODE that is considerably faster than the numerical solver (Liu et al., 2025), and TopoFlow is a topography-aware neural network for high-resolution airquality prediction (Kheder et al., 2026). Neural networks can also partially solve the estimation of rate coefficients in chemical ODEs (Peng et al., 2025; Liu et al., 2026). However, purely data-driven models often lack interpretability and the ability to predict anomalies or extreme events, and their training is data-intensive and becomes infeasible when observations are sparse. Neural networks are trained by gradient descent, which relies on automatic differentiation, and although this is possible in Fortran, the majority of neural-network models are developed in Julia or in Python frameworks such as PyTorch and JAX. Hybrid models that combine physically based process models with data-driven neural networks are expected to unite the flexibility of neural networks with the interpretability and physical constraints of process-based models, enabling both fast emulation and learning from data and an understanding of the underlying processes (Shen et al., 2023; Innes et al., 2019; Gelbrecht et al., 2023). For example, the hybrid weather and climate model NeuralGCM, which combines a differentiable dynamical core with neural networks, both implemented in JAX, achieves a performance comparable to the best machine-learning and process-based models (Kochkov et al., 2024). In ecohydrology, JAX-CanVeg integrates a land-surface model with neural-network components and improves the simulation (Jiang et al., 2025). Such hybrid models require the process model itself to be differentiable, so that gradients can be propagated through it to the network parameters. General-purpose differentiable modelling tools for reaction networks exist, such as Catalyst.jl in Julia (Loman et al., 2023), but to the best of our knowledge no tool generates differentiable code from mechanisms written in the KPP language, in which atmospheric mechanisms such as the MCM are distributed and used. Differentiability matters most for inverse problems, in which model parameters or inputs are inferred from observations of the model output. Estimating rate coefficients from chamber or flow-tube observations is such a problem, and it is often ill-posed because the observations constrain only some combinations of the rate coefficients, so that the identifiability of the parameters should be analysed before they are estimated (Guillaume et al., 2019). Another example is the inversion of emissions from concentration observations, which has traditionally relied on variational data assimilation with hand-derived adjoint models (Elbern et al., 2007; Henze et al., 2007). To learn how emissions depend on environmental drivers, the retrieved emission fields can be used to train a neural network offline, in the same way as data assimilation has been combined with machine learning to learn model dynamics from sparse and noisy observations (Brajard et al., 2020). Such a two-step approach, however, discards the information that the individual retrievals cannot resolve, whereas online training, in which the network is fitted to the observations through the chemical model, requires the gradient of the model output with respect to the network parameters. To address this, we present KPPax, a kinetic preprocessor implemented in the JAX framework that generates differentiable, just-in-time compiled, and vectorisable chemistry code from mechanisms written

2

in the KPP language, which is directly usable within the data-driven modelling ecosystem. 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 (Kidger, 2022), which enables inverse modelling, parameter estimation, and hybrid modelling of physical and neural models. Section 2 describes the implementation. Section 3 uses the sensitivities from KPPax to analyse the identifiability of the rate coefficients of a benchmark mechanism before estimating them from noisy synthetic observations. Section 4 trains a neural-network correction of an emission inventory and shows that online fine-tuning through the chemical model improves on the offline two-step approach. Section 5 discusses the limitations, and Sect. 6 concludes.

2

Implementation of differentiable kinetic pre-processing

2.1

The chemistry kinetic problem

An atmospheric chemical reaction system can be mathematically formulated by a coupled system of stiff ordinary differential equations (ODEs). Consider a reaction system of n species and m reactions. We denote the concentration of species i as yi and reaction j as rj , with 1 ≤ i ≤ n, 1 ≤ j ≤ m. Let y = [y1 , ..., yn ]T be a vector of concentration of all the species. The stoichiometric coefficients define the quantitative relationship of the reactants and the products involved in the reaction. We define the stoichiometric coefficient s− i,j as the number of molecules of species + i that react and consumed in the reaction j, and similarly si,j as the number of molecules produced. The reaction j can be generically formulated as: X X s− − s+ (1) i,j yi → i,j yi The kinetic reaction rate law define the reaction rate ωj of the reaction j as a proportion with the multiplication of the concentration of the reactants raised by the partial order for reactant s− i,j , with the proportionality constant defined as the rate coefficient (rate constant) kj : ωj (t, y) = kj (t)

n Y

s−

yi i,j

(2)

i=1

Depending on the reaction system, the rate coefficient may be either a constant or a time-dependent quantity that varies with environmental conditions. For instance, the photolysis rate of NO3 depends on the actinic flux (light intensity), while the rate of cross-reactions among intermediate peroxy radicals depends on the total concentration of peroxy radicals present in the system. The vector k consists of the rate coefficients kj from all j reactions. The concentration of species i change at a rate given by the cumulative effect of consumption and production in all reactions: m X dyi − (s+ (3) = i,j − si,j )ωj (t, y) dt j=1 + − We organise the stoichiometric coefficients s− and S + and the reaction rates in i,j and si,j in matrices S T + − a vector ω = [ω1 , ..., ωm ] , and let S = S − S , the rate equation (3) can be written as:

dy = (S + − S − )ω(t, y) = Sω(t, y) = f (t, y), dt

(4)

This forms an initial value problem and the solution y(t) ∈ Rn represents the time evolution of the concentration, starting from the initial concentration y(t0 ) over the time span (t0 , t): Z t y(t) = y(t0 ) +

f τ, y(τ ) dτ

(5)

t0

2.2

Reaction parsing

The proposed KPPax adopts the KPP language as a description syntax for kinetic reaction system. The KPP language, developed by the KPP project, consists of sections, commands, and inline codes. KPP sections contain the description of reaction schemes, including the species and reactions. KPP commands 3

control the code generation and the behaviour of the generated code. More custom behaviour can be implemented as inline code, and be inserted into a certain location in the generated code. KPPax uses only the sections that describe the mechanism, i.e. the species and the equations, because its objective is to generate the chemistry code alone, and it is the responsibility of the user to provide the initial concentrations and to construct the customised simulation. Following the taxonomy of reaction parsing in (Damian et al., 2002), KPPax parses the kinetic equations in a mixture of the totally integrated approach and the preprocessing approach. The former parses the chemical mechanism once into a set of coefficient arrays that are subsequently evaluated by generic array operations at run time, and the latter generates mechanism-specific source code ahead of time, as KPP itself does. We utilise the parsing module from pykpp, which defines and parse the KPP language using the pyparsing module and returns a list of reactions composed of reactants, products, and rate coefficient expression. The stoichiometric coefficients are converted into sparse array literals in the coordinate list (COO) format, with row/column indices and values stored in separate arrays (i, j, v). The time dependent rate coefficients are converted from Fortran expression to JAX expression and wrapped as a python function. To exploit the computational power of accelerators, we implement the rate law in fully sparse matrix operation. Rather than materialising S − and S as dense matrices, the rate law (Eq. 2-4) is evaluated directly from their COO representation using index-gather and segment-reduction primitives, as summarised in Algorithm 1. The reactant product term in Eq. 2 is computed as a segment product: partial s−

powers yi i,j are gathered at the non-zero row indices of S − and reduced onto the reaction axis j by scatter-multiplication. Algorithm 1 Sparse evaluation of the rate law Require: y; k; S − : (ir , jr , vr ); S: (is , js , vs ) 1: p ← y[ir ] vr {gather and raise to partial order} 2: ω ← k ⊙ segment_prod(p, jr ) {scatter-multiply onto reaction axis} 3: ẏ ← segment_sum(vs ⊙ ω[js ], is ) {scatter-add onto species axis} 4: return ẏ This power-multiply formulation is not differentiable in practice, because the reactant product is realised as a scatter-multiplication over repeated reaction indices, for which JAX does not implement a differentiation rule. We avoid this implementation limitation by reformulating the rate law in log space, X log ωj = log kj + s− (6) i,j log yi , i

which replaces the scatter-multiplication with a scatter-addition, segment_sum, before a single exponentiation recovers ω: Algorithm 2 Sparse evaluation of the rate law in log space Require: y; k; S − : (ir , jr , vr ); S: (is , js , vs ) 1: q ← vr ⊙ log y[ir ] {gather and scale log-concentrations by partial order} Q s− i,j 2: l ← segment_sum(q, jr ) {scatter-add onto reaction axis, i.e. log i yi } 3: ω ← k ⊙ exp(l) {recover reaction rate} 4: ẏ ← segment_sum(vs ⊙ ω[js ], is ) {scatter-add onto species axis} 5: return ẏ Since JAX defines a standard gradient rule for segment_sum, this substitution leaves the sparse indexing pattern and computational cost of Algorithm 1 unchanged while restoring full differentiability of the rate law with respect to both y and k, at the cost of requiring y to remain strictly positive throughout the integration. Atmospheric chemical mechanisms are numerically stiff, because reactive radical species relax on timescales of milliseconds to seconds, while longer-lived species evolve over hours, so the system’s Jacobian eigenvalues span many orders of magnitude. Explicit integration schemes are constrained by this stiffness to prohibitively small time steps for stability, rather than accuracy, and are therefore unsuitable, which is why KPP itself provides implicit and Rosenbrock-type solvers (Sandu et al., 2003). We integrate Eq. (4) using diffrax (Kidger, 2022), a JAX-native library of differentiable ODE solvers, selecting the implicit

4

Kvaerno family of solvers, which combine embedded error estimates with an implicit treatment of the stiff terms via a diagonally implicit Runge-Kutta scheme, avoiding the step-size collapse of explicit methods. Step size is controlled adaptively by a proportional-integral (PID) controller, which adjusts dt to meet a prescribed local error tolerance rather than relying on a fixed or manually tuned step size. Because concentrations must remain strictly positive for the log-space rate law (Eq. 6), we additionally clip y to a small positive floor at each evaluation of the right-hand side, preventing the solver from probing negative or exactly zero concentrations during error-controlled step adjustment. We perform all stiff ODE integration in 64-bit double precision.

2.3

Automatic differentiation

Automatic differentiation (AD) computes exact derivatives of a numerical program by decomposing it into elementary operations and applying the chain rule systematically, either forward alongside the primal computation (forward mode) or backward from the output (reverse mode). Once the rate law and its numerical integration are expressed as JAX-traceable array operations, the sensitivity of an output quantity, for instance an observed concentration at time tF , with respect to an input quantity, for instance an emission rate or a rate coefficient, can therefore be obtained automatically, without deriving the corresponding tangent-linear or adjoint equations by hand. Differentiating through an ODE solve is nonetheless a distinct problem from differentiating a single expression, since the numerical solution is produced by an iterative sequence of solver steps rather than a closed-form function. Two broad strategies exist for obtaining the required sensitivities, namely the optimize-then-discretize adjoint, which derives an adjoint ODE from the continuous forward equations and integrates it backwards independently of how the forward solve was discretised, and the discretizethen-optimize adjoint, which differentiates the forward numerical scheme itself, step by step, yielding the exact gradient of the numerical solution rather than an approximation to the gradient of the underlying continuous problem (Sandu et al., 2003). KPP implements both, but its discretize-then-optimize adjoint must be derived analytically for each integrator in turn. For example, Sandu et al. (2003) derive it explicitly for the linearly implicit Euler method, requiring not only the Jacobian but also its derivative with respect to concentration (the Hessian tensor) and dedicated sparse linear-algebra routines to evaluate it efficiently. They note this construction to be considerably more demanding than their optimize-thendiscretize adjoint, and correspondingly, KPP provides hand-derived discretize-then-optimize adjoints for only a subset of its integrators. diffrax (Kidger, 2022) instead obtains a discretize-then-optimize adjoint of the forward solve through reverse-mode AD applied directly to the solver’s internal steps, exposing the differentiation strategy as a configurable component of the solve rather than a property fixed by the choice of integrator. Naively differentiating through every accepted step this way would require storing each intermediate solver state, incurring memory cost proportional to the number of steps. The recursive checkpoint adjoint avoids this by checkpointing only a logarithmic subset of states and recomputing the remainder on demand during the backward pass, keeping memory cost sub-linear in step count while still returning the exact discrete gradient. Because this adjoint is generated automatically for any solver JAX can trace, KPPax removes the need to derive a Hessian or a solver-specific adjoint step by hand for every integrator, and with it the correctness risk that Sandu et al. (2003) themselves highlight for hand-written or generalpurpose adjoint-compiler code. The trade-off is computational rather than conceptual: checkpointed recomputation adds cost that a hand-derived discretize-then-optimize adjoint, once implemented, would not incur, particularly for very large mechanisms. We consider this an acceptable exchange for the substantial reduction in implementation and maintenance effort it affords when extending KPPax to new mechanisms or integrators.

3

Application I: identifiability and parameter estimation

Rate coefficients of atmospheric chemical mechanisms are uncertain, and many of them can only be constrained indirectly, through laboratory or chamber observations of a few species. A fit of all rate coefficients to such observations is often non-unique, because different combinations of parameter values reproduce the observations equally well. A parameter is non-identifiable when its value cannot be determined uniquely from the data, and this can arise from the structure of the model, the observed quantities and the objective function (structural non-identifiability), or from the experimental design and the observation error (practical non-identifiability) (Guillaume et al., 2019). Therefore identifiability is recomended to be assessed routinely before parameters are estimated, so that the estimation is 5

restricted to what the data can support. Derivative-based identifiability analysis requires the sensitivity of the model output with respect to every parameter, which is laborious to derive by hand for a chemical mechanism and inaccurate and expensive to approximate by finite differences. With KPPax, the exact sensitivities are obtained by automatic differentiation through the stiff solver at a cost comparable to a few forward integrations. In this section, we use these sensitivities first to predict, before any fit, which rate coefficients a given experimental design can constrain, and then to estimate the rate coefficients from noisy synthetic observations, which tests these predictions.

3.1

Method

We parameterise the rate coefficients as k = kref ⊙ exp(θ), where kref are the reference values of the mechanism and θ is the vector of log correction factors. The observation operator h(θ) stacks the logarithms of the simulated concentrations of all observed species at all observation times and in all experiments. Given observations ỹ with independent Gaussian errors of standard deviation σ in log space, the residuals and the least-squares cost are r(θ) =

h(θ) − ỹ , σ

χ2 (θ) = r(θ)T r(θ).

(7)

Because h and θ are both logarithmic, the Jacobian J = ∂h/∂θ contains the relative sensitivities ∂ ln y/∂ ln k, which are comparable across species and reactions whose magnitudes differ by many orders. Given the number of θ and y, we compute J by forward-mode automatic differentiation through the solver, which propagates one tangent per parameter. The local practical identifiability at the true parameters is characterised by the Fisher information matrix (FIM) 1 (8) F = 2 J T J = V ΛV T , σ which for noise-free observations equals the Gauss–Newton approximation of the Hessian of χ2 /2 at the truth. Each eigenvector vm (a column of V ) is a direction in parameter space, i.e. a combination of rate coefficients, and its eigenvalue λm measures how strongly the data constrain this combination. The Cramér–Rao bound gives the smallest covariance F −1 that any unbiased estimator can attain, which −1/2 along vm and to a standard deviation of corresponds to a standard deviation of λm σpCR =

h

2 1/2
i1/2  X Vpm F −1 pp = λm m

(9)

for the individual parameter θp . We classify a rate coefficient as identifiable if σpCR < 0.1, i.e. if its onestandard-deviation uncertainty is about 10% (a factor of e0.1 ≈ 1.1), which is small compared with the factor-of-two uncertainty (ln 2 ≈ 0.69) of the initial guesses used below. In the experiment of Sect. 3.2, the classification does not depend on the exact threshold, because the values of σpCR jump from 0.04 to 0.15 between the identifiable and the remaining parameters (Sect. 3.3). Directions with eigenvalues close to zero are combinations of rate coefficients that the observations cannot distinguish, and their loadings Vpm show which reactions are confounded. This analysis requires only the Jacobian at a single parameter point and is performed before any fit. To test the predictions, we estimate the parameters from noisy synthetic observations. In each of 20 trials, we draw a new noise realisation of ỹ at the truth and a new initial guess θ (0) ∼ N (0, (ln 2)2 I), i.e. initial rate coefficients that typically deviate from the truth by a factor of 2, and minimise χ2 with the Levenberg–Marquardt algorithm, i.e. Gauss–Newton iterations with adaptive damping, of the optimistix library (Rader et al., 2024), using the same forward-mode Jacobian and at most 200 iterations. All trials are solved together under jax native vectorisation jax.vmap. For each parameter, we compare the root-mean-square (RMS) deviation of the estimates from the truth across the trials with σpCR . For identifiable parameters, the two should agree. For non-identifiable parameters, the estimates are expected to remain scattered at least as widely as the initial guesses, because the data cannot pull them back, even though the fit to the observations is as good as at the truth.

3.2

Experimental design

We use the POLLU scheme (Verwer, 1994) with its reference rate coefficients, which are constant in time, and estimate all 25 rate coefficients. The design mimics a flow-tube experiment, in which the reaction 6

(a) 105

(c)

103

recovered (per trial)

FIM eigenvalue λ

101 10−1

R8 ALD + OH → C2O3

−3

R3 HO2 + NO → NO2 + OH

10

predicted ±2 sd

R4 HCHO → 2 HO2 + CO

R20 SO2 + OH → SO4 + HO2

10−5

R14 NO2 + OH → HNO3

10−7

R6 HCHO + OH → HO2 + CO

10−9

R9 C2O3 + NO → NO2 + MEO2 + CO2 R13 CH3O → HCHO + HO2

identifiability threshold (sd = 0.1)

10−11 0

5

10

15

20

25

R12 MEO2 + NO → CH3O + NO2 R10 C2O3 + NO2 → PAN

FIM eigen-direction

R11 PAN → C2O3 + NO2 R21 NO3 → NO

(b)

R5 HCHO → CO R23 NO2 + O3 → NO3

loading 1.00

R1 NO2 → NO + O(3P)

R7 ALD → MEO2 + HO2 + CO

R2 NO + O3 → NO2 R5 HCHO → CO

R22 NO3 → NO2 + O(3P)

0.75

R1 NO2 → NO + O(3P)

R7 ALD → MEO2 + HO2 + CO

R15 O(3P) → O3

0.50

R10 C2O3 + NO2 → PAN

R2 NO + O3 → NO2

R11 PAN → C2O3 + NO2

0.25

R15 O(3P) → O3 R16 O3 → O(1D)

R24 NO3 + NO2 → N2O5 R25 N2O5 → NO3 + NO2

0.00

R17 O3 → O(3P)

R18 O(1D) → 2 OH

R18 O(1D) → 2 OH

−0.25

R19 O(1D) → O(3P)

R17 O3 → O(3P)

R21 NO3 → NO

R16 O3 → O(1D)

−0.50

R22 NO3 → NO2 + O(3P)

R19 O(1D) → O(3P)

−0.75

R23 NO2 + O3 → NO3

−10

R24 NO3 + NO2 → N2O5

−1

0

1

10

̂ true) ln(θ/θ

−1.00

R25 N2O5 → NO3 + NO2

16 17 18 19 20 21 22 23 24 25

FIM eigen-direction (below threshold)

Figure 1: Identifiability analysis and estimation of the 25 POLLU rate coefficients. (a) Eigenvalues of the Fisher information matrix in descending order, where the dashed line marks λ = 100, which corresponds to a Cramér–Rao standard deviation of 0.1 along the eigenvector. (b) Loadings of the ten eigenvectors below this threshold on the rate coefficients that contribute to them. (c) Estimated log rate coefficients relative to the truth in the 20 trials (dots), compared with the interval of ±2σpCR predicted from the FIM before the fit (grey bars). The horizontal axis is linear between −1 and 1 and logarithmic outside. Reactions are numbered as in POLLU. time is fixed and the injected precursor concentrations are varied. POLLU initialises six species (NO, O3 , HCHO, CO, ALD, and SO2 ), and we generate eight initial conditions by multiplying each of these initial concentrations by an independent log-normal factor with a log-space standard deviation of ln 2. All other species start from zero. Each initial condition is integrated for 60 min, and all 20 species are observed once at the end, which gives 160 observations with a log-space noise of σ = 0.05. The ODEs are integrated with the Kvaerno3 solver with relative and absolute tolerances of 10−8 and 10−9 ppm. If all species were observed continuously in time, all 25 rate coefficients of POLLU would be globally structurally identifiable, as we verified with the StructuralIdentifiability.jl package (Dong et al., 2023). Any non-identifiability in this experiment is therefore practical and arises from the design, which observes a single time after a fixed reaction time. Such a design is outside the scope of the structural analysis, which assumes continuous-time outputs, whereas the FIM applies to any set of observations.

3.3

Results

The automatically differentiated Jacobian (160 × 25) agrees with the central finite-difference approximation to a relative error of 8.7 × 10−7 , and it takes 0.77 s, compared with 28.7 s for the 50 evaluations of the forward model required by the finite differences. The finite-difference approximation also depends on the choice of the step size, whereas automatic differentiation gives the exact derivative of the numerical solution. The eigenvalues of the FIM span 16 orders of magnitude (Fig. 1a). Fifteen eigen-directions are constrained to a standard deviation better than 0.1, and the remaining ten fall off steeply, with the four smallest eigenvalues below 1 and the smallest at 10−11 . Since the well-constrained directions are

7

combinations of rate coefficients, only nine individual rate coefficients have σpCR < 0.1 (Fig. 1c). These are the reactions of the HOx , HCHO, ALD, and SO2 chemistry (R3, R4, R6, R8, R9, R12, R13, R14, and R20). The loadings of the weakly constrained eigenvectors (Fig. 1b) group the non-identifiable rate coefficients into three clusters of fast reactions. These are the NO–NO2 –O3 photostationary cycle (R1, R2, and R15), the equilibrium between NO3 , NO2 , and N2 O5 (R24 and R25), and the fast interconversion of O3 , O(1 D), and O(3 P) (R16–R19). Within each cluster, the reactions are fast compared with the reaction time of 60 min, so the observed concentrations are set by the ratios of the rate coefficients, and their absolute values are not constrained. For example, the eigenvector that loads equally on R24 and R25 corresponds to a simultaneous increase of both rate coefficients, which leaves the N2 O5 equilibrium constant k24 /k25 unchanged. The noisy estimation confirms these predictions (Fig. 1c). For the nine identifiable rate coefficients, the RMS deviation of the estimates from the truth agrees with σpCR to within about 30%, which shows that the local, linearised prediction from a single Jacobian quantifies the achievable accuracy correctly. The seven rate coefficients with intermediate σpCR of 0.15–0.57 are estimated with errors of the predicted size, which is somewhat larger for R22 and R23. The non-identifiable rate coefficients drift far from the truth, by up to 15 in log units, i.e. by more than six orders of magnitude. For the cluster R1, R2, and R15, the RMS deviation of 6.3 is seven times the Cramér–Rao prediction, because the drift follows a curved valley of the cost function along which the linearisation at the truth is not valid. For the O3 –O(1 D)–O(3 P) cluster, the RMS deviation is much smaller than the Cramér– Rao value, which in this case only indicates that the cost is flat, so that the spread of the estimates is set by the initial guesses and the finite number of iterations. The drift is not random, because the estimates within a cluster move together and keep their ratios constant, which is visible in Fig. 1c as identical estimates of R1, R2, and R15, and of R24 and R25, in each trial. The FIM therefore identifies correctly which combinations of rate coefficients the design cannot constrain, but it underestimates how far the estimates can move along these combinations. All fits reach the noise level of the observations. The median of χ2 /n over the trials is 0.85 at the estimates and 1.01 at the truth, and the reduction is consistent with the number of parameters p that the fit can adjust to the noise, since 1 − p/n is 0.84–0.91 for p between the 15 well-constrained directions and all 25 parameters. Only two of the 20 trials met the convergence tolerance of the Levenberg–Marquardt algorithm, whereas the others reached the maximum of 200 iterations. These trials did not fail, but kept creeping along the flat valleys of the non-identifiable clusters, where each step reduces the cost only marginally, while their estimates of the identifiable rate coefficients agree with the Cramér–Rao prediction (Fig. 1c). This behaviour is itself a symptom of nonidentifiability, and it illustrates why the analysis should precede the fit: once the FIM has revealed the confounded clusters, their rate coefficients can be fixed at literature values or reparameterised by their ratios, and only the identifiable combinations are estimated.

4

Application II: hybrid neural emission inversion

Emission inventories are a major source of uncertainty in chemical transport modelling. Their dependence on environmental drivers, such as temperature, light, or population density, is usually prescribed by simple empirical functions. Concentration observations constrain emissions only indirectly, through the chemistry, so a correction to the inventory has to be inferred by inverting the chemical model, as in four-dimensional variational (4D-Var) emission inversion (Elbern et al., 2007). A neural network that maps environmental drivers to an emission correction can then be trained offline, by regressing it on the emission fields retrieved by such an inversion. However, the accuracy of the network is then limited by the retrieved fields themselves, because each retrieval is regularised towards the prior inventory and is constrained only by the observations of its own scenario, so that information that the observations do contain, but that a single retrieval cannot isolate from the prior, is lost before the network is trained. Here, we use KPPax to learn a neural-network correction of the emission inventory as a function of its environmental drivers. Because the network sits inside the kinetic ODE, its training requires gradients of the simulated concentrations with respect to the network parameters, propagated through the stiff chemistry, which is exactly what the differentiable solver provides. We compare the offline training with online training through the ODE in a synthetic experiment where the true correction is known.

8

4.1

Experimental setup

We use the POLLU scheme (Verwer, 1994) (20 species, 25 reactions) in a 0-dimensional box model with emission and dilution towards a background concentration cbg :
y − cbg dy = S ω t, y + E(t) − , dt τ

(10)

where the first term is the rate law of Eq. (4) generated by KPPax, and τ = 12 h is the dilution time scale. The rate coefficients of POLLU are constant, so we impose a diurnal cycle on its eight photolysis reactions by scaling them with a smooth daylight factor that equals 1 at noon and vanishes between 18:00 and 06:00. Concentrations are in ppm and time is in minutes, following POLLU. Emissions come from an anthropogenic sector (s = anth), which emits NO, NO2 , CO, SO2 , HCHO, and ALD, and a biogenic sector (s = bio), which emits ALD and HCHO. The total emission rate is X

 E(t) = es gs x(t) ps (t) exp fs x(t) . (11) s

Here, es is a fixed speciated base emission rate, and ps (t) is a diurnal profile known a priori and normalised to a daily mean of one, which follows a traffic double peak at 08:00 and 18:00 for the anthropogenic sector and the daylight factor for the biogenic sector. The covariate vector x = (T, R, log10 P, danth , dbio ) consists of the temperature T (in ◦ C), the daily maximum light level R ∈ [0, 1], the population density P , and two synthetic covariates, which can be interpreted as the share of fossil-fuel heating (danth ) and the vegetation type or thermal acclimation (dbio ). The inventory scaling gs is a deliberately crude function of the covariates, in which anthropogenic emissions scale linearly with population density, whereas biogenic emissions follow a fixed exponential temperature response and depend linearly on light:   ganth = 10log10 P −3 , gbio = exp 0.09 (T − 30) R. (12) The log correction factor fs is the unknown to be learned, and the inventory corresponds to fs = 0. The true correction depends on three covariates per sector and is constructed to be smooth and nonlinear, so that neither a linear nor an additive model can reproduce it:

2 true fanth = 0.6 e−[(log10 P −2.7)/0.35] + (0.3 + danth ) σ 8−T σ 3 (log10 P − 2.8) 2
+ 0.3 sin(2πdanth ) tanh 2 (log10 P − 2.75) − 0.3, (13)
2 true fbio = 0.07 (T − Topt ) − ln 1 + e0.3 (T −Topt ) + ln √ + 0.8 sin(2πdbio ) (R − 0.55) + 0.49, (14) 1 + 3R2 where σ is the logistic function and Topt = 23 + 6 dbio . The anthropogenic correction (Eq. 13) contains a non-monotonic dependence on population density, peaking in suburban densities, and a cold-day heating term in populated areas that grows with danth . Combined with the inventory, the biogenic correction (Eq. 14) gives a temperature response √ that rises at 0.16 ◦ C−1 at low temperature and has an optimum near Topt , and a light dependence 2R/ 1 + 3R2 that saturates instead of growing linearly. Each sector also includes a sinusoidal interaction with its synthetic covariate, and the constants centre both corrections at approximately zero over the training ranges (Fig. 2). Since biogenic emissions are small at low temperature, fbio there has little effect on any concentration and is only weakly constrained by the observations. In each scenario, the covariates are sampled uniformly from their training ranges. The daily mean temperature is chosen such that the diurnal temperature cycle, which has a range of 4–14 ◦ C and a maximum between 13:00 and 16:00, stays within 5–30 ◦ C, while R, log10 P , and danth , dbio are drawn from [0.2, 0.85], [2, 3.5], and [0.1, 0.9], respectively. Background O3 and CO are also sampled per scenario. Temperature varies over the day, so x and hence fs are functions of time. We generate 150 training, 30 validation, and 30 interpolation test scenarios. For each covariate, we additionally generate 20 extrapolation test scenarios in which that covariate lies in bands extending its training range by 12.5% of the range width on either side. Each scenario is spun up for one day with the true emissions, and then integrated over a one-day window from midnight, and the post-spin-up state is assumed to be known exactly and is used, without noise, as the initial state of the assimilation window. Observations are taken hourly from 01:00 to 24:00 and perturbed with multiplicative log-normal noise with a log-space standard deviation of 0.05. We consider two observation sets, namely all 20 species and only the five species routinely measured at air-quality monitoring stations (O3 , NO, NO2 , CO, and SO2 ). Each set is used with and without noise, giving four settings. 9

f_true

f_anth dummy_anth = 0.70

3.50

3.25

3.25

-0.25

3.25

3.00

0.00

3.00

0.25

2.75

2.75 2.50

-0.50

2.00 5

10 15 20 25 T (°C) f_bio dummy_bio = 0.30

30

5

-0.75

-0.50

25

30

-0.25

-1.00

0.00 0.25

0.2 5

10

15 20 T (°C)

25

30

1.0 0.5

0.5

0.0

0.4

−0.5

0.3 0.2

0.50

15 20 T (°C)

0.50

10

-0.50

0.25

0.00

0.3

-0.50

0.4

0.3

-0.75

0.4

5

-0.75

-0.50

-0.25

-0.75

0.6

30

-0.50

0.6

10 15 20 25 T (°C) f_bio dummy_bio = 0.70

-1.00 -0.75

0.7

0.5

0.2

0.8

0.7 -0.25

-1.00

-0.25

0.5

0.9 -1.00

0.25

R (daily max light)

10 15 20 25 T (°C) f_bio dummy_bio = 0.50

0.9

0.6

0.00

2.00 5

0.00

0.7

-0.25

30

0.8

0.25

2.25

2.00

0.9 0.8

2.50

0.00

2.25

0.7 5

-0.25

0.25

-0.5 0

0.00

2.25

0.25

0.0 0

2.50

0.50

0.2 5

0.25

-1.2 5

0.50

-1.0 0

2.75

-0.25

0.00

0.5 0

-0.7 5

0.25

0.6 0.4 0.2 0.0 −0.2 −0.4 −0.6

-0.50

-0.2 5

0.00

3.50

f_anth

3.50 3.00 log10 P

f_anth dummy_anth = 0.50

f_bio

f_anth dummy_anth = 0.30

5

10

15 20 T (°C)

25

−1.0

30

Figure 2: True log emission correction fstrue (Eqs. 13 and 14) for the anthropogenic (top) and biogenic (bottom) sectors, at three values of the synthetic covariate (columns). The dashed boxes mark the training ranges, and the regions outside them are the extrapolation bands.

4.2

Neural emission correction

The unknown correction fs is represented by one multilayer perceptron (MLP) per sector, which takes the three covariates of that sector as input, scaled linearly from their training ranges to [−1, 1]. Each MLP has two hidden layers of width 32 with softplus activations, and its output is bounded as fs = fmax tanh(MLPs /fmax ) with fmax = 3, so that an untrained network cannot push the chemistry into extreme states. The weights and bias of the last layer are initialised to zero, so the network reproduces the inventory (fs = 0) at the start of training. The two MLPs together have 2434 parameters, which we denote by θ. The correction varies in time through the temperature,
and we evaluate it at hourly knots tj (j = 0, . . . , 24) over the one-day window, as fs,j = fs x(tj ); θ , and interpolate between the knots with a cubic Hermite spline. This gives the same time representation of the correction as the 4D-Var control variable in Sect. 4.3, so that the offline and online approaches differ only in how the knot values are constrained. Both approaches measure the misfit to the observations y obs in log space, as in Eq. (7) but with a detection limit: obs ln max(yt,i , yDL ) − ln max(yt,i , yDL ) rt,i = , (15) σobs for every observation time t and observed species i, where yDL = 5 × 10−5 ppm is a detection limit that prevents the logarithm from amplifying differences between near-zero concentrations, and σobs = 0.05 in all settings. All ODEs are integrated with the Kvaerno3 solver with relative and absolute tolerances of 10−6 and 10−12 ppm.

4.3

Offline two-step training

The offline approach first retrieves the correction separately for each scenario n by 4D-Var and then trains the network to reproduce the retrievals. The control variable of scenario n is the vector δn of knot values of both sectors, and the retrieval minimises   1X T |tj − tl | 1X 2 rt,i + δn,s B −1 δn,s , Bjl = σb2 exp − , (16) Jn (δn ) = 2 t,i 2 s L

10

where the background term keeps the retrieval close to the inventory (δ = 0) with a prior standard deviation σb and a temporal correlation length L. We minimise Eq. (16) by Gauss–Newton iterations with Levenberg–Marquardt damping for 15 iterations (Courtier et al., 1994). The Jacobian of the residuals with respect to the 50 control variables is obtained by forward-mode automatic differentiation through the stiff solver because the number of controls is much smaller than the number of residuals. The posterior post standard deviation σn,j,s of each knot value is taken from the diagonal of the inverse Gauss–Newton Hessian Hn = JnT Jn + B −1 at the minimum. The prior hyperparameters (σb , L) are chosen from the data alone rather than from the truth. For each candidate pair, we sum the Laplace approximation of the log marginal likelihood, ln p(ynobs | B) ≈ −Jn (δn∗ ) +

1 1 ln det B −1 − ln det Hn + const, 2 2

(17)

over the training scenarios and select the maximum on a grid. Unlike the cost Jn itself, the evidence penalises prior freedom that the data cannot use, and it therefore discriminates between priors that fit the observations equally well. The selected prior has σb = 0.35 in three of the four settings and σb = 0.25 with noise-free station observations. The selected correlation lengths are 48 h and 96 h with all species observed, and 192 h, the largest value on the grid, with station observations, which means that these observations cannot resolve the variation of the correction within a day. In the second step, the network is trained to reproduce the retrieved knot values δ̂ from the covariates at the knots of all training scenarios, by minimising the weighted mean squared error h i2
1 X post
−2 wn,j,s fs xn (tj ); θ − δ̂n,j,s , wn,j,s ∝ σn,j,s , (18) Loff (θ) = N n,j,s where N is the number of terms and the weights are normalised to a mean of one per sector. Without weighting, retrievals that remain at the prior because the observations carry no information, such as biogenic emissions at night, would enter the regression as if they were genuine values of zero. We train with the Adam optimiser (Kingma and Ba, 2015) for 5000 full-batch steps with a learning rate decaying from 10−3 along a cosine schedule. No ODE is solved in this step. The weighted regression was more accurate than the unweighted one for the anthropogenic sector in all settings and for the biogenic sector when all species are observed, whereas the two performed similarly for the biogenic sector with station observations. We therefore use the weighted regression as the offline reference below.

4.4

Online fine-tuning through the ODE

In the online approach, the network parameters are fitted directly to the concentration observations of all training scenarios, with the network inside the ODE (Eqs. 10 and 11), by minimising h
2
i 1 X X J (θ) = rn,t,i δn (θ) , δn (θ) = fs xn (tj ); θ . (19) 2 j,s t,i n∈train

The data term is the same as in Eq. (16), but the knot values of every scenario are now produced by the shared network rather than being free control variables, and no background term is used, because the shared network already restricts the knot values to a smooth function of the covariates. This allows information from all scenarios to be combined, including information that no single scenario constrains well enough to be retrieved on its own. Training requires the gradient of the misfit with respect to θ, which passes through the stiff chemistry and is therefore only available with a differentiable ODE such as the one generated by KPPax. We initialise the network with the offline solution and fine-tune it by minimising Eq. (19) with Gauss– Newton iterations with Levenberg–Marquardt damping. By the chain rule, the Jacobian of the residuals of scenario n with respect to the parameters factorises as Jδ,n Dn , where Jδ,n is the residual Jacobian with respect to the 50 knot values used in 4D-Var, and Dn = ∂δn /∂θ is the Jacobian of the network at the knots, which requires no ODE solve. The Gauss–Newton Hessian and gradient are accumulated over the training scenarios as X X  −1 T T H= DnT Jδ,n Jδ,n Dn , g= DnT Jδ,n rn , ∆θ = − H + λ diag(H) g, (20) n

n −2

and the damping λ, initialised to 10 , is divided by 3 after an accepted step and multiplied by 4 after a rejected one. The ODE sensitivities Jδ,n are therefore the same as in one 4D-Var iteration, and their 11

Table 1: RMSE of the learned log emission correction with respect to the truth, for the offline two-step training (offline) and after the online Gauss–Newton fine-tuning (+ online). “Train” denotes the training ranges of the covariates, and “extrap.” the extrapolation bands outside them. anthropogenic

biogenic ◦

observations

method

train

extrap.

train

5–12 C

12–20 ◦ C

20–30 ◦ C

extrap.

all species, no noise

offline + online

0.063 0.056

0.144 0.144

0.162 0.041

0.290 0.063

0.062 0.025

0.033 0.028

0.315 0.128

all species, noise

offline + online

0.063 0.043

0.141 0.103

0.145 0.028

0.260 0.037

0.052 0.022

0.028 0.023

0.289 0.100

station, no noise

offline + online

0.063 0.053

0.132 0.125

0.287 0.107

0.479 0.167

0.145 0.056

0.151 0.075

0.544 0.265

station, noise

offline + online

0.055 0.036

0.121 0.106

0.309 0.128

0.528 0.212

0.155 0.073

0.127 0.062

0.554 0.264

cost does not grow with the number of network parameters. The parameters only determine the size of H, which has one row and one column per parameter (2434 × 2434), and the damped system in Eq. (20) is solved directly for the step ∆θ at a cost that is small compared with the ODE solves. We run 15 iterations, which take about 2 min for the 150 training scenarios on a single NVIDIA GH200 GPU. We use a second-order method because the directions in parameter space that control the biogenic correction are weakly constrained and have small gradients. In a test with station observations and noise, 400 steps of the first-order Adam optimiser from the same offline network barely changed the biogenic correction, whereas Gauss–Newton rescales these directions by the inverse curvature. Training the network online from scratch with Adam was also possible, but it took more than 20 times the computing time of the offline training followed by fine-tuning and remained less accurate.

4.5

Results

We evaluate the learned correction against the truth on a 60 × 60 grid over the two physical covariates of each sector, i.e. (T, log10 P ) for the anthropogenic sector and (T, R) for the biogenic sector, at three values of the synthetic covariate (0.3, 0.5, and 0.7). The grid covers the training ranges extended by the extrapolation bands. Table 1 lists the root-mean-square error (RMSE) of fs − fstrue inside the training ranges and in the extrapolation bands. For the biogenic sector, the RMSE inside the training ranges is also split into three temperature bands, because the information that the observations carry about fbio depends strongly on temperature. The online fine-tuning reduces the error in every setting and region, except for the anthropogenic extrapolation error with noise-free observations of all species, which is unchanged. The anthropogenic correction is already reproduced well by the offline training, with an RMSE of about 0.06 inside the training ranges, and the fine-tuning reduces it moderately to 0.04–0.06. The largest gains are in the biogenic sector, where the RMSE inside the training ranges decreases by a factor of 2.4–5. With all species observed, the fine-tuning reduces the error in the coldest band from 0.26–0.29 to 0.04–0.06, so the biogenic correction at low temperature is constrained by the observations even though the offline training could not recover it. Low temperatures occur mainly in the early morning and on cold days, when the biogenic emissions are small but not zero, and their effect on the concentrations is therefore weak in each single scenario. A per-scenario retrieval attributes such a weak signal mostly to the prior, whereas the shared network accumulates it over all scenarios that pass through the same temperature range. Figure 3 shows the error of the correction for the most realistic setting, with noisy station observations. The offline network overestimates the biogenic correction both below about 12 ◦ C and above about 27 ◦ C, which means that it stays close to the inventory at both ends of the temperature range and misses both the steeper increase at low temperature and the optimum of the true response. This is at least partly a consequence of the retrievals rather than of the regression. With station observations, the evidence selects a correlation length of 192 h (Sect. 4.3), so each retrieval is nearly constant over the day and the regression cannot attribute any variation within the day to the diurnal temperature cycle. After the online fine-tuning, which evaluates the network at the instantaneous temperature of every knot, the error is mostly below 0.25 inside the training ranges, and the remaining bias is concentrated below about 10 ◦ C, where biogenic emissions are smallest. The anthropogenic error is small in both cases.

12

With all species observed, the fine-tuned biogenic error is larger without noise than with noise (0.041 and 0.028, respectively). In the noise-free setting, the iterations moved the parameters far from the offline solution along flat directions of the parameter space, to a distance ∥θ − θoff ∥ ≈ 2 × 106 compared with 33–190 in the other settings, and the final iterate is not the one with the lowest validation loss. The correction remains accurate inside the training ranges, but a proximity term towards the offline parameters or early stopping would be advisable for noise-free data. Errors in the extrapolation bands are about 2–4 times larger than inside the training ranges for both approaches, which is expected for a network trained on a bounded range of covariates, and the fine-tuning reduces them in all cases but one.

5

Discussion

The applications in Sects. 3 and 4 use the sensitivities of the chemical model in three different ways, and the most efficient mode of automatic differentiation differs between them. When the number of parameters is small compared with the number of observations, as for the 25 rate coefficients in Sect. 3 and the 50 knot values of a 4D-Var retrieval in Sect. 4.3, forward-mode differentiation through the solver gives the full Jacobian at the cost of one tangent integration per parameter and requires no checkpointing. When a scalar loss is differentiated with respect to many parameters, as in first-order training of a neural network, reverse mode with recursive checkpointing is preferable. The online Gauss–Newton fine-tuning in Sect. 4.4 combines both, because the network enters the ODE only through a low-dimensional set of knot values, so that the expensive sensitivities are computed in forward mode with respect to these knots, and the network parameters only enter through the Jacobian of the network itself. KPPax inherits these choices from diffrax, so the differentiation mode is a configurable option of the solve rather than a property of the generated code, and no adjoint has to be derived by hand for any of the three uses. The computational cost of KPPax differs from that of KPP-generated code. Just-in-time compilation adds a fixed cost to the first call, which is amortised when the same mechanism is integrated many times, for example for many scenarios in parallel with jax.vmap or over many iterations of an optimisation, as in our experiments. The implicit solvers of diffrax solve dense linear systems with the full Jacobian of the right-hand side in each step, whereas KPP exploits the sparsity of the Jacobian through symbolic sparse LU decomposition. For POLLU, with 20 species, this difference is negligible, but the cost of dense factorisation grows with the cube of the number of species and will dominate for mechanisms with thousands of species such as the MCM. A sparse or matrix-free linear solver for the implicit stages is therefore the most important extension for large mechanisms. A systematic comparison of the accuracy and computational cost of KPPax with KPP-generated Fortran code for a large mechanism is ongoing and will be reported separately. The log-space formulation of the rate law (Eq. 6) keeps the rate law differentiable, but it requires strictly positive concentrations. We clip the concentrations to a small positive floor before the rate law is evaluated, which prevents the solver from probing zero or negative concentrations but also sets the gradient with respect to a clipped concentration to zero. This affects species that start from zero or are fully consumed, and applications that differentiate with respect to such near-zero concentrations should take it into account. The identifiability analysis in Sect. 3 is local, since it relies on the Jacobian at a single parameter point. It identified correctly which combinations of rate coefficients the design cannot constrain, but it underestimated how far the estimates drift along curved valleys of the cost function. Global methods, such as profile likelihoods or Markov chain Monte Carlo sampling, would quantify this drift and are straightforward to combine with KPPax, since every forward solve is differentiable and vectorisable. Because the FIM can be computed for any set of observations within seconds, it can also be used to design experiments, for example to choose the observed species, observation times, or precursor concentrations of a flow-tube experiment that make a given set of rate coefficients identifiable, before the experiment is carried out. The emission experiment in Sect. 4 is idealised in several respects. The true correction is a known synthetic function, the chemical mechanism of the model is the same as that of the truth, the box model has no transport, and the initial state of each scenario is known exactly. In particular, errors in the initial state could be aliased into the correction of the early-morning emissions, where we found the largest benefit of online fine-tuning, and a joint estimation of the initial state and the emissions, as in Elbern et al. (2007), would be needed in a real application. The Gauss–Newton fine-tuning also forms the Hessian with respect to all network parameters, which is feasible for the 2434 parameters used here but would require matrix-free or approximate second-order methods for larger networks. Nevertheless, the experiment shows that information which is present in the observations but lost in a per-scenario 13

f_NN - f_true (offline_w_station_noise, model_last.eqx) f_anth dummy_anth = 0.50

30

5

10

15 20 T (°C)

0.00

25

30

5

0.75

f_NN - f_true (finetune_station_noise, model_last.eqx) f_anth dummy_anth = 0.50

3.50

5 0.2

3.50

15 20 T (°C)

25

f_anth

−0.4

30

f_anth dummy_anth = 0.70 0.25

0.2 5

f_anth dummy_anth = 0.30 5 0.2

10

0.25

0.2

0.00

0.2 5

f_bio

−0.2

0.2

3.50

0.2

0.4 0.3

30

0.4

0.0

0.3 25

30

0.5

0.4

15 20 T (°C)

15 20 25 T (°C) f_bio dummy_bio = 0.70

0.6 0.25

0.00

0.25

0.50

0.7

0.3 10

10

0.8

0.4

5

0.4

3.25

3.25

3.00

3.00

3.00

2.75

2.75

2.75

0.0

2.50

2.50

2.50

−0.2

2.25

2.25

0.00

2.00

5

10

15 20 25 T (°C) f_bio dummy_bio = 0.30

5

10

15 20 25 T (°C) f_bio dummy_bio = 0.50

30

0.9

f_anth

0.00

5

10

15 20 25 T (°C) f_bio dummy_bio = 0.70

30

0.9

0.8

0.8

0.4

0.00

0.25

−0.4

0.6

0.6

0.6

0.5

0.5

0.5

0.0

−0.2

0.3

0.2

0.50

5

10

15 20 T (°C)

25

30

0.2 5

10

15 20 T (°C)

25

30

0.2

0.00

5

10

15 20 T (°C)

25

f_bio

0.25

0.2

0.25

0.4

0.3 0.25

0.4

0.3 0.2 5

0.4

0.50

0.7 0.2 5

0.7 0.00

R (daily max light)

0.2

2.00 0.00

5 0.2

0.7

30

5 0.2

0.9 0.8

0.00

2.00

0.00

2.25

0.00

3.25

0.00

log10 P

0.5

5

0.25

0.50

0.6 0.00

0.6 0.75

0.7 0.50

0.8

0.7

30 0.9

1.00

0.8

15 20 25 T (°C) f_bio dummy_bio = 0.50

0.50

0.9

10

0.00

0.9

−0.4

2.00

0.75

15 20 25 T (°C) f_bio dummy_bio = 0.30

0.25

R (daily max light)

10

−0.2

2.25

2.00 5

0.0

0.00

2.50

0.00

2.25

2.00

2.75

0.00

2.50 0 0.0

2.25

0.2

3.00

2.75

2.50

0.4

0.00

0.50

2.75

3.25

0.00

3.00

0 0.0

0.25

3.00

0.25

0.75

3.25

f_anth dummy_anth = 0.70

3.50

0.50

3.25

0.5

0.25

3.50 0.00

log10 P

3.50

1.00

f_anth dummy_anth = 0.30 0.25

−0.4

30

Figure 3: Error of the learned log emission correction, fs − fstrue , with noisy station observations, (top two rows) after the offline two-step training and (bottom two rows) after the online Gauss–Newton finetuning. In each pair of rows, the top row shows the anthropogenic and the bottom row the biogenic sector, at three values of the synthetic covariate (columns). The dashed boxes mark the training ranges.

14

retrieval can be recovered when the network is trained through the chemical model, which is only possible with a differentiable chemistry code.

6

Conclusions

We have presented KPPax, a kinetic preprocessor that generates differentiable chemistry code in JAX from mechanisms written in the KPP language. KPPax parses the mechanism with the pykpp parser, stores the stoichiometry as sparse coordinate arrays, transpiles the rate-coefficient expressions from Fortran to JAX, and evaluates the rate law with sparse operations in log space, which makes it differentiable with respect to both the concentrations and the rate coefficients. Combined with the implicit solvers and the configurable forward- and reverse-mode adjoints of diffrax, the generated code provides exact sensitivities of stiff chemical simulations without any hand-derived tangent-linear or adjoint code, it can be vectorised over many scenarios, and it runs on GPUs. We demonstrated two uses of these sensitivities with the POLLU benchmark mechanism. First, the Fisher information matrix computed from a single automatically differentiated Jacobian predicted, before any fit, which rate coefficients a flow-tube-like design can constrain. Nine of the 25 rate coefficients were identifiable, and their estimates from noisy synthetic observations scattered as predicted to within about 30%, whereas the non-identifiable rate coefficients formed clusters of fast reactions whose ratios, but not absolute values, were constrained. Second, a neural-network correction of an emission inventory, trained offline on per-scenario 4D-Var retrievals and then fine-tuned online through the chemical model by Gauss–Newton iterations, reduced the error of the biogenic correction by a factor of 2.4–5 compared with the offline training alone, at a cost of about two minutes on a single GPU. Online training recovered information on low-temperature emissions that is present in the observations but lost in the individual retrievals. KPPax thereby brings chemical mechanisms written in the KPP language into the differentiable modelling ecosystem, in which process-based chemistry and machine learning can be combined and trained against observations. Future work will address large mechanisms, which require sparse linear solvers for the implicit integration, and applications to real chamber and flow-tube observations.

Acknowledgements The authors wish to acknowledge CSC – IT Center for Science, Finland, for computational resources. WP acknowledges financial support from the Finnish Cultural Foundation, Päijät-Häme Regional Fund (grant 70251819). Large language models were used to assist with code development and with language editing of the manuscript. All code and text were reviewed by the authors, who take full responsibility for the content.

References Michail Alvanos and Theodoros Christoudias. Gpu-accelerated atmospheric chemical kinetics in the echam/messy (emac) earth system model (version 2.52). Geoscientific Model Development, 10(10): 3679–3693, 2017. Julien Brajard, Alberto Carrassi, Marc Bocquet, and Laurent Bertino. Combining data assimilation and machine learning to emulate a dynamical model from sparse and noisy observations: A case study with the Lorenz 96 model. Journal of Computational Science, 44:101171, 2020. doi: 10.1016/j.jocs.2020. 101171. Guy P. Brasseur and Daniel J. Jacob. Modeling of Atmospheric Chemistry. Cambridge University Press, 2017. Philippe Courtier, Jean-Noël Thépaut, and Anthony Hollingsworth. A strategy for operational implementation of 4D-Var, using an incremental approach. Quarterly Journal of the Royal Meteorological Society, 120(519):1367–1387, 1994. doi: 10.1002/qj.49712051912. Valeriu Damian, Adrian Sandu, Mirela Damian, Florian Potra, and Gregory R Carmichael. The kinetic preprocessor kpp-a software environment for solving chemical kinetics. Computers & Chemical Engineering, 26(11):1567–1579, 2002.

15

Oscar H Díaz-Ibarra, Samuel G Frederick, Jeffrey H Curtis, Zachary D’aquino, Peter A Bosler, Lekha Patel, Cosmin Safta, Matthew West, and Nicole Riemer. Tchem-atm (v2. 0.0): scalable performanceportable multiphase atmospheric chemistry. Geoscientific Model Development, 19(3):1281–1299, 2026. Ruiwen Dong, Christian Goodbrake, Heather A. Harrington, and Gleb Pogudin. Differential elimination for dynamical models via projections with applications to structural identifiability. SIAM Journal on Applied Algebra and Geometry, 7(1):194–235, 2023. doi: 10.1137/22M1469067. Hendrik Elbern, Achim Strunk, Hauke Schmidt, and Olivier Talagrand. Emission rate and chemical state estimation by 4-dimensional variational inversion. Atmospheric Chemistry and Physics, 7(14): 3749–3769, 2007. doi: 10.5194/acp-7-3749-2007. Maximilian Gelbrecht, Alistair White, Sebastian Bathiany, and Niklas Boers. Differentiable programming for earth system modeling. Geoscientific Model Development, 16(11):3123–3135, 2023. Joseph H. A. Guillaume, John D. Jakeman, Stefano Marsili-Libelli, Michael Asher, Philip Brunner, Barry Croke, Mary C. Hill, Anthony J. Jakeman, Karel J. Keesman, Saman Razavi, and Johannes D. Stigter. Introductory overview of identifiability analysis: A guide to evaluating whether you have the right type of data for your modeling purpose. Environmental Modelling & Software, 119:418–432, 2019. doi: 10.1016/j.envsoft.2019.07.007. Barron H. Henderson. pykpp: a KPP-like chemical mechanism parser for Python. https://github. com/barronh/pykpp. Accessed 27 October 2026. Daven K. Henze, Amir Hakami, and John H. Seinfeld. Development of the adjoint of GEOS-Chem. Atmospheric Chemistry and Physics, 7(9):2413–2433, 2007. doi: 10.5194/acp-7-2413-2007. Mike Innes, Alan Edelman, Keno Fischer, Chris Rackauckas, Elliot Saba, Viral B Shah, and Will Tebbutt. A differentiable programming system to bridge machine learning and scientific computing. arXiv preprint arXiv:1907.07587, 2019. Michael E Jenkin, Sandra M Saunders, and Michael J Pilling. The tropospheric degradation of volatile organic compounds: a protocol for mechanism development. Atmospheric Environment, 31(1):81–104, 1997. Peishi Jiang, Patrick Kidger, Toshiyuki Bandai, Dennis Baldocchi, Heping Liu, Yi Xiao, Qianyu Zhang, Carlos Tianxin Wang, Carl Steefel, and Xingyuan Chen. Jax-canveg: A differentiable land surface model. Water Resources Research, 61(3):e2024WR038116, 2025. Ammar Kheder, Helmi Toropainen, Wenqing Peng, Samuel Antão, Jia Chen, Michael Boy, and Zhi-Song Liu. Topoflow: topography-aware pollutant flow learning for high-resolution air quality prediction. npj Climate and Atmospheric Science, 2026. Patrick Kidger. On neural differential equations. arXiv preprint arXiv:2202.02435, 2022. Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In 3rd International Conference on Learning Representations (ICLR), 2015. arXiv:1412.6980. Dmitrii Kochkov, Janni Yuval, Ian Langmore, Peter Norgaard, Jamie Smith, Griffin Mooers, Milan Klöwer, James Lottes, Stephan Rasp, Peter Düben, et al. Neural general circulation models for weather and climate. Nature, 632(8027):1060–1066, 2024. Zhi-Song Liu, Petri Clusius, and Michael Boy. Neural network emulator for atmospheric chemical ode. Neural Networks, 184:107106, 2025. Zhi-Song Liu, Wenqing Peng, Helmi Toropainen, Ammar Kheder, Andreas Rupp, Holger Froning, Xiaojie Lin, and Michael Boy. Inverse neural operator for ode parameter optimization, 2026. Torkel E. Loman, Yingbo Ma, Vasily Ilin, Shashi Gowda, Niklas Korsbo, Nikhil Yewale, Chris Rackauckas, and Samuel A. Isaacson. Catalyst: Fast and flexible modeling of reaction networks. PLOS Computational Biology, 19(10):e1011530, 2023. doi: 10.1371/journal.pcbi.1011530. SP O’Meara, S Xu, D Topping, MR Alfarra, G Capes, D Lowe, Y Shao, and G McFiggans. Pycham (v2. 1.1): a python box model for simulating aerosol chambers, geosci. model dev., 14, 675–702, 10.5194, 2021. 16

Wenqing Peng, Zhi-Song Liu, and Michael Boy. Spin-ode: Stiff physics-informed neural ode for chemical reaction rate estimation. In European Conference on Artificial Intelligence, pages 2033 – 2040, 2025. Gabriele G. Pfister, Sebastian D. Eastham, Avelino F. Arellano, Bernard Aumont, Kelley C. Barsanti, Mary C. Barth, Andrew Conley, Nicholas A. Davis, Louisa K. Emmons, Jerome D. Fast, Arlene M. Fiore, Benjamin Gaubert, Steve Goldhaber, Claire Granier, Georg A. Grell, Marc Guevara, Daven K. Henze, Alma Hodzic, Xiaohong Liu, Daniel R. Marsh, John J. Orlando, John M. C. Plane, Lorenzo M. Polvani, Karen H. Rosenlof, Allison L. Steiner, Daniel J. Jacob, and Guy P. Brasseur. The multi-scale infrastructure for chemistry and aerosols (musica). Bulletin of the American Meteorological Society, 101(10):E1743 – E1760, 2020. doi: 10.1175/BAMS-D-19-0331.1. URL https://journals.ametsoc. org/view/journals/bams/101/10/bamsD190331.xml. Jason Rader, Terry Lyons, and Patrick Kidger. Optimistix: modular optimisation in JAX and Equinox. arXiv preprint arXiv:2402.09983, 2024. Adrian Sandu, Dacian N Daescu, and Gregory R Carmichael. Direct and adjoint sensitivity analysis of chemical kinetic systems with kpp: Part i—theory and software tools. Atmospheric Environment, 37 (36):5083–5096, 2003. Sandra M Saunders, Michael E Jenkin, Richard G Derwent, and Mike J Pilling. Protocol for the development of the master chemical mechanism, mcm v3 (part a): tropospheric degradation of non-aromatic volatile organic compounds. Atmospheric Chemistry and Physics, 3(1):161–180, 2003. John H. Seinfeld and Spyros N. Pandis. Atmospheric Chemistry and Physics: From Air Pollution to Climate Change. John Wiley & Sons, Hoboken, NJ, 3 edition, 2016. Chaopeng Shen, Alison P Appling, Pierre Gentine, Toshiyuki Bandai, Hoshin Gupta, Alexandre Tartakovsky, Marco Baity-Jesi, Fabrizio Fenicia, Daniel Kifer, Li Li, et al. Differentiable modelling to unify machine learning and physical models for geosciences. Nature Reviews Earth & Environment, 4 (8):552–567, 2023. David Topping, Paul Connolly, and Jonathan Reid. Pybox: An automated box-model generator for atmospheric chemistry and aerosol simulations. Journal of Open Source Software, 3(28):755, 2018. Jan G. Verwer. Gauss–Seidel iteration for stiff ODEs from chemical kinetics. SIAM Journal on Scientific Computing, 15(5):1243–1250, 1994. doi: 10.1137/0915076.

17

Related documents

• KPPax: preprocess differentiable kinetic ODE for machine learning in atmospheric chemistry#1469325
• Physics-Guided Machine Learning for Predictive Turbulent Combustion Simulation#793257
• Physics-Guided Machine Learning for Predictive Turbulent Combustion Simulation#789253
• Torch-DPhys: A Torch-Native, MOOSE-Compatible Framework for Differentiable Physics Simulation and Hybrid Machine Learning Integration#1388867
• FlowMLLab: An open-source framework for reproducible computational-fluid-dynamics and scientific-machine-learning experiments#1064519