KPPax
KPPax is a kinetic preprocessor that converts atmospheric chemical mechanisms described in the KPP language into JAX code that is differentiable, just-in-time compiled, and vectorizable, supporting parameter estimation and hybrid physics-machine learning modeling.
Timeline (last 90 days)
KPPax is a kinetics preprocessor that generates differentiable, just-in-time compiled and vectorizable chemical code in JAX from the KPP language.
KPPax evaluates rate laws in log space using sparse operations to avoid numerical underflow and non-differentiable operations.
KPPax calls the diffrax library's implicit solvers and forward and reverse mode adjoints to integrate stiff chemical systems.
KPPax leverages diffrax's built-in adjoint integration to obtain exact gradients without manually deriving and writing adjoint equation code, avoiding the memory explosion and numerical instability that occur when automatic differentiation passes directly through stiff solvers.
The authors demonstrate two applications using the POLLU benchmark mechanism: Fisher information matrix prior identifiability analysis and neural-network correction of emission inventories.
KPPax's demonstration is based only on the POLLU benchmark mechanism; its computational overhead and numerical stability on larger-scale mechanisms remain to be verified.
The effectiveness of KPPax's online fine-tuning stage is affected by the configuration and quality of the 4D-Var inversion.
All Facts (7)
KPPax is a kinetics preprocessor that generates differentiable, just-in-time compiled and vectorizable chemical code in JAX from the KPP language.
50%UnverifiedKPPax evaluates rate laws in log space using sparse operations to avoid numerical underflow and non-differentiable operations.
50%UnverifiedKPPax calls the diffrax library's implicit solvers and forward and reverse mode adjoints to integrate stiff chemical systems.
50%UnverifiedKPPax leverages diffrax's built-in adjoint integration to obtain exact gradients without manually deriving and writing adjoint equation code, avoiding the memory explosion and numerical instability that occur when automatic differentiation passes directly through stiff solvers.
50%UnverifiedThe authors demonstrate two applications using the POLLU benchmark mechanism: Fisher information matrix prior identifiability analysis and neural-network correction of emission inventories.
50%UnverifiedKPPax's demonstration is based only on the POLLU benchmark mechanism; its computational overhead and numerical stability on larger-scale mechanisms remain to be verified.
50%UnverifiedThe effectiveness of KPPax's online fine-tuning stage is affected by the configuration and quality of the 4D-Var inversion.
50%