STAMNet—A spatiotemporal attention module and network for upscaling reactive transport simulations of the hyporheic zone

Reactive transport (RT) simulations are important tools for understanding and predicting phenomena in the subsurface. However, RT is computationally intensive and complex simulations can be numerically unstable. Here, we present STAMNet, a low-parameter attention-based suite of neural nets that can upscale and upsample reactive transport simulations, applied to example problem of bioremediation in the hyporheic zone. We show that a simple MLP offers 30x speedup over standard multiphysics RT simulations and can accurately ($\approx 90\text{\%}$ $R^2$) predict the output of multiple variables of a 1x20 meter RT simulation by using the output from a 1 × 2 m simulation as input. We add efficient channel attention to our optimized MLP which significantly improves the mean average error but does not affect the $R^2$. We further develop a novel spatiotemporal attention module (STAM), which results in improvements both in mean square error and $R^2$ (92.5%). Finally, we present a network architecture that utilizes STAM to accurately (99.9% $R^2$) upsample simulations in two dimensions. Specifically, our model allows for the 2x upsampling of simulations in the $x$ and $y$ dimensions to convert a coarse-grained input into a fine-grained output. These models have potential use for Monte-Carlo-style investigations of bioremediation and the work presented serves as a proof-of-concept for accurate prediction of large sets of spatiotemporal outputs.