Time-Splitting Fourier Neural Operator with Coordinate Injection for Scalable Reservoir Simulation

Gabriel F. Barros, Amanda C N de Oliveira, Rômulo Montalvão Silva, Ezequiel Souza dos Santos, Rodolfo da Silva Machado de Freitas, Dakshina Murthy Valiveti, Xiao-Hui Wu, Fernando Alves Rochinha, Alvaro L. G. A. Coutinho

Keywords: Fourier Neural Operator, Time Splitting, High Performance Computing, Porous Media Flow, Reservoir Simulation, Scientific Machine Learning

TL;DR: We propose a time-splitting strategy with coordinate injection for Fourier Neural Operators that reduces GPU memory usage by 7x in reservoir simulations while maintaining accuracy on production metrics.

Abstract: In this work, we propose a Time-Splitting FNO with Coordinate Injection that partitions the temporal domain into manageable sub-windows, reducing training memory requirements regarding the time horizon. To mitigate the loss of global context inherent to splitting, we introduce an explicit time-coordinate injection mechanism that breaks the shift-invariance of the spectral operator, allowing the network to learn non-stationary dynamics. We validate our approach on the SPE10 benchmark. Results demonstrate that our method reduces peak GPU memory usage by approximately 7 times while maintaining competitive accuracy on key quantities of interest, such as oil production rates. The proposed time-splitting strategy is a promising way to reduce memory consumption in long-horizon spatio-temporal FNOs while preserving useful predictive quality on the tested benchmark.

Journal Opt In: Yes, I want to participate in the IOP focus collection submission

Submission Number: 93