Papers

KPPax:将大气化学动力学机制编译为可微分 JAX 代码

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

Why we picked it

做大气化学或物理-机器学习混合建模的可以看看,能把 KPP 机制直接转成 JAX 可微分代码,做参数估计不用手推伴随了。

KPPax 是一个动力学预处理器,能把 KPP 语言书写的大气化学机制编译成 JAX 中可微分、可即时编译、可向量化的化学代码。它以稀疏操作在对数空间计算速率定律,并通过 diffrax 库的隐式求解器与前向、反向伴随模式积分刚性系统,无需手写伴随代码即可获得精确敏感性。在 POLLU 基准机制上,基于自动微分雅可比矩阵计算的 Fisher 信息矩阵可在拟合前预测 25 个速率系数中哪些可识别,噪声合成估计验证了 9 个可识别系数的不确定度预测误差在 30% 以内。另一项实验中,神经网络对排放清单的修正经 4D-Var 反演离线训练后在线微调,将生物源修正误差降低 2.4-5 倍。

Original · arXiv cs.AI

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

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.