Differentiable hydrology · work in progress

A differentiable, globally-calibrated
parameter set for CREST

Reimplementing the CREST hydrologic model (after EF5) in an autodiff framework, then learning parameters by backpropagation against Earth-observation truth. Target: a global 0.1° parameter field — built CONUS-first, in three stages: snow → water balance → routing.

0.1°global target grid
3 + 1staged · then joint
Stage 1active: Snow17
autodiffAdam, not SCE-UA

Why this project

CREST is a distributed, grid-based water-balance model used for large-scale flash-flood forecasting. EF5 calibrates it with basin-uniform parameter multipliers via MCMC/random search. We want something different: spatially-distributed parameters that generalize globally, learned efficiently with gradients.

Gradients, not search

EF5 uses DREAM (MCMC) and ARS. We reimplement the per-grid physics so it is fully differentiable and optimize parameters with Adam — orders of magnitude more sample-efficient in high dimensions.

Grid-independent = parallel

Snow and water balance are column processes with no lateral coupling. Each 0.1° cell is solved independently — vmap over the grid, embarrassingly parallel. Only routing couples cells.

Observation-anchored

Force with reanalysis (ERA5-Land / MSWEP), but calibrate against independent observations (SNODAS, SNOTEL, fSCA, streamflow, ET, GRACE) — never model-to-model, to avoid circularity.

The approach

Build the differentiable model bottom-up, validating physics before scale. Calibrate stage-by-stage so each component is well-posed, freezing upstream outputs as inputs to the next stage — then optionally fine-tune everything end-to-end.

Staged, then joint

  • Stage 1 — Snow (Snow17). Direct gradient calibration of snow parameters against gridded SWE + snow-cover truth.
  • Stage 2 — Water balance (CREST). Snowmelt becomes input; calibrate runoff parameters against streamflow volume, ET, and total water storage.
  • Stage 3 — Routing. Grid runoff routed on the MERIT-Hydro network against daily hydrographs. approach TBD
  • ★ Joint. Warm-start from staged parameters; fine-tune all params end-to-end.

From per-grid calibration → dPL

For snow, SWE truth is near-global, so we calibrate parameters θ directly as trainable tensors (no neural net yet), with light spatial regularization for smooth maps.

The calibrated θ field is kept structured to warm-start differentiable parameter learning (dPL) later — a network θ = gφ(attributes) (incl. AlphaEarth embeddings) that extrapolates to ungauged cells for the runoff stage and the global rollout.

Stage 1 · Snow Stage 2 · Water balance Stage 3 · Routing Joint end-to-end

The stages in detail

Each stage lists its forcing inputs, the observational truth/constraints, the learnable output parameters, the modeling route, and the calibration scale. Parameter names follow the EF5 source.

STAGE 1

Snow — differentiable Snow17

current focus

A column snow-accumulation/ablation model. Parameters θ are trainable leaf tensors optimized by autodiff + Adam — replacing SCE-UA/DDS with backprop. No neural net at this stage.

Forcing
ERA5-Land / MSWEP — daily precipitation & air temperature
Truth
SNODAS (1 km SWE, primary) · MODIS/VIIRS fSCA (melt-out + ADC) · SNOTEL (held-out points) · UA-SWE / SWEML (cross-check). Never ERA5-Land SWE.
Route
Differentiable Snow17 → soil water balance (S→C), per-grid · spatially regularized
Scale
Per grid cell (0.1°), snow-affected cells masked
SCFMFMAXMFMIN UADJMBASENMF TIPMPLWHC PXTEMP*ADC*DAYGM*

* added as learnable beyond EF5, which hardcodes PXTEMP=1°C and implements no ADC / ground-melt.

Differentiable details

  • Time recursion via scan / BPTT over states (ice, liquid, ATI, cold deficit).
  • Soft rain/snow split: f_snow = σ((PXTEMP − Ta)/τ) — no hard threshold.
  • ADC = differentiable piecewise-linear with learnable tie-points.
  • Bounds via sigmoid reparam θ = lo+(hi−lo)·σ(θ_raw) (Anderson 2006 ranges).
  • Spatial regularization (neighbor TV / graph-Laplacian) for coherent maps.

Loss & splits

L = Σcell wcell · [ KGESWE + λ₁·meltout + λ₂·fSCACE ] + λsp·Rspatial

  • wcell by multi-product agreement — down-weight where SWE products disagree (deep mountain snow).
  • Splits: temporal holdout (train vs test years) + spatial holdout (leave-out SNOTEL sites).

Phase A — single-basin physics validation

Debugging rig (e.g. East River nr Gunnison, CO). All gates must pass before scaling up.

A1Forward sanity — default θ accumulates winter SWE, melts in spring; plausible melt-out vs SNOTEL/SNODAS.
A2Identical-twin recovery — recover known θ from synthetic SWE (SCF, MFMAX/MIN, UADJ, PXTEMP).
A3Gradient check — autodiff vs central finite differences, max rel. error < 1e-4 per parameter.
A4Mass-balance closure — Σ snowfall − Σ melt − ΔSWE ≈ 0 (~1e-6).
A5BPTT stability — 10–20 yr daily run, no NaN/Inf, bounded gradient norms, checkpointing if needed.

Phase B — Upper Colorado / Rockies cluster (0.1°)

θ as [N_cells, N_params], vmap forward. Acceptance: SWE KGE median > 0.7 (temporal holdout) · melt-out bias within ±7 days median · SNOTEL leave-out NSE/KGE reported · spatially-coherent θ maps · gradient health logged.

Deliverables: differentiable Snow17 module · Phase-A report (A1–A5 numbers/plots) · calibrated θ field (NetCDF/zarr, saved as dPL warm-start) · diagnostics (KGE / melt-out / per-param maps, loss & gradient curves) · README.
STAGE 2

Water balance — CREST runoff

next

CREST's variable-infiltration curve turns rain + snowmelt into fast (surface) and slow (interflow) runoff. Grid-independent → vmap over cells. This is where dPL becomes primary for global generalization.

Forcing
Rain + snowmelt (Stage-1 output, frozen) + PET + soil attributes
Truth
Caravan / GRDC-Caravan streamflow (volume) · ET (MOD16 / PML / FLUXCOM) · GRACE TWS
Route
Differentiable water balance (W→C), multi-basin · dPL: θ = gφ(attributes)
Scale
Basin-aggregated · volume / monthly
WMBIM KEFCIWU
STAGE 3

Routing — river network

approach TBD

The only network-coupled stage: cells drain downstream along the river network. Routing approach is still being decided — see tech routes.

Forcing
Grid runoff (Stage-2 output, frozen) + MERIT-Hydro river network
Truth
Gauge daily hydrographs
Route
Differentiable routing, or merged into the joint pass — TBD
Scale
River network · daily

Parameters follow EF5's actual routing schemes (the table's coeM/expM/coeR/coeS map onto these):

Linear: coemriverunder leakoleakith
Kinematic: alpha0alphabeta underleakith
★ JOINT

End-to-end fine-tuning

later
Forcing
Full chain (snow → water balance → routing)
Truth
SWE + Q + ET + TWS (multi-objective, mixed scales)
Route
End-to-end, multi-objective · warm-started from staged θ
Scale
Mixed (cell · basin · network)

Stack & data

Autodiff core, parallel over the grid; reproducible, config-driven runs on HPC.

Compute

JAX (lax.scan for time, vmap over grid) preferred; PyTorch acceptable (loop + gradient checkpointing). Optimizer: Adam.

Infrastructure

OSCER SLURM cluster · data stored as zarr · access via Google Earth Engine where possible · YAML config, seeded & reproducible.

Package (planned)

pyCREST: modular differentiable Snow17 / CREST / routing + dPL, data loaders, training, objectives. Built in a later round.

RoleDatasetsUsed in
Forcing (P, T)ERA5-Land · MSWEPAll stages
Snow truthSNODAS · SNOTEL · MODIS/VIIRS fSCA · UA-SWE · SWEMLStage 1
Runoff / flux truthCaravan / GRDC-Caravan (Q) · MOD16 / PML / FLUXCOM (ET) · GRACE (TWS)Stage 2
Hydrograph truthGauge daily dischargeStage 3 / Joint
River networkMERIT-HydroStage 3 / Joint
Attributes (for dPL)Soil · climate · topography · AlphaEarth embeddingsStage 2+ · global rollout

Technical-route options

Open design decisions, with the leading candidate marked. These will be resolved as the work proceeds — the routing approach in particular is still undecided.

Routing engine decision pending

  • (a) Differentiable linear-reservoir / Muskingum-Cunge on MERIT-Hydro as sparse scatter-add along the downstream index, in topological order — fully autodiff. leading
  • (b) Keep EF5 kinematic wave (implicit Newton solve, non-differentiable) and calibrate routing with DREAM / ARS as a separate step.
  • (c) Learned unit-hydrograph / Nash cascade per reach — light, flexible, trainable.

Other open choices

  • Framework: JAX (scan/vmap) vs PyTorch (loop + checkpointing).
  • Forcing: ERA5-Land vs MSWEP (or ensemble of both).
  • Snow regularization strength λsp: trade off map smoothness vs detail.
  • dPL onset: when to switch from direct-θ to θ = gφ(attributes) (Stage 2 vs global rollout).
  • Joint coupling: fully end-to-end vs frozen-upstream fine-tuning.

Stage-1 methods & configuration (as run)

The exact experimental setup for the full-CONUS snow calibration (WY2010–2024).

Model

  • Differentiable Snow-17 (PyTorch), transcribed from EF5 and verified vs the C++ kernel to single-precision round-off.
  • 9 learnable params / cell: uadj, mbase, mfmax, mfmin, tipm, nmf, plwhc, scf, pxtemp.
  • Soft rain/snow partition; sigmoid reparam bounds every param to Anderson-2006 ranges; branches via torch.where (subgradient).

Domain & resolution

  • Full CONUS, 0.1° (~11 km), grid 250×585 = 146,250 cells (lon −125…−66.5, lat 24…49).
  • Snow cells (peak>30 mm): 58,265; seasonal snowpack (≥100 mm): 26,335.

Time & splits

  • WY2010–2024 (15 water years), daily, 5,114 days. SNODAS truth every 8 days → 640 obs steps.
  • Temporal split (66/34): train 2010-10 … 2019-12 (422 steps); test 2019-12 … 2024-09 (218 steps).
  • Spatial split: SNOTEL cells held out from the loss for independent validation.

Inputs / truth

RoleData
TemperatureERA5-Land daily T2m (GEE) + per-cell warming (West from SNOTEL bias; East uncorrected)
PrecipitationERA5-Land daily total precip (GEE) — swappable: IMERG/MSWEP
ElevationSRTM (for atmospheric pressure)
Calibration targetSNODAS gridded SWE (1 km→0.1°), 8-daily
Independent val.SNOTEL 830 stations (West), point SWE

Loss & optimization

  • Loss = weighted MSE on SWE (melt-season AMJ ×2) + spatial-TV reg (λ=2), NaN-masked.
  • Adam lr 0.08; two-stage: regional-shared (50 it) → per-cell (70 it); grad-clip 2 + NaN-grad guard.
  • Compute: 2.5° tiles (138), memmap cubes, 46 workers / 48 cores, ~2.5 h.

Result (held-out test-KGE vs SNODAS)

Snowpack depthnmedian KGE
peak ≥ 100 mm26,2820.63 (64%>0.5)
peak ≥ 200 mm10,4690.66
peak ≥ 400 mm4,4100.66
SNOTEL independent (West)8040.36

Product: snow17_conus_era5.nc — 9 param maps + test-KGE + warming + snow/snowpack masks + peak-SWE.

Progress tracker

Live milestone board — flip the badges (todoin progdone) as the work advances. Last updated: 2026-07-01 — Stage-2 water balance done (424 CONUS basins, KGE 0.67 / holdout 0.56). See Stage-1 Snow → and Stage-2 Water →

done
Clone & map EF5 source
Snow17 / CREST / routing parameters, states, equations, and grid-independence all traced in code.
done
Project roadmap & this tracker site
Three-stage plan locked; static site deployed to HuggingFace.
done
Lossless Python port verified vs EF5 C++
NumPy ports of Snow17 & CREST reproduce the verbatim C++ kernels to single-precision round-off (SWE ~3e-7 rel; CREST near-exact).
done
Stage 1 · Step 0 — snow data access
ERA5-Land + MODIS via GEE (re-authed); SNODAS (NSIDC, open) + SNOTEL (USDA) downloaded. Phase-A site: Butte SNOTEL, East River CO.
done
Stage 1 · differentiable Snow17 (PyTorch)
Scan-based BPTT, soft rain/snow partition, sigmoid-bounded θ, learnable PXTEMP. Adam-trainable.
done
Stage 1 · Phase A gates (East River, real data)
A1 sanity ✓ · A3 gradient check (rel 9e-9) ✓ · A4 mass closure (float32 round-off) ✓ · A5 6-yr BPTT stable ✓. Calibrated to Butte SNOTEL SWE: KGE test 0.48→0.81 (2-yr temporal holdout).
done
Stage 1 · Phase B — SWE calibration + validation
240-cell 0.1° grid, ERA5-Land forcing + gridded SNODAS truth (3 WY), per-cell θ via BPTT+Adam + spatial reg. Median test-KGE 0.75; SNOTEL leave-out (cells excluded from loss) KGE 0.71 vs independent ground obs. Base θ field exported (NetCDF).
done
Stage 1 · multi-region calibration (efficient recipe converged)
6-round autonomous loop over 4 CONUS climates. Converged recipe: two-stage (regional→per-cell) + data-driven per-region temperature correction (from SNOTEL T-bias) + SCF cap + ablation-weight + light fSCA. ~33% faster, rescues data-sparse regions (Sierra 0.50→0.69). Held-out KGE: UC 0.71, N.Rockies 0.58(temporal), Sierra 0.68, Cascades 0.48(spatial). Diagnosis: melt-out lag = ERA5-Land cold bias (continental) — fixed by warming; maritime limited by 3-WY record. θ exported per region (NetCDF).
done
Stage 1 · Western-US CONUS, 15-year record (WY2010–2024)
Scaled to the whole Western US at 0.1° (24,857 cells, 59 tiles, 12 parallel workers). Data cubes: SNODAS + ERA5-Land (GEE) + 830 SNOTEL + SRTM; per-cell warming = interpolated SNOTEL T-bias. Seasonal snowpack (peak≥200mm): median test-KGE 0.62; deep snow ≥400mm: 0.67 across 12k+ cells — the recipe transfers West-wide. Key lesson: mask to seasonal snowpack (ephemeral <100mm cells are unfittable). Product: snow17_conus_west.nc.
done
Stage 1 · FULL CONUS, 15-year record (WY2010–2024)
Whole CONUS at 0.1° (585×250, 58,265 snow cells, 138 tiles) on all 48 cores via memmap cubes. Snowpack (peak≥100mm): median test-KGE 0.61 (n=21,581); ≥400mm 0.66. East (Great Lakes/Midwest/Northeast) calibrates well. Swappable-precip architecture (ERA5/IMERG/MSWEP) in place. Product: snow17_conus_era5.nc.
done
Stage 1 · FINALIZED — forcing verified + physical-prior parameters
Diagnostics: ERA5 melt-season temperature unbiased vs SNOTEL (−0.04 °C), grid-scale precip adequate. A bounds experiment showed melt/retention params are weakly identified → added a physical prior (λ=200) pulling them to literature values: bound-railing MFMAX 54%→0.8%, MFMIN 76%→3%, at only −0.036 KGE. Precip-swap tested (ERA5 vs IMERG) = low yield in maritime. Full Stage-1 results page →
todo
Stage 1 · optional — PRISM maritime precip + East (GHCN) warming
Only lever left for maritime deep-snow bias (sub-grid orography); East temperature correction needs GHCN/COOP (no SNOTEL east).
done
Stage 2 · differentiable CREST water balance (424 CONUS basins)
EF5-US a-priori grids × per-basin multipliers, fed by frozen Stage-1 snowmelt, closed with a groundwater baseflow reservoir. Monthly-runoff median KGE 0.67 (93% basins >0); temporal-holdout test KGE 0.56. Full Stage-2 results page →
todo
Stage 2b · dPL regionalization (ungauged) + multi-signal
Attributes→multipliers NN for a continuous 0.1° field & ungauged prediction; optional ET + GRACE constraints.
todo
Stage 3 · routing on MERIT-Hydro
Resolve routing engine choice; calibrate vs daily hydrographs.
todo
★ Joint end-to-end & global 0.1° rollout
Warm-start fine-tuning on SWE + Q + ET + TWS; extend to ungauged cells via dPL.
todo
pyCREST Python package
Package the differentiable modules + training harness for release.

Key references

Anderson, E. A. (1973). National Weather Service River Forecast System — Snow Accumulation and Ablation Model. NOAA Tech. Memo. NWS HYDRO-17.

Anderson, E. A. (2006). Snow Accumulation and Ablation Model — SNOW-17. NWS documentation (parameter ranges).

Wang, J., Hong, Y., et al. (2011). The Coupled Routing and Excess Storage (CREST) distributed hydrological model. Hydrol. Sci. J., 56(1).

Flamig, Z. L., Vergara, H., Gourley, J. J. (2020). The Ensemble Framework For Flash Flood Forecasting (EF5). Geosci. Model Dev.

Feng, D., Liu, J., Lawson, K., Shen, C. (2022). Differentiable, learnable, regionalized process-based models with multiphysical outputs. Water Resour. Res.

Kratzert, F., et al. (2023). Caravan — A global community dataset for large-sample hydrology. Scientific Data.

Yamazaki, D., et al. (2019). MERIT-Hydro: high-resolution global hydrography. Water Resour. Res.

Zeng, Z., et al. (UA-SWE) & NOHRSC SNODAS — gridded snow water equivalent products.