Temporal Extrapolation and Reliable Generalization via 2U-Nets Deep Operator Network (2U-DeepONet) for time-dependent PDEs

Recent advances in scientific machine learning have demonstrated the added value of machine learning for scientific applications. In particular, neural operator learning algorithms such as Fourier Neural Operator (FNO), Deep Operator Network (DeepONet), and their extensions have demonstrated great ability to generalize to new sets of input functions reliably. However, temporal extrapolation and real-time inference remains a major challenge for these algorithms. In this work, we introduce a novel approach that can simultaneously learn the time evolution operator of parametric time-dependent partial differential equations (PDEs) and multiple other operators that map input functions to the solution space beyond the temporal training horizon. The proposed method, named 2U-Nets Deep Operator Network (2U-DeepONet) is inspired by the U-Net enhanced DeepONet (U-DeepONet), which can generalize to new sets of input functions efficiently. We incorporate U-Net blocks in both the branch and the trunk of the DeepONet, combined with an autoregressive training technique for efficient temporal extrapolation. The proposed approach allows for the accurate prediction of state variables from initial conditions and a wide range of input functions far beyond the temporal training horizon. The 2U-DeepONet stands out from other architectures in addressing the intricate challenges of generalization across various operators and the time evolution operator. By mapping function spaces and learning a spectrum of PDE solutions corresponding to diverse initial conditions and independent variables, our approach transcends the current limitations of neural operator learning. This would allow for true PDE operator learning which opens the door for absolute generalization in physics problems. We evaluated the 2U-DeepONet on a CO2 sequestration dataset. The dataset contains 5500 realizations split as 9:1:1 for training, validation, and testing, respectively. Each realization is obtained by varying 10 scalar and field input variables. The outputs are the saturation and pressure buildup data containing 24 time steps spanning 30 years. We train the 2U-DeepONet on saturation and pressure buildup data from the first year (50% of the available time steps). We test the trained 2U-DeepONet model on 500 unique realizations to measure the generalization capabilities. Additionally, we test the model on the full temporal horizon of 30 years to measure its ability to extrapolate in time. The trained 2U-DeepONet demonstrates excellent generalization capabilities on the test dataset and can extrapolate well given a single initial time step.