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

TL;DR

This work tackles the computational bottleneck of stochastic hydrologic simulations by integrating a hybrid workflow in which a process-based model (ParFlow-CLM) seeds an ML surrogate that completes the remaining simulations. Seven diverse ML architectures (including CNNs, ViTs, U-FNOs, and PredRNN++) are trained on a 3D MAR case study and evaluated for accuracy, training stability, and runtime efficiency, with careful data preprocessing and downsampling to manage high-dimensional inputs. All models achieve $MAPE < 10\%$ and show substantial runtime gains over the process-based model, with the best performers (U-FNO4d and PredRNN++) delivering rapid end-to-end predictions after only a few training examples. The study also highlights the importance of min-max normalization and loss normalization for training stability and accuracy, and discusses interpretability and limitations, indicating strong potential for ML surrogates to inform MAR site assessment and water resources decision-making. Overall, the results support adopting ML surrogates to accelerate large ensembles of hydrologic simulations while preserving essential physical fidelity.

Abstract

Process-based hydrologic models are invaluable tools for understanding the terrestrial water cycle and addressing modern water resources problems. However, many hydrologic models are computationally expensive and, depending on the resolution and scale, simulations can take on the order of hours to days to complete. While techniques such as uncertainty quantification and optimization have become valuable tools for supporting management decisions, these analyses typically require hundreds of model simulations, which are too computationally expensive to perform with a process-based hydrologic model. To address this gap, we assess a hybrid modeling workflow in which a process-based model is used to generate an initial set of simulations and a machine learning (ML) surrogate model is then trained to perform the remaining simulations required for downstream analysis. As a case study, we apply this workflow to simulations of variably saturated groundwater flow at a prospective managed aquifer recharge site. We compare the accuracy and computational efficiency of several ML architectures, including deep convolutional networks, recurrent neural networks, vision transformers, and networks with Fourier transforms. Our results demonstrate that ML surrogate models can achieve under 10% mean absolute percentage error and yield order-of-magnitude runtime savings over process-based models. Building on these findings, we examine the impacts of key modeling choices on surrogate model accuracy and efficiency. Results show that a normalized loss function improves training stability, while min-max data normalization can significantly reduce error up to a factor of 10 when compared to other treatments such as Z-score and no normalization. Downsampling input features using an autoencoder also decreases memory requirements by training with tensors 4% their original size. By reducing computational costs and...

Figures (14)

• Figure 1: Diagram of the proposed hybrid framework. The process-based hydrologic model generates a training set ($\mathcal{D}_\text{train}$), which is then used to train and test an ML model. Subsequently, the trained ML model is used to generate the remaining hydrologic simulations.
• Figure 2: Illustration of the three types of ML model architectures compared in this study. Each blue cube represents a single 3D tensor. The non-sequential models (a) receive four-dimensional input (x, y, z and time) and produce all pressure fields for all time steps in a single forward pass of the model. The one-step models (b) simulate a single time step per forward pass and roll forward, using the output from previous time steps as input for subsequent steps. The recurrent models (c) use a similar procedure as the one-step models, except that they contain memory cells that retain hidden states between forward passes.
• Figure 3: Workflows to train an autoencoder (a) and subsequently use the autoencoder in the surrogate model's training (b).
• Figure 4: Snapshot of the 3D pressure field for a random test set example at a single time step in Stage 1, as produced by the process-based hydrologic model (a) and the PredRNN++ (b). The cell-by-cell absolute error (c) between the ground truth and predictions is greatest near the wetting front. Note that we omit the deepest layer, which has a height of 60 $m$, for visual clarity.
• Figure 5: Test set MAPEs for all ML models across all stages. (b) is a local enlargement of (a), indicated by the "Inset" label in (a). The ML surrogate models are grouped by model type (see § \ref{['subsec:ml_models']}), with the non-sequential models on the left (CNN4d, ViT4d and U-FNO4d), the one-step models in the middle (CNN3d, ViT3d and U-FNO3d) and the recurrent model on the right (PredRNN++). Exact MAPE values are listed in Table \ref{['tab:res_compile']}.