tnfr-streamflow-arizona

TNFR — Transformer for Streamflow Forecasting (Arizona)

A standard encoder-decoder Transformer trained to forecast daily streamflow for Arizona basins. This is the time-series-only baseline used in the TADA / Global-TADA study — it does not consume satellite imagery or ViT features. Use it as a strong text-of-the-paper baseline, or as the warm-start backbone for the ViT-G gated variant.

Quick start

from huggingface_hub import snapshot_download
import sys, json, torch
from safetensors.torch import load_file

local = snapshot_download(repo_id="kdahal/tnfr-streamflow-arizona")
sys.path.insert(0, local)  # so `model_code` is importable

from model_code import Model

cfg = json.load(open(f"{local}/config.json"))
model = Model(
    enc_in=cfg["enc_in"], dec_in=cfg["dec_in"], c_out=cfg["c_out"], out_len=cfg["out_len"],
    d_model=cfg["d_model"], n_heads=cfg["n_heads"],
    e_layers=cfg["e_layers"], d_layers=cfg["d_layers"], d_ff=cfg["d_ff"],
    dropout=cfg["dropout"], embed=cfg["embed"], freq=cfg["freq"],
    activation=cfg["activation"], output_attention=cfg["output_attention"], mix=cfg["mix"],
)
model.load_state_dict(load_file(f"{local}/model.safetensors"))
model.eval()

Or run the bundled demo (uses random inputs, just proves the load path works):

pip install -r requirements.txt
python inference.py

Inputs

Output: (B, 30, 1) — 30-day-ahead streamflow, in the same normalized units as training (undo with per-basin mean/std).

Architecture

Encoder-decoder Transformer with:

• Sinusoidal positional embedding + 1-D conv token embedding (no temporal embedding)
• 3 encoder blocks, 2 decoder blocks, 8 heads
• d_model = d_ff = 256, dropout = 0.1, GELU activations
• 5.26M parameters, 21 MB on disk

Input layout (250 features) = daily meteorology (PRISM) + static catchment attributes (CAMELS + Arizona-specific), all z-score normalized against the locked training statistics.

Training

• Dataset: Arizona USGS stations (az_streamflow split in the source repo), ~40 years of daily data
• Loss: Masked Huber (δ=1.0) on the 30-day forecast window
• Optimizer: AdamW, LR=1e-3, weight decay 1e-2, OneCycle schedule
• Seq. lengths: 365 / 180 / 30 (enc / label / pred)
• 5 random seeds; this checkpoint is seed 1

Performance

Reported as median over all Arizona test basins (see the paper for full tables):

Normalization

Inputs to this model must be z-score normalized with the frozen training statistics — not statistics computed from your own data. Those stats are shipped as locked_stats.json and cover all 315 Arizona USGS basins used in training.

import json, numpy as np
stats = json.load(open("locked_stats.json"))
site  = "09379025"                                        # USGS site number
means = stats["basin_dynamic_stats"][site]["means"]       # dict: feat -> mean
stds  = stats["basin_dynamic_stats"][site]["stds"]        # dict: feat -> std

# apply to a daily dataframe `df` with columns matching feature names
for col in means:
    df[col] = (df[col] - means[col]) / (stds[col] + 1e-6)

# un-normalize model output (streamflow is in the same per-basin z-space)
mean_q, std_q = means["streamflow"], stds["streamflow"]
y_real = np.maximum(y_pred.numpy() * std_q + mean_q, 0.0)

Top-level keys in locked_stats.json:

• basin_dynamic_stats: {site: {"means": {feat: µ}, "stds": {feat: σ}}} — per-basin
• global_static_stats: {"means": {...}, "stds": {...}} — for static catchment attributes
• static_data_lookup: {site: {feat: value}} — raw static attribute values per basin (already z-scored using global_static_stats when assembling x_enc)

Files

• model.safetensors — weights (state_dict)
• config.json — exact constructor kwargs
• locked_stats.json — frozen per-basin normalization statistics (315 Arizona basins, 40 dynamic features)
• model_code/ — self-contained PyTorch source (no external repo needed)
• inference.py — end-to-end demo

License

MIT.

Citation

If you use this model, please cite the accompanying paper (dissertation / forthcoming). See the project repo for the canonical BibTeX.

Related repos

• kdahal/vitg-streamflow-arizona — ViT-Gated variant (learnable gate over satellite ViT features)
• kdahal/cnn-streamflow-arizona — CNN hybrid (ResNet18 on raster tiles + Transformer time-series)