TrunX: A massively parallel, differentiable implementation of the 3-PG forest growth model in JAX
TL;DR
An implementation of the Physiological Processes Predicting Growth (3-PG) model in JAX that uses just-in-time compilation, vectorization, and GPU acceleration to reduce execution time and enables efficient gradient-based optimization and gradient-informed Bayesian calibration.
Abstract
Process-based forest models are widely used to simulate forest growth and responses to environmental change, but their calibration and application often require many computationally expensive model evaluations. We present an implementation of the Physiological Processes Predicting Growth (3-PG) model in JAX that uses just-in-time compilation, vectorization, and GPU acceleration to reduce execution time. The implementation also supports automatic differentiation, providing gradients of model outputs and calibration objectives with respect to model parameters. This enables efficient gradient-based optimization and gradient-informed Bayesian calibration, extending 3-PG beyond conventional gradient-free approaches. The implementation produced results numerically consistent with r3PG for the evaluated configuration. Overall, the JAX implementation provides a faster and differentiable framework for calibrating and applying the 3-PG model.