Multivariate spatio-temporal modeling of drought prediction using graph neural network

Droughts develop slowly and affect vast regions over time, causing severe social and economic damage. Droughts in China have become more frequent, severe, and widespread in the last six decades due to global climate change. From 1950 to 2015, the cumulative disaster area in Henan Province was 113,796,627 acres, and the average disaster area over the years was 1,724,191 acres. As the main grain-producing area in China, drought not only impacts the economic development of Henan Province but also affects the national food security when the drought is severe. Therefore, we choose Henan Province as the focus of this study to improve the ability of drought prediction and help decision-makers make informed decisions, which is of great significance in reducing agricultural economic loss.

Previous studies have developed various indices to capture the different features of drought. Palmer (1965) developed a drought index called the Palmer drought severity index (PDSI) using precipitation and temperature data, which is still in use today. McKee et al. (1993) developed the standard precipitation index (SPI), and in 2009, the World Meteorological Organization recommended the index to other countries around the world. Currently, the SPI is widely used globally. The standardized precipitation evapotranspiration index (SPEI) was developed by Vicente-Serrano et al. (2010). Compared to SPI, SPEI includes temperature information by introducing potential evapotranspiration. Moreover, its different time scales can characterize different types of drought. SPEI is extensively used in drought assessment and forecasting. According to the National Center for Environmental Information (NCEI), drought classification converts a large amount of drought index into a category representing drought severity. Interpreting the drought category is simpler than interpreting drought index values and can help stakeholders easily understand. Therefore, decision-makers are most concerned with the category of drought rather than the values of the drought indices (Bazrkar & Chu 2022).

Drought prediction research can be broadly categorized into physical models, statistical models, and deep learning models. The physical processes of the land, ocean, and atmosphere are simulated by physical models. However, their accuracy in forecasting precipitation on a monthly or seasonal scale limits their ability to predict drought (Deo & Şahin 2015). Statistical models use various influencing factors as predictors to analyze relationships between historical records. This includes various techniques like regression (Li et al. 2020b), time series analysis (Han et al. 2010), and machine learning approaches (Komasi et al. 2018; Fung et al. 2020; Ma et al. 2022). These models are widely used because of their simplicity in structure, small data requirements, and low computational costs. However, they fail to adapt to the non-stationary nature of drought estimation, and they are prone to overfitting because of the lagged terms in the time series data. Therefore, we should explore the potential of more advanced deep learning methods.

Deep learning models are increasingly used for drought forecasting. One of the most popular models is the recurrent neural network (RNN) (Le et al. 2017), which can handle sequential data. However, RNNs struggle to capture long-range dependencies in time series data, which are important for drought prediction. To overcome this limitation, long short-term memory (LSTM) (Poornima & Pushpalatha 2019) and gated recurrent unit (GRU) (Zhou et al. 2022; Yan et al. 2023) models have been developed, which can retain information over longer periods. However, these models may face the problem of gradient vanishing or gradient explosion. Transformer (Vaswani et al. 2017) solves their limitations in handling long-term dependencies and capturing long sequences. Another type of model that can be useful for time series classification is the convolutional neural network (CNN) (Ham et al. 2019), which can extract features from different dimensions and reduce model complexity. Fully convolutional networks (FCNs) (Wang et al. 2017) are a variant of CNNs that achieve high performance in time series classification tasks. A recent model that combines LSTM and FCN with a Squeeze-and-Excitation Block is MLSTM-FCN (Karim et al. 2019), which can learn the relationships between different features at each time step. However, these models all ignore the spatial structure of large regions, which can affect drought patterns. Graph neural networks (GNNs) (Scarselli et al. 2008) are a class of models that can exploit the spatial structure within a region and have shown success in time series prediction tasks such as crop yield prediction (Fan et al. 2022) and COVID forecasting (Kapoor et al. 2020), traffic flow prediction (Lan et al. 2022). Common graph neural network (GNN) models include graph convolutional network (GCN) (Bhatti et al. 2023), graph attention network (GAT) (Chen et al. 2023), GraphSAGE (Liu et al. 2023), etc. To enhance the performance of these models, the attention mechanism (Li et al. 2020a) was introduced. It is a technique that enables a neural network to selectively attend to the most significant portions of the input and output data, while disregarding less relevant parts. Vaswani et al. (2017) introduced a compelling enhancement to this mechanism known as multi-head self-attention, which allows the network to further refine its focus by attending to multiple subspaces of the input and output sequences simultaneously. It has several benefits over single-head attention, such as capturing various dependencies and facilitating interpretability.

The dataset used in this study is the climatic research unit (CRU) dataset, which was developed by the University of East Anglia. The dataset features a spatial resolution of 0.5° × 0.5° and comprises 11 variables, including cloud cover, frost day frequency, potential evapotranspiration, rainfall, diurnal temperature range, relative humidity, daily mean temperature, monthly average daily maximum temperature, vapor pressure, monthly average daily minimum temperature, and wet day frequency. This dataset has been utilized for a broad range of applications, such as climate variability, agronomic research (Renard & Tilman 2019), and paleo-climatic studies (Nagavciuc et al. 2019). We use ArcMap (Wadwekar & Kapshe 2023) to filter the dataset for Henan Province's meteorological data, based on its boundary vector file. The resulting dataset contains meteorological data for 69 grid points, which is illustrated by the red dots in Figure 2.

In this study, we use the standardized precipitation evapotranspiration index (SPEI) as the drought index for prediction. This index takes into account not only the statistical distribution of precipitation but also potential surface evapotranspiration, providing a more comprehensive reflection of regional drought conditions. The specific calculation process is described in Vicente-Serrano et al. (2010). The SPEI values are available at multiple time scales, including 1, 3, 6, 9, 12, and 24 months, and different time scales can characterize different types of drought. Generally, shorter time scales are typically used to assess meteorological drought, while medium time scales are commonly utilized to evaluate agricultural drought. Longer time scales are more appropriate for describing hydrological drought. The CRU dataset provides SPEI values at different time scales globally. The drought categories, which are divided into drought categories according to the grades of meteorological drought¹, are presented in Table 1.

In drought prediction, we denote each grid point's by and ground-truth drought category by ⁠, where c, t represent grid point and month, respectively. Each contains 11 weather features, detailed symbol explanations are shown in Table 2.

Table 2

List of symbols

Symbols .	Meanings .
	the c-th grid point at time t,
A	symmetric adjacency matrix
N,M	number of grid points, number of samples
	the aggregated embedding of c-th grid point at time t
	GNN embedding of c-th grid point at time t
⁠,⁠,⁠,	the values of the reset gate, the values of the update gate, the current candidate state, the current hidden state
	the activation function
,,	the learnable weight matrices for the i-th attention head
Q,K,V	query matrix, key matrix, value matrix
⁠,⁠,	the real drought category of the c-th grid point at time t, the final output of the TEM module, the predicted distribution of the c-th grid point at t-th month
GNN	Graph neural network
GRU	Gated recurrent unit
SPEI	Standardized precipitation evapotranspiration index
CRU	Climate research unit
LSTM	Long short-term memory
FCN	Fully convolutional network
CNN	Convolutional neural network
ROC-AUC	Receiver operating characteristic – area under the curve

Symbols .	Meanings .
	the c-th grid point at time t,
A	symmetric adjacency matrix
N,M	number of grid points, number of samples
	the aggregated embedding of c-th grid point at time t
	GNN embedding of c-th grid point at time t
⁠,⁠,⁠,	the values of the reset gate, the values of the update gate, the current candidate state, the current hidden state
	the activation function
,,	the learnable weight matrices for the i-th attention head
Q,K,V	query matrix, key matrix, value matrix
⁠,⁠,	the real drought category of the c-th grid point at time t, the final output of the TEM module, the predicted distribution of the c-th grid point at t-th month
GNN	Graph neural network
GRU	Gated recurrent unit
SPEI	Standardized precipitation evapotranspiration index
CRU	Climate research unit
LSTM	Long short-term memory
FCN	Fully convolutional network
CNN	Convolutional neural network
ROC-AUC	Receiver operating characteristic – area under the curve

View Large

As shown in Figure 1, it can be seen that geographically adjacent grid points have a strong correlation in drought. Intuitively, if some grid points have experienced severe drought, nearby grid points tend to have similar situations. Incorporating relevant information from the neighboring grid points can potentially enhance the accuracy of the prediction if appropriately integrated. Previous studies have used convolutional neural networks (CNNs) to extract spatial features from structured images, but they are not suitable for our problem, where the grid locations form an unstructured graph with irregular node arrangements. Graph neural networks (GNNs) are a recent class of neural networks designed to handle complex dependencies that exist within graph-structured data sources. With GNNs, there is greater flexibility and a broader range of representation space available to encode node and edge information from the graph, thereby facilitating more effective inference. Formally, a graph can be represented as G=(V, E), where V denotes the collection of nodes and E represents the connections between them. In this drought prediction task, each node is a grid point. E is represented as a symmetric adjacency matrix⁠, where if two grid points adjacent and otherwise. N is the total number of grid points. For each month, there is an associated value of for each node.

GraphSAGE is a widely used GNN model that employs node feature information to learn node embeddings via neighborhood aggregation. Unlike other techniques that rely on matrix factorization and normalization, GraphSAGE simply aggregates features from a node's local neighborhood, resulting in lower computational requirements. The model is highly adaptable since it can use features from different numbers of hops or search depths, leading to better generalization. As the adjacency matrix is sparse due to the fact that a majority of grid points have only a few neighboring points, GraphSAGE is a suitable approach for drought prediction.

To analyze drought patterns and cycles, it is essential to extract temporal features from historical knowledge. We propose a TEM that captures the temporal dynamics of the output embedding from GNN. This module consists of two components: (1) GRU models long-term dependencies and trends over time and (2) multi-head self-attention mechanism enhances its effect by assigning a weight to each month.

In this task, the meteorological characteristics of the input months have varying degrees of influence on the prediction results. For instance, the impact of adjacent months' meteorological factors on the prediction results is likely more significant than that of non-adjacent months. Therefore, by incorporating a multi-head self-attention mechanism, the model can focus on the characteristics of important moments.

The output of each attention head is concatenated and transformed linearly again to obtain the final output of the multi-head self-attention mechanism. This allows the model to attend to different aspects of the input in parallel and capture different types of information.

To evaluate the predictive performance of each model, we utilize four evaluation metrics: precision, recall, F1 score, and accuracy. Additionally, we employ a multi-class receiver operating characteristic–area under the curve (ROC-AUC) for each classifier, providing a better illustration of their performance.

For binary classification, we use a 2 × 2 confusion matrix to show the prediction results of a classifier, and Table 3 displays true positive (TP), true negative (TN), false negative (FN), and false positive (FP).

Given that this study involves multiple categories, the performance metrics can be computed independently for each category. Specifically, samples from a specific category are treated as positive, while the remaining categories are considered negative. The final result is obtained by averaging the metrics across all categories. To account for the imbalanced distribution of the dataset, a macro approach is utilized to calculate the average, which does not take into account the proportion of each category in the dataset. This approach results in a greater penalty when the model performs poorly in minority categories. All results presented in this study are macro-averaged.

The performance of a classifier is typically evaluated using the AUC-ROC curve, which takes into account the true positive rate (TPR) and false positive rate (FPR) at various thresholds. A higher AUC–ROC value corresponds to better classification performance, as it indicates a superior trade-off between TPR and FPR. Therefore, this curve is an important metric for evaluating the performance of a classifier and can aid in selecting the best classifier and optimizing its parameters.

To normalize the data and reduce errors from data magnitude differences, we use min–max scaling to map the values to the [0, 1] range. Then, we divide the data into training and testing sets: 1961–2014 for training and 2015–2018 for testing. We randomly select 20% of the data from the training set as a validation set.

To prevent overfitting, a dropout layer with a 20% rate is incorporated into the model. The Adam optimization algorithm (Kingma & Ba 2015) is utilized to iteratively update the weight parameter matrix, with a learning rate of 0.0001. We apply early stopping with a patience of 10 to stop training when the validation loss does not decrease for 10 consecutive rounds but increases instead. After multiple tests, we found that a batch size of 32 is optimal for the model. We set the time step between 5 and 15 and choose the best time step based on the prediction results of each model at each time scale. Through experiments, we set the timestep of 1, 3, 6, 12, and 24 months to 5, 8, 10, 12, and 15, respectively, that is, we use the time steps of the past 5, 8, 10, 12, and 15 months of data to predict the drought category for the next month at each scale.

To practically validate the prediction accuracy of MSTSN for drought category prediction, we compare MSTSN with classical deep learning models, including LSTM, fully convolutional network (FCN), MLSTM-FCN, and GNN-RNN. These four models are briefly introduced as follows:

LSTM (Dikshit et al. 2021): a variant of recurrent neural networks is designed to capture long-term dependencies in sequential data.

FCN (Wang et al. 2017): a neural network architecture that relies on convolutional layers to extract features from the input time series data. Subsequently, global pooling and fully connected layers are utilized for classification.

MLSTM-FCN (Karim et al. 2019): a hybrid architecture that integrates both LSTM and FCN designs to effectively capture both short-term and long-term temporal dependencies within time series data.

GNN-RNN (Fan et al. 2022): the combination of GNN and RNN is a powerful technique that synergistically blends spatio-temporal information, enabling accurate predictions that incorporate both geospatial and temporal dependencies.

To ensure a fairer comparison and demonstrate the superiority of our model, we selected the best-performing hyperparameters for each model. Table 4 shows the chosen hyperparameters for MSTSN. The detailed results of the model are presented in Table 5, which displays the performance of five models: LSTM, FCN, MLSTM-FCN, GNN-RNN, and MSTSN in predicting drought categories at different time scales.

From the perspective of accuracy, MSTSN performs the best at all five time scales, followed by GNN-RNN, LSTM, and FCN, while MLSTM-FCN has the lowest accuracy. For example, at the 6-month time scale, MSTSN's accuracy is 15% higher than MLSTM-FCN's. The use of graph neural networks to extract spatial characteristics, as seen in MSTSN and GNN-RNN, proved to be more accurate than methods that only consider temporal characteristics. The accuracy rate of MSTSN is higher than that of GNN-RNN, which can be attributed to the effectiveness of the TEM module. This module uses GRU to extract the long-term dependencies of data and enhance its effect by assigning a weight to each month through multi-head self-attention. From the perspective of F1 score, our model scores the highest, except for the FCN model, which performs slightly better than MSTSN at the 3-month time scale. A high macro-F1 score means the model can classify drought categories accurately and fairly without bias toward the majority class. Therefore, MSTSN is superior to other methods, particularly for small-scale, high-intensity droughts.

From the overall perspective, the five models perform worst at the 1-month time scale. However, as the time scale increased, the predictive accuracy of the five models improved. This improvement may be attributed to the reduction in data and the tendency of data sequences to become more stable with increased time scales, leading to better predictive performance.

To evaluate the contributions of the SAM module and the components of the TEM Module: GRU layer, and multi-head self-attention layer to our model's performance on drought prediction, we perform ablation experiments on five time scales. Table 6 shows the average evaluation metrics for each ablation setting. ‘MSTSN-SAM’ means removing the SAM module, while ‘MSTSN-TEM-GRU’ and ‘MSTSN-TEM-ATT’ mean removing the corresponding layer from the TEM block.

Table 6 shows that MSTSN outperforms the model without the SAM module, which confirms the correlation of drought conditions among adjacent grid points and the effectiveness of graph neural networks for spatial feature extraction. By handling long-term dependencies and capturing the temporal relationship with the GRU layer, the model can better predict the future drought development trend based on the historical drought situation, which further improves the accuracy by 2.6%. Since the data of each month have different impacts on drought, the multi-head self-attention layer assigns a weight to each month, allowing the model to focus on key months, which improves the accuracy by 1.8%. Generally, SAM has the greatest impact on the model. Removing the SAM module reduces the model's accuracy by 3.8%, which shows that introducing graph neural networks has an important influence on drought prediction.

From Figure 8, we can see that MSTSN has the closest prediction to the actual spatial distribution of drought and can predict the approximate time and range of drought, followed by GNN-RNN. LSTM fails to predict different categories of drought and predicts all drought categories as 0, which proves that it cannot predict categories with small sample sizes well. FCN and MLSTM-FCN have lower prediction accuracy, especially in predicting the first half of 2015 and 2018. When small-scale drought occurs, all five models have some deviations.

In this study, we have developed a novel MSTSN for predicting drought categories. The model consists of two modules: the SAM and the TEM. The SAM module aggregates information from neighboring grid points to incorporate spatial information and generate a new feature matrix. This matrix is then inputted into the TEM module, which extracts temporal features. Finally, the extracted features are fed into a softmax layer to obtain the final drought category predictions.

We compare MSTSN with common deep learning models and find that MSTSN has the highest prediction accuracy and F1 score at most time scales, indicating its ability to predict small-scale and high-intensity droughts. Additionally, we discover a positive correlation between the drought category and the SPEI time scale. Ablation experiments are conducted to analyze the contributions of each module. The results show that the SAM module improves the model's performance by leveraging graph neural networks to extract spatial features. The TEM module captures long-term dependencies through GRU and enhances its effectiveness by incorporating multi-head self-attention with weighted monthly assignments.

Although this model demonstrates improved accuracy in drought prediction by considering the correlation and temporal dependence between adjacent regions, it has some limitations. The study solely relies on distance to determine adjacency between grid points, disregarding important geographical features. Additionally, the GRU model used struggles to capture long-term temporal dependencies, leading to prediction errors in long-term tasks and an inability to accurately reflect real change trends.

In summary, this study proposes a novel drought category prediction model, MSTSN, and showcases its effectiveness and superiority through extensive experiments. By introducing a GNN and considering geographical knowledge, the model surpasses the limitations of treating each area as independent, resulting in improved prediction accuracy. This has significant implications for industries, such as agriculture, water resource management, and climate monitoring, providing more accurate drought monitoring and forecasting information for decision-makers.

However, future research should incorporate additional factors that influence drought, such as human factors, terrain, and other natural factors. Moreover, geographic features can be utilized to determine adjacency between grid points more accurately. To predict long-term drought conditions, advanced deep learning methods like transformers can be explored to capture temporal dependency relationships effectively.

This work was supported in part by the National Key Research and Development Program of China (No.2021YFE014400). This work was supported in part by the National Natural Science Foundation of China (No.62102187). This work was supported in part by the Science and Technology Development Fund of Egypt (No.43088).

Name of the code/library: MSTSN

Contact: e-mail and phone number

Hardware requirements: CPU: Intel Xeon Gold 6330; GPU: NVIDIA GeForce RTX A5000; Memory: 30 G.

Program language: Python 3.7

Software required: Pycharm, Anaconda3, Pandas, Numpy

Program size: 37KB

The source codes are available for download at the link: https://github.com/nuist-yjx/MSTSN.

All relevant data are available from an online repository or repositories: https://crudata.uea.ac.uk/cru/data/hrg/.

The authors declare there is no conflict.