Water Flow Prediction Based on Improved Spatiotemporal Attention Mechanism of Long Short-Term Memory Network

Abstract

The prediction of water plant flow should establish relationships between upstream and downstream hydrological stations, which is crucial for the early detection of flow anomalies. Long Short-Term Memory Networks (LSTMs) have been widely applied in hydrological time series forecasting. However, due to the highly nonlinear and dynamic nature of hydrological time series, as well as the intertwined coupling of data between multiple hydrological stations, the original LSTM models fail to simultaneously consider the spatiotemporal correlations among input sequences for flow prediction. To address this issue, we propose a novel flow prediction method based on the Spatiotemporal Attention LSTM (STA-LSTM) model. This model, based on an encoder–decoder architecture, integrates spatial attention mechanisms in the encoder to adaptively capture hydrological variables relevant to prediction. The decoder combines temporal attention mechanisms to better propagate gradient information and dynamically discover key encoder hidden states from all time steps within a window. Additionally, we construct an extended dataset, which preprocesses meteorological data with forward filling and rainfall encoding, and combines hydrological data from multiple neighboring pumping stations with external meteorological data to enhance the modeling capability of spatiotemporal relationships. In this paper, the actual production data of pumping stations and water plants along the East-to-West Water Diversion Project are taken as examples to verify the effectiveness of the model. Experimental results demonstrate that our STA-LSTM model can better capture spatiotemporal relationships, yielding improved prediction performance with a mean absolute error (MAE) of 3.57, a root mean square error (RMSE) of 4.61, and a mean absolute percentage error (MAPE) of 0.001. Additionally, our model achieved a 3.96% increase in R² compared to the baseline model.

1. Introduction

An essential aspect of urban water resource management [1] is the prediction of pipeline flow in water treatment plants, which plays a crucial role in formulating policies such as water resource allocation, reservoir operation, flood prevention measures, and drought relief strategies. Accurately predicting pipeline flow not only helps optimize the use of water resources but also ensures downstream demand, flood prevention, and promotes water balance in reservoir operation [1,2].

Traditional hydrological forecasting methods rely on researchers’ in-depth understanding of water resource management and hydrological knowledge [3,4]. These methods require combining hydraulic process mathematical models with computer simulation techniques and mathematical equations to numerically simulate water conveyance systems [5,6]. However, such methods are limited by the difficulty of manual modeling and fixed engineering physical models [7] and boundary conditions, resulting in poor usability and accuracy. The introduction of machine learning methods has brought new perspectives to flow prediction tasks [8,9]. Since machine learning models learn the complex relationships of systems based on historical data without relying on prior engineering physical models, they help overcome problems caused by imperfect or inaccurate engineering models [10,11,12]. Farshed Fathian analyzed the stochastic and deterministic components of data time series and constructed a hybrid model RF-SETAR to forecast river flow [13]. Ahmad Khazaee Poul combined methods such as MLR, ANN, ANFIS, and KNN with wavelet transform [14] to improve the accuracy of monthly flow forecasts for the St. Clair River. Although there are some achievements for machine learning methods and their improved algorithms in the field of water flow prediction, their ability to extract and process multidimensional features still needs improvement [15].

The introduction of deep learning methods can effectively utilize multidimensional data. The Recurrent Neural Networks (RNNs) proposed by Zaremba et al. [16] have shown significant advantages in terms of time complexity and accuracy. Zhou et al. [17] improved the RNN with Unscented Kalman Filtering post-processing to predict the inflow sequence of the Three Gorges Reservoir in China. However, the increasing difficulty of data fitting and the increase in the number of network layers often lead to the problems of vanishing gradients and exploding gradients during the optimization process. The emergence of Long Short-Term Memory Networks (LSTMs) [18] effectively alleviates these issues and demonstrates more powerful capabilities in time series analysis tasks. Le’s study in 2019 indicated that the ability of LSTM in watershed hydrological modeling has been confirmed in multiple fields such as river flow and flood prediction [19]. Hu et al. utilized an LSTM deep learning model to predict the runoff data of the Tunxi hydrological station for the next six hours [20]. Compared to Support Vector Regression (SVR) and Multilayer Perceptron (MLP) models, LSTM exhibited superior performance. Although LSTM is very helpful in capturing long-term dependencies, it cannot focus on different variables at different time steps. To overcome these challenges, an attention-based encoder–decoder network was proposed in [21]. It converts the input sequence into a fixed-length vector through the encoder part [22], which is then converted into an output sequence by the decoder.

Although the deep learning method LSTM performs well in hydrological prediction models, the current research still faces two challenges in practical applications [23,24]. Firstly, flow prediction tasks are often not only directly related to upstream flow, but also to various factors such as geographical location, environmental conditions, and weather [25]. Studies have shown that models based on spatial and temporal features have better predictive capabilities [26,27,28,29]. Wang et al. accurately predicted the water level of the Hanjiang River using an LSTM network augmented with the spatiotemporal attention mechanism [30]. Li et al. utilize a spatiotemporal enhanced network for predicting machine Remaining Useful Life [31]. The lack of ability to capture spatiotemporal features in the prediction model limits the accuracy of the flow prediction task and the adaptability to dynamic scenes [32,33]. Secondly, deep learning models often lack intuitive interpretability [34,35]. Although LSTM has good accuracy in predicting flow data over time, the inability to verify the predicted results hinders the practical application of deep learning methods [36,37].

To address the two challenges mentioned earlier and capture key information from multidimensional spatiotemporal input attributes relevant to flow prediction objectives, this paper introduces an improved spatiotemporal attention mechanism on top of the standard LSTM network. Our research innovatively utilizes data from upstream and downstream pumping stations to predict current pump station flows, rather than relying solely on the physical parameters of the current station. By leveraging data from multiple geographic locations, STA-LSTM captures complex spatiotemporal relationships that traditional models often ignore. In addition, the attention mechanism provides a visual representation of the attention score, which can provide a deeper understanding of the model’s decision-making process and further promote the model’s application in the real world.

The architecture of this model consists of two main parts: the encoder and the decoder. The encoder consists of a spatial attention mechanism module based on LSTM units, which can adaptively identify and capture important spatial attributes closely related to the prediction target at each time step. The decoder contains another set of LSTM unit-based time attention mechanism modules, whose main function is to dynamically identify key hidden states related to the prediction target in all time steps within the window. With the help of the attention mechanism, the STA-LSTM model can not only adaptively select input spatial attributes most relevant to the prediction but also detect long-term dependencies in the time series. The method is validated using actual data from the East-to-West Water Diversion Project in Beijing, demonstrating the effectiveness and prediction accuracy of the approach.

The main work of this study can be summarized as follows:

• Integration of LSTM with Spatiotemporal Attention Mechanism: Construct a water plant pipeline flow prediction model that combines Long Short-Term Memory (LSTM) networks with spatiotemporal attention mechanisms, improving the model’s interpretability of flow changes. Achieve the visual representations of spatial and temporal attention weights, and preliminarily explain their relationship with hydrological prediction.
• Fusion of Hydrological Station Parameters and Meteorological Conditions: Innovatively utilize surrounding hydrological station parameters in conjunction with external meteorological conditions for water plant flow prediction and to improve prediction accuracy.

The structure of the article is as follows: Section 2 provides a detailed description of the construction of the dataset, the preprocessing process, and the formulation of flow prediction task objectives. Section 3 specifically introduces the STA-LSTM model, including the design of the backbone LSTM model and the spatiotemporal attention mechanism. Section 4 elaborates on the experimental part, including the settings of the experimental environment, experimental results, and analysis. Finally, Section 5 summarizes the research findings and proposes prospects for future work.

2. Materials

2.1. Study Area

The East-to-West Water Diversion Project is a large-scale water conservancy project designed to address the water supply needs of residents in the western part of Beijing, China. Since its supply began in 1997, the project has delivered over 300 million cubic meters of water to the Chengzi Water Plant in the Mentougou District by the year 2020. This has effectively ensured water safety for the residents and production needs in the Shijingshan and Mentougou districts. The first phase of the project primarily includes the water transmission pipeline from the Tuancheng Lake to the Mentougou Chengzi Water Plant, with a lift exceeding 60 m. After the transformation, the project unified water intake from the Tuancheng Lake, transported through the Yuquanshan, Xingshikou, and Mayu Pump Stations to the Chengzi Water Plant in the western part of Beijing, with a daily water supply capacity of 80,000 cubic meters. Additionally, the project branches off between the Yuquanshan Pump Station and the Xingshikou Pump Station to provide water source security to the newly built water plant in the Shijingshan District. The water transmission capacity of this branch is 0.58 cubic meters per second.

2.2. Research Data

The sensors utilized for data collection in this study are positioned at Yuquanshan, Xingshikou, Mayu Pump Stations, and Shijingshan Water Plant. Figure 1 illustrates the spatial relationships among these four monitoring points: points A, B, and D are sensors on the water transmission line in sequential order, whereas point C monitors the water quantity of the Shijingshan Water Plant, the focal point of the prediction task. Analyzing the spatial relationships among these monitoring points can improve prediction accuracy. Prior research indicates that meteorological factors significantly affect the accuracy of hydrological prediction models. To enhance model accuracy, this study incorporates external meteorological data for a multi-attribute analysis of the target.

The specific data are shown in

Table 1. Units of input attributes.

Input Attribute	Description	Unit
PIPE_PRESS (P)	Pipeline pressure	Mpa
Water Level (H)	Water level height	m
Water Flow (Q)	Water flow rate	m3/h
Rainfall (R)	Local rainfall	mm
Temperature (T)	Local temperatures	°C

, P, H, and Q represent the pipeline pressure, water level, and discharge flow rate of each pump station, respectively. Because the valve opening value of the water plant has little fluctuation in the data and contributes little to the prediction, the prediction is made by combining the upstream and downstream hydrology station correlation information. R denotes the rainfall in meteorological conditions, T represents daily temperature, and the water flow attribute of point C is the prediction target. Other pump stations’ attributes and spatial information are used as the input for water flow prediction.

The sensor data collection started from 26 September 2022, at 0:00:00 AM to 1 August 2023, at 0:00:00 AM, totaling 309 days and including 395,072 samples. During the data collection process, pump station data were sampled at 1 min intervals. However, to consider realistic prediction needs, this study extracted fixed 30 min interval samples from the full dataset. Downsampling is the process of reducing high-frequency data sampling to a lower frequency. Assuming the original data sequence is

$$X=\left[x_1,x_2,x_3,···,x_n\right]$$

(1)

where the $x_i$ represents the sample at the i-th time point in the original data. When conducting 30 min downsampling, the original data are first divided into consecutive 30 min time intervals. Then, the first sample of each time interval is selected as the data point after downsampling. The resulting data sequence is

$$U=\left[u_1,u_2,u_3,···,u_n\right]$$

(2)

The $u_i$ represents the i-th data point after downsampling and $m=\frac{n}{30 }$ is the number of all data points in the new sampled dataset. To quantify the impact of precipitation, the textual descriptions in the weather data are converted into numerical representations. In this study, a coding scheme is applied based on the magnitude of precipitation from small to large: light snow is coded as 1, light rain as 1.5, light to moderate rain as 2.5, thunderstorm as 3, heavy snow as 4, heavy rain as 4, and torrential rain as 6. This approach integrates the textual weather information into the algorithm. Rainfall and temperature data are recorded at two-hour intervals. To maintain consistency with the pump station data, a forward-filling method is employed to process the weather data, ensuring its alignment with the time granularity of the pump station data, resulting in the creation of an extended dataset containing 13,197 records.

The complete dataset is constructed by concatenating the data from the four sensor monitoring points, following the formula

$$\left\{\begin{array}{c}{x}^A=\left(x_{P,1}^A,\text{⋯}{x}_{P,T}^A;x_{H,1}^A,\text{⋯}{x}_{H,T}^A;x_{Q,1}^A,\text{⋯}{x}_{Q,T}^A)\right.\\ x^B=\left(x_{P,1}^B,\text{⋯}{x}_{P,T}^B;x_{H,1}^B,\text{⋯}{x}_{H,T}^B;x_{Q,1}^B,\text{⋯}{x}_{Q,T}^B)\right.\\ x^C=\left(x_{Q,1}^C,\text{⋯}{x}_{Q,T}^C\right.) \\ x^D=\left(x_{P,1}^B,\text{⋯}{x}_{P,T}^B;x_{H,1}^B,\text{⋯}{x}_{H,T}^B;x_{Q,1}^B,\text{⋯}{x}_{Q,T}^B)\right.\\ x^M=\left(x_{T,1}^M,\text{⋯}{x}_{T,T}^M;x_{R,1}^M,\text{⋯}{x}_{R,T}^M)\right. \end{array}\right.$$

(3)

where $x^A$, $x^B$, $x^C$, and $x^D$ represent all data from sensor monitoring points A, B, C, and D, respectively. The datasets $x^A$, $x^B$, $x^D$ contain three attributes, P, H, and Q, with time information marked in collection order as {1, 2,..., T}, corresponding to the period from Beijing Time 26 September 2022, 0:00:00 AM to 1 August 2023, 0:00:00 AM. The dataset $x^C$ represents the flow rate of the Shijingshan Water Plant, which is the forecast target, and contains only attribute Q. Data $x^M$ represent external meteorological data including attributes R and T by time information supplemented using the aforementioned forward-filling method. Symbol “;” denotes the data concatenation operation. Using the aforementioned methods, the extended dataset is constructed as follows:

$$D_{all}=\left\{x^A,x^B,{x^C,x}^D,x^M\right\}$$

(4)

The dataset $D_{all}$ is the extended dataset we constructed, combining weather and hydrological data. To facilitate its application in flow prediction tasks, we conduct a series of data preprocessing steps. We employ an interpolation strategy for handling missing and null values, replacing them with either the median of each respective attribute or interpolation of neighboring valid values. Furthermore, to reduce data fluctuations caused by measurement noise, we apply an improved exponential smoothing filter to remove random fluctuations in the dataset, thereby enhancing the accuracy of subsequent analyses. Finally, to ensure the fair treatment of data at different scales, we normalize the data using Z-scores based on the mean and standard deviation, which ensures that all variables in the dataset have a consistent scale and distribution range for subsequent analysis and modeling, providing a balanced and standardized data foundation for further analysis.

3. Methodology

3.1. Framework

To address the gap in runoff prediction research based on Long Short-Term Memory (LSTM) attention mechanisms and to enhance LSTM’s capability to extract spatiotemporal features from hydrological sequences, this paper proposes an improved STA-LSTM method that combines spatial and temporal attention for predicting water plant flow.

The overall structure of the STA-LSTM model is illustrated in Figure 2, which is based on an encoder–decoder network architecture with LSTM units as the backbone. The encoder includes a spatial attention mechanism module to adaptively identify important variables related to the target (i.e., water flow). Firstly, the spatial attention module calculates the weighted scores for different hydrological variables’ contributions to the target prediction, and the weighted variables are then input into the encoder LSTM network to update the hidden states. The decoder contains a temporal attention mechanism module to dynamically identify important hidden states related to the prediction variable. The hidden vectors at different time steps obtained from the encoder layer are input into the temporal attention mechanism module, which distinguishes crucial information in a weighted manner and forms a context vector for inputting into the decoder. Subsequently, the decoder LSTM network decodes the context vector and predicts the target flow data through a fully connected layer. Details of the STA-LSTM model will be elaborated on later in the paper.

3.2. LSTM

The Long Short-Term Memory (LSTM), proposed by Hochreiter and Schmidhuber [18] and further developed and popularized by Alex Graves [38], is now widely used in the field of machine learning. The basic structure of its recursive module is illustrated in Figure 3a. For intuitive representation, the output of the hidden module is set to be feedback to the input. Therefore, based on the current input $x_{\left(t\right)}$ and the previous hidden state $h_{(\mathrm{t}-1)}$, the hidden state $h_{\left(t\right)}$ and the output $r_{\left(t\right)}$ at step t are obtained. The computation formula of the basic recursive unit is as follows:

$$r_{\left(t\right)}=g\left(V{h}_{\left(t\right)}\right)$$

(5)

$$h_{\left(t\right)}=f\left(U{x}_{\left(t\right)}+W{h}_{\left(t-1\right)}\right)$$

(6)

The functions $f$ and $g$ are nonlinear activation functions. $U$, $W,$ and $V$ are the connection weights, which are shared across different time steps. The network specifies the hidden layer through the designed LSTM module, which performs well in handling the issues of gradient vanishing and exploding. The specific structure of each basic recursive unit in the LSTM network is shown in Figure 3b, which includes three inputs $c_{(t-1)}$, $h_{(t-1)}$, and $x_{\left(t\right)}$, representing the unit state at the previous time step, the hidden state at the previous time step, and the current input, respectively. The LSTM contains three gate controllers: the forget gate, input gate, and output gate, used to determine what information should be forgotten, remembered, and outputted, respectively. Their computation formulas are as follows:

$$f_{\left(t\right)}=\sigma \left(W_{fx}{x}_{\left(t\right)}+W_{fh}{h}_{\left(t-1\right)}+b_f\right)$$

(7)

$$i_{\left(t\right)}=\sigma \left(W_{ix}{\tilde{x}}_{\left(t\right)}+W_{ih}{h}_{\left(t-1\right)}+b_i\right)$$

(8)

$$o_{\left(t\right)}=\sigma \left(W_{ox}{x}_{\left(t\right)}+W_{oh}{h}_{\left(t-1\right)}+b_o\right)$$

(9)

In the subsequent sections, the values computed from spatial and temporal attention serve as inputs to the encoder and decoder-side LSTMs to obtain the prediction results.

3.3. Encoder Decoder with Self-Attention

The encoder–decoder framework is a commonly used framework in deep learning to address sequence-to-sequence modeling problems [39]. By using RNN or other recursive neural networks, the encoder–decoder can effectively capture long-term dependencies in sequences, making it perform exceptionally well in many sequence modeling tasks such as language modeling and time series prediction. Compared to traditional RNN models, LSTM can effectively alleviate the accumulated prediction errors, showing significant advantages in handling data with long-term sequential features. Therefore, we utilize LSTM networks as the backbone network in the encoder–decoder layers of the flow prediction model.

However, as the length of the input sequence increases, long-term information has less impact on the present, limiting the performance of predictions [40,41]. To address this issue, the attention mechanism is introduced into the encoder–decoder framework [42,43]. The structure of the attention-based encoder–decoder model is shown in Figure 4. Different vectors $c_{\left(t\right)}$ are designed instead of fixed ones for $y_{\left(t\right)}$ and can be calculated as

$$c_{\left(t\right)}=\sum _{i=1}^T{w_{i}h}_i$$

(10)

For the prediction at time step t, the attention mechanism assigns separate attention weights $w_{\left(i\right)}$ to each encoded hidden feature state $h_{\left(i\right)}$ based on its relationship with the previous hidden state in the decoder. In this study, the encoder transforms the input sequences, such as meteorological factors, into a sequence of hidden states, which is further transformed into a fixed-length vector and the decoder then utilizes it to predict the target output sequence. The encoder and decoder are trained end-to-end, typically executed in a sliding window manner, where the input and output sequences are progressively processed within a given window. Below, we specifically introduce the spatial attention and temporal attention mechanisms adopted in the encoder and decoder, respectively.

3.4. Spatial Attention with Encoder

The encoder, serving as a feature extractor, represents the input data sequence as a fixed-length encoding vector, capturing the intrinsic correlations among the input auxiliary variable sequences. However, as the input sequence length increases, the encoder can only compute intermediate values to represent the original input sequence and struggles to individually capture the key features related to the prediction target from a large number of input sequences. Therefore, we introduce a spatial attention mechanism at the encoder layer to adaptively acquire the importance of pump station data and weather data for the flow prediction results relative to each geographical location, which is illustrated in Figure 5.

${(x}_t^1,x_t^2\cdots x_t^n)$ represents the given feature at time step t and n is the number of auxiliary prediction variables. We capture the relationship between each input variable and the water flow by utilizing the previous decoder hidden state $s_{\left(t-1\right)}$, which is computed by measuring the similarity between the current raw input $x_{\left(t\right)}$ and the previous hidden state. Generally, a two-layer network structure is employed to extract the spatial attention between different variables. The spatial attention values $e_{\left(t\right)}^i$ for each variable at each time step t are calculated as follows:

$$e_{\left(t\right)}^i=V_1^i\mathrm{L}\mathrm{e}\mathrm{a}\mathrm{k}\mathrm{y}\mathrm{R}\mathrm{e}\mathrm{L}\mathrm{u}\left(W_1^i{s}_{\left(t-1\right)}+U_1^i{x}_{\left(t\right)}^i+b_1^i\right) 1\le i\le n$$

(11)

where $W_1^i$, $U_1^i$, and $b_1^i$ represent the fully connected layer weights and bias terms of a multi-layer perceptron, and LeakyReLU denotes a nonlinear activation function. $e_{\left(t\right)}^i$ signifies the spatial attention for each variable, indicating the importance of the i-th input for flow prediction at time t. To ensure the sum of attention values for all input variables equals 1 at each sampling time t, the values of $e_{\left(t\right)}^i$ are normalized as follows:

$${\alpha }_{\left(t\right)}^i=\frac{\left|e_{\left(t\right)}^i\right|}{{\sum }_{j=1}^n\left|e_{\left(t\right)}^j\right|} 1\le i\le n$$

(12)

The normalized spatial attention weights ${\alpha }_{\left(t\right)}^i$ are utilized to adjust the importance of each feature vector, reflecting its relevance with the water flow prediction target. Therefore, the samples weighted by spatial attention can be represented as follows:

$${\tilde{\mathrm{X}}}_t=\left({\alpha }_t^1{x}_t^1,{\alpha }_t^2{x}_t^2,\cdots ,{\alpha }_t^n{x}_t^n\right)$$

(13)

Then, we feed the attention-weighted input data into the encoder LSTM network to compute the context hidden states, which assist the model in memorizing crucial information to enhance prediction accuracy.

3.5. Temporal Attention with Decoder

Data are sequentially input into the encoder, where information carried by the previous timestep interacts with that of the subsequent timestep to contribute better prediction. However, due to interference from data noise, the accuracy of information carried by different timesteps is unbalanced. This situation may lead to relevant information being obscured by irrelevant data, and this issue becomes more severe with an increase in the length of the input time series. To alleviate this problem, we introduce a module of time attention mechanism at the decoder layer. This module dynamically identifies important hidden states similar to the prediction target across all timesteps, i.e., how the water plant flow at the prediction timestep depends on hydrological information at each timestep, which is shown in Figure 6.

In this study, the time feature represents hydrological data at different time points along the Eastern Route. After processing the input adjusted by spatial attention, the time attention module calculates the task-related weight for each timestep by analyzing the hidden states of the encoder. These weights indicate the importance of each timestep for the current prediction task. For each timestep within a sliding window before timestep t, the time attention of the hidden state is calculated as

$$g_{\left(t\right)}^k=V_2^k\mathrm{t}\mathrm{a}\mathrm{n}\mathrm{h}\left(W_2^k{s}_{\left(t-1\right)}+U_2^k{h}_{\left(t-T+k\right)}+b_2^k\right) 1\le k\le T$$

(14)

where $V_2^k$, $W_2^k$, $U_2^k$, and $b_2^k$ are learnable parameters related to the sample in the window, $s_{\left(t-1\right)}$ is the previous timestep’s decoder hidden states, k is the sub-window size, and tanh is the hyperbolic tangent nonlinear activation function. The temporal attention is then normalized as follows:

$${\beta }_{\left(t\right)}^k=\frac{\left|g_{\left(t\right)}^k\right|}{{\sum }_{m=1}^T\left|g_{\left(t\right)}^m\right|} 1\le k\le T$$

(15)

This ensures that the sum of attention weights for all timesteps is equal to 1. The weighted context vector is calculated using the time attention weights, and this vector serves as the input to the decoder LSTM:

$$c_{\left(t\right)}={\sum }_{k=1}^T{h}_{\left(t-T+k\right)}{\beta }_{\left(t\right)}^k$$

(16)

The time context vector $c_{\left(t\right)}$ varies at each timestep, and is then combined with one of the historical target data $y_{\left(t\right)}$ and input into the LSTM unit to update the decoder hidden state for the next timestep:

$$s_{\left(t\right)}=f_l\left(y_{\left(t\right)},c_{\left(t\right)},{\mathrm{s}}_{\left(t-1\right)}\right)$$

(17)

Here, the function $f_l$ represents the calculations of the LSTM unit, $s_{\left(t\right)}$ represents the updated decoder hidden state at timestep t and $s_{\left(t-1\right)}$ is the previous decoder hidden state. This process iterates for each timestep of the sequence, allowing the model to capture the temporal dependencies for accurate prediction. By iteratively computing this process until $s_{\left(T-1\right)}$, the final LSTM unit in the decoder is able to generate the predicted target flow:

$$\begin{gathered} {\widehat{y}}_{\left(T+1\right)}=F(y_{\left(1\right)},y_{\left(2\right)},\dots ,y_{\left(T\right)},x_{\left(1\right)},x_{\left(2\right)},{\dots ,x}_{\left(T\right)}) \\ =V\left[f_l\left(y_{\left(T\right)},c_{\left(T\right)},{\mathrm{s}}_{\left(T-1\right)}\right)\right]+b_v \end{gathered}$$

(18)

F() is our deep learning network to predict the target, and the data of x and y in a sliding window of the previous 16 h are used to predict the data of the next half hour. We employ the Adam optimizer to train the model parameters, which offers superior computational time and less memory requirements compared to the traditional gradient descent with momentum. The model parameters are learned through standard backpropagation through time to minimize the Mean Squared Error (MSE) loss. The MSE measures the average of the squared errors, which is the average squared difference between the predicted values and the ground truth. The MSE is computed as follows:

$$MSE=\frac{1}{n}\sum _{i=1}^n{\left(\widetilde{y_i}-y_i\right)}^2$$

(19)

where $\widetilde{y_i}$ denotes the predicted flow values, $y_i$ represents the ground truth flow values, and n is the number of test samples.

4. Results and Discussion

4.1. Evaluation Metrics

To evaluate the performance of the model, we use several performance metrics including Mean Absolute Error (MAE), Root Mean Square Error (RMSE), Mean Absolute Percentage Error (MAPE), and the coefficient of determination (R²). MAE is the average of the absolute differences between the predicted values and the actual values. The formula is as follows:

$$MAE=\frac{1}{n}\sum _{i=1}^n\left|\widetilde{y_i}-y_i\right|$$

(20)

where $\widetilde{y_i}$ represents the predicted flow values, and $y_i$ represents the actual flow values. RMSE evaluates the accuracy of the regression results:

$$RMSE=\sqrt{\frac{1}{n}\sum _{i=1}^n{\left(\widetilde{y_i}-y_i\right)}^2}$$

(21)

MAPE is the Mean Absolute Percentage Error:

$$MAPE=\frac{100\mathrm{\%}}{n}\sum _{i=1}^n\left|\frac{\widetilde{y_i}-y_i}{y_i}\right|$$

(22)

$R^2$ represents the correlation between two random variables and is calculated as

$$R^2=1-\frac{{\sum }_{i=1}^n{\left(\widetilde{y_i}-y_i\right)}^2}{{\sum }_{i=1}^n{\left(y_i-\overline{y_i}\right)}^2}$$

(23)

where $\overline{{\mathrm{y}}_{\mathrm{i}}}$ represents the mean of the true value sequence.

4.2. Hyperparameter Analysis

In this section, we analyze the model’s performance in relation to key hyperparameters, including the number of hidden neurons p and q in the encoder and decoder LSTM layers, and the window size T. We tune p and q using a grid search over {40, 60, 80, 100, 120}. Figure 7a illustrates RMSE for different hidden neuron numbers and window sizes. We select 40 hidden neurons to be optimal, which offer sufficient capacity without excessive complexity compared to 120.

The window size is searched in {8, 16, 32, 48, 96, 120}, as shown in Figure 7b. A higher or lower window size leads to a waste of performance and a decrease in accuracy. As sequence length increased, initial prediction errors decreased due to more historical data. However, longer sequences slow performance improvement as LSTM struggled with extensive historical dependencies. We choose a window size of 32 and obtain the best results. We also adjust other hyperparameters, such as neural network layers, to optimize performance and maintain robustness across conditions.

In studying the impact of the attention module on the STA-LSTM model, we specifically investigate the effects of different activation combinations. LeakyReLU is used in the spatial attention part for its ability to handle shallow features concentrated in a small range. Its advantage lies in maintaining gradients even for negative inputs, preventing neuron inactivity, and preserving feature distinctiveness for clearer weight distribution in subsequent Softmax operations. For the time attention part, we choose tanh due to its balanced treatment of positive and negative vector inputs, essential for timestep features. The results in

Table 2. Comparison of model performance using different activation function combinations in STA-LSTM attention mechanisms.

Activation Combination	MAE	RMSE	R2
tanh–LeakyReLU	3.95	5.74	0.974
tanh–ReLU	3.93	5.70	0.976
ReLU–LeakyReLU	3.95	5.74	0.974
LeakyReLU–tanh	3.31	4.37	0.991

show that models using the LeakyReLU–tanh activation combination perform best and demonstrate stable performance. To ensure model robustness, we select the LeakyReLU–tanh combination as the final activation setting.

4.3. Experiment Settings

For all experimental data, 80% of the data are allocated for the training phase, with the remaining 20% used for testing. The model parameters are set with input feature dimensions fixed at 12, processing 128 samples per batch, a learning rate of 0.001, and a maximum training epoch of 200. We define a look-back window of 32 timesteps. We conduct all experiments with five different random seeds and average the results to obtain a fair comparison. At the end of each training epoch, we compute MAE, RMSE, MAPE, and R² on the test set to monitor changes in model performance. All experiments are conducted on a computer equipped with an Intel(R) Xeon(R) CPU i5-13500H and NVIDIA GeForce RTX™ 3050 GPU.

4.4. Results and Analysis

4.4.1. Performance Comparison

We compare the proposed STA-LSTM model with the following baselines:

LSTM (Long Short-Term Memory network): a variant of Recurrent Neural Networks (RNNs) capable of learning long-term dependencies, commonly used in time series prediction.

GRU [44] (Gated Recurrent Unit): a variant of LSTM networks, it enhances computational efficiency by simplifying the temporal unit structure.

CNN (Convolutional Neural Network): a type of neural network with characteristics such as local connections and weight sharing, commonly used in deep feed-forward neural networks.

CNN-LSTM [45] (Combination of Convolutional Neural Network (CNN) and LSTM): a combination where CNN is used to extract spatial features from time series data and LSTM is used for temporal sequence processing.

The experimental parameter settings are shown in

Table 3. Parameter settings of baseline models and our STA-LSTM model.

Models	Parameters
	Learning Rate	Run Units	Activation
LSTM	0.001	128	tanh
GRU	0.001	128	tanh
CNN	0.001	128	tanh
CNN-LSTM	0.001	128	ReLU
STA-LSTM	0.001	p = 60 * q = 60 *	LeakyReLU–tanh

Note: * p and q are the number of hidden neurons in the encoder and decoder LSTM layers, respectively.

. A detailed prediction of water plant flow is shown in Figure 8, and experimental performance results are shown in

Table 4. Comparison of MAE, RMSE, MAPE, and R² for the test set of each model in the test process.

Models	MAE	RMSE	MAPE	R2
LSTM	26.72	35.97	0.009	0.9298
GRU	29.07	36.73	0.010	0.9148
CNN	41.75	53.54	0.015	0.8142
CNN-LSTM	23.73	29.45	0.008	0.9516
STA-LSTM	3.57	4.61	0.001	0.9912

.

Quantitative experiments demonstrate that the STA-LSTM model outperforms other models significantly in key metrics such as MAE, RMSE, and MAPE. Additionally, the R² value approaching 1 indicates its outstanding accuracy and generalization ability. Furthermore, the predictions of each model on the test set are shown in Figure 8. As indicated in Table 4, the STA-LSTM model exhibits the lowest MAE and RMSE, corresponding to prediction curves that closely align with the truth water flow curve, demonstrating higher accuracy compared to all baselines.

4.4.2. Ablation Study

In this section, we examine the influence of meteorological condition variables on prediction accuracy. The results presented in

Table 5. Comparison of Evaluation Metrics for Test Sets of Different Models on Hydrological and Hydro-meteorological Datasets.

Dataset	MAE	RMSE	MAPE	R2
Hydrological data	5.25	7.32	0.0023	0.9668
Hydrological and meteorological data	3.57	4.61	0.0012	0.9912

reveal that models incorporating meteorological data achieve higher prediction accuracy. This observation is consistent with real-world dynamics, where increased rainfall correlates with greater upstream water volume, thereby impacting the flow along the pipeline.

Subsequently, we conduct a systematic analysis by sequentially removing the Spatial Attention (SA) and Temporal Attention (TA) modules from the STA-LSTM model to assess their respective contributions. SA-LSTM integrates a spatial attention mechanism within the encoder to emphasize critical spatial features in the input data. In contrast, TA-LSTM incorporates a temporal attention mechanism within the decoder to assign varying levels of significance to different time steps.

The results, as presented in

Table 6. Comparison of Evaluation Metrics for Test Sets Under Different Attention Mechanisms.

Model	MAE	RMSE	MAPE	R2
TA-LSTM	23.43	32.01	0.0081	0.9428
SA-LSTM	19.56	28.81	0.0069	0.9507
STA-LSTM	3.57	4.61	0.0012	0.9912

, indicate a significant decrease in prediction accuracy when either module is removed. SA-LSTM could not effectively address the issue of memory decay that can occur in long sequences, leading to increased prediction errors. Conversely, TA-LSTM may suffer from limitations in handling multivariate spatial features due to the lack of spatial attention mechanisms.

4.4.3. Visual Attention Analysis

In this section, we visualize the attention weights of the proposed STA-LSTM model and analyze the trend. To better analyze the correlation between pumping station data and water plant flows, a hydrologic dataset without weather is used here. It contains 9-dimensional pumping station data and the variables are in the same order as the full dataset. The attention weight distribution on the hydrological dataset is shown in Figure 9 when predicting the water flow at the current time step t of the water plant. We observe that the first and fifth spatial attributes (i.e., the pressure at point A pump station and the water level in front of point B pump station), as well as the second spatial attribute (i.e., the water level in front of point A pump station), exhibit darker colors than other spatial attributes across most historical timesteps. This implies their higher importance throughout the entire prediction process and thus, they are assigned greater weights.

4.5. Discussion

Among the experimental results mentioned above, STA-LSTM demonstrates higher accuracy compared to other methods. These findings are supported by the data presented in Table 4, where STA-LSTM consistently achieves significantly lower values across all error metrics. This indicates its superiority in capturing complex spatiotemporal relationships and delivering precise predictions. The average absolute error (MAE) for STA-LSTM is 3.57, the root mean square error (RMSE) is 4.61, and the mean absolute percentage error (MAPE) is 0.001, demonstrating enhanced predictive performance. We find that the performance differences among the models mainly stem from their ability to handle spatiotemporal relationships and sequential data. While both GRU and LSTM can maintain long-term dependencies in time series, the simplicity of GRU’s structure may lead to the loss of important information, resulting in a lower performance compared to LSTM. However, neither of these models effectively captures spatial and temporal information, leading to poor predictive performance. SLSTM, due to its simplistic feature extraction and lack of attention mechanisms, shows an inferior performance. Although CNN-LSTM extracts spatial information, it still has limitations in handling complex spatiotemporal data, hence its subpar performance.

Ablation study suggests that both SA and TA modules play vital roles in capturing spatiotemporal dependencies. STA-LSTM proves the crucial significance of its strategy in comprehensively leveraging spatiotemporal information for performance enhancement. The results of the visual experiment analysis are consistent with the logic of the hydrological system. When the opening of Shijingshan gate is certain, point A (Yuquanshan pump Station) is located in the upstream area, and its pipeline pressure directly affects the flow of the downstream Shijingshan gate. The variation of water level in the lower reaches of Xingshikou implies the inherent constraint of the data of hydrographic stations in the upper and lower reaches, which plays a key role in predicting the water quantity of Shijingshan.

Our STA-LSTM model, with its unique spatiotemporal attention mechanism, effectively utilizes both spatial and temporal information, enhancing its ability to capture complex spatiotemporal patterns and demonstrating exceptional performance in the task of predicting water flow.

5. Conclusions

To address the strong nonlinearity and dynamics of hydrological problems in flow prediction tasks, we propose the improved STA-LSTM model. Building upon the LSTM backbone network, our model incorporates a spatial attention mechanism in the encoding layer and introduces a time attention mechanism in the decoding layer. This allows the model to adaptively determine the information relevant to the predicted flow from the input hydrological data, and dynamically distinguish important hidden states related to flow velocity from the output of the encoder.

To validate the effectiveness of the proposed STA-LSTM, we construct a real dataset containing data from three adjacent pumping stations and external meteorological data. Extensive experiments show that the STA-LSTM model outperforms all baselines. Specifically, it achieved reductions of 20.16 in MAE, 24.84 in RMSE, and 0.007 in MAPE, and a 3.96% increase in R² compared to the baseline model. STA-LSTM provides valuable insights for the future application and promotion of deep learning methods in fields such as water supply system management.

The STA-LSTM model proposed in this study demonstrates efficient training and accurate prediction on conventional computing platforms. However, in the context of automated pumping stations, the issue of delayed prediction result generation remains to be addressed. Future research will focus on developing a lightweight multi-step flow prediction model to tackle these challenges.