JAX‐CanVeg: A Differentiable Land Surface Model
Skip to main content
eScholarship
Open Access Publications from the University of California

UC Berkeley

UC Berkeley Previously Published Works bannerUC Berkeley

JAX‐CanVeg: A Differentiable Land Surface Model

Abstract

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 capture these complex interactions by seamlessly hybridizing process‐based models with deep neural networks (DNNs), benefiting both worlds, that is, the physical interpretation of process‐based models and the learning power of DNNs. Here, we developed a differentiable land model, JAX‐CanVeg. The new model builds on the legacy CanVeg by incorporating advanced functionalities through JAX in the graphic processing unit support, automatic differentiation, and integration with DNNs. We demonstrated JAX‐CanVeg's hybrid modeling capability by applying the model at four flux tower sites with varying aridity. To this end, we developed a hybrid version of the Ball‐Berry equation that emulates the water stress impact on stomatal closure to explore the capability of the hybrid model in (a) improving the simulations of latent heat fluxes and net ecosystem exchange , (b) improving the optimization trade‐off when learning observations of both and , and (c) benefiting a multi‐layer canopy model setup. Our results show that the proposed hybrid model improved the simulations of and at all sites, with an improved optimization trade‐off over the process‐based model. Additionally, the multi‐layer canopy set benefited hybrid modeling at some sites. Anchored in differentiable modeling, our study provides a new avenue for modeling land‐atmosphere interactions by leveraging the benefits of both data‐driven learning and process‐based modeling. Plain Language Summary Land‐atmosphere interactions involve flux exchanges of carbon, water, and energy. They are important terrestrial ecosystem components that are being gradually modified by the warming climate. Despite the progress in the land surface model development, accurately modeling these interactions still remains a challenge owing to multiple complicated processes going from the canopy top to the soil system. In this paper, we developed a new land surface model that uses a novel modeling approach called differentiable programming to seamlessly integrate process‐based models and deep neural networks. The new model, JAX‐CanVeg, is consistent with the known ecohydrological processes while flexible to be coupled with neural networks to improve the simulations of water and carbon fluxes. We demonstrated the hybrid modeling capability of JAX‐CanVeg by coupling the equation to calculate stomatal conductance with a neural network that quantifies the water stress impact through soil moisture observations. As a proof of concept, applying the hybrid JAX‐CanVeg in four different ecosystems improves simulations of latent heat fluxes and net ecosystem exchanges. The improvement shows promise in using the new model to enhance the simulations of terrestrial water and carbon cycling and better facilitate answering research questions related to climate change. Key Points We developed a differentiable land surface model using a high‐performance machine learning package JAX We proposed a hybrid Ball‐Berry model that accounts for the influence of water stress on stomatal conductance through a deep neural network We showed that the hybrid JAX‐CanVeg improved the simulations of water and carbon fluxes at four ecosystems with varying aridity

Many UC-authored scholarly publications are freely available on this site because of the UC's open access policies. Let us know how this access is important for you.