Distributed Hydrological Modeling in the Feature Space
Organizations: Max Planck Institute for Biogeochemistry, Germany · European Centre for Medium-Range Weather Forecasts (ECMWF), United Kingdom · Institute of Bio- and Geosciences Agrosphere (IBG-3), Research Centre Jülich, Germany · University of Bonn, Germany · Lamarr Institute for Machine Learning and Artificial Intelligence, Germany
Abstract
Accurate forecasting of river discharge and floods is very challenging. River dynamics are affected by storage, meteorological forcing, and flow propagation at different spatial and temporal scales. Forecasting thus requires a framework that considers the upstream-to-downstream flow through river networks across grid cells and catchments. This modeling is known in hydrology as distributed modeling and routing. Existing deep learning approaches either ignore this topology, operate on lumped catchments, or route predicted physical quantities through a separate graph or physical routing model. We instead introduce feature-space routing: a topology-aware state-space operator embedded directly in the forecasting dynamics. At every forecast step, the operator gathers latent states from upstream grid cells and causally updates the downstream state according to the known river network. This preserves the physical connectivity of the river system while allowing the propagated state itself to be learned end-to-end and allows the model to predict river discharge considering both local dynamics and neighboring upstream contributions. To address uncertainty and provide probabilistic forecasts, we minimize the fair continuous ranked probability score (fCRPS) as a training objective. Our experiments on the European Flood Awareness System (EFAS) and observational data for river discharge forecasting demonstrate that encoding the physical structure of river networks explicitly in the feature space substantially improves the forecasting skill, particularly in an ungauged setting. Our approach achieves state-of-the-art results on both reanalysis and observational data and is able to forecast maps of river discharge at 1 arcminute and 6-hourly resolution up to 10 days lead time.
Figures & tables
| Validation (2022-2023) | Test (2024-2025) | ||||||||
|---|---|---|---|---|---|---|---|---|---|
| Model | R2 r ( ) | KGE r ( ) | CRPS r ( ) | R2 r ( ) | KGE r ( ) | CRPS r ( ) | |||
| Probabilistic | DRUM (mean) | 0.1496 | 0.1505 | 2.0872 | 0.0984 | 0.1352 | 2.3581 | ||
| ED-LSTM (mean) | 0.1726 | 0.2431 | 1.9121 | 0.1347 | 0.2333 | 2.1971 | |||
| ME-LSTM (mean) | 0.2093 | 0.3071 | 1.8395 | 0.1210 | 0.2640 | 2.2249 | |||
| Ours (mean) | 0.4515 | 0.5751 | 1.4720 | 0.3907 | 0.5474 | 1.7619 | |||
| (a) Routing and Warmup | (b) Routing Module | (c) Upstream Aggregation | ||||||
| Routing | Warmup | KGE r | KGE o | F1 r ( ) | Module | KGE r | KGE o | F1 r ( ) | Type | KGE r | KGE o | F1 r ( ) | ||
| ✗ | ✗ | 0.666 | 0.692 | 0.096 | GS-mean | 0.680 | 0.721 | 0.310 | Summation | 0.738 | 0.682 | 0.335 | ||
| ✓ | ✗ | 0.672 | 0.739 | 0.114 | GS-sum | 0.702 | 0.702 | 0.285 | Mean | 0.730 | 0.739 | 0.332 | ||
| ✗ | ✓ | 0.687 | 0.721 | 0.162 | GAT | 0.714 | 0.633 | 0.268 | GraphMamba | 0.722 | 0.720 | 0.276 | ||
| ✓ | ✓ | 0.740 | 0.747 | 0.339 | SSM | 0.740 | 0.747 | 0.339 | Ours | 0.740 | 0.747 | 0.339 | ||
| Val | Test | |
|---|---|---|
| RiverMamba | 0.141 0.01 | 0.138 0.01 |
| Ours | 0.334 0.02 | 0.314 0.01 |
| Validation (2022-2023) | Test (Jun 2024- Dec 2025) | ||||||||
| Model | R2 o ( ) | KGE o ( ) | CRPS o ( ) | R2 o ( ) | KGE o ( ) | CRPS o ( ) | |||
| EFAS (control run) | - | - | - | 0.0410 | 0.3964 | - | |||
| Probabilistic | DRUM (mean) | 0.1750 | 0.2834 | 4.5375 | -0.0089 | 0.2355 | 4.3474 | ||
| ED-LSTM (mean) | 0.2841 | 0.3957 | 3.9024 | 0.1214 | 0.3788 | 3.6818 | |||
| ME-LSTM (mean) | 0.3339 | 0.4771 | 3.4928 | 0.0913 | 0.3843 | 3.7962 | |||
| Ours (mean) | 0.6042 | 0.6689 | 2.4592 | 0.3693 | 0.5520 | 2.9708 | |||
Appendix figures & tables109 assets
Supplementary material from the paper’s appendix.
Appendix
| Results on EFAS Reanalysis | ||||||
| Static | CPC | IMERG | EMO | Analysis | IFS-HRES | R2 | KGE | R2 | KGE | F1 r ( ) |
| ✗ | ✓ | ✓ | ✓ | ✓ | ✓ | 0.6963 | 0.7338 | 0.6919 | 0.7443 | 0.2734 |
| ✓ | ✗ | ✓ | ✓ | ✓ | ✓ | 0.6741 | 0.6972 | 0.6674 | 0.7145 | 0.3053 |
| ✓ | ✓ | ✗ | ✓ | ✓ | ✓ | 0.7021 | 0.7303 | 0.7041 | 0.7411 | 0.3100 |
| ✓ | ✓ | ✓ | ✗ | ✓ | ✓ | 0.6863 | 0.7474 | 0.6637 | 0.7116 | 0.3101 |
| ✓ | ✓ | ✓ | ✓ | ✗ | ✓ | 0.6869 | 0.6979 | 0.6835 | 0.7107 | 0.3069 |
| Results on EFAS Reanalysis | ||
|---|---|---|
| LISFLOOD Static | AlphaEarth | R2 | KGE | R2 | KGE | F1 r ( ) |
| ✗ | ✗ | 0.6963 | 0.7338 | 0.6919 | 0.7443 | 0.2734 |
| ✓ | ✗ | 0.7045 | 0.7404 | 0.6943 | 0.7462 | 0.3360 |
| ✗ | ✓ | 0.7061 | 0.7396 | 0.7043 | 0.7493 | 0.3394 |
| Results on CAMELS observational data | ||
| LISFLOOD Static | AlphaEarth | R2 | KGE | F1 ( ) |
| EFAS Reanalysis | CAMELS Observations | |
| Bi-SSD | R2 | KGE | R2 | KGE | F1 r ( ) | R2 | KGE | F1 ( ) |
| ✗ | 0.6715 | 0.7057 | 0.6307 | 0.6915 | 0.3225 | 0.6793 | 0.7324 | 0.2947 |
| ✓ | 0.7061 | 0.7396 | 0.7043 | 0.7493 | 0.3394 | 0.6990 | 0.7472 | 0.3109 |
| Results on EFAS Reanalysis | |
|---|---|
| # Routing modules | R2 | KGE | R2 | KGE | F1 r ( ) |
| 0.7061 | 0.7396 | 0.7043 | 0.7493 | 0.3394 | |
| 0.6547 | 0.7057 | 0.6289 | 0.6650 | 0.2839 | |
| EFAS Reanalysis | CAMELS Observations | |
| Type | R2 | KGE | R2 | KGE | F1 r ( ) | R2 | KGE | F1 ( ) |
| Connection per Neighbor | 0.7027 | 0.7264 | 0.7064 | 0.7456 | 0.2950 | 0.6936 | 0.7318 | 0.2461 |
| Binary Mask | 0.6622 | 0.6816 | 0.6404 | 0.6751 | 0.1122 | 0.6245 | 0.7158 | 0.2954 |
| All Sort(area) | 0.7061 | 0.7396 | 0.7043 | 0.7493 | 0.3394 | 0.6990 | 0.7472 | 0.3109 |
| Results on EFAS Reanalysis | ||
|---|---|---|
| Bi-directionality | Training time* ( ) | Inference time* ( ) |
| 1 epoch 16 steps | 280,820 points | |
| Mamba (RiverMamba) | 159 min | 6.20 0.005 sec |
| Mamba2 (Ours) | 49 min | 4.10 0.008 sec |
| ∗ tested on NVIDIA A40 GPU. | ||
| Overall | Lead Time | ||||
|---|---|---|---|---|---|
| Source | Mean | ||||
| EFAS Reanalysis (GT) | 99.97% | - | - | - | - |
| ✗ routing | 99.09% | 99.19% | 99.34% | 99.37% | 99.35% |
| ✓ routing | 99.62% | 99.73% | 99.87% | 99.89% | 99.89% |
| Benefit of adding routing | +0.53% | +0.54% | +0.53% | +0.52% | +0.54% |
| Class (Uparea km 2 ) | MAE ( ) | RMSE ( ) | R2 ( ) | KGE ( ) | F1@(1.5-500) ( ) |
|---|---|---|---|---|---|
| 1: 1.95-3.04 | 0.0143 0.0078 | 0.0213 0.0133 | 0.016 0.417 | 0.236 0.532 | 0.1218 0.1279 |
| 2: 3.04-5.28 | 0.0180 0.0144 | 0.0306 0.0244 | 0.287 0.447 | 0.429 0.554 | 0.0687 0.1249 |
| 3: 5.28-12.88 | 0.0349 0.0257 | 0.0589 0.0437 | 0.243 0.476 | 0.399 0.581 | 0.0471 0.1215 |
| 4: 12.88-84.87 | 0.0960 0.0860 | 0.1681 0.1542 | 0.422 0.521 | 0.540 0.626 | 0.0578 0.1071 |
| 5: 84.87-559.26 | 0.5460 0.5327 | 1.0069 0.9718 | 0.559 0.588 | 0.630 0.671 | 0.0503 0.0851 |
| 6: 559.26-160002 | 7.9160 7.9300 | 14.408 14.003 | 0.557 0.617 | 0.641 0.693 | 0.0410 0.0840 |
| Class (Uparea km 2 ) | MAE ( ) | RMSE ( ) | R2 ( ) | KGE ( ) | F1@(1.5-500) ( ) |
|---|---|---|---|---|---|
| 1: 1.95-3.04 | 0.0346 0.0126 | 0.0536 0.0183 | -6.758 0.089 | -1.159 0.237 | 0.0539 0.0543 |
| 2: 3.04-5.28 | 0.0502 0.0228 | 0.0835 0.0338 | -2.085 0.105 | -0.349 0.268 | 0.0391 0.0519 |
| 3: 5.28-12.88 | 0.0843 0.0409 | 0.1348 0.0616 | -1.034 0.109 | -0.098 0.291 | 0.0334 0.0531 |
| 4: 12.88-84.87 | 0.2624 0.1469 | 0.4032 0.2264 | -0.633 0.119 | -0.031 0.300 | 0.0288 0.0454 |
| 5: 84.87-559.26 | 1.4830 0.9720 | 2.1890 1.4710 | -0.602 0.143 | -0.044 0.319 | 0.0266 0.0390 |
| 6: 559.26-160002 | 20.564 14.658 | 28.943 22.410 | -0.918 -0.144 | -0.036 0.239 | 0.0186 0.0223 |
| KGE ( ) | R2 ( ) | F1@(1.5-500) ( ) | |
|---|---|---|---|
| mean | 0.563 | 0.316 | 0.185 |
| median | 0.602 | 0.581 | 0.159 |
| KGE ( ) | R2 ( ) | F1@(1.5-500) ( ) | |
|---|---|---|---|
| mean | 0.542 | 0.202 | 0.200 |
| median | 0.595 | 0.545 | 0.175 |
| Model | # Param. | auto-regressive | Probabilistic | Nowcasting |
|---|---|---|---|---|
| RiverMamba | 4.4 million | ✗ | ✗ | ✗ |
| Ours | 2.3 million | ✓ | ✓ | ✓ |
| Validation (2022-2023) | Test (2024-2025) | ||||||||
|---|---|---|---|---|---|---|---|---|---|
| Model | R2 ( ) | KGE ( ) | F1@(1.5-20) ( ) | R2 ( ) | KGE ( ) | F1@(1.5-20) ( ) | |||
| Determ | RiverMamba | 0.7075 0.0138 | 0.7255 0.0077 | 0.1242 0.0160 | 0.5806 0.0221 | 0.6580 0.0118 | 0.0611 0.0104 | ||
| Ours | 0.7075 0.0056 | 0.7410 0.0134 | 0.1627 0.0143 | 0.5676 0.0097 | 0.6589 0.0162 | 0.0841 0.0037 | |||
| Prob | Ours (mean) | 0.7247 | 0.7695 | 0.2129 | 0.6280 | 0.7284 | 0.2266 | ||
| Validation (2022-2023) | Test (2024-2025) | ||||||||
|---|---|---|---|---|---|---|---|---|---|
| Model | R2 ( ) | KGE ( ) | F1@(1.5-20) ( ) | R2 ( ) | KGE ( ) | F1@(1.5-20) ( ) | |||
| Determ | RiverMamba | 0.7326 0.0099 | 0.7579 0.0052 | 0.1106 0.0125 | 0.6161 0.0222 | 0.6927 0.0141 | 0.0990 0.0151 | ||
| Ours | 0.6957 0.0273 | 0.7507 0.0273 | 0.1326 0.0272 | 0.5631 0.0237 | 0.6645 0.0241 | 0.1132 0.0136 | |||
| Prob | Ours (mean) | 0.7359 | 0.7876 | 0.1998 | 0.6505 | 0.7552 | 0.2125 | ||
| Validation (2022-2023) | Test (2024-2025) | ||||||
|---|---|---|---|---|---|---|---|
| Model | F1@(1.5-20) ( ) | F1@(1.5-500) ( ) | F1@(1.5-20) ( ) | F1@(1.5-500) ( ) | |||
| Determ | RiverMamba | 0.1670 0.0122 | 0.1407 0.0099 | 0.1713 0.0134 | 0.1379 0.0108 | ||
| Ours | 0.3542 0.0199 | 0.3337 0.0176 | 0.3183 0.0191 | 0.3143 0.0129 | |||
| Validation (2022-2023) | Test (2024) | ||||||||
|---|---|---|---|---|---|---|---|---|---|
| Model | R2 ( ) | KGE ( ) | F1@(1.5-20) ( ) | R2 ( ) | KGE ( ) | F1@(1.5-20) ( ) | |||
| Determ | RiverMamba | 0.7236 0.0000 | 0.7563 0.0113 | 0.3181 0.0041 | 0.5428 0.0151 | 0.6257 0.0235 | 0.2864 0.0675 | ||
| Ours | 0.7286 0.0011 | 0.7716 0.0065 | 0.2883 0.0008 | 0.6216 0.0313 | 0.6761 0.0256 | 0.2805 0.0107 | |||
| Prob | Ours (mean) | 0.7175 | 0.7744 | 0.3151 | 0.6134 | 0.6941 | 0.3580 | ||
| Test set (2024-05-30 to 2025-12-21) | |||||||
|---|---|---|---|---|---|---|---|
| Model | RMSE ( ) | R2 ( ) | KGE ( ) | F1@(1.5-20) ( ) | F1@(1.5-500) ( ) | ||
| Determ | EFAS (control run) | 9.8394 | 0.0481 | 0.3967 | 0.0945 | 0.0525 | |
| Ours | 7.2889 0.1628 | 0.3413 0.0135 | 0.5180 0.0073 | 0.0703 0.0159 | 0.0512 0.0175 | ||
| Prob | Ours (mean) | 7.2860 | 0.3673 | 0.5515 | 0.1015 | 0.1515 | |
| Test set (2024-06-01 to 2024-08-31) | ||||||||
|---|---|---|---|---|---|---|---|---|
| Model | RMSE ( ) | R2 ( ) | KGE ( ) | CRPS ( ) | CRPS mae ( ) | CRPS spread ( ) | ||
| Prob | EFAS (mean of 51 ens members) | 6.8122 | -0.8517 | 0.1667 | 3.8862 | 5.4043 | 1.5180 | |
| Ours (mean of 12 ens members) | 4.5627 | 0.0519 | 0.3718 | 2.1599 | 4.1280 | 1.9682 | ||
| Validation (2022-2023) | Test (2024-2025) | ||||||||
| Model | RMSE ( ) | R2 ( ) | KGE ( ) | RMSE ( ) | R2 ( ) | KGE ( ) | |||
| Deterministic | Persistence | 4.949 | 0.086 | 0.533 | 6.013 | -0.118 | 0.456 | ||
| ED-LSTM | 4.622 0.016 | 0.229 0.007 | 0.300 0.009 | 5.316 0.010 | 0.157 0.007 | 0.264 0.013 | |||
| ME-LSTM | 4.830 0.034 | 0.181 0.011 | 0.286 0.018 | 5.584 0.042 | 0.095 0.014 | 0.235 0.020 | |||
| Ours | 3.972 0.052 | 0.428 0.015 | 0.534 0.011 | 4.747 0.040 | 0.333 0.011 | 0.481 0.008 | |||
| Probabilistic | Clima (median) | 5.382 | -0.005 | -0.035 | 6.025 | -0.068 | -0.074 | ||
| Validation (2022-2023) | Test (2024-2025) | ||||||||
| Model | RMSE ( ) | R2 ( ) | KGE ( ) | RMSE ( ) | R2 ( ) | KGE ( ) | |||
| Deterministic | Persistence | 10.417 | 0.096 | 0.542 | 12.056 | -0.076 | 0.468 | ||
| ED-LSTM | 9.739 0.065 | 0.239 0.009 | 0.318 0.009 | 10.975 0.005 | 0.151 0.004 | 0.269 0.013 | |||
| ME-LSTM | 10.242 0.067 | 0.186 0.010 | 0.321 0.020 | 11.515 0.062 | 0.085 0.013 | 0.249 0.021 | |||
| Ours | 8.494 0.138 | 0.420 0.019 | 0.531 0.014 | 9.851 0.074 | 0.314 0.010 | 0.474 0.006 | |||
| Probabilistic | Clima (median) | 11.014 | 0.002 | -0.019 | 12.169 | -0.062 | -0.058 | ||
| Results on the Validation set (2022-2023) | ||||||
| Clima | DRUM | ED-LSTM | ME-LSTM | Ours | ||
| CRPS ( ) | 4.3328 | 4.5104 | 4.1379 | 3.9776 | 3.1837 | |
| Results on the Test set (2024-2025) | ||||||
| Clima | DRUM | ED-LSTM | ME-LSTM | Ours | ||
| CRPS ( ) | 4.6667 | 4.9376 | 4.6384 | 4.6492 | 3.7604 | |
| Validation (2022-2023) | Test (2024-2025) | ||||||||
| Model | RMSE ( ) | R2 ( ) | KGE ( ) | RMSE ( ) | R2 ( ) | KGE ( ) | |||
| Deterministic | ED-LSTM | 6.0542 0.0082 | 0.2705 0.0012 | 0.3491 0.0214 | 11.3616 0.1632 | 0.2366 0.0164 | 0.4501 0.0097 | ||
| ME-LSTM | 5.8810 0.0206 | 0.2993 0.0044 | 0.3550 0.0083 | 11,4754 0.1683 | 0.2384 0.0246 | 0.4330 0.0138 | |||
| Ours | 4.6795 0.0681 | 0.5603 0.0095 | 0.5890 0.0038 | 9.0163 0.0952 | 0.4576 0.0066 | 0.5839 0.0076 | |||
| Probabilistic | DRUM (mean) | 7.0236 | 0.1514 | 0.2177 | 12.6358 | 0.1153 | 0.3059 | ||
| ED-LSTM (mean) | 6.4217 | 0.2004 | 0.2780 | 11.5629 | 0.2308 | 0.4266 | |||
| Results on the Validation set (2022-2023) | |||||
| DRUM | ED-LSTM | ME-LSTM | Ours | ||
| CRPS ( ) | 4.5375 | 3.9024 | 3.4928 | 2.4592 | |
| Results on the Test set (2024-2025) | |||||
| DRUM | ED-LSTM | ME-LSTM | Ours | ||
| CRPS ( ) | 5.6934 | 5.0091 | 5.0361 | 3.5194 | |
| Validation (2022-2023) | Test (2024) | ||||||||
| Model | RMSE ( ) | R2 ( ) | KGE ( ) | RMSE ( ) | R2 ( ) | KGE ( ) | |||
| Deterministic | ED-LSTM | 6.0542 0.0082 | 0.2705 0.0012 | 0.3401 0.0214 | 6.8658 0.1326 | 0.0544 0.0072 | 0.0906 0.0136 | ||
| ME-LSTM | 5.8810 0.0206 | 0.2993 0.0044 | 0.3550 0.0083 | 6.9621 0.0874 | 0.0500 0.0100 | 0.1186 0.0123 | |||
| Ours | 4.6795 0.0681 | 0.5603 0.0095 | 0.5890 0.0038 | 5.8406 0.1311 | 0.3485 0.0202 | 0.4162 0.0284 | |||
| Probabilistic | DRUM (mean) | 7.0236 | 0.1514 | 0.2177 | 7.3433 | 0.0115 | 0.0576 | ||
| ED-LSTM (mean) | 6.4217 | 0.2004 | 0.2780 | 7.0896 | 0.0452 | 0.0831 | |||
| Results on the Validation set (2022-2023) | |||||
| DRUM | ED-LSTM | ME-LSTM | Ours | ||
| CRPS ( ) | 2.7239 | 2.3110 | 2.1237 | 1.4840 | |
| Results on the Test set (2024) | |||||
| DRUM | ED-LSTM | ME-LSTM | Ours | ||
| CRPS ( ) | 2.3685 | 2.2683 | 2.2683 | 1.5678 | |
| Configuration | Pre-training | Fine-tuning |
| (EFAS Reanalysis) | (GRDC/CAMELS) | |
| Optimizer | AdamW | AdamW |
| Learning rate | 0.0003 | 0.0001 |
| Minimum learning rate | 0.0001 | 0.0001 |
| Batch size ( ) | 1 | 1 |
| Learning rate scheduler | Cosine annealing | Cosine annealing |
| 10K | 20K | 40K | 80K | 160K | 300K | 600K |
| 0.0154 | 0.301 | 0.591 | 1.172 | 2.347 | 4.395 | 8.913 |
| Dataset | Source | Country | Start | End | # Stations |
|---|---|---|---|---|---|
| CAMELS-CH | Höge et al. [2023] | Switzerland | 1981 | 2020 | 136 |
| CAMELS-CZ | Institute et al. [2025] | Czechia | 1978 | 2024 | 173 |
| CAMELS-DE | Loritz et al. [2024] | Germany | 1951 | 2020 | 605 |
| CAMELS-DK | Liu et al. [2025] | Denmark | 1989 | 2019 | 54 |
| CAMELS-ES | Casado-Rodríguez et al. [2026] | Spain | 1991 | 2020 | 244 |
| CAMELS-FI | Seppä et al. [2026] | Finland | 1961 | 2023 | 234 |
| Variable | Long name | Unit | Height | Surface parameters |
|---|---|---|---|---|
| d2m | 2m dewpoint temperature | K | 2m | instantaneous |
| e | total evaporation | m of water equivalent | surface | accumulated |
| ro | runoff | m | surface & subsurface | accumulated |
| sd | snow depth | m of water equivalent | surface | accumulated |
| sf | snowfall | m of water equivalent | surface | accumulated |
| smlt | snowmelt | m of water equivalent | surface | accumulated |
| Variable | Long name | Unit | Height | Surface parameters |
|---|---|---|---|---|
| pd | water vapour pressure | hPa | - | instantaneous |
| pr | total precipitation | mm | surface | accumulated |
| rg | downward surface solar radiation | J/m 2 | surface | accumulated |
| tn | minimum air temperature | C ∘ | 2m | instantaneous |
| tx | maximum air temperature | C ∘ | 2m | instantaneous |
| ws | wind speed | m/s | 10m | instantaneous |