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 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 (LE) and net ecosystem exchange (NEE), (b) improving the optimization trade-off when learning observations of both LE and NEE, and (c) benefiting a multi-layer canopy model setup. Our results show that the proposed hybrid model improved the simulations of LE and NEE 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.