dmg 1.4.0

PyTorch Differentiable Modeling Framework

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

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. (2024).

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

Setup	Installation options (PyPI, pip, UV, Conda)
How to Run	CLI usage and custom model development
Configuration	Config file system and full settings glossary
API Reference	Public API — models, loss functions, NNs, utilities
Examples	Jupyter notebook tutorials
Changelog	Release history

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., et al. (2023). Differentiable modelling to unify machine learning and physical models for geosciences. Nature Reviews Earth & Environment, 4(8), 552–567. https://doi.org/10.1038/s43017-023-00450-9

BibTeX

@article{shen2023differentiable,
  title={Differentiable modelling to unify machine learning and physical models for geosciences},
  author={Shen, Chaopeng and others},
  journal={Nature Reviews Earth \& Environment},
  volume={4},
  number={8},
  pages={552--567},
  year={2023},
  publisher={Nature Publishing Group},
  doi={10.1038/s43017-023-00450-9}
}

Contributing

We welcome contributions! See CONTRIBUTING.md for details.

---

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