δSnowSac

A neural network that learns to calibrate NOAA's operational snowmelt and soil-moisture models across many river basins at once, including ones it has never seen, while leaving NOAA's original Fortran code untouched.

Every day, NOAA forecasts river flow across the U.S. using two models written in the 1970s: one for snowmelt, one for soil moisture. Before these models can be used on a river, hydrologists usually have to tune these models with respect to each of the basins, creating a need for workflows that enable automated calibrations across multiple basins.

Machine learning has already shown potential to learn calibration across many basins at once. However, pure data driven models (like rainfall-runoff LSTMs) don't expose familiar hydrologic states or parameters that forecasters use for diagnosis and turn the problem into a typical ML black box. Another class of models called the differentiable hydrological models provide a better alternative by having a neural network learn parameters for physics model, and optimize end-to-end based on final loss function. However, this requires rewriting the physics model in a framework like Jax or PyTorch. I didn't do this.

By leveraging Tesseract, the original Fortran NOAA models (Snow-17 and SAC-SMA) that run in production are left untouched, and the neural network is trained through them by gradient descent. What the network learns are the parameters of the real operational models: melt rates, snowfall correction, soil-water storage capacities, drainage rates, among 27 in total.

Trained on 35 snowy basins and tested on 10 it never saw, the model reaches a median NSE of 0.70 on the unseen basins (NSE: 1 = perfect, 0 = no better than predicting average flow every day; ~0.7+ is usually considered usable). A pure-ML LSTM does better on accuracy (0.795), but what you get in exchange is a model a forecaster can actually open up: every prediction comes with physical parameters and states (snowpack, soil moisture) that can be inspected, audited, and compared against how NOAA already calibrates these basins.

The network learning to calibrate NOAA's Fortran models, shown on a basin it never trained on. Every frame is real model output from a saved checkpoint, run through both unmodified Fortran models. The untrained network roughly doubles the snowmelt peak (NSE -0.42, worse than predicting the average flow). By the end of training it tracks the observed flow (NSE +0.82). This basin was chosen because the learning is most visible there, and it ends among the better fits. The curves underneath are the honest summary: median NSE across all 10 held-out basins (0.23 untrained to 0.70) and across the 35 training basins (0.38 to 0.84). Regenerate with results/animate_training.py.

How it works

Snow-17 and SAC-SMA are compiled Fortran, so PyTorch's autograd cannot see inside them and backpropagation stops at their boundary. This project does not rewrite them. Each model is wrapped, unmodified, in its own Tesseract container. Each container exposes a forward run and a finite-difference gradient, and tesseract-torch splices both into the PyTorch graph as ordinary differentiable layers. A neural network can then learn the parameters of NOAA's operational code directly, rather than the parameters of a reimplementation of it.

Solid arrows are the forward pass; the dashed arrow is the gradient. It is computed by forward-mode automatic differentiation over the two Tesseracts: a tangent is seeded on each parameter, Snow-17's jacobian_vector_product turns it into a tangent on RAIM, and SAC-SMA's jacobian_vector_product carries that tangent through to a tangent on runoff. tesseract-torch chains the two endpoints automatically — the RAIM tangent is handed from one container to the next as a single vector, never as a matrix. Snow-17 produces RAIM (rain-plus-melt) — the same coupling flux NOAA runs operationally into SAC-SMA — so the container boundary sits at a real, existing operational seam.

Why Tesseract

PyTorch's .backward() walks a recorded graph — Snow-17 and SAC-SMA are compiled Fortran, so nothing is recorded and autodiff stops cold. Each model is wrapped as its own Tesseract exposing apply() and both finite-difference derivative endpoints, jacobian_vector_product() (forward mode) and vector_jacobian_product() (reverse mode); tesseract-torch splices both into the autograd graph as ordinary differentiable layers. Two Tesseracts are composed here: NOAA maintains Snow-17 and SAC-SMA as separate modules, and a standalone Snow-17 Tesseract is reusable with any downstream rainfall-runoff model, not just this one.

Both containers are built and gradient-checked end-to-end — against autograd ground truth and an independent brute-force check on cheap stand-ins first (tests/test_coupling_toy.py), then against the real Tesseracts (tests/test_pipeline_hhwm8.py, tests/test_gradients.py). tesseract build runs in CI on every push, building both containers from scratch and smoke-testing apply() against the built images (see .github/workflows/ci.yml).

Results

NSE (Nash-Sutcliffe Efficiency) is the standard skill metric for streamflow models. 1.0 is a perfect match to observed flow. 0 means the model is no better than guessing the historical average every day, and a negative value is worse than that. What counts as good depends on the domain and the basin, but 0.7 or higher is generally read as a solid, usable model.

ParamNet predicts all 27 learnable parameters (11 Snow-17 + 16 SAC-SMA) from each basin's static CAMELS attributes plus a climatology sequence. It is trained end to end across 35 snow-dominated CAMELS basins, and 10 more basins are held out (WY1991-1999). The held-out basins test spatial generalization, which is the same problem as predicting flow in ungauged basins.

	median train NSE	median held-out NSE
epoch 1	+0.38	+0.28
epoch 150 (final)	+0.84	+0.70

Held-out skill rises quickly and then plateaus. Most of the held-out gain comes in the first ~10 epochs, while training NSE keeps improving, so the train/held-out gap widens to about 0.14 by epoch 150. Held-out NSE never declines, so this is limited generalization from 35 basins rather than overfitting. Full numbers and reproduction commands are in results/README.md.

Daily streamflow in two held-out basins after training. The black line is the USGS gauge and the blue line is the hybrid model. The top basin is one of the best held-out fits (NSE 0.82). The bottom one is near the median (NSE 0.69): it starts spring melt a little late and overshoots the 1997 peak. Regenerate with results/plot_hydrograph.py.

These numbers come from the published run, trained with an earlier, hand-written cross-container coupling. Retrained from scratch with the current forward-mode code (same config and seed, results/runs/model_9yrs_spatial_fwdmode/), the model reaches median NSE 0.84 on training basins and 0.73 on held-out basins. Basin by basin, 5 of the 10 held-out basins improve and 5 get worse (mean held-out NSE 0.66 → 0.62), so this reproduces the result within run-to-run variation rather than improving on it. See results/README.md.

For comparison, a properly engineered LSTM (NeuralHydrology) was trained and tested on the exact same 35/10 basin split:

model	median held-out NSE
NeuralHydrology LSTM	0.795
this hybrid model	0.70

Basin by basin, the gap is smaller than the medians suggest. The LSTM leads on 6 of 10 basins and the hybrid model wins on 4. The gap ranges from essentially tied (09035900, 0.822 vs. 0.809) to wide (11230500, 0.408 vs. 0.827). Regenerate with results/plot_basin_comparison.py.

On raw NSE, the hybrid model currently trails a competent LSTM (see results/external/neuralhydrology_lstm_pub/). Beating an LSTM was never the goal. The goal was to learn the parameters of NOAA's actual operational models end to end, without rewriting the physics.

Reproduce

Requires uv and gfortran.

1. Build and test

git submodule update --init --recursive   # vendors NOAA-OWP/snow17 + sac-sma, pinned commits
make test                                  # creates .venv, builds Fortran shims, runs pytest

make env (run by make test) installs PyTorch's CUDA 12.6 build, which works with NVIDIA driver 525 or newer. For other hardware, see the comment at the top of the Makefile.

2. Get the CAMELS data (once)

About 3.4 GB. make test does not fetch it.

data/download_camels.sh
.venv/bin/python data/select_basins.py
.venv/bin/python data/build_attributes.py
.venv/bin/python data/build_pet.py
.venv/bin/python data/build_climatology.py

These build the default set: the 45 most snow-dominated CAMELS basins. To train on other basins, see step 5; it downloads CAMELS by itself if needed.

3. Train

.venv/bin/python src/train.py

By default this trains on the 45 snow-dominated basins with the temporal split: every basin trains on water years 1991–1993 and is tested on 1994–1996. It runs 150 epochs. Each epoch prints a line like epoch 12 train_nse=+0.41 test_nse=+0.38, where test_nse is the median NSE over the test set. Use train.n_epochs=5 for a quick check that everything runs.

Choose the basins with data= and the evaluation with split=:

data=camels_snow35	45 most snow-dominated basins (default)
data=camels_531	the standard 531-basin CAMELS benchmark subset
data=camels_671	all 671 CAMELS basins
data=camels_list data.basin_list=<file>	your own list (step 5)
split=temporal	same basins, train and test on different years (default)
split=spatial	80% of basins train, 20% held out, same years (WY1991–1999)

The 531 and 671 sets are built automatically on first use (about a minute). The saved run results/runs/model_9yrs_spatial/ is reproduced by src/train.py split=spatial device=cpu.

How long it takes: the Fortran runs are spread over n_workers processes (default 32). On a shared 104-core server:

Data, split	Time per epoch	150 epochs
45 basins, spatial (9 years)	about 2 s (55 s with n_workers=0)	about 6 min
531 basins, temporal (3 years)	about 35 s	about 1.5 h

Everything is written to results/runs/hybrid_<split>_<timestamp>/:

File	Contents
checkpoint.pt	final network weights, the input to inference
normalization.npz	the feature scaling the network was trained with
checkpoints/epoch_NNNN.pt	weights and optimizer state every 10 epochs
history.json	per-epoch train/test NSE
test_predictions.json	simulated streamflow and NSE for each test basin
config.yaml	the full config the run used

Common overrides (any config value can be set this way):

.venv/bin/python src/train.py device=cpu                # network + loss on CPU (default: cuda)
.venv/bin/python src/train.py n_workers=64              # more processes for the Fortran runs
.venv/bin/python src/train.py train.batch_size=64       # minibatches of 64 basins per gradient step
.venv/bin/python src/train.py seed=1 train.n_epochs=50 train.lr=1e-3
.venv/bin/python src/train.py output_dir=results/runs/my_run

For a long run on a remote machine, detach it and keep a log:

nohup .venv/bin/python src/train.py data=camels_531 > train.log 2>&1 &
tail -f train.log

4. Run inference

Point checkpoint= at a checkpoint.pt from step 3, and pass the same data= you trained with:

.venv/bin/python src/infer.py checkpoint=results/runs/hybrid_temporal_<timestamp>/checkpoint.pt
.venv/bin/python src/infer.py data=camels_531 checkpoint=<path>
.venv/bin/python src/infer.py split=spatial checkpoint=results/runs/model_9yrs_spatial/checkpoint.pt   # saved run

This runs every basin in both the train and the test set, prints the median NSE of each, and writes predictions.json to results/predictions/hybrid_<split>_<timestamp>/. In that file, predictions.train and predictions.test map each gauge ID to its time window, simulated streamflow (mm/day) and NSE.

It is much faster than training: without gradients, each basin is a single Snow17 → SAC-SMA run (about 10 s for the 45 basins). Overrides work the same way; device= and n_workers= work here too. To score a trained model on another period, change the windows, e.g. split.test_window.start=1999-10-01 split.test_window.end=2004-09-30 (or split.window.* with split=spatial).

model must match what the checkpoint was trained with; the checkpoint stores only weights, not the architecture. data can be a different basin list (see step 5).

5. Use your own basin list

Any set of CAMELS basins works. Write their gauge IDs in a text file, one per line (leading zeros optional, # starts a comment):

# my_basins.txt
01013500
06623800
12145500

Then train on it:

.venv/bin/python src/train.py data=camels_list data.basin_list=my_basins.txt

On the first run this:

1. downloads CAMELS if it isn't there yet (one ~3.4 GB archive with all 671 basins; CAMELS has no per-basin download),
2. checks that every ID is a CAMELS basin,
3. holds out 20% of the basins for testing (data.heldout_fraction, chosen with data.split_seed),
4. builds the list's attributes, climatology and PET under data/camels/datasets/my_basins/.

Later runs reuse those files, and rebuild them automatically if you edit the list. To check a list before training, run .venv/bin/python data/prepare_dataset.py my_basins.txt. It prints the basins with their names, snow fraction and train/held-out split.

To choose the held-out basins yourself, use a CSV with a split column:

gauge_id,split
01013500,train
06623800,train
12145500,heldout

Things to know:

• Observations in the window. Every basin needs observed streamflow in its train and test windows. If some don't, the run stops and lists them. All 671 CAMELS basins have observations in the default windows.
• Held-out basins. data.heldout_fraction only matters for split=spatial; with the default temporal split every basin is in both sets, with different years.
• Inference on other basins. Inference can use a different list than training: src/infer.py data=camels_list data.basin_list=other.txt checkpoint=.... The network's input features are scaled with the training list's statistics, which training saves as normalization.npz next to the checkpoint.
• Larger lists. Time per epoch grows with the number of basins: each basin costs 27 model runs per gradient step, spread across n_workers processes. For hundreds of basins, raise n_workers if the machine has cores to spare, and consider train.batch_size so the network is updated more than once per epoch.

Notes

Training and inference are driven by Hydra configs under configs/ (data / split / model / train), not hardcoded constants. device=cuda (or auto) puts the parameter network and loss on a GPU; the Fortran physics always runs on CPU, and src/coupling.py moves tensors across that seam. The bottleneck is the Fortran/Tesseract calls (finite-difference gradients), not model size, so with the current small network a GPU does not speed training up. Those calls are what n_workers parallelizes: each basin needs 27 independent gradient passes, and src/physics_pool.py spreads every (basin, pass) pair across worker processes. Results are identical for any n_workers. See results/README.md for saved runs and results/compare_runs.py for comparing them.

Docker note: day-to-day apply() / jacobian_vector_product() development runs through tesseract_core.Tesseract.from_tesseract_api() directly (no container needed); actual tesseract build runs in CI, where Docker is available.

Layout

external/snow17/, external/sac-sma/   git submodules, pinned commits (Apache-2.0, unmodified)
patches/                              disclosed, minimal, build-time-only patch to vendored source
fortran/                              bind(C) shims threading each model's state explicitly
tesseracts/snow17/, tesseracts/sacsma/  the two Tesseract containers: apply() + finite-difference JVP/VJP
src/coupling.py                       forward-mode bridge: physics differentiation -> network autograd
src/pipeline.py                       chains the two Tesseracts (apply_tesseract) into coupling.py
src/physics_pool.py                   runs each basin's gradient passes in parallel worker processes
src/paramnet.py                       LSTM + MLP: attributes/climatology -> 27 bounded parameters
src/train.py, src/infer.py            Hydra-driven training / checkpoint scoring CLIs
configs/                              Hydra config groups (data/split/model/train)
data/                                 CAMELS download + basin selection + attribute/PET/climatology prep
data/prepare_dataset.py               any CAMELS basin list -> training dataset (run automatically)
data/basin_lists/                     the 531-basin benchmark subset and all 671 CAMELS basins
tests/                                shim determinism/mass-balance, JVP/VJP checks, coupled-chain regression,
                                      parallel-vs-serial equality, CPU/GPU gradient agreement, dataset prep
notes/NOTES.md                        upstream Fortran findings, with a before/after proof
notes/logs.md                         design-decision rationale log
results/                              saved, seeded, reproducible run directories + external comparisons

Status and what's next

• Status: research prototype. The core pipeline works and is tested end to end. Development is to be continued.
• Contributions: issues are very welcome (bug reports, questions, ideas, basins where it fails). If you'd like to contribute code, please open an issue first so we can agree on the approach; the codebase is still moving.
• Contact: Kamlesh Sawadekar on LinkedIn, or email kas7897 [at] psu [dot] edu
• Citation: if you use this work, please cite it using the "Cite this repository" button on GitHub (from CITATION.cff).

Origin

This project started at the Pasteur Labs Tesseract Hackathon 2026 (Track 03: Hybrid ML + mechanistic models), where it won Second Prize 🥈.

License

Apache-2.0. See LICENSE and NOTICE. Snow-17 is vendored and linked against unmodified. SAC-SMA is vendored unmodified as a pinned submodule. One disclosed, minimal patch is applied to a build-time copy only to fix a confirmed upstream defect; external/sac-sma itself is never modified. "Original work" applies to this project, not its dependency tree: the shims, patches, Tesseract wrappers, gradient endpoints, and training pipeline are original work written during the hackathon period (Aug 3-31, 2026).