GRACE TWS Decomposition with Deep Learning — pretrained checkpoints
Decompose GRACE / GRACE-FO total water storage anomalies into 5 storage compartments (groundwater, soil moisture, surface water, snow water equivalent, residual) using a 2-D U-Net trained on global hydrology simulations (WGHM / LISFLOOD) conditioned on real GravIS observations.
This repository hosts the four canonical model checkpoints. Training code, preprocessing, and the evaluation pipeline live on GitHub.
Paris Basin: standardized groundwater anomaly — in-situ wells (dashed) vs G3P vs our hybrid model. GRACE→GRACE-FO gap shaded.
Checkpoints
| File | Training input | Supervisor | Best at |
|---|---|---|---|
05-22_wghm_w5e5_quick-volcano/05-22_wghm_w5e5_quick-volcano.pt |
WGHM synthetic + residual-pool aug | WGHM | closed-loop & global wells |
05-22_wghm_hybrid_wise-sandstone/05-22_wghm_hybrid_wise-sandstone.pt |
WGHM synthetic pre-2002 + real GravIS paired post-2002 | WGHM | aquifer-mean r — the headline |
05-22_wghm_gravis_only_steady-delta/05-22_wghm_gravis_only_steady-delta.pt |
real GravIS only | WGHM | HPA (0.83), strong global |
05-22_lisflood_gravis_only_swift-glacier/05-22_lisflood_gravis_only_swift-glacier.pt |
real GravIS only | LISFLOOD | open-loop SM/SW; supervisor-bias check |
All four are 2-D U-Nets, ~31.1 M params, 5-channel input [tws, month_sin, month_cos, land_mask, dtws_12], 5-channel output [GW, SM, SW, Snow, residual].
Results — Pearson r vs in-situ wells
| Aquifer | quick-volcano | wise-sandstone | steady-delta | swift-glacier | G3P (our eval) | Paper (G3P) |
|---|---|---|---|---|---|---|
| Paris Basin | 0.76 | 0.74 | 0.66 | 0.59 | 0.54 | 0.63 |
| High Plains | 0.52 | 0.84 | 0.83 | 0.44 | 0.14 | 0.82 |
| Guarani | 0.63 | 0.67 | 0.51 | 0.59 | 0.81 | 0.81 |
Across 2,863 quality-gated GGMN wells worldwide, every model beats G3P on per-well median r (up to 0.48 vs 0.38). Full numbers in the GitHub README.
How to use
# 1. Get the GitHub repo (training/inference/eval code)
git clone git@github.com:NB11/grace-tws-decomp.git
cd grace-tws-decomp
python3.12 -m venv .venv
.venv/bin/pip install -r requirements.txt huggingface_hub
# 2. Download these checkpoints into the expected runs/{name}/{name}.pt layout
.venv/bin/huggingface-cli download Noe-B/gravis-tws-decomposition \
--local-dir runs/ --include "*.pt"
# 3. End-to-end open-loop + wells evaluation (~10 min per run)
.venv/bin/python scripts/evaluation/run_all.py \
eval.checkpoint=runs/05-22_wghm_hybrid_wise-sandstone/05-22_wghm_hybrid_wise-sandstone.pt
# → runs/{run}/eval_plots/{summary.csv, model_vs_g3p/, wells/, _utils/}
The data side (preprocessed caches at /scratch/grace-data/processed/...) is described in the GitHub README. See src/evaluation/inference.py::load_model for the loader the eval scripts use.
License
MIT. See the GitHub repo for the full license text.