Simulation and Driving Factor Analysis of Satellite-Observed Terrestrial Water Storage Anomaly in the Pearl River Basin Using Deep Learning

Abstract

Accurate estimation of terrestrial water storage (TWS) and understanding its driving factors are crucial for effective hydrological assessment and water resource management. The launches of the Gravity Recovery and Climate Experiment (GRACE) satellites and their successor, GRACE Follow-On (GRACE-FO), combined with deep learning algorithms, have opened new avenues for such investigations. In this study, we employed a long short-term memory (LSTM) neural network model to simulate TWS anomaly (TWSA) in the Pearl River Basin (PRB) from 2003 to 2020, using precipitation, temperature, runoff, evapotranspiration, and leaf area index (LAI) data. The performance of the LSTM model was rigorously evaluated, achieving a high average correlation coefficient (r) of 0.967 and an average Nash–Sutcliffe efficiency (NSE) coefficient of 0.912 on the testing set. To unravel the relative importance of each driving factor and assess the impact of different lead times, we employed the SHapley Additive exPlanations (SHAP) method. Our results revealed that precipitation exerted the most significant influence on TWSA in the PRB, with a one-month lead time exhibiting the greatest impact. Evapotranspiration, runoff, temperature, and LAI also played important roles, with interactive effects among these factors. Moreover, we observed an accumulation effect of precipitation and evapotranspiration on TWSA, particularly with shorter lead times. Overall, the SHAP method provides an alternative approach for the quantitative analysis of natural driving factors at the basin scale, shedding light on the natural dominant influences on TWSA in the PRB. The combination of satellite observations and deep learning techniques holds promise for advancing our understanding of TWS dynamics and enhancing water resource management strategies.

1. Introduction

Terrestrial water storage (TWS) refers to the total amount of water stored on or within the land surface, mainly including surface runoff, soil moisture, groundwater, snow, and water in lakes and reservoirs. It serves as a comprehensive indicator of the global hydrological cycle [1] and plays an important role in measuring water resource availability [2]. Furthermore, TWS reflects vegetation growth [3], climate change [4], and drought [5]. With an increasing water demand and the more frequent occurrence of drought and flood disasters under climate change [6], investigating the major drivers of TWS and quantifying their relative contributions within the Pearl River Basin (PRB) are essential for understanding extreme hydrological disasters, food security, and water resource management in the PRB [7].

The analysis of driving factors has been challenging in the past due to the difficulty in obtaining TWS measurements with high accuracy. Recently, a number of other optical and microwave-based satellite remote sensing techniques have been proposed to detect individual components of terrestrial water reserves. For example, Soil Moisture Active and Passive (SMAP) satellites and Fengyun-3 satellites are used to monitor soil water, Water Cycle Observation Mission (WCOM) satellites are used to detect snow water content, Google Earth Engine (GEE) is used to detect water volume change information, and Surface Water and Ocean Topography (SWOT) microsensors are used to monitor surface and groundwater. The launches of the Gravity Recovery and Climate Experiment (GRACE) satellites in 2002 and the GRACE Follow-On (GRACE-FO) in 2018 have revolutionized our ability to understand and monitor global water resources. Unlike other satellite sensor products, GRACE, which utilizes remote sensing technology, has delivered invaluable and highly accurate observations of the global total water storage anomalies with all components [8], enabling us to assess changes in water availability on a large scale. GRACE’s data have been instrumental in studying hydrological processes, tracking water availability, detecting droughts, and understanding the impacts of climate change on water resources. Its contributions have significantly advanced our knowledge and have become a cornerstone for effective water resource management strategies worldwide.

Various approaches have been employed to analyze the driving factors of TWS anomaly (TWSA), including land surface models, hydrological models, and statistical analysis. Pokhrel et al. [6] used ensemble hydrological simulations and found that the projected changes in global TWS are driven primarily by climate forcing rather than land and water management activities. Deng et al. [9] calculated the climate resilience coefficient to analyze the impact of climate factors on hydrological systems and concluded that in the process of water storage reduction in 2012, potential evapotranspiration accounted for 90%, precipitation accounted for 7%, and surface properties accounted for 3%. But due to the uncertainty of the climate forcing data and the absence of relevant process modules (e.g., human activity and vegetation dynamics), TWS simulated by land surface models and hydrological models—only a subset of TWS components—often does not match the data observed by GRACE [5,6]. Additionally, the PRB exhibits a highly heterogeneous geological environment and intense human activities. These factors pose significant challenges for accurately estimating TWS using hydrological models [5,10,11]. Furthermore, as the driving factors are interdependent, solely relying on linear relationships to investigate the primary driving factors of TWS may lead to misleading results.

Fortunately, the rapid advancement of deep learning techniques provides a viable solution to the aforementioned challenges in hydrology and related fields [12,13,14,15]. Deep learning modeling, a nonparametric technique, proves effective in approximating the intricate nonlinear relationship between input and output variables, devoid of any prior assumptions regarding the underlying physical system [16,17,18]. In fact, deep learning models have gained widespread adoption in hydrological research. On the one hand, these models are used to improve the simulation and prediction of hydrological variables, thereby improving the quality of data products [19,20,21]. On the other hand, they are gradually being used for hydrological process understanding and data mining [22,23,24,25]. Among the various deep learning models, the LSTM model has been well used in hydrological time series simulation and forecasting research, including flood forecasting, water temperature forecasting, and soil water data extension [17,19,24]. Further, explainable machine learning methods have been developed to elucidate patterns captured by machine learning and deep learning models. Techniques such as accumulated local effects (ALE) [22], global to local analysis (local interpretable model-agnostic explanations, LIME) [26,27], and SHapley Additive exPlanations (SHAP) techniques [28,29] have been widely used in the field of hydrology and meteorology. Particularly, the SHAP method is a novel explainable machine learning method that provides an intuitive understanding of the relationship between the model’s input variables and the predicted results [30,31].

So far, there have been a large number of studies using deep learning models to reconstruct TWSA sequences and achieve favorable results [1,18,32]. However, few studies have focused on the driving factor analysis of TWS with deep learning methods. Some studies have applied soil moisture, which is a component of TWS and highly correlated with TWS, as an input factor for models to achieve better simulation results, which hinders the examination of causal relationships among drivers. In this study, we avoid using direct TWS components as input factors. The SHAP method is combined with an LSTM model to uncover patterns of TWS captured by the LSTM network.

Therefore, the primary objectives of this study are as follows: (1) Employ GRACE data, meteorological data, and underlying surface data to train an LSTM model for simulating the dynamics of TWS in the PRB and evaluate its performance; (2) quantify the relative importance of driving factors of TWS in the PRB using an explainable machine learning framework; and (3) analyze the influencing mechanism of TWS driving factors in the PRB.

2. Materials and Methods

2.1. Study Area

The PRB encompasses a geographical area between the longitudes 102°14′E–115°53′E and the latitudes 21°31′N–26°49′N (Figure 1). With a basin area of about 450,000 km² and an annual water discharge of 3.36 × 10¹¹ m³, the Pearl River is China’s second largest river by discharge volume and third largest river in terms of basin area [33]. The PRB stretches across the Yunnan–Guizhou Plateau and the Guangdong–Guangxi Hills, primarily situated in the subtropical and tropical monsoon climate regions characterized by diverse landscapes and abundant precipitation. Notably, from April to September, precipitation during these months accounts for 80% of the total annual precipitation [34]. The Pearl River Delta region is densely populated, exhibits high water consumption, and experiences frequent droughts and floods [35,36].

2.2. Data Sources

2.2.1. GRACE Data

The GRACE data used in this study were obtained from the NASA Goddard Space Flight Center (GSFC). The estimation of TWS was derived using the mass concentration blocks (mascons) method [37]. This method facilitates the implementation of geophysical constraints and effectively filters noise from GRACE/GRACE-FO observations, providing a more rigorous alternative to the standard spherical harmonic method [38]. The TWSA values in this data product are relative to the average from 2004 to 2010, with a spatial resolution of 0.5° and a temporal resolution of one month. Our analysis focuses on data spanning from April 2002 to December 2020. However, due to technical issues with satellite equipment, there are gaps in the data for several months, and these periods of missing data are directly excluded from the analysis.

2.2.2. Input Data

In this study, the analysis of driving factors was carried out by excluding TWS components such as soil moisture to ensure the integrity of the analysis. However, it is worth noting that explainable machine learning methods are most effective when the model’s performance is satisfactory [39]. Therefore, to ensure the performance of the LSTM model and the validity of the driving factor analysis, five driving factors including precipitation, evapotranspiration, temperature, runoff, and LAI were considered as input features.

The precipitation and mean temperature data used in this study were obtained from the widely used Climate Research Unit (CRU) Time Series Climate dataset, specifically version 4.03, released by the University of East Anglia. This dataset employs an angle–distance weighting method to interpolate in situ climate data [40]. It provides a spatial resolution of 0.5° and a monthly temporal resolution. Additionally, the monthly actual evapotranspiration (AET) data were extracted from the Global Land Evaporation Amsterdam Model (GLEAM), which incorporates satellite observations of climate and environmental variables [41,42].

Furthermore, the influence of runoff and leaf area index (LAI) was taken into account in this study. Since we needed gridded format data, which are not available for runoff from in situ measurement, we utilized the European Centre for Medium-range Weather Forecasts ReAnalysis 5-Land (ERA5-Land) reanalysis dataset, which provides assimilated runoff data. This dataset combines model data with global observations and is considered one of the most state-of-the-art datasets in land surface applications. It has a spatial resolution of 0.25° and a monthly temporal resolution.

LAI data can be downloaded from the Land-Atmosphere Interaction Research Group at Sun Yat-sen University (SYSU). This dataset is an improved version derived from the post-processing of the MODerate-resolution Imaging Spectroradiometer (MODIS) version 6 LAI product. It utilizes a two-step integration method, resulting in enhanced data quality compared to the raw MODIS data [43,44]. This ensures a continuous and consistent dataset for analysis [44]. The LAI dataset has a spatial resolution of 0.5° and a temporal resolution of monthly intervals.

This study selected the above data from April 2002 to December 2020 in the PRB. To ensure consistency, the data were unified into the grid of the GRACE data with the nearest neighbor interpolation method. More details are listed in

Table 1. Description of datasets used.

Data	Data Sources	Spatial Resolution (°)	Temporal Resolution	Date
Precipitation	CRU	0.5	Monthly	April 2002–December 2020
Temperature	CRU	0.5	Monthly	April 2002–December 2020
AET	GLEAM	0.25	Monthly	April 2002–December 2020
Runoff	ERA5-Land	0.25	Monthly	April 2002–December 2020
LAI	SYSU	0.5	Monthly	April 2002–December 2020
TWSA	NASA GSFC	0.5	Monthly	April 2002–December 2020

.

2.3. Methods

2.3.1. LSTM Model

We used LSTM models (Figure 2b) to simulate the TWSA dynamics at the setting time scale in this study. As a special recurrent neural network (RNN), LSTM is a typical sequence learning model [45]. Its notable advantage lies in its ability to capture information from distant time steps, effectively addressing the challenge of vanishing gradients commonly encountered in standard RNNs. This is achieved through LSTM’s incorporation of two types of states and a specialized gated structure (Figure 2a). During the training process, the RNN updates the weights in the recurrent unit through the backpropagation-through-time algorithm [46,47]. When the sequence is long, the weight-updated signal decays rapidly during propagation along the chain, preventing the model from learning information from distant time steps. This phenomenon is known as vanishing gradients.

Under the premise of considering the effect of the model and applying the explainable machine learning method, we designed the model for this study as shown in Figure 2b. The model architecture consists of an input layer, an LSTM layer, a hidden layer, and an output layer. The LSTM layer and the dense layer have 256 neurons each, while the output layer has 1 neuron. To ensure that each simulation of the TWSA value incorporates information of the same length, we utilized only the output from the last time step of the LSTM layer. We then mapped the output from 256 hidden cells to 1 output cell through the fully connected layer.

Previous studies have shown that there is a time lag of one to two months between climate forcing data and the TWSA observed by GRACE satellites [46]. At the same time, in order to verify whether the LSTM model can effectively capture information from the previous time step, we set the LSTM time step to 11, that is, the precipitation, air temperature, evapotranspiration, runoff, and LAI in the current month and the simulated values from the preceding 10 months were used as input series to simulate the TWSA of the current month (Figure 2c). The model can be expressed by the following formula:

$$y_t=f\left(x_t, x_{t-1}, \dots ,x_{t-10}\right)$$

(1)

where $y_t$ refers to the simulation of TWSA in month $t$; $x_t$ refers to the input variable in month $t$; $x_{t-1}, \dots , x_{t-10}$ is the input vector for the corresponding months from $t-1$ to $t-10$.

To fetch the best model, we employed a strategy to increase the number of samples by treating each time grid point in the basin as an individual sample. We utilized the five-fold cross-validation method to train the models and select the optimal one. The dataset was divided evenly into five consecutive folds in periods from 2002 to 2020 (

Table 2. Date of training and testing sets in five groups.

Group ID	Training Set	Testing Set
1	October 2005–December 2020	April 2002–September 2005
2	April 2002–September 2005, April 2009–December 2020	October 2005–March 2009
3	April 2002–March 2009, October 2012–December 2020	April 2009–September 2012
4	April 2002–September 2012, April 2016–December 2020	October 2012–March 2016
5	April 2002–March 2016	April 2016–December 2020

). During each iteration, one of the five folds was designated as the testing set, while the remaining four folds were used to train the model. This process resulted in the generation of five different models. For each model, we used the Adam optimization algorithm [48] to continuously update the model parameters to gradually reduce the root mean square error of the TWSA simulated values and the observed values of the training set. When the Nash–Sutcliffe efficiency (NSE) coefficient between the simulated values and the observed values of TWSA of the testing set no longer grew, the model corresponding to the maximum NSE coefficient was chosen as the model of the experiment in this group. To prevent overfitting, we applied dropout regularization [49] with an empirical setting of 0.5. Finally, we reconstructed the TWSA sequences from 2002 to 2020 using five different models and utilized the best model to analyze TWSA’s response to precipitation, temperature, evapotranspiration, runoff, and vegetation.

2.3.2. Model Evaluation Metrics

In this study, the correlation coefficient (r), normalized root mean square error (NRMSE), NSE, and bias are used to evaluate the simulation performance of the model. These evaluation indices can be calculated using the following formulas:

$$r=\frac{{{\displaystyle \sum }}_{i=1}^n\left(y_{si}-\overline{y_{si}}\right)\left(y_{oi}-\overline{y_{oi}}\right)}{\sqrt{{{\displaystyle \sum }}_{i=1}^n{\left(y_{si}-\overline{y_{si}}\right)}^2\times {{\displaystyle \sum }}_{i=1}^n{\left(y_{oi}-\overline{y_{oi}}\right)}^2}}$$

(2)

$$NRMSE=\sqrt{\frac{{{\displaystyle \sum }}_{i=1}^n{\left(y_{si}-y_{oi}\right)}^2}{{{\displaystyle \sum }}_{i=1}^n{\left(y_{oi}\right)}^2}}$$

(3)

$$NSE=1-\frac{{{\displaystyle \sum }}_{i=1}^n{\left(y_{si}-y_{oi}\right)}^2}{{{\displaystyle \sum }}_{i=1}^n{\left(y_{oi}-\overline{y_{oi}}\right)}^2}$$

(4)

$$bias=\frac{{{\displaystyle \sum }}_{i=1}^n\left(y_{si}-y_{oi}\right)}{{{\displaystyle \sum }}_{i=1}^n{y}_{oi}}$$

(5)

where $n$ denotes the number of samples; $y_{si}$, $y_{oi}$ are the simulated values and GRACE observed values, respectively. The closer the bias and NRMSE are to 0, the better the model performance, and the closer the correlation coefficient and the NSE are to 1, the better the model performance.

2.3.3. Interpretation Methods

SHAP is a game theory-based machine learning interpretation technique that unifies common additive feature attribution methods with better computational performance and better alignment with human intuition [30,50].

Additive feature attribution methods assign an effect ${\varphi }_i$ to each feature, and by summing the effects of all feature attributions, we can approximate the output $f\left(x\right)$ of the original model. The process can be expressed using the following formula:

$$f\left(x\right)=g\left(x^{\prime }\right)={\varphi }_0+{\displaystyle \sum }_{i=1}^M{\varphi }_i{x}_i^{\prime }$$

(6)

$${\varphi }_i={\displaystyle \sum }_{S\subseteq \left\{N/\left\{i\right\}\right\}}\frac{\left|S\right|!\left(M-\left|S\right|-1\right)!}{M!}\left(f_x\left(S\cup \left\{i\right\}\right)-f_x\left(S\right)\right)$$

(7)

where $x^{\prime }\in {\left\{0, 1\right\}}^M$, $M$ is the number of simplified input features, $S$ is the set of nonzero indexes, and ${\varphi }_i\in R$; $f\left(x\right)$ is the output of the original model; $h\left(x\right)$ is a mapping function which explanation models used to turn original inputs $x$ into simplified inputs $x^{\prime }$ through $x=h_x\left(x^{\prime }\right)$; the explanation model function $g\left(x^{\prime }\right)$ matches the original model $f\left(x\right)$ when $x=h_x\left(x^{\prime }\right)$; ${\varphi }_0=f\left(h_x\left(0\right)\right)$ represents the model output with all simplified inputs toggled off.

SHAP values serve as a comprehensive measure of feature importance and provide a solution to Equation (7). They represent the Shapley values of a conditional expectation function of the original model. The SHAP values indicate the changes in the model’s predicted values when considering a specific feature. They explain how the features contribute to the predicted values compared to the baseline value without any feature.

DeepSHAP is a computational method designed to approximate the SHAP values efficiently. It is an improvement on the traditional additive interpretation method called Deeplift [30]. In this study, DeepSHAP was employed to reveal the impact of driving factors at different time steps and quantify their individual contributions.

3. Results

3.1. Performance of the LSTM Models

In our study, the LSTM-based models exhibited satisfactory performance overall, effectively capturing the trends in TWSA. Additionally, there was minimal difference in performance measures between the training and testing sets, indicating that the data have a relatively good representativeness and the models have been sufficiently trained. The evaluation indicators’ standard deviation (Std) across the five groups was small, suggesting a certain degree of stability in the models.

From the perspective of spatiotemporal average, the specific results are shown in

Table 3. Evaluation metrics between the simulated values and GRACE observations from the perspective of spatiotemporal average.

Group ID	Training Period				Testing Period
	r	NRMSE	NSE	bias	r	NRMSE	NSE	bias
1	0.898	0.044	0.803	−0.166	0.901	0.058	0.749	−1.241
2	0.917	0.04	0.841	−0.062	0.909	0.046	0.823	0.279
3	0.895	0.044	0.801	−0.067	0.855	0.064	0.714	−0.997
4	0.949	0.031	0.896	0.412	0.826	0.06	0.665	−0.22
5	0.918	0.038	0.837	0.417	0.777	0.068	0.586	−0.222
Mean	0.915	0.039	0.836	0.107	0.854	0.059	0.707	−0.48
Std	0.022	0.005	0.039	0.284	0.055	0.008	0.089	0.624

. The scatter density plots (Figure 3) show that the LSTM models performed well in simulating TWSA, with correlation coefficients ranging from 0.777 to 0.909 and an average of 0.854. The NRMSE ranges from 0.046 to 0.068, with an average of 0.059. The bias ranges from −1.241 to 0.279, with an average of −0.48. The NSE is relatively high, ranging from 0.586 to 0.823, with an average of 0.707.

From the perspective of basin average, specific results are shown in

Table 4. Evaluation metrics between the simulated values and GRACE observations from the perspective of basin average.

Group ID	Training Period				Testing Period
	r	NRMSE	NSE	bias	r	NRMSE	NSE	bias
1	0.963	0.068	0.924	0.039	0.984	0.078	0.946	−0.848
2	0.969	0.060	0.937	−0.062	0.962	0.093	0.909	0.279
3	0.966	0.068	0.930	0.057	0.964	0.073	0.925	−0.534
4	0.966	0.069	0.925	−0.310	0.966	0.081	0.909	−0.165
5	0.969	0.068	0.929	−0.115	0.957	0.081	0.871	−0.244
Mean	0.967	0.067	0.929	−0.078	0.967	0.081	0.912	−0.302
Std	0.003	0.004	0.005	0.148	0.010	0.007	0.027	0.422

. On the testing set, five groups of correlation coefficients between the simulated values and GRACE observations range from 0.957 and 0.984, with an average of 0.967. Additionally, the time series plots (Figure 4) show that the LSTM models can effectively capture the seasonal variations of TWSA. However, there is a certain gap between the simulated values and the observations at the peak values. This discrepancy can be attributed to the limited number of samples during peak periods, resulting in the model being relatively undertrained and making conservative predictions at the peaks. Overall, the LSTM model exhibits excellent performance in simulating the average TWSA of the entire basin. Therefore, analyzing the average TWSA reconstructed by LSTM can help mitigate the impact of missing months on driving factor analysis.

From a temporal average perspective, Figure 5 illustrates the spatial distribution of each evaluation metric within the basin, while specific results for five groups are provided in

Table 5. Evaluation metrics between LSTM-based grid series TWSA and GRACE-based grid series TWSA.

Group ID	Training Period				Testing Period
	r	NRMSE	NSE	bias	r	NRMSE	NSE	bias
1	0.914	0.096	0.805	−0.063	0.926	0.152	0.710	−0.640
2	0.927	0.087	0.837	−0.179	0.933	0.110	0.819	0.645
3	0.905	0.104	0.790	1.796	0.890	0.130	0.744	0.108
4	0.952	0.074	0.888	0.338	0.841	0.160	0.630	−0.246
5	0.936	0.089	0.845	0.274	0.851	0.168	0.433	0.205
Mean	0.927	0.090	0.833	0.433	0.888	0.144	0.667	0.014
Std	0.018	0.011	0.038	0.793	0.042	0.024	0.147	0.484

. In general, the correlation coefficient (r) between GRACE observations of equivalent water height and the reconstructed values for each grid point is consistently above 0.8, with more than 88% of the grid cells having a correlation coefficient greater than 0.85. Additionally, NSE exceeds 0.6 for more than 80% of the grid cells, NRMSE remains below 0.17 for over 85% of them, and bias is less than 0.9 for more than 78%. The correlation near the boundary is strong, although the degree of fitting is low. This can be attributed to the limited training of LSTM near the boundary due to the absence of data beyond it.

3.2. Correlation and Importance Analysis of TWSA Driving Factors Based on SHAP

TWS is influenced by precipitation, temperature, evapotranspiration, runoff, and vegetation. These driving factors interact with each other simultaneously. Thus, the linear correlation coefficient and lag correlation analysis between a single factor and TWS do not yield the desired results. For instance, according to the water balance equation, an increase in evapotranspiration and runoff can lead to a decrease in TWS, but the correlation coefficient shows that both evapotranspiration and runoff have a positive correlation with TWS (Figure 6). Therefore, accurately reflecting the correlation between driving factors and TWS requires considering the coupling effect of multiple factors. Here, we utilized a LSTM model (G2) to simulate the coupling effect compared with direct analysis between a single factor and TWSA. Additionally, we employed the SHAP method to analyze the role of each driving factor in this process.

Figure 6 shows the temporal variations of various driving factors and GRACE TWSA from a basin average perspective. Generally, all driving factors have a positive correlation with TWSA with a lead time no more than 2 months between driving factors and TWSA under the coupling effect. Furthermore, the seasonal fluctuations and correlation coefficients of each driving factor are comparable to those of TWSA (

Table 6. Trend and lagged correlation analysis of driving factors of TWSA.

Variables	Unit	Rate of Change	Coincident Correlation Coefficient	Lag Time/Month	Lagged Correlation Coefficient
Precipitation	mm	0.101	0.50 *	2	0.79 *
Runoff	mm	0	0.54 *	2	0.76 *
AET	mm	0.024	0.67 *	1	0.77 *
Temperature	°C	0.003	0.67 *	1	0.75 *
LAI	-	0.002	0.79 *	0	0.79 *

Note: “*” indicates that the corresponding correlation coefficient results passed the significance test with p < 0.01.

). They tend to reach their maximum and minimum values almost simultaneously with TWSA, although slight variations exist among different driving factors, indicating their different influences on TWSA.

Precipitation is the main water source for TWS and plays a key role in the hydrological water cycle, making it a crucial driving factor for TWS. In the PRB, characterized by subtropical monsoonal climate zones, precipitation exhibits distinct seasonal variations, with precipitation in summer (June–August) accounting for about 50–70% of the annual total. Hence, TWSA in the PRB displays a similar seasonal trend to precipitation. Notably, both precipitation and TWSA in the PRB show a slight upward trend between 2003 and 2020, with a substantial increase during the summer months. Furthermore, the analysis reveals a two-month time lag between peak values of TWSA and precipitation (Table 6), and the lagged correlation coefficient between precipitation and TWSA is 0.76, which is consistent with previous studies. Moreover, the lagged correlation coefficient between TWSA and runoff is smaller than that of precipitation, indicating that the response of runoff to TWSA is faster than that of precipitation, which delays the rise in the soil moisture.

The effects of evapotranspiration and temperature on TWSA are similar, as the increase in temperature leads to increased evapotranspiration, with temperature having a smaller lagged correlation coefficient. Additionally, both evapotranspiration and temperature display a lag time of one month in relation to TWSA, indicating that the replenishment of atmospheric precipitation through evapotranspiration delays the increase in TWSA.

The seasonal variations of LAI exhibit similarity to that of TWSA, with a coincident correlation coefficient of 0.79. Both LAI and TWSA show a slight upward trend, highlighting the significance of vegetation as an additional important driving factor alongside precipitation. Vegetation plays a role in intercepting precipitation, reducing runoff, and subsequently decreasing regional water exchange, which contributes to an increase in TWSA. This process occurs within a relatively short timeframe.

Figure 7 shows the importance ranking of driving factors. Generally, precipitation emerges as the most crucial driving factor, followed by evapotranspiration, runoff, temperature, and LAI. The influence of precipitation is much greater than that of any other factor, emphasizing its paramount importance.

Considering that the hydrological process represented by runoff is more complex and interacts more frequently with other driving factors, dependence plots between runoff and other driving factors are drawn to analyze the interaction between driving factors further (Figure 8). In general, runoff shows a negative correction with TWS. The impact is less pronounced when runoff amounts are small, but becomes more evident with larger runoff volumes, indicating a stronger inhibitory effect. Furthermore, the dependence plots demonstrate an obvious vertical distribution pattern, which reflects its interaction with other driving factors. Especially, larger precipitation, higher temperatures, and greater evapotranspiration contribute to both an increase in TWS and an enhanced inhibitory effect of runoff on TWS. However, LAI does not exhibit a significant impact on this interaction. These findings underscore the intricate nature of the coupling effects among driving factors.

3.3. Lead Time Analysis between GRACE TWSA and Driving Factors

Figure 9 showcases the influence of different lead times for each driving factor on TWS revealed through SHAP analysis. Generally, the influence of a driving factor under the coupling effect becomes stronger as the lead time approaches the current month. The coupling effect of driving factors reaches the maximum when the lead time is one month, followed by a two-month lead time, and the current month (namely, zero lead time). Specifically, precipitation from the previous month exerts the greatest impact, while the effect of evapotranspiration is strongest with a two-month lead time. LAI and runoff have the most significant effect in the current month, whereas temperature has a relatively small and uniform impact.

Figure 10 illustrates the cumulative local effects of partial dependence plots with three different lead times on TWS in the PRB. As mentioned in the Methods section, a partial dependence plot presents the SHAP values of a driving factor at a certain lead time in each sample, representing the contribution’s change when the value of the driving factor changes. In general, the relationship between the SHAP value and its driving factors exhibits nonlinearity and nonmonotonicity on the time scale.

Figure 10 shows that the SHAP value of precipitation with longer lead times exhibits strong nonlinearity with the value of precipitation. When the precipitation is high, the SHAP value has a positive correlation with the precipitation and the SHAP value is large. However, when the precipitation is low, the SHAP value has no obvious correlation with the precipitation and the range of SHAP values is smaller. As the lead time approaches the current month, even when the precipitation is low, the SHAP value also gradually exhibits strong linearity with the value of precipitation and the range of SHAP values also becomes larger. This indicates that higher precipitation with different lead times has a significant effect on increasing the TWS in the current month, while lower precipitation has little effect. Notably, when the lead time is closer to the current month, the effect of precipitation on TWS in the current month becomes more pronounced.

In the case of evapotranspiration and temperature, they exhibit similar patterns. When the lead time is long, the range of SHAP values is small, showing no obvious correlation with the amount of evapotranspiration. As the lead time gets closer to the current month, the SHAP value gradually shows a negative correlation with evapotranspiration. This suggests that evapotranspiration with a longer lead time has less effect on TWS. As the lead time gets closer to the current month, the inhibitory effect of evaporation on the TWS in the current month is more obvious. Additionally, it is worth noting that the negative correlation cannot be drawn directly by conducting a separate correlation analysis between evapotranspiration and TWSA (Figure 6). However, the deep learning model that integrates all five driving factors can effectively identify the negative correlation of evapotranspiration on TWS.

As for LAI, it exhibits a slightly different trend compared to evapotranspiration, and its correlation with the value of LAI is significantly influenced by both the magnitude of LAI and the lead time, which is more complex. When the lead time is short, lower LAI suppresses the increase in TWS, while higher LAI promotes the increase in TWS. However, for longer lead times, this relationship is reversed.

4. Discussion

4.1. The Uncertainty of the LSTM Model

Deep learning methods have proven to be highly effective in tackling complex problems, and LSTM, in particular, excels in handling long time sequence tasks. Wang et al. [51] constructed an LSTM model to simulate GRACE TWSA in the Tarim River Basin, achieving a correlation coefficient of 0.922 and NRMSE of 0.107. In this study, the average of correlation coefficient and NSE of TWSA, evaluated through 5-fold cross-validation in the PRB, are 0.854 and 0.707, respectively, from a spatiotemporal perspective. Moreover, when examining the basin average from a temporal perspective, the correlation coefficient and NSE reach high values of 0.967 and 0.912, respectively. These results highlight the outstanding capability of LSTM models in simulating time series data.

While LSTM can effectively simulate TWSA, it still introduces a certain level of uncertainty. This uncertainty primarily stems from inherent model structure limitations and potential shortcomings in training. Regarding the model structure, researchers need to carefully adjust it based on specific application scenarios. Although LSTM has good performance in temporal simulation, it may struggle to capture spatial patterns adequately, as evidenced by average correlation coefficients of 0.927 and 0.967 for spatial and temporal dimensions, respectively. However, TWS is often influenced by factors in neighboring areas, which makes spatial considerations important.

Shortcomings in training can be attributed to both the training data and model optimization. First, the limited time scale of data available (only 20 years since the launch of the GRACE satellites) poses challenges in effectively training the model. Treating each time grid point in the basin as a separate sample increased the sample size but introduced the uncertainty in the spatial distribution of TWSA patterns. Moreover, factors such as human activities and geographical conditions, which are not considered in the model, can also impact TWS. Additionally, the data themselves also carry inherent uncertainty.

As for optimization, as can be seen from Table 3, Table 4 and Table 5, the LSTM model always performs better on the training set than on the testing set. Ideally, both the training and testing sets should obey the same data distribution. The LSTM model learns the joint probability distribution of the training set and applies it to the testing set. However, due to the inability to consider all factors, variations in other factors may lead to the invalidity of this assumption, especially when the training and testing data come from different time periods. Consequently, the performance of the model on the testing set may not be as strong as on the training set [47]. As can be seen from the scatter plots (Figure 3), the model tends to be conservative in its simulation of peak values. This may be due to the uneven distribution of the samples and inherent model characteristics. The proportion of peak values in periodic data is small, making it challenging to train the model effectively. Additionally, the spatial distribution figure (Figure 5) indicates a scarcity of data around the edges of the basin, which partially reduces the effectiveness of the simulation.

4.2. Analysis of Natural Driving Factors

Figure 9 illustrates that these driving factors exert distinct hysteresis effects on TWS, a phenomenon corroborated by previous studies. Wang et al. [52] found the hysteresis effect from precipitation, explaining that the TWS for a given month only incorporates precipitation information from the preceding month. This is due to the time required for precipitation infiltration and evaporation processes. We thus delve into the hysteresis of different driving factors in the GRACE-derived monthly TWS to examine their temporal impacts on TWS.

From Figure 9, it is evident that runoff exhibits its greatest impact in the current month. Meanwhile, precipitation from the previous month and evapotranspiration from two months prior exert stronger influences. We postulate that this is due to GRACE’s measurement techniques and the hydrological cycle. GRACE likely exhibits heightened sensitivity to substantial water body quality fluctuations, encompassing rivers, lakes, and groundwater. This leads to spatial and temporal hysteresis variations. Figure 9 highlights the most immediate TWS responses to the runoff since runoff aptly represents rivers and lakes. Some of the runoff seeps into the soil, needing time to reach groundwater, thereby influencing the subsequent month’s TWS.

The situation of the precipitation is different from runoff. Upon reaching the land surface, precipitation transforms into major hydrological components, including runoff, soil moisture, groundwater, and evapotranspiration. We observe that precipitation has the highest impact on TWS at a one-month lead time (Figure 9). This indicates that the combined amount of runoff, soil moisture, groundwater, and evapotranspiration is predominantly determined by the precipitation from the preceding month. In other words, the TWS response to precipitation is most significant with a one-month lag, highlighting the time required for these hydrological processes to take effect.

Temperature affects plant growth and evapotranspiration. Its contribution pattern over the initial three months mirrors that of evaporation. However, since temperature can also influence precipitation and plant growth, it adds complexity to the relationship between TWS and temperature.

Compared to temperature and precipitation, the hysteresis effect of evapotranspiration is more pronounced. Apart from its impact on lakes or rivers, evapotranspiration mainly influences TWS by depleting soil moisture and groundwater. These processes take longer to manifest, explaining why evapotranspiration from the previous two months has a more significant impact on TWS. Moreover, evapotranspiration is also affected by temperature and soil moisture, indicating a potential reciprocal influence on TWS.

Lastly, vegetation’s influence on TWS is multifaceted. It can modify soil moisture, significantly impact both runoff and evapotranspiration, and consume varied water amounts at different stages of plant growth. Figure 9 suggests that in the PRB, vegetation shows a similar trend to runoff, indicating that vegetation primarily influences TWS by altering runoff patterns.

4.3. Uncertainty of Explainable Machine Learning

Machine learning models have been widely used in Earth science for their excellent performance in stimulating complex nonlinear processes, even without prior assumptions about the underlying physical systems. However, one challenge associated with machine learning models, particularly deep learning models, is their lack of interpretability. To address this issue, explainable machine learning was introduced by the Defense Advanced Research Projects Agency in 2017, and it has gained popularity across various fields, especially in scientific research [53]. Several machine learning interpretation techniques, such as LIME and ALE, have been developed and applied to elucidate the patterns captured by these models. Previous studies have demonstrated the significant capabilities of explainable machine learning in revealing complex natural processes. In this study, we utilized the SHAP method to quantify the relative contributions of precipitation, temperature, vegetation, runoff, and evapotranspiration to TWS. Notably, the results align with the existing understanding of hydrological processes from previous studies, indicating a certain level of explanatory power for hydrological phenomena.

However, the SHAP method used in this study introduces its own uncertainty. Firstly, SHAP is a method of additive interpretation that assigns an additive contribution value to each variable. However, this method can be affected by multicollinearity between variables, which reduces the reliability of the model when there is high interaction among driving factors. Additionally, the learning ability of the deep learning model is a crucial factor in the interpretability of the results obtained through explainable techniques. It is essential for the deep learning model to capture the physical laws of the research objects rather than solely relying on correlations for prediction. The deep learning model structure itself also introduces uncertainty in the attribution allocation. While LSTM can mitigate the issue of gradient vanishing, it cannot completely balance the difficulty of gradient updating at different time steps. Some scholars believe that due to the existence of forget gates in LSTM, the model tends to consider the influence of similar factors over time [24].

Despite the inherent uncertainties and limitations associated with data-driven methods such as deep learning and interpretable machine learning methods, they offer researchers valuable modeling techniques to accurately estimate TWS and comprehend water cycles [54]. Such methods have the potential to unveil complex relationships that have not yet been observed or are challenging to observe through traditional approaches.

For future work, we aim to incorporate additional driving factors into our models, including human activities, soil texture, and other relevant variables. Furthermore, we aim to explore models that incorporate both temporal and spatial aspects, such as convolutional neural networks and spatiotemporal convolutional neural networks. By comparing and utilizing datasets with less uncertainty, we can enhance the robustness of our simulations. Additionally, we will incorporate spatial interpretation methods, such as gradient-weighted class activation mapping (Grad-CAM), to improve the interpretability of our models and make the simulation and interpretation more reasonable.

5. Conclusions

In this study, a long short-term memory (LSTM) neural network model, based on the data of precipitation, runoff, evapotranspiration, temperature, and leaf area index (LAI), was applied to simulate terrestrial water storage anomaly (TWSA) combined with GRACE gravity satellite products. We employed the SHapley Additive exPlanations (SHAP) method to quantitatively analyze the contributions of different driving factors to terrestrial water storage (TWS) in the Pearl River Basin (PRB) and explored the impacts of different lead times. Our results show that the LSTM model performs well in stimulating TWSA. From the perspective of the basin average, the basin-average r is 0.967 and NSE is 0.912. Moreover, our modeled findings demonstrated that TWSA in the PRB is predominately influenced by precipitation, followed by evapotranspiration, runoff, temperature, and LAI. There is also a clear accumulation effect of precipitation and evapotranspiration on TWS, which becomes more obvious with a shorter lead time in general. Finally, the impact of driving factors on TWS is related to lead time. Precipitation exerts the most significant effect with a lead time of one month, while that of evapotranspiration reaches its peak impact when the lead time is one to two months. Runoff has the greatest effect in the current month, while the effect of the lead time of temperature on TWS is relatively minor.