Machine learning-based simulation of monthly water quality in the Santa Lucía Chico river basin

Abstract

This study aims to develop a data-driven tool for monthly water quality simulation using machine learning techniques. The study focuses on the upper basin of the Santa Lucía Chico River in Uruguay, utilizing data from two water quality monitoring stations (XSLH010 and XSLH020). The variables considered include dissolved oxygen (DO), temperature (T), total nitrogen (NT), and phos-phate (PO₄³⁻). The time series data were split into training (80%) and testing (20%) sets, with separate min-max normalization applied to ensure consistent scaling across variables. The prediction models were trained using Extra Trees Regressor (ET) and Histogram-based Gradient Boosting Regressor (HGB), eval-uated with Mean Absolute Error (MAE) and Mean Squared Error (MSE). This resulted in four models trained per variable. Nash-Sutcliffe Efficiency (NSE) was also calculated for model performance evaluation. Optimal hyperparameters were identified using a 5-fold cross-validation process and optimized with Op-tuna. The input dataset integrates domain knowledge by incorporating spatial de-pendencies, spatial correlations, physical dependencies, and temporal variability. Additionally, SHapley Additive exPlanations (SHAP) values were used to refine model inputs by removing low-importance variables. The models operate at a monthly time step, allowing for the assessment of long-term water quality trends. The results were highly satisfactory, with NSE values exceeding 0.6 for all vari-ables across both stations, except for PO₄³⁻ at XSLH010. These findings demon-strate the potential of machine learning models for water quality prediction and provide a valuable tool for improving water resource management. Future efforts will focus on refining the model, incorporating additional data sources, and ex-tending its applicability to other basins.