TrunX: A massively parallel, differentiable implementation of the 3-PG forest growth model in JAX
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.
Related
- A Massively Parallel Hybridizable Discontinuous Galerkin Solver for Direct Numerical Simulation of Compressible Flows on GPUs
- A Massively Parallel Three-Grid Preconditioner for the High-Frequency Helmholtz Equation
- Massively parallel Schwarz methods for the high-frequency Helmholtz equation with O(k) convergence
- Counterfactual Probing for Parallel Unmasking with Hidden Forest Structure
- Polynomial Spectral Gap Estimates for Finitely Differentiable Quasi-Periodic Schrödinger Operators
BibTeX
@misc{givi2026trunx,
title = {{TrunX: A massively parallel, differentiable implementation of the 3-PG forest growth model in JAX}},
author = {Glory Mary Givi and Cédric Travelletti and Grégory Mermoud},
howpublished = {arXiv},
year = {2026},
month = sep,
doi = {10.48550/arxiv.2609.02557},
eprint = {2609.02557},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2609.02557},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - TrunX: A massively parallel, differentiable implementation of the 3-PG forest growth model in JAX
AU - Glory Mary Givi
AU - Cédric Travelletti
AU - Grégory Mermoud
PY - 2026
DA - 2026/09/09
PB - arXiv
DO - 10.48550/arxiv.2609.02557
UR - https://arxiv.org/abs/2609.02557
AB - 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.
ER -