April 10, 2025
Journal Article

JAX-CanVeg: A Differentiable Land Surface Model

Abstract

Land surface models consider the exchange of water, energy, and carbon along the soil-canopy-atmosphere continuum, which is challenging to model due to their complex interdependency and associated challenges in representing and parameterizing them. Differentiable modeling provides a new opportunity to explore the parameter space and capture these complex interactions by seamlessly hybridizing process-based models with deep neural networks. The new modeling paradigm thus has the benefits of both worlds, i.e., the physical interpretation of process-based models and the learning power of DNNs. Here, we developed a differentiable land surface model (LSM), JAX-CanVeg. The new model builds on a legacy LSM, CanVeg, by incorporating advanced functionalities through JAX in the graphic processing unit (GPU) support, automatic differentiation, and notably integration with deep neural networks. We demonstrated the hybrid modeling capability of JAX-CanVeg by applying the model to simulate the water and carbon fluxes at two flux tower sites with varying aridity. To account for the influence of water stress on stomatal closure, we developed a hybrid version of the Ball-Berry equation that emulates the impact of water stress on stomatal conductance calculation through a deep neural network parameterized on the observed soil water content. Trained against the latent heat flux observations, the hybrid model improves the water flux simulations over the pure process-based model at both sites, as a result of a better representation of soil moisture impacts on stomatal conductance. The updated stomatal conductance further alters the model prediction on canopy carbon fluxes, such as photosynthesis. At both sites, the GPU-enabled JAX-CanVeg model executes several hundred times faster than the Matlab-based CanVeg. Our study provides a new avenue for modeling land-atmospheric interactions by leveraging the benefits of both data-driven learning and process-based modeling.

Published: April 10, 2025

Citation

Jiang P., P. Kidger, T. Bandai, D. Baldocchi, H. Liu, Y. Xiao, and Q. Zhang, et al. 2025. JAX-CanVeg: A Differentiable Land Surface Model. Water Resources Research 61, no. 3:Art. No. e2024WR038116. PNNL-SA-199036. doi:10.1029/2024WR038116

Research topics