Large-scale parametric PDE approximations with model-parallel Fourier neural operators

Solving PDEs to simulate two-phase flow is expensive since it involves the inversion of large ill-conditioned matrices. FNOs represent a special type of neural network capable of approximating solutions to two-phase flow equations. In order to speed up this computation, we develop a high-level software abstraction tool to exploit the linearly separable property of Fourier Transforms via Kronecker products. We perform a series of all-to-all operations where we partition the data to apply the operations in a distributed fashion. We apply these FNOs to predict the evolution of CO2-plumes in subsurface environments. Our model takes an input permeability model, and outputs time-varying CO2 saturations in a quick and cost-effective manner. Additionally, our research involves developing a distributed matrix-free abstraction library that can be used to represent any generic linear and nonlinear operator. This library is scalable and auto-differentiable thanks to the hand-written customized AD rules, allowing us to represent and train any network. We provide an evaluation on the accuracy of FNO simulations compared to traditional PDE simulations in solving various classes of PDEs.