gravis-tws-decomposition

Noe-B/gravis-tws-decomposition

gravis-tws-decomposition is a machine learning model from Noe-B. Use it for the machine learning task on the model card, and read the license before you ship it in a product. It is set up for pytorch. The card lists the license as mit.

[](https://github.com/NB11/grace-tws-decomp)

Downloads · 30 days

Access

Public

Updated Jun 12, 2026

Repo size

1.6 GB

Likes

Public

Hugging Face

Repo makeup

Click a slice to open those files.

.pt1.6 GB · 100%

At a glance

Library: pytorch
License: mit
Access: Public
Created: May 21, 2026
Updated: Jun 12, 2026
SHA: 4474f767

Model card

From the Hugging Face model README

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 — 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

# 1. Get the GitHub repo (training/inference/eval code)
git clone [email protected]: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.

Library: pytorch
License: mit
Created: May 21, 2026
Updated: Jun 12, 2026