Abstract: Water pollution issues in China are becoming increasingly prominent, making improvements in the accuracy of water quality prediction models crucial for effective water resource management and ecological protection. This study addresses the challenges of complex nonlinear relationship modeling and computational efficiency in multi-step water quality time series prediction by proposing an improved PatchTST model. The model incorporates three key module optimizations: 1) a lightweight CMixer encoder replacing the traditional Transformer encoder, which efficiently extracts temporal features through one-dimensional convolution and residual connections while reducing computational burden; 2) an Adaptive Mid-Frequency Energy Optimizer (AMEO) that enhances mid-frequency spectral information, improving the model’s ability to detect periodic changes in water quality parameters; and 3) a CKAHead prediction module based on Chebyshev polynomials and the Kolmogorov-Arnold representation theorem, strengthening the modeling of complex nonlinear relationships. In dissolved oxygen prediction at the Shimenzi section, the improved model achieves an MSE reduction of 12.9% compared with PatchTST and 14.0% compared with iTransformer, while maintaining a balance between computational efficiency and resource consumption. Furthermore, in generalization tests across five different monitoring sections, the model reduces MSE by approximately 10% compared with the next-best model for 48-hour forecasting tasks. Experimental results demonstrate that the improved model effectively enhances the accuracy and computational efficiency of multi-step water quality prediction, offering reference value for environmental time series analysis and water quality prediction research.

Key words: water quality prediction, time series forecasting, PatchTST, deep learning, water quality monitoring