Machine learning surrogates for efficient hydrologic modeling: Insights from stochastic simulations of managed aquifer recharge

The goal of this hybrid workflow is to harness the fast output rate of ML models to efficiently generate a large volume of simulations for downstream hydrologic analysis. Our task is to execute $N$ simulations, generating $N$ simulation outputs from $N$ distinct sets of simulation inputs. As illustrated in Figure 1, we employ a traditional processed-based hydrologic model (ParFlow-CLM) to generate a training set $\mathcal{D}_{\text{train}}$ comprising simulation input–output pairs. Subsequently, an ML surrogate model is trained on these examples and used to generate simulation outputs for the remaining $N-|\mathcal{D}_{\text{train}}|$ simulations, where $|\mathcal{D}_{\text{train}}|$ is the number of examples in $\mathcal{D}_{\text{train}}$. If the ML surrogate provides comparable accuracy to the process-based model, the combined group of $N$ simulations can then be used for further analyses, such as uncertainty quantification, parameter estimation and resource optimization.

We apply this hybrid workflow to a single case study in which hydrologic simulations are used to evaluate a recharge efficiency at a prospective managed aquifer recharge (MAR) site. MAR is a water management strategy whereby land managers deliberately flood parcels of land or inject water into wells to replenish underlying aquifers. Recent work has shown that MAR is a valuable tool for mitigating groundwater depletion in many regions of the globe (Dillon et al., 2019; Sprenger et al., 2017), but evaluating prospective recharge sites remains challenging as subsurface heterogeneity can significantly impact recharge rates (Knight et al., 2022; Perzan et al., 2023). Hydrologic and reactive transport simulations are effective tools for quantifying outcomes of interest at prospective MAR sites, such as changes in groundwater storage (e.g., Ganot et al., 2018) or the fate of subsurface contaminants (e.g., Perzan and Maher, 2024). However, generating these simulations comes with a notable computational expense, as the time required for a single simulation can range from hours to days depending on model scale and resolution. Thus, MAR site selection is a logical case study that could benefit from a computationally efficient hybrid modeling workflow.

In this case study, we focus on a prospective MAR site in Tulare County, California, that measures $$800\text{\,}\mathrm{m}$\times$400\text{\,}\mathrm{m}$$. The site is underlain by a $45\text{\,}\mathrm{m}$ thick vadose zone that consists of interbedded layers of sand, silt and clay. The unconfined aquifer extends to $150\text{\,}\mathrm{m}$ depth, below which the Corcoran Clay, a regional aquitard, limits vertical flow (Mid-Kaweah GSA, 2019). The goal of the simulations at this site is to quantify the increase in saturated zone storage that will occur within two years of inundating the site with a $0.8\text{\,}\mathrm{m}$ recharge event.

To evaluate recharge efficiency at this site, Perzan et al. (2023) performed several hundred stochastic simulations of MAR using ParFlow-CLM (§2.3). The model domain is discretized into a $120\times 80\times 25$ rectilinear grid ($$1200\text{\,}\mathrm{m}$\times$800\text{\,}\mathrm{m}$\times$150\text{\,}\mathrm{m}$$ in the x, y and z directions, respectively), utilizing a uniform horizontal cell size of $10\text{\,}\mathrm{m}$ and variable cell thicknesses in the $z$ direction ($1.02\text{\,}\mathrm{m}$ at the surface and $9\text{\,}\mathrm{m}$ at $90\text{\,}\mathrm{m}$ depth). Specified head boundaries corresponding to the depth to the water table ($45\text{\,}\mathrm{m}$) are applied on the sides of the domain. A no-flow boundary is implemented on the bottom of the domain and variable forcing is applied to the top of the domain based on the modeling stage, as explained at the end of this section.

The domain is parameterized through a stochastic geophysical-geostatistical workflow (Perzan et al., 2023). Because subsurface heterogeneity can exert a dominant control on recharge efficiency, we parameterize the domain using a subsurface electrical resistivity model derived from a towed transient electromagnetic (tTEM) survey of the site. The tTEM system is a hydrogeophysical tool that can map meter-scale variations in electrical resistivity, a property directly related to sediment lithology. We first generate 600 gridded realizations of the tTEM-derived resistivity model using sequential Gaussian simulation, a stochastic technique for populating a rectilinear grid with a random variable. We then compare vertical profiles of electrical resistivity to collocated profiles of sediment type collected through cone penetration tests. Using the bootstrapping procedure of Knight et al. (2018), we transform each resistivity realization to a realization of coarse fraction, a metric that describes the fraction of coarse-grained (sand and gravel) sediment within each grid cell. The remainder of sediment within the cell is assumed to be fine-grained (silt and clay). For example, a cell with a coarse fraction of 0.5 is 50% coarse-grained material (either 100% sand, 100% gravel, or some combination of the two) and 50% fine-grained material (some combination of silt and clay). To assign hydrogeologic parameters (hydraulic conductivity, porosity, etc.) to each coarse fraction realization, we assume that only two types of sediment exist within a single realization — coarse-grained and fine-grained — and that each cell is a mix of the two types. We then sample from broad distributions to assign parameter values to each end member and use a weighted averaging procedure to calculate values for cells that are a mix of the two end members. The input parameter distributions are derived from a mix of literature and site-specific data. For more information about these parameter values, we refer the reader to the more complete description provided by Perzan et al. (2023).

Each MAR simulation consists of three consecutive stages. In Stage 1, a constant recharge forcing equal to the long-term natural recharge rate is applied to the top of the domain for 160,000 days ($\sim$440 years), which is enough time for each simulation to reach steady state. In Stage 2, we apply one year of meteorological and irrigation forcing using hourly measurements from the southern Central Valley. These two spin-up stages initialize water content in the vadose zone and establish realistic soil moisture profiles prior to implementing recharge. Finally, Stage 3 simulates recharge by inundating the orchard with $0.8\text{\,}\mathrm{m}$ of water in a single winter season (applied in either 1, 4, 8 or 16 individual recharge events), alternating between periods of flooding and no flooding. In between inundation events, we apply meteorologic and irrigation forcing to the top of the domain.

Once each simulation has completed all three stages, we quantify the increase in saturated zone storage induced by MAR. Each simulation produces spatiotemporally variable output (transient pressure fields, represented as a time series of $120\times 80\times 25$ rectilinear grids) from a combination of temporally static, spatially variable input (e.g., 3D fields of porosity or saturated hydraulic conductivity, represented as $120\times 80\times 25$ rectilinear grids) and temporally variable, spatially homogeneous input (e.g., ambient air temperature). The number of output time steps varies depending on the modeling stage, with a finer temporal resolution in later modeling stages. We output pressure fields every 1000 days in Stage 1 (160 output time steps), every 2 days in Stage 2 (183 output time steps) and every 2 days in Stage 3 (366 output time steps). Note that the number of time step outputs here simply reflects the frequency at which the model outputs 3D pressure fields and not the temporal resolution of the underlying solver. From these 3D transient pressure fields, we calculate other system states (e.g., water content in each grid cell) and quantify the change in saturated zone storage.

In the first phase of our proposed framework (§2.1), a process-based hydrologic model is used to generate an initial set of simulations, which are then used to train an ML surrogate model. In this case study, we perform these process-based simulations using ParFlow-CLM. ParFlow-CLM couples the integrated hydrologic code ParFlow with the Common Land Model (CLM) to solve partial differential equations that describe water and energy fluxes over and beneath the land surface. ParFlow solves the Richards’ equation for variably saturated subsurface flow using a Newton-Krylov approach and the shallow-wave equations for overland flow using a kinematic wave approximation. CLM, meanwhile, uses a mass transfer approach to compute water and energy fluxes at the land surface (Maxwell and Miller, 2005; Kollet and Maxwell, 2008), including irrigation, evapotranspiration and precipitation. To accurately simulate the impacts of diurnal temperature and radiation fluctuations on water fluxes across the land surface, we run ParFlow-CLM using an hourly time step in Stages 2 and 3. For a detailed description of the governing equations, we refer readers to previous work describing ParFlow-CLM (Jones and Woodward, 2001; Ashby and Falgout, 1996; Kollet and Maxwell, 2006; Maxwell, 2013; Maxwell and Miller, 2005; Kollet and Maxwell, 2008).

In the second phase of this hybrid framework (§2.1), we train an ML surrogate on the initial set of simulations generated by the process-based model. In total, we evaluate the performance of seven different machine learning surrogate model architectures. The architectures include recurrent and highway networks developed for spatiotemporal forecasting (Wang et al., 2018), vision transformers (Dosovitskiy et al., 2021), convolutional nerual networks and U-shaped networks paired with Fourier transforms (Wen et al., 2022; Zhou et al., 2022; Kang et al., 2023). We selected these architectures because some of them have been successfully applied in previous hydrologic surrogate modeling studies (Tran et al., 2021; Maxwell et al., 2021; Wen et al., 2022), while others are some of the most common architectures used for spatiotemporal forecasting in related fields like computer vision.

The ML surrogate architectures can be categorized into three types based on their method of prediction (Table 1, Figure 2). First, the Non-sequential models emulate the entire simulation in a single forward pass (i.e., the emulator produces 3D pressure fields for all output time steps in one run of the model). By contrast, the One-step models emulate one output time step (producing a single 3D pressure field) per forward pass. These models then use the output from the previous time step to perform another forward pass, repeating until all 3D pressure fields have been generated. Finally, the Recurrent models are similar to the one-step models in that they emulates one output time step per forward pass. However, they use memory cells to retain hidden states between forward passes, so that information can be stored between output time steps. As discussed above, the output frequency varies between model stages (see 2.2), which results in a different shape input tensor for each model stage. Because most ML architectures require a consistent shape input tensor across all training examples, we train a separate surrogate for each model stage. The ML models are summarized below, though complete details are provided in the Supplementary material (Section S1.1).

To train each ML surrogate model, we designate the inputs of the ParFlow-CLM simulations (e.g., 3D porosity fields, transient meteorological forcing, etc.) as our training features and the outputs of the ParFlow-CLM simulations (transient, 3D pressure fields) as our training targets. However, prior to training, we first perform several steps of preprocessing. In particular, we perform feature selection to curate a list of salient input features, we broadcast input features across the $x$, $y$, $z$ and $time$ dimensions to standardize input and output formats, and we downsample certain inputs to reduce memory requirements for training.

In this section, we discuss the training setup used for all surrogate models, the hardware used to calculate computational efficiency and the evaluation criteria used to quantify model performance.

In summary, we trained six different groups of surrogate models, each for a different numerical experiment:

1. 1.

   We train all seven surrogate architectures — the CNN4d, ViT4d, U-FNO4d, CNN3d, ViT3d, U-FNO3d, and PredRNN++ (§2.4) — on Stage 1, 2, and 3 ParFlow-CLM simulations, amounting to 21 trained models. We also train an autoencoder for all Stage 3 models. This totals 22 trained models.

2. 2.

   We train two additional Stage 1 CNN3d models using alternative data normalization techniques (§2.6.1). One of the additional models does not use any data normalization, and the other uses Z-score normalization. The Stage 1 CNN3d model is chosen over other stages and models for its efficiency in training and evaluation.

3. 3.

   We train six additional Stage 1 CNN3d models using different combinations of loss function and data normalization (§2.6.1). Three CNN3d models use the normalized L2 loss with min-max, Z-score, and no data normalization, and three CNN3d models use the mean squared error loss with min-max, Z-score and no data normalization.

4. 4.

   To evaluate the effect of training set size on model accuracy, we train all seven ML models with nine Stage 1 training sets of decreasing sizes, as described in §2.6.3. This experiment includes 63 trained models in total.

5. 5.

   To evaluate the ML models’ runtime efficiency, we re-train all seven ML models with variable Stage 1 training set sizes, using a binary search technique to find the minimum training set size that achieves $<$10% test MAPE. The number of models trained at this stage varies according to the efficiency of the search algorithm and the test MAPE results from previous experiments but ranged from 2 to 7 across the seven architectures. We then compare the runtimes of the seven models trained on their respective minimum Stage 1 training sets with the procedure described in §2.6.4.

Each ML surrogate model produces pressure fields that are visually similar to those produced by the process-based hydrologic model. For example, a comparison of ParFlow and the PredRNN++ shows that both models produce similarly-shaped heterogeneous wetting fronts during Stage 1 (Figure 4), with monotonically increasing pressure values below the front. Plotting the cell-by-cell absolute error between the two models reveals that differences are greatest near the wetting front (Figure 4c). However, absolute pressure differences are generally less than a few meters, which is of a reasonable range given the sharp hydraulic gradients that arise near infiltration wetting fronts.

All seven ML models achieve $<$10% MAPE on the test set (Figure 5 and Table S8), which is the threshold for “highly accurate forecasting” (Lewis, 1982). The PredRNN++ displays the lowest error for Stage 1 and end-to-end (E2E) predictions, while the CNN3d and ViT3d perform the best on Stage 2 and Stage 3, respectively. Even though the non-sequential models have more parameters than their one-step counterparts (Table 1), they do not exhibit an appreciable increase in predictive performance. The lone recurrent model (the PredRNN++), however, outperforms the others and is better able to learn the temporal dependencies that develop as each hydrologic simulation progresses through time. The large memory cells retained between time steps likely provide ample degrees of freedom for the model to store spatially heterogeneous information over time.

Overall, prediction error is the highest for Stage 1, with a mean MAPE of 5.6% across all ML models, as opposed to 0.13% and 0.42% for Stages 2 and 3, respectively. Previous modeling work at this site has shown that managed aquifer recharge efficiency is most strongly controlled by vadose zone water content prior to inundation (Perzan et al., 2023). The low model error for Stages 2 and 3 makes sense, then, given that Stage 2 and 3 are provided the ground-truth pressure field output by ParFlow-CLM at the end of Stage 1 and 2, respectively, as initial pressure head fields during training.

Alternatively, the higher Stage 1 MAPE could result from its large time scale (§2.2 and 2.5.3). While Stage 2 and 3 span 1 and 2 years, respectively, Stage 1 spans over 400 years, with more dramatic changes in pressure between outputs. Pressure head changes across large time scales may be more difficult to learn than smaller, more incremental changes, potentially contributing to the poorer performance of Stage 1 models.

End-to-end (E2E) predictions, on the other hand, do not receive ground-truth pressure fields between stages; the predicted pressure at the end of one stage is used as input for the ML surrogate in the next stage. Overall, the MAPE for E2E predictions is typically lower than that of Stage 1 and similar in magnitude to that of Stage 3 (Figure 5). The similarity between Stage 3 and E2E error likely arises because E2E error is averaged equally over the entire time series and Stage 3 has the most time steps. However, not all models follow this trend; the CNN3d and CNN4d both exhibit E2E MAPE higher than 1%, with a $3.1\times$ and $6.6\times$ respective increases compared to their Stage 3 MAPEs. Plotting the E2E MAPE per time step in Figure 6 reveals that the CNN3d and CNN4d exhibit the highest MAPE per time step among other models in Stage 3. This heightened difference between Stage 3 MAPE and Stage 3 MAPE in E2E evaluation could indicate that the CNN3d and CNN4d both exhibit a strong dependence on the initial pressure head at the start of a model stage. The curves for all other models follow the same general shape, suggesting that these models may have learned similar patterns in the Stage 3 data that have resulted in similar accuracy.

Collectively, the results in this section suggest that ML surrogate models can simulate variably saturated groundwater flow with a similar level of accuracy as process-based hydrologic models, supporting the feasibility of applying a hybrid modeling framework to this case study.

While each ML surrogate model achieves $<$$10\%$ MAPE, the models vary widely in the amount of time and memory they require to produce simulations. One attribute they share, however, is that all ML surrogate models offer significant runtime advantages over the process-based hydrologic model (Figure 9). Figure 9 reveals positive slopes across all lines. Recall that the slope for these plots are calculated as the difference in time taken by ParFlow-CLM and the ML model in generating a single Stage 1 simulation (Equation 6). All lines appear to exhibit identical slopes because each ML surrogate outputs simulations orders of magnitude faster than the process-based hydrologic model ($t_{\text{PF}}=$ 15,017 CPU-seconds/simulation), so the difference in $t_{\text{PF}}-t_{\text{ML}}$ across the ML models is miniscule. Nonetheless, positive slopes for all lines signify that all seven ML models produce simulations at a faster rate than ParFlow-CLM. The ViT4d technically exhibits the steepest slope, producing simulations the fastest ($t_{\text{ML}}=0.115$ GPU-seconds/simulation), but, due to the large training set size required for the ViT4d, time savings only appear for the ViT4d after 246 simulations. Closer examination of the $x$-intercepts reveals that 6 out of 7 of the ML models achieve a runtime advantage over ParFlow-CLM when the desired simulation output count is as low as 17. We provide a detailed overview of all time measurements in Table 2, including distinct values of $\mathcal{D}_{\text{train}}$ (the minimum size training set required to achieve $<10\%$ MAPE, which strongly impacts the computational efficiency of each model.

When accounting for both training time and forward simulation time, the U-FNO4d is the most efficient ML model. To perform 500 simulations, the U-FNO4d requires a total of 13.1 processor-hours (the sum of Equations 4 and 5), compared to the second-most efficient ML model, the PredRNN++, which requires 13.3 processor-hours. These computational times are orders of magnitude smaller than the time it would take to generate 500 simulations exclusively with ParFlow-CLM (2090 processor-hours). These findings, expressed in Figure 9, affirm the tangible benefits of this hybrid framework, which outpaces a naive approach of solely relying on ParFlow-CLM. We believe this computational speedup arises from two factors. First, ParFlow-CLM uses an iterative Newton-Krylov approach to solve the coupled system of nonlinear partial differential equations at each time step. This iterative scheme requires many Jacobian evaluations and floating point operations per time step, while the ML surrogates only require forward passes, consisting of fixed calculations determined by the model architecture and model weights. Second, the ML surrogate models can take arbitrarily large time steps without compromising convergence due to their data-driven nature, while ParFlow-CLM must use an hourly time step to accurately represent meteorology-driven water fluxes across the land surface. Other process-based models are subject to similar constraints — such as the Courant–Friedrichs–Lewy condition — which limit their computational efficiency relative to ML surrogates. Since ML surrogates can learn temporal patterns in training data, they are unbound by time step size and can thereby complete simulations at faster rates.

In addition to runtime, the memory required to train an ML model is an important component of computational efficiency. While the memory cells in the PredRNN++ allow the model to retain information from one time step to the next, these large tensors require significant memory allocation (Table 2). Using a batch size of 1, the PredRNN++ requires 15.1$\times$ more GPU memory to train (31.7 GB) than the ViT3d (2.1 GB), the least memory-intensive architecture. At this allocation size, the PredRNN++ may be too memory-intensive to train on some modern GPUs. Similarly, all three non-sequential models (CNN4d, ViT4d and U-FNO4d) require more memory than their one-step counterparts (CNN3d, ViT3d and U-FNO3d), though this increase in memory allocation does not necessarily translate to an increase in accuracy (Figure 5). Ultimately, even though some architectures, such as the PredRNN++, can achieve low MAPE with only a few training examples, they may not be flexible enough to adapt to all compute environments due to their high memory requirements.

Another crucial factor to consider when selecting an ML architecture within a hybrid modeling framework is model interpretability. Error recognition is more straightforward with an interpretable surrogate model and can enhance trust among practitioners. The PredRNN++ architecture, for example, contains $\mathcal{C}$ and $\mathcal{M}$ memory cells (Figure S2) that are intended to serve distinct roles, with the former facilitating the transfer of temporal information and the latter facilitating the transfer of spatial information (Wang et al., 2018). However, in practice, deciphering the values stored in these cells and validating their role in each model has proven challenging (Figure 10). Each memory cell is the same size as the model domain ($120\times 80\times 25$), with two memory cells for each of the four causal LSTM layers in the PredRNN++. Given that these memory cells change between each time step and between each evaluation example, manually inspecting and interpreting $\mathcal{C}$ and $\mathcal{M}$ quickly becomes cumbersome. In other words, the hidden states of the PredRNN++ model and, consequently, the internal mechanisms of the model, prove challenging to interpret.

In addition to the PredRNN++, other ML architectures have been shown to be interpretable when applied in different domains and may prove useful in hydrologic surrogate modeling. For example, previous studies have indicated that hierarchical convolutional layers, as employed in CNN3d, can learn hierarchical image information (Zhou et al., 2014; Yang et al., 2019). Furthermore, research on transformers like those employed by the ViT3d has demonstrated that attention mechanisms within these models may also learn hierarchical information (Jain and Wallace, 2019; Jawahar et al., 2019; Wiegreffe and Pinter, 2019). When combined with conventional ML interpretability methods such as saliency maps, deconvolutional networks, guided backpropagation, and feature occlusion, each of these mechanisms could serve as a check to ensure that a surrogate model is behaving as expected. Future work is needed to explore the interpretablility of each of these architectures when applied to hydrologic simulations.

While each ML model can reliably reproduce transient pressure fields generated by a process-based model, future researchers should carefully tune models to limit spatial and temporal variability in model performance. For example, when loss is calculated evenly across the entirety of the domain, ML models may not accurately simulate changes in pressure head in the most important portions of the domain. In Figure 4, which visualizes one time step from a PredRNN++ emulation of Stage 1, model errors are predominantly concentrated at and above the wetting front. Even though the loss function is weighted to more severely penalize pressure changes within the vadose zone (Figure S3), the high MAPE in this area suggests that the ML model still cannot fully resolve sharp hydraulic gradients near the wetting front. Future applications could employ more complex weighted loss functions to ensure that each surrogate most accurately simulates the processes of interest.

The full results that display ML model performance on all stages in Figure 5 reveal a similar challenge in capturing temporal changes in pressure head. The performance in Stage 1 for all models consistently lags behind that of Stage 2 and Stage 3. Because these later stages simulate shorter time periods (one year in Stage 2 and two years in Stage 3, as opposed to 440 years in Stage 1), pressure head changes over the model stage are less pronounced and their temporal relationships are thereby easier to learn. Thus, while this case study achieves low MAPE, future applications should carefully adapt the implementation strategies in this study (for example, the choice of loss function, data normalization technique and input features) to their problem of interest, especially when simulating model domains characterized by abrupt variations in pressure head across either time or space.

The autoencoder is a convolutional neural network containing four layers and around 676,000 parameters, which makes the autoencoder much smaller than each surrogate model. Two of the four layers are encoding convolutional layers, which squeezes the input into compact spatial dimensions, and the remaining two are decoding transposed convolutional layers, which recovers the original spatial resolution. All convolutional layers are 2D, using the $z$ dimension as the channel dimension. They employ a (7, 7)-sized convolutional kernel, a stride of (2, 2), and a hidden size of 128. The convolutional layers are separated by ReLU activation functions and batch normalization layers.

The autoencoder, though a vital component in training Stage 3 models, may introduce inaccuracies. We attempt to disentangle the errors from the surrogate model and the errors from the autoencoder in Figure S4, focusing on the CNN3d for its efficiency in training. Computing MAPE between the predictions $\hat{y}$ taken directly from the autoencoder, i.e., $\hat{y}\leftarrow h_{\text{dec}}(h_{\text{enc}}(y))$, and the simulation targets reveals that the autoencoder can reproduce its own input with 0.28% error on the held-out test set (Figure S4, Config A). We expect that this autoencoder error contributes to Stage 3 error, since we use the autoencoder in training Stage 3 models; indeed, when computing the MAPE between decoded Stage 3 model predictions and simulation targets, this error, at 0.29%, is greater than the autoencoder error alone. It is surprising, then, when we use the autoencoder to encode and decode the targets and use these encoded and decoded targets as our new evaluation references, our percent error decreases to 0.11%. This shift indicates that the autoencoder introduces error by compressing the space of the target encodings. In other words, our autoencoder encodes targets in a way that makes it more difficult to distinguish between unique targets once they have been encoded. Consequently, during the decoding step, essential information distinguishing unique targets is not fully recovered. Exploring improved autoencoder architectures may help to mitigate these challenges and enhance the overall performance our hybrid framework in future applications.