You can deploy this model and start earning money today!

Noe-B/
gravis-tws-decomposition

Name: gravis-tws-decomposition
Availability: InStock

Views

No views yet

GRACE TWS Decomposition with Deep Learning — pretrained checkpoints

GitHub

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 — hybrid model vs G3P vs in-situ wells (standardized monthly anomalies)

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

bash

# 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.