Enhancing water quality prediction: a machine learning approach across diverse water environments

WWTP

wastewater treatment plant

RE

nutrient removal efficiency

ML

machine learning

AI

artificial intelligence

BOD

biochemical oxygen demand

COD

chemical oxygen demand

DO

dissolved oxygen

L_at

sampling location latitude

L_ng

sampling location longitude

Y

year

M

month

D

day

H

hour

SAD

site actual depth

SS

sea state, degree of turbulence at sea

WS

wind speed

WDO

Winkler method dissolved oxygen

C

sample temperature

SAL

salinity

TC

total coliform

SD

water light penetration

ACA

active chlorophyll-A

PCB

polychlorinated biphenyl plate count

NN

nitrate/nitrite

OP

ortho-phosphorus

AMN

ammonium

TP

total phosphorus

PH

potential of hydrogen

TSS

total suspended solid

CDT

conductivity

SD

sample depth

WD

water density

TSP

transparency of water

NH₃-N

ammonia-nitrogen

TOC

total organic carbon

DTP

dissolved total phosphorus

PO₄-P

phosphate

TN

total nitrogen

DTN

dissolved total nitrogen

NO₃-N

nitrate nitrogen

EC

electroconductivity

DMF

deep matrix factorization

DNN

deep neural network

GB

gradient boosting

RF

random forest

WQI

water quality index

ANFIS

adaptive neuro-fuzzy inference system

CNN

convolutional neural networks

ANN

artificial neural network

LSTM

long short-term memory

IA-LSTM

input-attention LSTM

FPN

feature pyramid network

PSPNet

pyramid scene parsing network

ResNet

residual networks

PSO

particle swarm optimization

NBC

naive bayes classifier

SVM

support vector machine

GCWs

groundwater circulation wells

MLR

multiple linear regression

NGR

natural gamma ray

TH

total hardness

HCO₃⁻

bicarbonate

Cl⁻

chloride

SO₄²⁻

sulfate

nitrate

Ca²⁺

calcium

Mg²⁺

magnesium

F⁻

iron

Na⁺

sodium

K⁺

potassium

K_V

vertical hydraulic conductivity

K_H

horizontal conductivity

M

aquifer thickness

Μ

specific yield

K_H/K_V

vertical heterogeneity

N

porosity

I

hydraulic gradient

P_r

ratio of particle recovery

R

radius of influence

TWI

topographic wetness index

TPI

topographic position index

FA

factor analysis

ORP

oxidation-reduction potential

SWRO

seawater reverse osmosis

PET

polyethylene terephthalate

PS

polystyrene

PP

polypropylene

PE

polyethylene

PA

polyamide

PUR

polyurethane

PC

polycarbonate

PVC

polyvinyl chloride

PMMA

poly (methyl methacrylate)

CA

cellulose acetate

T_p

permeate temperature

T

time

p_f

pressure

Q_f

permeate flow

Cond_f

feed conductivity

Cond_p

permeate conductivity

T_b

brine temperature

R

flow recovery

B_B

boronpermeability coefficient

LMT

logistic model tree

CMOS

complementary metal-oxide-semiconductor

BiGRU

bidirectional gated recurrent unit

IoU

intersection over union

AAO

anaerobic–anoxic–oxic

Q_r

internal recirculation flow rate

Q_sr

sludge internal recycle flow rate

YOLO5

You Only Look Once

VGG

intense convolutional networks

GRUs

gated recurrent neural network units

AP

atmospheric pressure

WS_avg

average wind speed

P_r

total rainfall or snow melt

VIS

average visibility

WS_max

maximum wind speed

Power_1h

power consumption per hour

F1_run

fan one running status

F2_run

fan two running status

M_LSS

mixed liquor suspended solids

T_Coli

total coliform

F_Coli

fecal coliform

P

precipitation

Chl-a

chlorophyll-a

D_d

drainage density

RH

humidity

Water is essential for sustaining life, ecosystems, and economic activities (Silva 2023). Water quality is critical for sustainable development, public health, and environmental balance. However, predicting and managing water quality remains challenging, particularly across diverse aquatic environments such as rivers, lakes, reservoirs, coastal waters, and groundwater systems (Oyedotun & Ally 2021). These environments are influenced by a range of natural and anthropogenic factors, which include industrial discharge, agricultural runoff, urbanization, and climate variability (Datta et al. 2021; Urban Wastewater Scenario in India 2022). In this context, conventional water quality monitoring techniques that depend on field sampling and laboratory analysis are indispensable but limited. Their spatial-temporal coverage is often time-consuming, expensive, and constrained (Korostynska et al. 2013). Consequently, these methodologies face challenges in delivering timely and comprehensive data analysis, which is vital for efficient water resource management.

Machine learning (ML) is one of the promising solutions to these challenges (Sheng et al. 2020). ML techniques excel at analyzing large, complex, diverse datasets and uncovering patterns to generate accurate predictions. ML models provide near real-time insights into water quality dynamics by integrating source data such as remote sensing, sensor networks, and historical records (Sagan et al. 2020; Mucheye et al. 2022). These capabilities complement traditional monitoring techniques, enhancing predictive accuracy and scalability across multiple aquatic environments. This paper provides an overview of the potential of ML in addressing water quality challenges, focusing on heterogeneous water systems. The key areas of investigation include:

• 1. Complexity of diverse water environments: Examining how varying hydrological, geological, and ecological conditions influence water quality dynamics.

• 2. The efficiency of ML techniques: Evaluating the performance of ML and ensemble models on predicting water quality parameters.

• 3. Data integration: Leveraging spatial, temporal, and environmental datasets to improve model accuracy and address data limitations.

• 4. Real-world applications: Demonstrating the role of ML-driven insights in supporting decision-making, policy formulation, and sustainable water resource management.

This paper attempts to bridge the gap between data-driven technologies and water quality management. The study demonstrates how ML can revolutionize monitoring practices by facilitating proactive and informed interventions to protect water resources in the face of increasing environmental challenges.

A four-step process is employed to filter the literature systematically:

• i. Identification: Initially, articles are identified by selecting subject areas such as ‘Engineering’ and ‘Computer Science.’

• ii. Screening: Exclusion criteria are applied to remove systematic reviews, meta-analyses, and conference proceedings. The articles published only in a single language (English) during the previous 5 years are included.

• iii. Eligibility: Titles and abstracts are carefully assessed to select articles relevant to the research scope.

• iv. Inclusion: Full-text articles deemed relevant and available are included in the final dataset. Key details, including input parameters, methodologies, datasets, and target parameters, are analyzed and summarized.

Network analysis is applied to examine the relationships between keywords and themes within the selected publications. Tools such as VOSviewer are utilized to generate co-occurrence networks based on data retrieved from 2020 to 2024. Keywords included ‘water treatment,’ ‘machine learning,’ ‘deep learning,’ and ‘artificial intelligence.’

Figures 2 and 3 collectively represent the importance of ML and artificial intelligence (AI) in enhancing water quality forecasting and management. Both are directly associated with main streams, which include predictive modeling, wastewater treatment, and environmental monitoring, emphasizing significant trends and directions for further studies.

Surface water refers to water found in streams, rivers, and lakes and represents a crucial freshwater resource. However, its quality is increasingly compromised by sewage discharge, industrial effluents, and agricultural runoff. These pollutants pose significant challenges to maintaining surface water quality, particularly in urban and rural areas (Oyedotun & Ally 2021). Geological and anthropogenic factors majorly impact global riverine zones, influencing the natural environment (Hoang et al. 2020; Islam et al. 2020). In the last 10 years, attention has increased toward water quality degradation in different regions of the world (Setia et al. 2020; Uddin et al. 2021). Table 1 highlights the various algorithms utilized in surface water research for predicting water quality parameters. Measuring water quality and recording parameters such as pH, dissolved oxygen (DO), and temperature is relatively straightforward and can often be done quickly using handheld detectors (Zhi et al. 2021). However, more complex analyses, such as detecting pesticides, nutrients, metals, or bacterial/algae concentrations require more resources, laboratory facilities, and time-intensive efforts (Ooi et al. 2022). Traditional statistical methods are generally ineffective in addressing temporal inefficiencies associated with biological oxygen demand (BOD) prediction. They are typically limited to sparse and incomplete datasets, which may arise due to sampling errors or facility constraints. Ma et al. (2020) resolved this issue by integrating deep matrix factorization (DMF) and deep neural network (DNN) data-driven models, significantly improving the accuracy of BOD prediction. Algorithms and datasets used in water quality models play a crucial role in determining the performance of prediction systems. For example, Chen et al. (2020) presented highly specified input features by comparing ten learning models, including seven traditional ML models and three ensemble models. Mohd Zebaral Hoque et al. (2022) applied the water quality index (WQI) to measure surface water quality. Sidek et al. (2024) improved WQI monitoring by integrating random forest (RF) and gradient boosting (GB) models, offering more reliable and accurate ML predictions with minimal effort.

Table 1

ML applications on different water environments

Water environment .	Task .	Dataset .	Algorithm .	Independent variables .	Target variables .
Surface water	BOD Prediction (Ma et al. 2020)	32,323 samples from New York City Open Data up to 2018	DMF and DNN	Lat, Lng, Y, M, D, H, SAD, Sea state, SS, WS, WDO, C, SAL, TC, SD, ACA, PCB, NN, OP, AMN, TP, PH, TSS, CDT, SD, WD, and TSP	BOD5 Prediction
	Water quality prediction (Chen et al. 2020)	2012–2018 data from the CNEMC stations	Three ensembles and seven traditional ML models	DO, NH3-N, COD, and pH	COD, NH3-N, and DO prediction
	WQI prediction (Sidek et al. 2024)	1,637 samples from 2008 to 2018 in Johor River, Malaysia	RF and GB	BOD, COD, and the DO (%)	WQI prediction
	Algal bloom prediction (Ly et al. 2021)	2011–2020 samples in Han River, China	ANFIS	NH3-N, NO3-N, COD, TOC, BOD, TSS, TP, PO4-P, DTP, TN, DTN, Chl-a P, Q, °C, DO, EC, pH, FColi and TColi	Chlorophyll-a prediction
Ground water	WQI (Agrawal et al. 2021)	One-year samples from hand pumps and bore wells in the Pindarwan tank	NBC, PSO, and SVM	pH, total dissolved solids (TDS), EC, TH, alkalinity, HCO3−, Cl−, ⁠, ⁠, Ca2+, Mg2+, F−, Na+ and K+	Predict water quality index
	Groundwater optimal design circulation wells (GCWs) (Fang et al. 2024)	3,000 samples of media types, aquifer thicknesses, and hydrogeological parameters in Xi'an City, Shaanxi	MLR, ANN, and SVM	KV, KH, M, μ, KH/KV, n, and I	Ratio of Pr and R
	Groundwater potential zones (Sarkar et al. (2024)	200 samples of groundwater and metrological data from Bangladesh Water Development (BWD)	LMT, ANN, and LR	Soil types, Dd, curvature, rainfall, slope, °C, geology, RH, roughness, lineament density, land use and land cover, geomorphology, TWI, and TPI	Groundwater potential zones prediction
	Predicting groundwater availability (Hussein et al. 2020)	174 monthly groundwater satellite images from 2002–2019	MLP, extreme GB, MLR, RF, and SVM	Local and global spatiotemporal features	Groundwater prediction
Seawater	Trapping and identifying small-sized microplastics (Gong et al. 2023)	Dataset using microfluidics and Raman spectroscopy devices	SVM, RF, CNN, and ResNet34	Eleven plastic types included PET, PS, PP, PE, PA, Nylon, polyester, PUR, PC, PVC, PMMA, and CA	Identify the type of plastic
	Forecasting boron coefficient values (Ajali-Hernández et al. 2024)	Data from 18 coastal wells in Spain with depths of 50 m and 100 m	Ensemble-based ML	Tp, t, pf, Qf, Condf, Condp, Tb, and R	Forecast BB.
	Classify saline particles (Alshehri et al. 2021)	The saline solution is observed using a CMOS camera	CNN with transfer learning	A Raspberry Pi device captured images scattered across various salt salinity concentrations	Classification for 10 salt concentration particles
	Marin fauna detection (Colefax et al. 2023)	Drone-based fauna Red Green Blue (RGB) images from 2021–2022	RetinaNet single-shot detector with ResNet-50	Dolphin groups video data from nine flight captures	Detecting submerged fauna
Wastewater	Energy consumption prediction (Harrou et al. 2023; Silva 2023)	2014–2019 energy consumption, climate, and wastewater data from Melbourne's eastern WWTP	LSTM and BiGRU	TN, COD, BOD, °C, AP, WSavg, Pr, VIS, WSmax, year, month, and day	Energy consumption prediction
	Predict water quality parameters (Wei et al. 2023)	3,912 h of data from February to June 2022 at a WWTP in southern China	LSTM with MLP network	COD, pH, NH3-N Power1h, °C, instantaneous flow rate at the effluent outflow, DO, F1run, and F2run	COD, pH, NH3-N prediction
	Real-time control of AAO (Liu et al. 2023b)	2019–2021 data from WWTP	LSTM model predictive control	DO, ORP, NH4+-N, ⁠, MLSS, Qin, COD, NH4+-N, TN, TP, pH, SS, °C Qr, and Qsr	Aeration volumes, Qr, and Qsr controlling
	Blockage detection (Patil et al. 2023)	14,765 sewer blockage frames ‘S-BIRD’ dataset	Transfer learning and YOLOv5	Blockage images	Blockage detection

Water environment .	Task .	Dataset .	Algorithm .	Independent variables .	Target variables .
Surface water	BOD Prediction (Ma et al. 2020)	32,323 samples from New York City Open Data up to 2018	DMF and DNN	Lat, Lng, Y, M, D, H, SAD, Sea state, SS, WS, WDO, C, SAL, TC, SD, ACA, PCB, NN, OP, AMN, TP, PH, TSS, CDT, SD, WD, and TSP	BOD5 Prediction
	Water quality prediction (Chen et al. 2020)	2012–2018 data from the CNEMC stations	Three ensembles and seven traditional ML models	DO, NH3-N, COD, and pH	COD, NH3-N, and DO prediction
	WQI prediction (Sidek et al. 2024)	1,637 samples from 2008 to 2018 in Johor River, Malaysia	RF and GB	BOD, COD, and the DO (%)	WQI prediction
	Algal bloom prediction (Ly et al. 2021)	2011–2020 samples in Han River, China	ANFIS	NH3-N, NO3-N, COD, TOC, BOD, TSS, TP, PO4-P, DTP, TN, DTN, Chl-a P, Q, °C, DO, EC, pH, FColi and TColi	Chlorophyll-a prediction
Ground water	WQI (Agrawal et al. 2021)	One-year samples from hand pumps and bore wells in the Pindarwan tank	NBC, PSO, and SVM	pH, total dissolved solids (TDS), EC, TH, alkalinity, HCO3−, Cl−, ⁠, ⁠, Ca2+, Mg2+, F−, Na+ and K+	Predict water quality index
	Groundwater optimal design circulation wells (GCWs) (Fang et al. 2024)	3,000 samples of media types, aquifer thicknesses, and hydrogeological parameters in Xi'an City, Shaanxi	MLR, ANN, and SVM	KV, KH, M, μ, KH/KV, n, and I	Ratio of Pr and R
	Groundwater potential zones (Sarkar et al. (2024)	200 samples of groundwater and metrological data from Bangladesh Water Development (BWD)	LMT, ANN, and LR	Soil types, Dd, curvature, rainfall, slope, °C, geology, RH, roughness, lineament density, land use and land cover, geomorphology, TWI, and TPI	Groundwater potential zones prediction
	Predicting groundwater availability (Hussein et al. 2020)	174 monthly groundwater satellite images from 2002–2019	MLP, extreme GB, MLR, RF, and SVM	Local and global spatiotemporal features	Groundwater prediction
Seawater	Trapping and identifying small-sized microplastics (Gong et al. 2023)	Dataset using microfluidics and Raman spectroscopy devices	SVM, RF, CNN, and ResNet34	Eleven plastic types included PET, PS, PP, PE, PA, Nylon, polyester, PUR, PC, PVC, PMMA, and CA	Identify the type of plastic
	Forecasting boron coefficient values (Ajali-Hernández et al. 2024)	Data from 18 coastal wells in Spain with depths of 50 m and 100 m	Ensemble-based ML	Tp, t, pf, Qf, Condf, Condp, Tb, and R	Forecast BB.
	Classify saline particles (Alshehri et al. 2021)	The saline solution is observed using a CMOS camera	CNN with transfer learning	A Raspberry Pi device captured images scattered across various salt salinity concentrations	Classification for 10 salt concentration particles
	Marin fauna detection (Colefax et al. 2023)	Drone-based fauna Red Green Blue (RGB) images from 2021–2022	RetinaNet single-shot detector with ResNet-50	Dolphin groups video data from nine flight captures	Detecting submerged fauna
Wastewater	Energy consumption prediction (Harrou et al. 2023; Silva 2023)	2014–2019 energy consumption, climate, and wastewater data from Melbourne's eastern WWTP	LSTM and BiGRU	TN, COD, BOD, °C, AP, WSavg, Pr, VIS, WSmax, year, month, and day	Energy consumption prediction
	Predict water quality parameters (Wei et al. 2023)	3,912 h of data from February to June 2022 at a WWTP in southern China	LSTM with MLP network	COD, pH, NH3-N Power1h, °C, instantaneous flow rate at the effluent outflow, DO, F1run, and F2run	COD, pH, NH3-N prediction
	Real-time control of AAO (Liu et al. 2023b)	2019–2021 data from WWTP	LSTM model predictive control	DO, ORP, NH4+-N, ⁠, MLSS, Qin, COD, NH4+-N, TN, TP, pH, SS, °C Qr, and Qsr	Aeration volumes, Qr, and Qsr controlling
	Blockage detection (Patil et al. 2023)	14,765 sewer blockage frames ‘S-BIRD’ dataset	Transfer learning and YOLOv5	Blockage images	Blockage detection

View Large

Harmful algal blooms (HAB) present a significant challenge in surface water due to their anoxic potential and toxin release (Barrientos-Espillco et al. 2023). Baek et al. (2022) explored deep-learning techniques to predict HAB in surface water. The researchers detected HAB using remote sensing and in situ monitoring; however, this approach relies entirely on algal biomass. For instance, weather conditions can influence physicochemical factors at spatial and temporal scales, with chlorophyll-a as a significant indicator of phytoplankton abundance. Ly et al. (2021) observed the impact of weather conditions to enhance the accuracy of HAB prediction, enabling proactive water management strategies.

As per a report by the US EPA (2025) on ‘Climate Change Impacts on Freshwater Resources,’ the growing population and changing weather patterns have increasingly contributed to a higher demand for freshwater resources. Groundwater accounts for 99% of the Earth's liquid freshwater reserves (Lall et al. 2020). Groundwater treatment and quality prediction are achieved by ML algorithms, as summarized in Table 1.

The WQI is used to evaluate the groundwater quality and aggregate the various water quality parameters. The traditional approach to WQI computation relies on complex mathematical models, is time-intensive, and struggles to handle large datasets. Advances in ML techniques, such as support vector machines (SVM), naive bayes classifiers (NBC), and particle swarm optimization (PSO), offer more efficient and accurate predictive models, enhancing water resource management (Agrawal et al. 2021). Groundwater circulation wells (GCWs) are one of the popular water purification techniques. Fang et al. (2024) presented artificial neural networks (ANN), multiple linear regression (MLR), and SVM to predict the influence radius and particle recovery ratio for optimum GCW design. Sarkar et al. (2024) predicted groundwater potential zones in Bangladesh using climatic data through ML models ANN, logistic tree models, and linear regression. Similarly, Hussein et al. (2020) utilized further advanced groundwater availability prediction by using GB, LR, multilayer perceptron (MLP), RF, and support vector regression (SVR) by satellite imagery.

Seawater, which constitutes 97% of Earth's water, is unsuitable for domestic use due to its high salinity (National Ocean Services 2024). Desalination has traditionally been used to convert seawater into freshwater (Qasim et al. 2019; Zapata-Sierra et al. 2021). Recently, ML models have emerged as promising tools for enhancing the efficiency of desalination processes, as detailed in Table 1.

The increasing contamination of seawater with microplastics poses a critical threat to marine ecosystems, aquatic life, and human health and emphasizes the importance of their identification, sampling, and contamination control (Cutroneo et al. 2020). Gong et al. (2023) proposed a data-driven model using SVM, RF, and deep-learning models, such as convolutional neural networks (CNN) and Residual Networks (ResNet34), for classifying and detecting 11 types of microplastics. Researchers utilize Raman spectroscopy to analyze microplastic samples obtained through microfluidics, facilitating high-precision identification.

The seawater reverse osmosis (SWRO) technique is one of the most reliable desalination methods, though extracting boron ions remains challenging (Abba et al. 2023). Ajali-Hernández et al. (2024) addressed this issue by using ensemble-based ML models to predict the permeability coefficient of boron over 1,500 days of operational data from a large-scale SWRO plant. Deep learning has also been applied to improve desalination efficiency. Alshehri et al. (2021) presented an ensemble-based model to classify salt particle concentrations into ten categories from the dataset of salt particle images captured with a Raspberry Pi device.

Monitoring marine fauna is essential to understanding and mitigating the negative impacts of human activities on aquatic ecosystems. The spectral filtering ML technique has been used to detect and track marine animals. Colefax et al. (2023) have applied this approach to monitor dolphin populations in aquatic systems, demonstrating the potential of ML for ecological conservation.

Wastewater treatment involves the collection, processing, and reuse of wastewater for safe disposal into natural water bodies. The treatment process typically involves three stages: primary, secondary, and tertiary stages (Adrados et al. 2014). A treatment plant typically includes the following stages (Garg 2012):

• (a) Screening: Passing sewage through different types of screens to trap and remove the floating matter.

• (b) Grit removal: Removal of non-putrescible material such as sand, bones, eggshells.

• (c) Sedimentation: Particles in the water settle out of the liquid due to gravity, followed by processes like detention and flocculation.

• (d) Aeration/activated sludge: Aeration tank of long detention, where the activated sludge mixed with segmented sewage is agitated and aerated.

• (e) Secondary settling: Coagulated suspended mass settles down by gravity.

• (f) Filtration: Removal of suspended and dissolved solids through physical and chemical processes.

• (g) Disinfection: Removal of pathogens in effluents for safe disposal in water bodies.

In the past few years, ML has optimized processes, enabled predictive capabilities, and thus revolutionized wastewater treatment, as summarized in Table 1. Most of the ML-based studies concentrated on filtration, which is a critical energy-intensive step in wastewater treatment. Harrou et al. (2023) utilized a predictive model using long short-term memory (LSTM) and bidirectional gated recurrent unit (BiGRU) algorithms for estimating energy consumption on datasets, including energy usage, climate data, and wastewater properties. Wei et al. (2023) presented a deep-learning algorithm to optimize energy consumption using sensor data from wastewater treatment plants (WWTPs).

One widely used technique in WWTPs for nitrogen and phosphorus removal is the anaerobic–anoxic–oxic (AAO) treatment (Sang et al. 2021). The AAO system contains three interconnected tanks: anaerobic, aerobic, and anoxic. Liu et al. (2023a) demonstrated the use of an ML model to optimize the AAO process by maximizing the internal recirculation flow rate (Q_r), sludge internal recycle flow rate (Q_sr), and aeration volumes. This approach significantly improves energy efficiency, enhances water quality, and ensures effective process regulation.

WWTP's performance is significantly impacted by non-biodegradable products in the sewer system. These result in blockages that lead to overflows and adverse effects on water quality, which affects human health and safety. Deep-learning techniques offer promising solutions to improve wastewater management processes (Alvi et al. 2023). Integrating the ML methods is essential to avoid flooding due to sewer system defects. Patil et al. (2023) worked on blockage detection using image processing techniques, including the You Only Look Once version 5 (YOLO5) deep-learning algorithm and transfer learning on the ‘S-BIRD’ images dataset in sewer blockage detection.

There are two major WWTPs in Melbourne City run by Melbourne Water: the Eastern Treatment Plant (ETP) and the Western Treatment Plant (WTP). The ETP resolves congestion at WTP, providing for Melbourne's expanding population in its southeastern areas. The records used in this study are open-access from Melbourne Water and the Melbourne Airport Weather Station. The data on hydraulic, biological, climate, and energy consumption are merged into a unified dataset with a daily time resolution (Table 2).

Table 2

Summary of values of measured water quality parameters

Parameter (Abbreviation) .	Unit .	Max .	Min .	Mean .
Average inflow (Qin)	/s	19	2.6	4.5
Average outflow (Qout)	/s	7.9	0.1	3.9
Energy consumption (EC)	MWh	398	116	275
Ammonium (NH4-N)	mg/L	93	13	39
BOD	mg/L	850	140	382
Chemical oxygen demand (COD)	mg/L	1,700	360	846
Total nitrogen (TN)	mg/L	92	40	62
Average temperature (Tavg)	C	35.5	0	15
Maximum temperature (Tmax)	C	43.5	0	20.5
Minimum temperature (Tmin)	C	28	−2	10
AP	(hPa)	10.22	0	3.7
Average humidity (H)	%	97	0	63
Total rainfall or snowmelt (Pr)	Mm	18	0	0.2
VIS	Km	512	0	9
Average wind speed (WSavg)	Km/h	49	0	19
Maximum wind speed (WSmax)	Km/h	83.5	0	35.4
Year (year)	-	2019	2014	-
Month (month)	-	12	1	-
Day (day)	-	31	1	-

Parameter (Abbreviation) .	Unit .	Max .	Min .	Mean .
Average inflow (Qin)	/s	19	2.6	4.5
Average outflow (Qout)	/s	7.9	0.1	3.9
Energy consumption (EC)	MWh	398	116	275
Ammonium (NH4-N)	mg/L	93	13	39
BOD	mg/L	850	140	382
Chemical oxygen demand (COD)	mg/L	1,700	360	846
Total nitrogen (TN)	mg/L	92	40	62
Average temperature (Tavg)	C	35.5	0	15
Maximum temperature (Tmax)	C	43.5	0	20.5
Minimum temperature (Tmin)	C	28	−2	10
AP	(hPa)	10.22	0	3.7
Average humidity (H)	%	97	0	63
Total rainfall or snowmelt (Pr)	Mm	18	0	0.2
VIS	Km	512	0	9
Average wind speed (WSavg)	Km/h	49	0	19
Maximum wind speed (WSmax)	Km/h	83.5	0	35.4
Year (year)	-	2019	2014	-
Month (month)	-	12	1	-
Day (day)	-	31	1	-

View Large

The Melbourne Water dataset includes records from 2014 to 2019 and is available under Victoria's open data policy (Bagherzadeh et al. 2021b). The study obtains wastewater characteristics through sampling and sensor readings and sources climate data from the Melbourne Airport Weather Station (Melbourne Airport Weather Station 2021). Table 2 summarizes the water quality parameters, presenting the ranges of measured water quality components. As shown in the table, measurement parameters include average inflow (Q_in), average outflow (Q_out), energy consumption (EC), ammonium (NH₄-N), total nitrogen (TN), BOD, average temperature (T_avg), chemical oxygen demand (COD), maximum temperature (T_max), minimum temperature (T_min), average humidity (H), total rainfall or snowmelt (P_r), average wind speed (WS_avg), atmospheric pressure (AP), maximum wind speed (WS_max), year (year), average visibility (VIS), month (month), and day (day).

Data engineering processes included cleaning, encoding, and outlier detection to enhance the quality of the datasets. Linear interpolation is used to address missing values. Categorical data are converted to numerical form using encoding techniques, while Z-score outlier detection identifies anomalous data points that deviate significantly from the dataset's normal distribution (Aggarwal et al. 2019).

FS methods are employed to evaluate and identify the most relevant features for modeling based on specific criteria. This research uses mutual information (MI) as the primary FS method, which captures linear and nonlinear relationships between features and the target variable (Zegaar et al. 2024). Similarly, methods like LASSO regression and recursive feature elimination are commonly used to identify the model features (Zhang et al. 2021; Julian et al. 2023). MI is adopted for its simplicity and effectiveness in predicting water quality, particularly in addressing the nonlinear and complex interactions commonly observed between environmental and water quality factors (Moeinzadeh et al. 2023).

ML models, including linear regression (LR), K-nearest neighbor regression (KNN), SVR, and decision tree regression (DT), are applied to predict water quality. In addition, the ensemble model technique combines multiple models to produce more accurate predictions than any single model. The ensemble models are RF, GB regression, and AdaBoost Regression (AR), which are employed to improve predictive performance.

LR models predict target values by mapping input features to target values through line fitting. This method is straightforward for linear patterns but limited to complex, nonlinear relationships. RF is an ensemble learning method that creates multiple decision trees and aggregates their output. It is appropriate for large-scale classification and regression tasks due to its improved accuracy and reduction of model overfitting. The SVR is an extension of SVMs on regression tasks. It seeks the hyperplane that maximizes the margin between the data points, best fits the data, and concerns itself with minimizing prediction errors.

KNN predicts target values by averaging the outcomes of the k closest points in the feature space. This method works on the assumption that similar inputs produce similar outputs. DT models divide data into subsets by evaluating feature values. Each node in the tree serves as a decision point, informed by these features, forming branches that lead to leaf nodes where the model makes the final prediction. Although decision trees are straightforward to interpret and visualize, they are particularly susceptible to overfitting, especially when dealing with noisy data. Pruning techniques are frequently applied to enhance generalization and reduce overfitting.

GB is an ensemble technique that builds models linearly, with each new model correcting the errors of the previous model. It iteratively enhances prediction accuracy by minimizing a specified loss function. While powerful and capable of modeling complex relationships, GB is sensitive to overfitting and requires careful tuning of parameters like the learning rate and tree depth. AR combines multiple weak learners, typically shallow decision trees, by focusing on the most challenging data points to predict. It assigns higher weights to misclassified instances, iteratively refining the model. The final prediction is a weighted sum of these weak learners. AdaBoost effectively boosts model performance, especially with weak base models, though it can be sensitive to noisy data.

Predicting water quality parameters like NH₄-N, BOD, TN, and COD is vital for environmental monitoring and management.

• (a) NH₄-N (Ammonium): Excess NH₄-N in water can lead to eutrophication, causing algal blooms that deplete oxygen levels and harm aquatic life.

• (b) BOD: BOD of sewage gives the amount of biologically active organic matter in sewage. Aerobic bacteria flourish in wastewater when sufficient oxygen is available, causing aerobic biological decomposition of sewage until the oxidation is complete. The amount of oxygen consumed in this process is the BOD (Garg 2012). A high BOD level indicates organic pollution, which reduces the available oxygen for fish and other marine organisms to breathe (Ooi et al. 2022).

• (c) TN: Total nitrogen includes all forms of nitrogen (ammonium, nitrate, nitrite, and organic nitrogen). A high level of TN indicates nutrient pollution in water bodies, worsening eutrophication and leading to ‘dead zones’ (Bagherzadeh et al. 2021a).

• (d) COD: The laboratory test determines the COD of wastewater using a strong oxidant like dichromate solution. Theoretical computations of COD are performed only on water solutions prepared with a known amount of a specific organic compound in a laboratory to compare the theoretical and test results (Garg 2012).

These parameters are crucial for understanding and addressing water quality problems, enabling targeted interventions to preserve aquatic ecosystems.

Table 2 summarizes the basic statistical properties of all features in the combined dataset. This research focuses on predicting four water quality parameters, including NH₄-N, BOD, TN, and COD, individually. Each iteration uses one parameter as the dependent or target variable; the rest are independent variables for prediction. For example, in the first iteration, NH₄-N is predicted, and all other features are predictors. This process is repeated for BOD, TN, and COD. This approach iteratively analyzes the relationships between variables and each water quality parameter.

Table 3 shows the ML model's performance metrics for predicting NH₄-N, BOD, TN, and COD. The models are evaluated using two key metrics: R² and RMSE. The results indicate that the ensemble learning models, including GB, RF, and AR, outperformed the rest significantly.

• The GB model indicates the highest overall accuracy with an RMSE value of 61.39 and R² value of 0.75, which shows its efficiency in dealing with complex relationships.

• The RF model is more accurate in TN and COD prediction than the others.

• Lower RMSE values across ensemble models indicate their robustness in predicting water quality parameters.

These results show that ensemble learning has the potential to predict water quality, providing reliable environmental monitoring and management tools.

GB is a highly effective technique in water quality prediction due to its capability to capture complex and nonlinear relationships between input features and target variables. The methodology behind its ensemble technique involves combining several weak learners into one robust predictive model. In achieving better accuracy with successive iterations, it refines errors in previously created models through iteration. GB is particularly well-suited to handling unstructured and noisy data, which are some of the common characteristics of water quality datasets. Effective hyperparameter optimization improves the generalization capability of GB, enhancing its robustness against overfitting. This makes GB particularly suitable for environmental prediction tasks, where input–output relationships are often complex, nonlinear, and subject to high variability.

As shown in Table 3, with ensemble models like GB, better predictions for WWTP water quality parameters would be obtained than solely relying on individual ML methods. The reduction of RMSE and the improved R² values indicate the possibility of including temporal dependencies and non-linearity using the ensemble method. Such capabilities are more valuable in WWTP management, thus making it more straightforward to make informed decisions and properly allocate resources.

Table 4 shows the comparison of the proposed model with the existing model. The R² value of the proposed model looks for a significant increase, indicating that the model provides a better overall fit and demonstrating that the water quality prediction model outperforms the existing model.

The proposed gradient boosting regression (GBR) model outperforms the traditional gradient boosting machine (GBM) model on the Melbourne WWTP dataset because of its enhanced ability to address inherent complexities, such as high-dimensional data, missing values, and noisy measurements. The GBR framework is designed to capture nonlinear relationships and feature interactions critical for wastewater treatment processes. Furthermore, modern boosting techniques natively handle missing data, manage categorical variables efficiently, and incorporate robust loss functions such as Huber loss to minimize the impact of outliers and noise.

Although the proposed model has high predictive power, its generalization to other regions and treatment plants is challenging. The key limitations are as follows:

• 1. Regional and Environmental Variability

Differences in climate, hydrology, land use, and pollution sources between regions may influence the model's performance when applied to datasets from other sites. Lack of representation can reduce the universality of the model.

• 2. Variability in Water Treatment Plant Operations

Different water treatment plants' treatment technologies, protocols, and pollutant characteristics further provide challenges. For instance, differences in sampling frequencies, types of equipment, maintenance schedules, and the nature of treatment (biological versus chemical) can cause differences within the input data and challenge adaptability in the model.

The study's findings have practical implications for water quality management by enabling more proactive, efficient, and data-driven decision-making. Predictive models enhance real-time monitoring and provide early warnings of water quality issues, allowing managers to address pollution events or treatment inefficiencies before they escalate, resulting in optimization of operations in water treatment plants and reducing costs while staying within regulatory standards. Through the identification of trends in such key parameters as NH₄-N, BOD, TN, and COD, these tools can point to pollution sources, guide targeted interventions, and support sustainable agriculture and industry practices. The resilience of water systems against climate variability is enhanced when it can predict water quality under changing conditions, whether seasonal or extreme. Water resource managers can create flexible, scalable, and transparent frameworks to protect ecosystems, keep people healthy, and ensure water resources.

This paper reviews the recent advances in ML for water monitoring, such as forecasting water quality, optimizing resource allocation, and managing water shortages. ML algorithms are promising, but challenges persist in handling data complexity, ensuring model interpretability, and adapting algorithms to diverse geographic and environmental contexts.

This study assessed seven ML algorithms to predict water quality parameters and demonstrated their potential in supporting decision-making and enhancing energy efficiency in WWTP operations. Key findings of this study are:

• The average, maximum, and minimum temperatures and humidity significantly influence the water quality parameters at Melbourne East WWTP.

• The GBR algorithm demonstrated superior predictive performance compared to other algorithms, effectively capturing nonlinear and irregular patterns in the data.

ML models have demonstrated their efficacy in water quality prediction, with accuracy rates often exceeding 70% in various studies, enabling more precise and timely interventions to safeguard water resources.

The future scope of this work involves enhancing the model by integrating real-time data feeds from sensor-coupled systems. This will ensure robust model performance, significantly improving wastewater quality management through accurate, continuous monitoring and proactive intervention. Additionally, explainable AI techniques can enhance the model's transparency and interpretability. In the context of water quality monitoring and prediction, ML algorithms can be optimized to improve practicality, scalability, and reliability.

We extend our sincere gratitude to Dr Vaishnavi Dabir, Principal Consultant at Green Cube LLC, USA, for her invaluable technical assistance and expert guidance.

This work was supported by the Research Support Fund (RSF) of Symbiosis International (Deemed University), Pune, India.

S.P. conceptualized the process, developed the methodology, rendered support in formal analysis, wrote the original draft, P.K. conceptualized the process, developed the methodology, wrote and reviewed and edited the article, support in resources, supervised the work.

All relevant data are available from an online repository or repositories: https://data.mendeley.com/datasets/pprkvz3vbd/1.

The authors declare there is no conflict.