PyTorch Differentiable Modeling Framework

𝛿MG is a generic framework for building differentiable models that seamlessly couple neural networks with process-based equations, leveraging PyTorch's auto-differentiation for efficient, GPU-accelerated optimization. This is the fast, spiritual successor to HydroDL.

• 🤝 Hybrid Modeling — Combine NNs with process-based equations; learn physical model parameters directly from data.
• 🔁 PyTorch Integration — Efficient training, modern ML tooling, and numerical solver compatibility.
• 🧩 Modular Architecture — Swap in domain-specific components (models, loss functions, data loaders) with ease.
• ⚡ Benchmarking — Rapid deployment and replication of published MHPI results.
• 🌊 Operations-ready — CSDMS BMI compliant for NOAA-OWP's NextGen Framework and AWI's NGIAB.

Installation

uv pip install dmg

Optional extras:

uv pip install "dmg[hydrodl2]"    # MHPI hydrologic models (δHBV, etc.)
uv pip install "dmg[logging]"     # TensorBoard and W&B
uv pip install "dmg[tune]"        # Hyperparameter tuning (Optuna/Ray)

For development installs, see setup .

Quick Start

Before Running:

• Environment: See setup.md for ENV setup. δMG must be installed with dependencies + hydrodl2 to run this notebook.

• Data: Download our CAMELS (details here) data extraction from AWS. Then, update the data configs:

  • In camels_531.yaml and camels_671.yaml, update...

    1. data_path with path to camels_daymetv2,
    2. gage_info with path to gage_id.npy,
    3. subset_path with path to 531sub_id.txt (camels_531 only).

  • The full 671-basin or 531-basin CAMELS datasets can be selected by setting observations: camels_671 or camels_531 in the model config, respectively.

• Hardware: The NNs used in this model require CUDA support only available with Nvidia GPUs. For those without access, T4 GPUs can be used when running this notebook with δMG on Google Colab.

Running:

Use an LSTM to learn parameters for the HBV hydrologic model:

from hydrodl2.models.hbv.hbv import Hbv
from dmg.core.data.loaders import HydroLoader
from dmg.core.utils import load_nn_model
from dmg.models.delta_models import DplModel
from example import load_config, take_data_sample

config = load_config('../example/conf/config_dhbv.yaml')

# Build differentiable model: NN learns parameters for physics model.
phy_model = Hbv(config['model']['phy'])
nn = load_nn_model(config['model'], phy_model)
dpl_model = DplModel(phy_model=phy_model, nn_model=nn)

# Load data and forward.
dataset = HydroLoader(config).dataset
sample = take_data_sample(config, dataset, days=730, basins=100)
output = dpl_model(sample)

Internally, DplModel composes the NN and physics model — the NN generates parameters, the physics model produces predictions:

parameters = self.nn_model(dataset_sample['xc_nn_norm'])
predictions = self.phy_model(dataset_sample, parameters)

We recommend starting with the δHBV 1.0 tutorial (Colab), then exploring the full example notebooks. See how to run for CLI usage.

Use Cases

1. Lumped Hydrology

Lumped differentiable rainfall-runoff models 𝛿HBV 1.0 and improved 𝛿HBV 1.1p.

2. Unseen Extreme Events Test with 𝛿HBV 1.1p

In the unseen extreme events 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 basins in the CAMELS dataset. This application has been benchmarked against LSTM and demonstrates better extrapolation abilities. Find more details and results in Song, Sawadekar, et al. (2026). The δHBV 1.1p tutorial walks through the training and evaluation of this model.

3. National- and Global-scale Distributed Modeling

A national-scale water modeling study on approximately 180,000 river reaches (with a median length of 7 km) across CONUS using the high-resolution, multiscale, differentiable water model 𝛿HBV 2.0. This model is also operating at global scales (Ji, Song, et al., 2025) and has been used to generate high-quality, seamless simulations for both CONUS and the globe. Find more details and results in Song, Bindas, et al. (2025).

4. Global-scale Photosynthesis Modeling

Differentiable modeling has also been applied to parameterize global-scale sapflow simulations. This work is currently in development; see Aboelyazeed et al. (2024) for more details.

Documentation

Architecture

src/dmg/
├── core/
│   ├── calc/                   # Metrics and calculation utilities
│   ├── data/                   # Data loaders and samplers
│   ├── logging/                # TensorBoard and W&B logging
│   ├── post/                   # Post-processing and plotting
│   └── utils/                  # Factory functions and helpers
├── models/
│   ├── criterion/              # Loss functions (MSE, NSE, KGE, ...)
│   ├── delta_models/           # Differentiable model types (DplModel, ...)
│   ├── neural_networks/        # NN architectures (LSTM, ANN, MLP, ...)
│   ├── phy_models/             # Physical model wrappers
│   └── model_handler.py        # High-level model manager
└── trainers/                   # Training orchestration

Ecosystem

• hydrodl2 — MHPI's suite of process-based hydrology models (lumped + distributed).
• diffEcosys — Physics-informed ML for ecosystem modeling (photosynthesis via FATES).
• In development — Numerical PDE solvers, adjoint sensitivity, surrogate models, data assimilation, and more.

Citation

This work is maintained by MHPI and advised by Dr. Chaopeng Shen. If you find it useful, please cite:

Shen, C., Appling, A.P., Gentine, P. et al. Differentiable modelling to unify machine learning and physical models for geosciences. Nat Rev Earth Environ 4, 552–567 (2023). https://doi.org/10.1038/s43017-023-00450-9

@article{shen_differentiable_2023,
    title = {Differentiable modelling to unify machine learning and physical models for geosciences},
    volume = {4},
    issn = {2662-138X},
    url = {https://doi.org/10.1038/s43017-023-00450-9},
    doi = {10.1038/s43017-023-00450-9},
    pages = {552--567},
    number = {8},
    journaltitle = {Nature Reviews Earth \& Environment},
    author = {Shen, Chaopeng and Appling, Alison P. and Gentine, Pierre and Bandai, Toshiyuki and Gupta, Hoshin and Tartakovsky, Alexandre and Baity-Jesi, Marco and Fenicia, Fabrizio and Kifer, Daniel and Li, Li and Liu, Xiaofeng and Ren, Wei and Zheng, Yi and Harman, Ciaran J. and Clark, Martyn and Farthing, Matthew and Feng, Dapeng and Kumar, Praveen and Aboelyazeed, Doaa and Rahmani, Farshid and Song, Yalan and Beck, Hylke E. and Bindas, Tadd and Dwivedi, Dipankar and Fang, Kuai and Höge, Marvin and Rackauckas, Chris and Mohanty, Binayak and Roy, Tirthankar and Xu, Chonggang and Lawson, Kathryn},
    date = {2023-08-01},
}

Contributing

We welcome contributions! See CONTRIBUTING.md for details.

Please submit an issue to report any questions, concerns, bugs, etc.