Machine learning framework for energy consumption prediction in wastewater treatment plants combining resampling and...

This study aims to enhance energy consumption prediction in wastewater treatment plants (WWTPs) by developing a robust machine learning framework. To address data imbalance and poor interpretability, the framework integrates three resampling methods with cohort Shapley additive explanations (SHAP). Specifically, stochastic under-sampling with sampling weights to the power of 3 (SUS-3) is combined with a random forest model, which achieves the best performance, yielding a determination coefficient of 0.928 and a root-mean-square error of 4.255 in predicting total energy consumption. In unit-energy-consumption classification, SUS-3 improves accuracy and precision by over 30% and reduces model uncertainty by 35%. Cohort-SHAP analyses reveal stable ammonia-nitrogen and biochemical oxygen demand aeration energy impacts, with season/load-specific energy patterns of chemical oxygen demand, suspended solids, total phosphorus, total nitrogen, and inflow rate, peaking in spring/summer and emphasizing phosphorus in winter. This framework offers a transferable, interpretable solution necessitating adaptive strategies to optimize energy use in WWTPs.