A Ranked Contamination Diagnosis Method for Water Distribution Systems Using Spatio-Temporal Graph Neural Networks

Contamination diagnosis in water distribution systems (WDS) is critical for ensuring public health and maintaining water quality. However, exact localization of contamination sources is often highly challenging due to sparse sensor deployment, uncertain water flow dynamics, and the ambiguity of contamination signatures. To address this challenge, we propose a novel end-to-end deep learning framework for ranked contamination diagnosis in WDS. Instead of relying on a single-source prediction, the deep learning model outputs a probabilistic ranking of likely contamination source nodes, enabling top- $k$ diagnosis strategies under uncertainty. A unified classification layer jointly performs contamination detection and source localization, allowing the model to operate in an end-to-end fashion. Extensive experiments on a benchmark WDS demonstrate that the proposed framework achieves strong top- $k$ diagnosis accuracy and remains robust across varying sensor configurations. This work offers a practical and scalable approach to contamination diagnosis in WDS, supporting more informed and effective decision-making in real-world operations. The code is available at: https://github.com/Xiaohan-Chen/ST-GAT.