Source-linked AI summary
A Dataset and Model for Imputing Water Surface Elevation on a Large and Extremely Sparse Spatiotemporal Graph
Ruben Cartuyvels, Karim Douch, Gabriele Bertoli, Mounia El Baz, Artemis Vrettou, Sébastien Lefèvre, Diego Fernandez Prieto
TL;DR
AmazonSWE addresses spatiotemporal water-surface-elevation imputation on a 19K-node Amazon river network where fewer than 1% of nodes are observed daily. It introduces a topology-aware bidirectional selective state space model using connected subgraphs and flattened space-time token sequences, outperforming existing methods and a state-of-the-art SWOT densification baseline while providing denser reconstructions.
Problem
AmazonSWE targets spatiotemporal graph imputation under a 19K-node river-network topology with fewer than 1% of nodes observed per day, a setting existing methods are not adapted to.
Method
The paper introduces a bidirectional selective state space model that samples connected subgraphs and flattens spatial and temporal observations into topology-aware token sequences.
Results
The model outperforms prior spatiotemporal graph imputation methods and a state-of-the-art SWOT-based densification baseline, producing denser reconstructions.
Takeaways & Limitations
AmazonSWE provides a benchmark for evaluating imputation methods on extremely sparse, large-scale, directed acyclic river graphs.
Takeaways & Limitations
In situ gauge data from 2022-07 to 2023-07 are reserved for evaluation, so the model does not use those observations during training.
Abstract
from arXiv · showhide
Continuous monitoring of water surface elevation across river networks is critical for flood forecasting, water resource management, and understanding the global water cycle. Yet, the scarcity of in situ gauges across much of the globe constrains the development of reliable modeling frameworks. Satellite altimetry has the potential to alleviate this problem but its use is currently hindered by sparse temporal coverage. To this end, we introduce AmazonSWE, a dataset for training and evaluating large-scale spatiotemporal graph imputation methods that integrates processed satellite altimetry measurements from a range of sources, including the recent wide-swath SWOT sensor. The dataset covers over 19K river sections and 10 years (2016-2026) in the Amazon river basin, with in situ gauges held out for evaluation. Besides contributing a novel real-world use case with the potential for societal impact, AmazonSWE introduces significant technical challenges: with fewer than 1% of sections observed per day, the dataset is far sparser than existing imputation benchmarks, and its directed acyclic river topology is both structurally different from and larger than graphs in existing datasets. We show that prior spatiotemporal graph imputation methods are not adapted to this topology, scale and sparsity, and propose a simple bidirectional selective state space model that outperforms them by sampling connected subgraphs and flattening space and time into a single token sequence with topology-aware positional encodings. Compared to the state-of-the-art published method for SWOT-based WSE densification, which integrates statistics with physical modeling, our model reduces RMSE against in situ gauges by 18-39%, while producing predictions for every river section rather than only those with sufficient nearby satellite coverage.
Introduction
AmazonSWE addresses the need for daily water-surface-elevation imputation across a very large, extremely sparse Amazon river graph. It combines heterogeneous satellite and gauge observations and introduces a sequence model tailored to river topology and sparsity.
- Motivation: Satellite altimetry improves spatial coverage but remains temporally sparse, while in situ gauges provide denser measurements at only a small, unevenly covered fraction of the network.SWOT adds wide-swath coverage with a 21-day repeat cycle but has a short temporal record.
- Motivation: Fewer than 1% of over 19K river nodes are observed on any day, creating a scale and sparsity challenge for existing graph-imputation methods.Full-graph STGNNs are resource intensive, while same-day neighborhoods contain too few observations for effective message passing.
- Dataset challenge: The dataset integrates observations from multiple sources whose measurements can differ when they observe the same location on the same day.Figure 1 illustrates locations observed by zero, one, or multiple sensors and the differing reference levels used by satellite and in situ measurements.
- Approach: The proposed model flattens space and time into one token sequence, using connected subgraph sampling and topology-aware encodings rather than explicit graph convolutions.Its inductive bias follows upstream-to-downstream and temporal order in river systems.
- Contributions: The model outperforms interpolation and existing transductive and inductive neural methods while improving coverage and accuracy over Reach-Reg.The paper presents AmazonWSE as a benchmark for daily imputation under extreme sparsity.
Related Work
Prior datasets and methods generally provide denser observations or smaller and structurally different graphs than AmazonWSE. Related hydrology work also offers either topology-free aggregates, gauge-only data, or mixed evidence about the value of river topology.
- Satellite Altimetry: Classical altimetry supplies long-running point-crossing series, whereas SWOT observes contiguous river segments and enables a different densification setting.Reach-Reg exploits SWOT’s measurement geometry for spatiotemporal WSE densification.
- Spatiotemporal Imputation: Spatiotemporal imputation methods include temporal models, transductive methods that see target nodes during training, and inductive methods that transfer to unseen nodes.The paper evaluates this methodological landscape on the larger, sparser AmazonWSE graph.
- Existing Datasets: Existing imputation benchmarks are substantially smaller or denser than AmazonWSE, limiting comparability under its extreme sparsity.Traffic, air-quality, and energy datasets range from tens to thousands of nodes, while many use artificial sparsification scenarios.
- Existing Datasets: Hydrology datasets often erase river topology through catchment aggregation or provide only in situ gauge data without satellite altimetry.AmazonWSE combines satellite and gauge-related data on an explicit reach graph.
- Rivers and GNNs: Prior findings report limited benefit from adjacency-based GNN topology for discharge forecasting, whereas this paper finds topology improves WSE prediction when encoded through its model design.The comparison motivates testing topology-aware alternatives to explicit graph convolutions.
Dataset
AmazonWSE is a multi-source daily WSE dataset built on the 19,172-reach SWORD graph of the Amazon basin. It aligns and quality-filters SWOT, HydroWeb, ICESat-2, and ANA observations while reserving ANA gauges for evaluation.
- Graph structure: The dataset uses SWORD’s 19,172 connected river reaches as nodes in an acyclic directed graph that is approximately, but not strictly, a tree.Edges point from an immediately upstream reach to its downstream neighbor, and some reaches have multiple parents.
- Data sources: SWOT contributes wide-swath observations mapped to SWORD reaches, with approximately 17K reaches observed and 10K retained after quality filtering.The SWOT science-cycle data begin in July 2023 and have a 21-day repeat cycle, with some reaches observed multiple times per cycle.
- Data sources: HydroWeb provides 3.7K retained virtual stations from altimetry missions spanning 2016–2026, with measurements every 10–35 days.The original database contains 7.3K Amazon crossings, of which locations matching nearby SWORD reaches are retained.
- Data sources: The dataset includes 375 quality-filtered ANA gauges and ICESat-2 transects spanning October 2018–March 2026, with ICESat-2 series remaining for over 90% of SWORD reaches.ANA gauges provide dense daily observations during operation and serve as ground truth.
- Processing: Source locations are matched to SWORD reaches within 10 km, while elevations are harmonized across geoid-referenced satellite products and locally referenced ANA records.Source-specific quality screening and expert hydrology processing are applied before export.
- Evaluation: The task reconstructs daily WSE for every SWORD reach in SWOT-era and pre-SWOT hindcast tracks, with ANA gauges withheld from model inputs and training targets.The withheld gauges are reserved solely for evaluation.
Model
The model represents sparse water-surface-elevation observations as topology- and metadata-aware tokens, processes them in a bidirectional selective state-space sequence, and predicts masked values. It samples connected river subgraphs and queries only the target reach, avoiding dense reach–time grids during inference.
- Sample Construction: The model samples a connected local SWORD neighborhood around an anchor and orders observed tokens by day and downstream river-flow rank.Each sample is defined by an anchor, time interval, and connected upstream/downstream neighborhood.
- Sample Construction: Sparse inputs create tokens only for available measurements, so missing values do not require a dense reach–time grid.Masked tokens are used for training, while query tokens represent targets during inference.
- Sample Construction: Topology-aware encodings represent each reach through branch choices and shared path segments, without node-index lookup, enabling induction to unseen reaches with available metadata.The encoding combines local branch paths with temporal, source, geographic, and scalar WSE metadata.
- Bidirectional Mamba: Bidirectional Mamba scans integrate information from upstream and downstream sequence directions before producing source-specific normalized WSE predictions.The model is a non-causal reconstruction model because it uses observations on both sides of a query.
- Inference: Inference adds query tokens only for the target reach, preserving a favorable observed-context-to-query ratio under extreme sparsity.Each target interval is reconstructed independently and overlapping windows are averaged.
Experiments
The experiments use filtered, multi-source altimetry data and evaluate imputation under mixed transductive and inductive conditions.
- Experiments: Training combines SWOT, HydroWeb, and ICESat-2 observations after filtering short time series, leaving approximately 75% of reaches observed.The model has 3.3M parameters and predicts all 19K nodes after training.
- Experiments: Figure 6 compares the model and baselines as neighborhood observations vary, targeting performance under changing contextual sparsity.
1. SWOT
The SWOT-era setting benefits from wide-swath observations, while evaluation preserves gauge data exclusively for held-out assessment.
- 1. SWOT: The SWOT-era period contains 2.7% observed day–reach slots because SWOT measurements are available during training and inference.
- 1. SWOT: Gauge data from 2022-07 to 2023-07 is reserved for ablation evaluation and is not used for early stopping or training.This prevents leakage even when training and inference periods coincide.
- 1. SWOT: Baselines span temporal, non-parametric, recurrent graph, attention-based, and inductive reconstruction methods.
Results
The proposed model outperforms existing machine-learning and hydrology baselines across sparse settings, while ablations identify benefits from topology, metadata, and multiple data sources.
- Results: The model outperforms all existing machine-learning methods in every reported setting, with sequence models degrading less than baselines as context observations decrease.Subgraph sampling generally outperforms full-graph processing, except for GRIN.
- Results: Metadata encodings, temporal or flow ordering, subgraph context, and all three data sources improve predictions over corresponding ablations.The Base subgraph model significantly outperforms the Isolated model, and 300–1000 km neighborhoods perform best.
- Results: Figure 7 evaluates temporal-window and spatial-neighborhood limits, with 3 months yielding the lowest RMSE and performance robust to temporal-window size.
- Results: KGE jointly evaluates correlation, variability, and mean agreement, with a perfect score of 1 and higher values indicating better performance.
- Results: The model outperforms Reach-Reg and ImputeFormer on both evaluation periods while providing greater coverage than Reach-Reg.Reach-Reg is evaluated for hindcasts despite being developed for the SWOT era.
Conclusion
AmazonWSE provides a large, highly sparse river-network benchmark, and the proposed sequence model improves imputation by combining connected-subgraph sampling with flattened space–time sequences.
- Conclusion: AmazonWSE covers 19K river reaches with 99% sparsity and a directed acyclic graph topology for spatiotemporal graph imputation.
- Conclusion: The bidirectional selective state space model outperforms prior graph-imputation methods and a state-of-the-art non-neural densification baseline while producing denser reconstructions.
A Appendix
The appendix provides supplementary technical materials, media, code, data, and examples supporting the dataset and model.
- The technical supplement covers dataset sources, preprocessing, storage format, hyperparameters, baselines, and predictions.
- Media supplements include animations of sparse inputs and daily WSE imputations for all SWORD reaches.
- The code and data supplement provides per-source NetCDF files, preprocessing and export scripts, and training and evaluation code.
A.1 Dataset Sources
AmazonWSE combines multiple satellite and gauge-derived water-surface-elevation sources with the SWORD river-reach graph.
- SWORD supplies a topological graph of 19,172 Amazon river reaches for organizing the dataset.
- ICESat-2 contributes measurements from three pairs of narrow laser beams, with beam pairs separated by approximately 3 km.
- SWOT RiverSP provides wide-swath reach-level observations from the science cycle beginning in July 2023.
- HydroWeb.next contributes operational and research time series from SWOT nadir and several other satellite altimeters.
A.2 Data Preprocessing Details
Preprocessing harmonizes observations onto daily UTC bins, applies source-specific quality filters, and creates interpolated SWOT statistics for prediction de-normalization.
- All timestamps are converted to UTC and measurements are assigned to daily bins on a shared temporal grid.
- SWOT preprocessing aggregates RiverSP records into 24-hour reach-day bins and applies extensive product-level and measurement-level quality filters.The filters remove about 50% of measurements.
- The pipeline filters SWOT observations using precision support, harmonic residuals, rating-curve consistency, and quality criteria.
- Pass-specific bias correction produces negligible differences in the resulting observations.
- SWOT mean and standard-deviation fields are spatially interpolated with inverse-distance weighting and used to de-normalize predictions at unobserved locations.
HydroWeb.
HydroWeb and other source records are cleaned, matched to SWORD reaches, stored in an indexed-ragged format, and used with source-aware prediction decoding.
- HydroWeb measurements are already cleaned and homogenized, while source files preserve accepted observations and quality metadata in a shared schema.
- ICESat-2 filtering removes about 4% of measurements among those matched to SWORD reaches.
- ANA gauge processing retains 375 series after automatic rejection, reach matching, quality filtering, and manual inspection.
- HydroWeb, ICESat-2, and ANA locations are matched to the nearest SWORD reach within 10 km, while SWOT uses its reach identifier directly.
- The model decodes predictions as HydroWeb for gauge comparison and de-normalizes them with interpolated SWOT statistics, the best-performing tested combination.
- Sampled subgraphs and source representations are evaluated alongside fixed-graph settings for existing baselines.