ProbEns

Constrained probabilistic modelling for rainfall–runoff — predict distributions, not points.

📄 Paper under review — please cite (see Citation)

ProbEns turns any PyTorch module into a model that predicts the full conditional distribution of streamflow, using distribution families that respect the physics of discharge: strictly positive, right-skewed, heavy-tailed. Across three CAMELS-DE catchments spanning the regime range it has the best CRPS on every one, and is never beaten by any published uncertainty method we could implement — including the hydrology SOTA — under a single fair protocol.

🏆 Benchmark — best CRPS on all three catchments, never beaten

Three CAMELS-DE catchments chosen a priori by flashiness (Richards–Baker index) to span the regime range — including the smoothest, the case least favourable to us. Test period 2000–2020, exactly the protocol of the dataset paper (Loritz et al., 2024, ESSD).

Every method gets the same 30 features, the same backbone, the same 5-seed ensembling, the same training budget; unbounded heads are zero-clipped exactly as Klotz et al. (2022) specify; nothing is tuned on the test period.

CRPS (lower is better) — ProbEns is best on every catchment:

Method	Reference	DE210860 (flashiest)	DEA10950	DE212390 (smoothest)
ProbEns (ours)	this work	0.361	0.186	0.121
MC dropout	Gal & Ghahramani 2016	0.365 ⁼	0.246 *	0.172 *
Quantile regression (25q)	pinball loss	0.367 ⁼	0.194 *	0.127 *
CMAL (3× asym. Laplace)	Klotz+ 2022 (SOTA)	0.370 *	0.192 *	0.135 *
DeepAR-style Gamma head	Salinas+ 2020 / Xie+ 2025	0.372 *	0.203 *	0.143 *
Deep ensemble (Gaussian)	Lakshminarayanan+ 2017	0.384 *	0.227 *	0.174 *
MDN / GMM (3× Gaussian)	Bishop 1994 / Klotz+ 2022	0.385 *	0.194 *	0.139 *
NGBoost (LogNormal)	Duan+ 2020 (ICML)	0.420 *	0.193 *	0.138 *
GBM quantile (25q)	tabular baseline	0.438 *	0.209 *	0.132 *
NGBoost (Normal)	Duan+ 2020 (ICML)	0.499 *	0.218 *	0.132 *

* ours wins (paired moving-block bootstrap, 30-day blocks, 2,000 resamples, p<0.05) · ⁼ tie

25 wins · 2 ties · 0 losses over 27 comparisons

Catchment	RB index	ProbEns CRPS	NSE	P(y<0)	Record
DE210860 (flashiest)	0.365	0.361	0.796	0.0000	7 W · 2 T · 0 L
DEA10950	0.302	0.186	0.781	0.0000	9 W · 0 T · 0 L
DE212390 (smoothest)	0.020	0.121	0.747	0.0000	9 W · 0 T · 0 L

It beats CMAL — the hydrology SOTA — on all three (+2.4%, +3.0%, +10.4%; p ≤ 0.006), and never assigns probability to physically impossible negative discharge.

🔬 Why it wins — and what we tried that failed

Ingredient	What it does
Constrained families	Gamma / LogNormal / Weibull — strictly positive and right-skewed, matching the physics of discharge. Gaussian and Laplacian heads (deep ensembles, MC dropout, MDN, CMAL) place probability on negative discharge; ours cannot, by construction (P(y<0) = 0.0000 everywhere).
Regime-gated experts	An input-dependent gate routes each day to a specialist — recession, wet-season body, flood peak — learned from likelihood alone, with mutual-information gating: confident routing, no dead experts, and no assumption that regimes are equally frequent (floods are ~5% of days, so a uniform load-balancing penalty actively hurts).
Warm start + cross-family pooling	Experts are pre-trained on flow bands so the mixture does not collapse to K copies of one model; the final predictive law pools LogNormal, Gamma and Weibull mixtures over 5 seeds.

Negative results we report rather than hide:

• PIT recalibration is not robust. It improved CRPS by 0.023 on DE210860 but made it worse by 0.090 and 0.031 on the other two catchments — the 5-year validation period is too short and too different from the 21-year test period for the map to transfer. It is therefore not part of the reported model (common.recalibrate, kept only to reproduce the finding).
• Heterogeneous mixtures (a different family per expert) did not beat a homogeneous mixture of the right family (p<0.001 the other way).
• NGBoost's natural gradient inside a neural head was consistently worse than plain backprop. Fisher preconditioning compensates for boosted trees' lack of an adaptive step size — Adam already does that job (probens/natural.py).

🚀 Quick start

pip install git+https://github.com/Bluerrror/probens

import torch, torch.nn as nn
from probens import ProbEns, train_model
from probens.uncertainty import build_predictive, decompose_uncertainty

class MLP(nn.Module):
    def __init__(self, input_size, num_params, hidden=192):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(input_size, hidden), nn.ReLU(),
                                 nn.Linear(hidden, num_params))
    def forward(self, x):
        return self.net(x)

model = ProbEns(base_learner_class=MLP, base_learner_params={'input_size': 30},
                dist_name='LogNormal',      # a family that respects the support
                n_estimators=3)
model, val_nll, val_crps = train_model(model, train_loader, val_loader,
                                       dist_name='LogNormal', loss_type='combined')

with torch.no_grad():
    dist = build_predictive('LogNormal', model(x_test))   # full predictive law
    u = decompose_uncertainty('LogNormal', model(x_test))
    u['aleatoric_std'], u['epistemic_std']                # noise vs model doubt

• 18 distribution families — swap 'LogNormal' for 'Gamma', 'Weibull', 'Gumbel', … with one argument.
• Any base learner — MLP, LSTM, CNN, Transformer: anything with forward(x) -> (batch, num_params).
• Aleatoric + epistemic, decomposed by the law of total variance (decompose_uncertainty).

🔁 Reproduce / extend the benchmark

Three catchments (CPU, ~20 min):

python -m venv .venv-boost && .venv-boost/bin/pip install numpy scipy scikit-learn ngboost
.venv-boost/bin/python experiments/env_b_ngboost.py <cache.npz> experiments/envB_samples_DE210860.npz
python experiments/final_benchmark.py DE210860   # then DEA10950, DE212390
python experiments/make_figures.py

Full study — 30 catchments × 2 backbones × 10 methods (Kaggle GPU, ~4–6 h):

experiments/camels_de_30stations_kaggle.ipynb — import it at kaggle.com (Create → Notebook → File → Import notebook), set GPU T4 x2 and Internet ON, Run All. It fetches the 30 catchments straight from Zenodo by HTTP range request, evaluates every method on both an MLP and the study's LSTM (hidden 128, dropout 0.4, 365-day sequences), bootstraps each catchment, and checkpoints after every station so an interrupted session simply resumes.

⚠️ Scope — what is and is not shown

• Results are for three catchments with an MLP backbone on lag features, under the CAMELS-DE protocol. A 30-catchment, LSTM-backbone study is in progress.
• The three stations were selected before any model was run, by a training-period flashiness index, to span the regime range (including the smoothest catchment, which is the least favourable case for positive-support families).
• The two ties are on the flashiest catchment against MC dropout (p=0.39) and quantile regression (p=0.058); ProbEns still has the lowest CRPS there.
• predict(return_quantiles=True) relies on torch.distributions.icdf, implemented only for some families — draw samples and use torch.quantile otherwise.

📄 Citation

This work is described in a manuscript currently under review:

ProbEns: A Constrained Probabilistic Model for Rainfall–Runoff Prediction. (under review, 2026).

If you use ProbEns, the benchmark, or the results in this repository, please cite the paper (BibTeX below) and link back to this repository.

@article{probens2026,
  title   = {ProbEns: A Constrained Probabilistic Model for Rainfall--Runoff Prediction},
  author  = {Bluerrror},
  year    = {2026},
  note    = {Manuscript under review},
  url     = {https://github.com/Bluerrror/probens}
}

🔗 References

CAMELS-DE · Loritz et al. 2024, ESSD · Klotz et al. 2022, HESS · Lakshminarayanan et al. 2017 · Duan et al. 2020, NGBoost · Salinas et al. 2020, DeepAR

---

Released under the MIT License — see LICENSE.