Adaptive Graph Learning with Transformer for Multi-Reservoir Inflow Prediction

TL;DR

AdaTrip tackles multi-reservoir inflow forecasting by learning adaptive, time-varying graphs and integrating spatial and temporal modeling with a Transformer-based encoder–decoder. The framework combines a feature extractor, adaptive graph learning via a Graph Attention Network, a temporal encoder–decoder, and a semi-supervised pretraining scheme to initialize representations. It demonstrates NSE improvements on 30 reservoirs in the Upper Colorado River Basin and provides interpretable edge-level and time-step attention maps that align with hydrological processes. The approach enhances robustness to topology changes and data sparsity, offering practical benefits for reservoir operation and planning.

Abstract

Reservoir inflow prediction is crucial for water resource management, yet existing approaches mainly focus on single-reservoir models that ignore spatial dependencies among interconnected reservoirs. We introduce AdaTrip as an adaptive, time-varying graph learning framework for multi-reservoir inflow forecasting. AdaTrip constructs dynamic graphs where reservoirs are nodes with directed edges reflecting hydrological connections, employing attention mechanisms to automatically identify crucial spatial and temporal dependencies. Evaluation on thirty reservoirs in the Upper Colorado River Basin demonstrates superiority over existing baselines, with improved performance for reservoirs with limited records through parameter sharing. Additionally, AdaTrip provides interpretable attention maps at edge and time-step levels, offering insights into hydrological controls to support operational decision-making. Our code is available at https://github.com/humphreyhuu/AdaTrip.

Figures (6)

• Figure 1: Overall Framework of AdaTrip. Starting with historical reservoir inflow observations $\{\mathbf{x}_{t}\}_{t=1}^{T}$, precipitation, and temperature, a feature extractor $\mathcal{F}_{\theta}$ maps them to embeddings $\mathbf{h}_{t}=\operatorname{MLP}(\mathbf{x}_{t})$; based on pairwise reservoir distances $d_{ij}$, an initial static graph $\mathcal{G}_{0}\!=\!(\mathcal{V},\mathcal{E},\mathbf{A})$ with adjacency $a_{ij}$ is formed; a graph attention network $\mathcal{G}_{\lambda}$ assigns time‑varying edge weights and removes edges with $\alpha_{ij,t}<\tau$ to obtain the adaptive graph $\mathcal{G}_{t}$; an encoder–decoder transformer $(\mathcal{T}_{e},\mathcal{T}_{d})$ summarizes $\mathbf{M}_{i}$ into latent states $\mathbf{z}_{i}$, and a linear layer yields $k\!\in\![1,7]$‑day ahead inflow predictions $\hat{y}_{i,t+k}$ for every reservoir.
• Figure 2: Comparison of the prediction performance between the ED--LSTM, Transformer, GCN+LSTM, and AdaTrip models for the 30-day forecasting horizons using the NSE metric (Very good: $\text{NSE} > 0.75$; Good: $0.65 < \text{NSE} \le 0.75$; Satisfactory: $0.5 < \text{NSE} \le 0.65$; Acceptable: $0.4 < \text{NSE} \le 0.5$; and Unsatisfactory: $\text{NSE} \le 0.4$). Results are shown for all three baselines and our AdaTrip across the seven-day forecast horizon.
• Figure 3: NSE scores and Mean Square Errors are shown across different graph constructions. We can observe that GCN+LSTM and AdaTrip without adaptation are more sensitive to the topology structure, which demonstrates the problem 1) in section \ref{['sec:intro']} also exists in the hydrology domain.
• Figure 4: The demonstration of the graph adaptation on day $t=30$ after the $4$-th epoch ends in AdaTrip. Self-loop edges are also involved in $G_{30}$.
• Figure 5: Reservoir inflow observations in reservoirs ECH, ECR, and LCR.