Geo-Aware Models for Stream Temperature Prediction across Different Spatial Regions and Scales

TL;DR

Geo-STARS addresses the generalization gap in environmental stream temperature prediction across spatial regions and scales by introducing a geo-aware embedding $z$ that encodes geographic patterns from meteorological data, static stream characteristics, and scale information. It integrates $z$ into a gated spatio-temporal graph neural network to adapt spatial interactions via adaptive filters and to model temporal dynamics with delayed effects, enabling zero-shot and few-shot transfer to new watersheds. Across 37 years of data from eastern US watersheds at multiple scales, Geo-STARS demonstrates superior cross-task generalization and robustness to missing characteristics, outperforming baselines. The approach offers a data-efficient, scalable framework for environmental monitoring and decision-making in watershed management.

Abstract

Understanding environmental ecosystems is vital for the sustainable management of our planet. However,existing physics-based and data-driven models often fail to generalize to varying spatial regions and scales due to the inherent data heterogeneity presented in real environmental ecosystems. This generalization issue is further exacerbated by the limited observation samples available for model training. To address these issues, we propose Geo-STARS, a geo-aware spatio-temporal modeling framework for predicting stream water temperature across different watersheds and spatial scales. The major innovation of Geo-STARS is the introduction of geo-aware embedding, which leverages geographic information to explicitly capture shared principles and patterns across spatial regions and scales. We further integrate the geo-aware embedding into a gated spatio-temporal graph neural network. This design enables the model to learn complex spatial and temporal patterns guided by geographic and hydrological context, even with sparse or no observational data. We evaluate Geo-STARS's efficacy in predicting stream water temperature, which is a master factor for water quality. Using real-world datasets spanning 37 years across multiple watersheds along the eastern coast of the United States, Geo-STARS demonstrates its superior generalization performance across both regions and scales, outperforming state-of-the-art baselines. These results highlight the promise of Geo-STARS for scalable, data-efficient environmental monitoring and decision-making.

Figures (8)

• Figure 1: Overview of geographic consistency across different watersheds and at various spatial scales. Different spatial regions on the regional side and both coarse and fine scales on the scale side have shared principles and patterns.
• Figure 2: Overview of Geo-STARS framework. The architecture consists of two main components: (a) Generation of Geo-Aware Embedding and (b) Gated Spatio-temporal Graph Neural Network. In (a), input meteorological features $x$, stream characteristics $c$, and stream adjacency matrix $\mathbf{A}$ are processed respectively to produce a geo-aware embedding $z$, which captures the shared underlying geographic principles of stream temperature by embedding overall information from these inputs for each task. In (b), the embedding $z$ and stream characteristics are used to compute adaptive influence filters that modulate the adjacency matrix at each time step. These filters enable intelligent information aggregation from current and past neighboring segments, allowing the model to capture non-uniform spatial dependencies and temporal lags in water temperature dynamics.
• Figure 3: Map of four distinct and ecologically varied watersheds in the eastern United States.
• Figure 4: Time-series comparison of stream water temperature predictions by Geo-STARS, EA-LSTM and RGRN.
• Figure 5: Comparison of average RMSE for stream water temperature prediction in zero-shot setting.