CauSTream: Causal Spatio-Temporal Representation Learning for Streamflow Forecasting

TL;DR

CauSTream is a novel causal spatiotemporal learning framework for streamflow forecasting that jointly learns dynamic causal graphs, improving accuracy and interpretability over existing methods across multiple river basins.

Contribution

It introduces a unified model that learns adaptive causal structures in spatiotemporal data, with proven identifiability and superior forecasting performance.

Findings

Outperforms state-of-the-art models on U.S. river basins

Improves accuracy at longer forecast horizons

Produces interpretable causal graphs aligned with domain knowledge

Abstract

Streamflow forecasting is crucial for water resource management and risk mitigation. While deep learning models have achieved strong predictive performance, they often overlook underlying physical processes, limiting interpretability and generalization. Recent causal learning approaches address these issues by integrating domain knowledge, yet they typically rely on fixed causal graphs that fail to adapt to data. We propose CauStream, a unified framework for causal spatiotemporal streamflow forecasting. CauSTream jointly learns (i) a runoff causal graph among meteorological forcings and (ii) a routing graph capturing dynamic dependencies across stations. We further establish identifiability conditions for these causal structures under a nonparametric setting. We evaluate CauSTream on three major U.S. river basins across three forecasting horizons. The model consistently outperforms prior state-of-the-art methods, with performance gaps widening at longer forecast windows, indicating stronger generalization to unseen conditions. Beyond forecasting, CauSTream also learns causal graphs that capture relationships among hydrological factors and stations. The inferred structures align closely with established domain knowledge, offering interpretable insights into watershed dynamics. CauSTream offers a principled foundation for causal spatiotemporal modeling, with the potential to extend to a wide range of scientific and environmental applications.