Differentiable modeling package release and seeking feedback.
Chaopeng Shen
A domain-agnostic, PyTorch-based framework for developing trainable differentiable models that merge neural networks with process-based equations. "Differentiable" means that gradient calculations can be achieved efficiently at large scale throughout the model, so process-based equations can be trained together with NNs on big data, on GPU. Following as a generalization of HydroDL, 𝛿MG (generic_deltaModel) aims to expand differentiable modeling and learning capabilities to a wide variety of domains where prior equations can bring in benefits.
𝛿MG is not a partcularly model. Rather, it is a generic framework that support many models across various domains (some are from HydroDL2.0) in a uniform way, while integrating lots of ecosystem tools. While the packages contains some basic examples for learners' convenience, the deployment models are supposed to exit in separate repositories and couple to the 𝛿MG framework. It is generalized and formalized after years of experiences working with various differentiable models across domains. Most of the differentiable modeling efforts in Shen's group will be using 𝛿MG. 𝛿MG can be configured to run through a configuration file and it should be easy and clear for new starters to learn. We even include a Graphical User Interface that allows easy job customization. The framework will closely synergize with advanced deep learning tools like foundation models and the scale advantage of PyTorch.
Maintained by the MHPI group advised by Dr. Chaopeng Shen. If this work is of use to you, please cite this paper for now, but we will have more dedicated citations later: Shen et al., 2023, Differentiable modelling to unify machine learning and physical models for geosciences, Nature Reviews Earth & Environment, https://www.nature.com/articles/s43017-023-00450-9
𝛿MG seamlessly integrates with:
- HydroDL2.0 (
hydroDL2): Home to MHPI's suite of physics-based hydrology models, and differentiable model augmentations (think variational data assimilation, model coupling, and other tools designed for hydrology). - HydroData (
hydro_data_dev): Data extraction, processing, and management tools optimized for geospatial datasets. (In development) - Config GUI (
GUI-Config-builder)(Source): An intuitive, user-friendly tool designed to simplify the creation and editing of configuration files for model setup and development. - Differentiable Ecosystem modeling (
diffEcosys (dev version only)): Data extraction, processing, and management tools optimized for geospatial datasets. - Concurrent development activities: We are working on these efforts connected to 𝛿MG: (i) numerical PDE solvers on torch; (ii) adjoint sensitivity; (iii) extremely efficient and highly accurate surrogate models; (iv) data assimilation; (v) downscaled and bias corrected climate data; (vi) mysteriously powerful neural networks, and more ...
Hybrid Modeling: Combines neural networks with physical process equations for enhanced interpretability and generalizability. For example, skip manually tuning model parameters by using neural networks to feed robust and interpretable parameter predictions directly.
PyTorch Integration: Easily scales with PyTorch, enabling efficient training and compatibility with modern deep learning tools, trained foundation models, and differentiable numerical solvers.
Domain-agnostic and Flexible: Extends differentiable modeling to any field where physics-guided learning can add value, with modularity to meet the diversity of needs along the way.
Benchmarking: All in one place. 𝛿MG + hydroDL2 will enable rapid deployment and replication of key published MHPI results.
NextGen-ready: 𝛿MG is designed to be CSDMS BMI-compliant, and our differentiable hydrology models in hydroDL2 come with a prebuilt BMI allowing seamless compatibility with NOAA-OWP's NextGen National Water Modelling Framework. Incidentally, this capability also lends to 𝛿MG being easily interfaced for other applications.
This package includes the differential lumped rainfall-runoff model 𝛿HBV1.0, improved 𝛿HBV1.1p, 𝛿PRMS, and 𝛿SAC-SMA. This package powers the global- and national-scale water model that provide high-quality seamless hydrologic simulations over US and the world. It also hosts global-scale photosynthesis learning and simulations. Many other use cases are being developed concurrently.
This package powers the global- and national-scale water model that provide high-quality seamless hydrologic simulations over US and the world. It also hosts global-scale photosynthesis learning and simulations. Many other use cases are being developed concurrently.
In the unseen extreme event spatial test, we used water years with a 5-year or lower return period peak flow from 1990/10/01 to 2014/09/30 for training and held out the water years with greater than a 5-year return period peak flow for testing. The spatial test was conducted using a 5-fold cross-validation approach for CAMELS dataset. This application is benchmarked against LSTM and demonstrates better extrapolation ability. More details and results can be found in Song, Sawadekar et al., 2024.
This is a national-scale water modeling study on approximately 180,000 river reaches (with a median length of 7 km) across CONUS, using a high-resolution, differentiable, national-scale water model. More details and results can be found in Song, Bindas et al., 2024, 𝛿HBV2.0.
Characterized by the combination of process-based equations with neural networks (NNs), differentiable models train these components together, enabling parameter inputs for the equations to be effectively and efficiently learned at scale by the NNs. There are many possibilities for how such models are built.
In 𝛿MG, we define a differentiable model with the class DeltaModel that can couple one or more NNs with a process-based model (itself potentially a collection of models). This class holds nn and a phy_model objects, respectively, as attributes internally and describes how they interface with each other. The DeltaModel object can be trained and forwarded just as any other PyTorch model (nn.Module).
We also define DataLoader and DataSampler classes to handle datasets, a Trainer class for running train/test experiments, and a ModelHandler class for multimodel handling, multi-GPU training, data assimilation and streaming in a uniform and modular way. All model, training, and simulation settings are be collected in a configuration file that can be adapted to custom applications. According to this schema, we define these core classes, from bottom up:
- nn: PyTorch neural networks that can learn and provide either parameters, missing process representations, corrections, or other forms of enhancements to physical models.
- phy_model: The physical model written in PyTorch (or potentially another interoperable differentiable platform) that takes learnable outputs from the
nnmodel(s) and returns a prediction of some target variable(s). This can also be a wrapper holding several physical models. - DeltaModel: Holds (one or multiple)
nnobjects and aphy_modelobject, and describes how they are coupled; connection to ODE packages. - ModelHandler: Manages multimodeling, multi-GPU compute, and data assimilation or streaming. Can contain its own optimizers. Acts as an interface to CSDMS BMI or other interfaces.
- DataSampler: Samples data according to data format and training/testing requirements.
- Trainer: Manages model training and testing, and connects data to models.
- DataLoader: Preprocesses data to be used in training, testing, and simulation.
The Pennsylvania State University
University Park, PA 16802
Office: 814-863-5844