Authors: Peishi Jiang – University of Alabama
Title: Learning Watershed Memory: Coupling LSTMs with Differentiable SAS Functions for Transport Process Discovery
Presentation Type: Poster Presentation
Abstract: Uncovering transit time distributions (TTDs) remains a long-standing challenge in understanding watershed transport processes. The StorAge Selection (SAS) function provides a numerically tractable way to derive TTDs from tracer data, but existing SAS models rely on predefined probability distributions, which require prior knowledge and limit flexibility. To address this, we propose a generic SAS formulation using a Mixed Density Network (MDN), representing the SAS function as a weighted combination of three base distributions: Gamma (skewed), Gaussian (symmetric), and Uniform (flat). This approach is implemented in a fully differentiable SAS-based transport model, \texttt{DiffSAS}, leveraging the JAX library for efficient training via automatic differentiation. \texttt{DiffSAS} was applied to two watersheds: a coastal site with observed flow and tracer data, and a mountainous site with synthetic data from an integrated hydrological model. The MDN-based SAS model achieved strong predictive performance, with Nash–Sutcliffe Efficiency values above 0.6 at both sites, substantially outperforming SAS models using fixed Gamma distributions.
The derived TTDs captured the increased contribution of young water during high-flow periods, and the time-varying weights of the MDN distributions revealed how TTD shapes adapted to changing flow conditions. We further coupled \texttt{DiffSAS} with an LSTM-based flow model, enabling simultaneous emulation of flow and tracer dynamics using only meteorological forcing as input. The coupled framework demonstrated strong operational performance under the assumption that observed streamflow is unavailable, highlighting its potential for transport prediction in ungauged or data-scarce basins. Overall, this study demonstrates that MDN-based SAS modeling provides a flexible, data-adaptive framework for uncovering time-varying transport behavior, offering a valuable tool for advancing process-based hydrological understanding in both well-monitored and data-scarce watersheds.