JAX-CanVeg: A Differentiable Land Surface Model

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.