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 opp...

Full description

Saved in:
Bibliographic Details
Main Authors: Peishi Jiang, Patrick Kidger, Toshiyuki Bandai, Dennis Baldocchi, Heping Liu, Yi Xiao, Qianyu Zhang, Carlos Tianxin Wang, Carl Steefel, Xingyuan Chen
Format: Article
Language:English
Published: Wiley 2025-03-01
Series:Water Resources Research
Subjects:
Online Access:https://doi.org/10.1029/2024WR038116
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1849687892437237760
author Peishi Jiang
Patrick Kidger
Toshiyuki Bandai
Dennis Baldocchi
Heping Liu
Yi Xiao
Qianyu Zhang
Carlos Tianxin Wang
Carl Steefel
Xingyuan Chen
author_facet Peishi Jiang
Patrick Kidger
Toshiyuki Bandai
Dennis Baldocchi
Heping Liu
Yi Xiao
Qianyu Zhang
Carlos Tianxin Wang
Carl Steefel
Xingyuan Chen
author_sort Peishi Jiang
collection DOAJ
description 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.
format Article
id doaj-art-b9b7b70821a741a091f9b107f8451e2b
institution DOAJ
issn 0043-1397
1944-7973
language English
publishDate 2025-03-01
publisher Wiley
record_format Article
series Water Resources Research
spelling doaj-art-b9b7b70821a741a091f9b107f8451e2b2025-08-20T03:22:12ZengWileyWater Resources Research0043-13971944-79732025-03-01613n/an/a10.1029/2024WR038116JAX‐CanVeg: A Differentiable Land Surface ModelPeishi Jiang0Patrick Kidger1Toshiyuki Bandai2Dennis Baldocchi3Heping Liu4Yi Xiao5Qianyu Zhang6Carlos Tianxin Wang7Carl Steefel8Xingyuan Chen9Atmospheric, Climate, and Earth Sciences Division Pacific Northwest National Laboratory Richland WA USACradle Bio Zurich SwitzerlandLawrence Berkeley National Laboratory Earth and Environmental Sciences Area Berkeley CA USADepartment of Environmental Science, Policy, and Management University of California Berkeley CA USADepartment of Civil and Environmental Engineering Washington State University Pullman WA USAAtmospheric, Climate, and Earth Sciences Division Pacific Northwest National Laboratory Richland WA USADepartment of Civil and Environmental Engineering Washington State University Pullman WA USADepartment of Environmental Science, Policy, and Management University of California Berkeley CA USALawrence Berkeley National Laboratory Earth and Environmental Sciences Area Berkeley CA USAAtmospheric, Climate, and Earth Sciences Division Pacific Northwest National Laboratory Richland WA USAAbstract 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.https://doi.org/10.1029/2024WR038116differentiable land surface modelinghybrid modelingstomatal conductancewater stressoptimization trade‐off
spellingShingle Peishi Jiang
Patrick Kidger
Toshiyuki Bandai
Dennis Baldocchi
Heping Liu
Yi Xiao
Qianyu Zhang
Carlos Tianxin Wang
Carl Steefel
Xingyuan Chen
JAX‐CanVeg: A Differentiable Land Surface Model
Water Resources Research
differentiable land surface modeling
hybrid modeling
stomatal conductance
water stress
optimization trade‐off
title JAX‐CanVeg: A Differentiable Land Surface Model
title_full JAX‐CanVeg: A Differentiable Land Surface Model
title_fullStr JAX‐CanVeg: A Differentiable Land Surface Model
title_full_unstemmed JAX‐CanVeg: A Differentiable Land Surface Model
title_short JAX‐CanVeg: A Differentiable Land Surface Model
title_sort jax canveg a differentiable land surface model
topic differentiable land surface modeling
hybrid modeling
stomatal conductance
water stress
optimization trade‐off
url https://doi.org/10.1029/2024WR038116
work_keys_str_mv AT peishijiang jaxcanvegadifferentiablelandsurfacemodel
AT patrickkidger jaxcanvegadifferentiablelandsurfacemodel
AT toshiyukibandai jaxcanvegadifferentiablelandsurfacemodel
AT dennisbaldocchi jaxcanvegadifferentiablelandsurfacemodel
AT hepingliu jaxcanvegadifferentiablelandsurfacemodel
AT yixiao jaxcanvegadifferentiablelandsurfacemodel
AT qianyuzhang jaxcanvegadifferentiablelandsurfacemodel
AT carlostianxinwang jaxcanvegadifferentiablelandsurfacemodel
AT carlsteefel jaxcanvegadifferentiablelandsurfacemodel
AT xingyuanchen jaxcanvegadifferentiablelandsurfacemodel