Neural Operator for Scientific Computing Citation Li, Zongyi (2025) Neural Operator for Scientific Computing. Dissertation (Ph.D.), California Institute of Technology. doi:10.7907/fz9s-fq86. https://resolver.caltech.edu/CaltechTHESIS:06032025-033640769 Abstract Scientific computing, which aims to accurately simulate complex physical phenomena, often requires substantial computational resources. By viewing data as continuous functions, we leverage the smoothness structures of function spaces to enable efficient large-scale simulations. We introduce the neural operator, a universal machine learning framework designed to approximate solution operators in infinite-dimensional spaces, achieving scalable physical simulations. The thesis begins with the introduction and definition of neural operators. Chapters 2-4 discuss architecture designs of neural operators including graph neural operator, multipole neural operator, and Fourier neural operator. Chapters 5-7 discuss physics-based learning techniques such as dissipative loss, physics-informed loss, and scale consistency loss. Chapters 8-10 discuss geometric neural operators with various boundary shapes, including latent space embedding, learned deformation, and optimal transport. Chapters 11-12 discuss further applications of neural operator in weather forecast and carbon capture storage. Item Type: Thesis (Dissertation (Ph.D.)) Subject Keywords: Machine Learning, Scientific Computing, Partial Differential Equations Degree Grantor: California Institute of Technology Division: Engineering and Applied Science Major Option: Computing and Mathematical Sciences Awards: Milton and Francis Clauser Doctoral Prize, 2025. Thesis Availability: Public (worldwide access) Research Advisor(s): Anandkumar, Anima Thesis Committee: Hou, Thomas Y. (chair) Anandkumar, Anima Bhattacharya, Kaushik Bruno, Oscar P. Hassanzadeh, Pedram Defense Date: 9 May 2025 Funders: Funding Agency Grant Number Nvidia Fellowship UNSPECIFIED Amazon AI4Science Fellowship UNSPECIFIED PIMCO Fellowship UNSPECIFIED Kortschak Scholarship UNSPECIFIED Record Number: CaltechTHESIS:06032025-033640769 Persistent URL: https://resolver.caltech.edu/CaltechTHESIS:06032025-033640769 DOI: 10.7907/fz9s-fq86 Related URLs: URL URL Type Description https://arxiv.org/abs/2003.03485 arXiv Article adapted for Chapter 2. https://dl.acm.org/doi/abs/10.5555/3495724.3496291 DOI Article adapted for Chapter 3. https://iclr.cc/virtual/2021/poster/3281 Publisher Article adapted for Chapter 4. https://dl.acm.org/doi/abs/10.5555/3600270.3601490 DOI Article adapted for Chapter 5. https://dl.acm.org/doi/10.1145/3648506 DOI Article adapted for Chapter 6. https://dl.acm.org/doi/abs/10.5555/3666122.3667678 DOI Article adapted for Chapter 8. https://dl.acm.org/doi/10.5555/3648699.3649087 DOI Article adapted for Chapter 9. https://arxiv.org/abs/2202.11214 arXiv Article adapted for Chapter 11. https://pubs.rsc.org/en/content/articlelanding/2023/ee/d2ee04204e Publisher Article adapted for Chapter 12. ORCID: Author ORCID Li, Zongyi 0000-0003-2081-9665 Default Usage Policy: No commercial reproduction, distribution, display or performance rights in this work are provided. ID Code: 17396 Collection: CaltechTHESIS Deposited By: Zongyi Li Deposited On: 03 Jun 2025 23:42 Last Modified: 17 Jun 2025 17:41 Thesis Files PDF - Final Version See Usage Policy. 92MB Repository Staff Only: item control page

Neural Operator for Scientific Computing Thesis by Zongyi Li In Partial Fulfillment of the Requirements for the Degree of Doctor of Philosophy CALIFORNIA INSTITUTE OF TECHNOLOGY Pasadena, California 2025 Defended May 9th, 2025 ii © 2025 Zongyi Li ORCID: 0000-0003-2081-9665 All rights reserved except where otherwise noted iii ACKNOWLEDGEMENTS First and foremost, I would like to express my deepest gratitude to my advisor, Anima Anandkumar, for her groundbreaking research vision, unwavering support, and encouragement. Thank you for introducing me to this field since my undergraduate years, and most importantly, teaching me the precious courage to pursue the difficult path. I am equally indebted to Andrew Stuart for being an inspiring mentor whose vast knowledge and incredible kindness have helped me grow both academically and personally. I extend my sincere appreciation to the members of my thesis committee, Kaushik Bhattacharya, Thomas Yizhao Hou, Oscar Bruno, and Pedram Hassanzadeh, for teaching me applied mathematics, collaborating with me, and supporting me throughout my doctoral journey. Special thanks go to my close collaborators, Nikola Kovachki and Samuel Lanthaler, who have been like second advisors to me. I have learned immensely from their expertise and approach to research. I am grateful to my early mentors, Kamyar Azizzadenasheli, Rose Yu, and Jean Kossaifi, for their guidance during the formative years of my PhD. I want to thank my senior collaborators: Tapio Schneider, Gianmarco Pinton, H. Jane Bae, Arvind Ramanathan, Morteza Gharib, Chiara Daraio who greatly extended my vision into science. I also want to thank my brilliant collaborators Burigede Liu, Daniel Zhengyu Huang, Yuanyuan Shi, Gege Wen, Vignesh Gopakumar, Di Zhou, Daniel Leibovici, and Yixuan Wang. It has been a privilege to work alongside such talented researchers. In addition, I am thankful to Karthik Duraisamy, Paris Perdikaris, Andrej Risteski, Johannes Brandstetter, Haoxing Ren, and Sanjay Choudhry for helping me. My academic journey has been enriched by the camaraderie of an amazing community: Sahin Lale, Florian Schaefer, Jiawei Zhao, Yujia Huang, Hongkai Zheng, Guanzhi Wang, Chuwei Wang, Robert Joseph George, Beom Seok Kang, David Pitt, and Jiachen Yao. I am equally thankful for the postdoctoral members of our group: Anqi Liu, Pan Xu, Huaizu Jiang, Rafal Kocielnik, Peter Jiayun Wang, Julius Berner, Kaiyu Yang, Bahareh Tolooshams, Valentin Duruisseaux, Colin White, and Luo Cheng. I have been fortunate to work with exceptional undergraduate students: David Jin, Derek Qin, Miguel Liu-Schiaffini, Kimia Hassibi, Haydn Maust, Zelin Zhao, Catherine Deng, Vansh Tibrewal, Xinyi Li, Reva Dhillon, and Michael Chen. Together, we have formed an inspiring academic family whose companionship I have truly cherished. iv I am grateful to Yisong Yue, Adam Wierman, Katie Bouman, and Georgia Gkioxari for fostering an inclusive and vibrant departmental environment. I would like to thank my fellow graduate students: Jiajie Chen, Yifan Chen, Berthy Feng, Angela Gao, Ivan Jimenez Rodriguez, Dohyeon Kim, Matthew Levine, Ziqi Ma, Nicholas Nelsen, Manuel Santana, Sabera Talukder, and Siki Wang. They have created a wonderful community that has supported me throughout this journey. Special thanks to my friends Xin Tong, Yue Liu, Han Liu, Shuming Zhang, Ran Gao, Zihui Wu, Jiaye Wu, Guandao Yang, and Tanya Marwah for the enjoyable moments and fun times we have shared. I also extend my thanks to the outstanding administrative staff in the CMS department and at Caltech, especially Alma Rangel-Fuentes, Maria Lopez, Jolene Brink, and Carmen Nemer-Sirois. It has truly been a privilege to be a student at Caltech. Finally, I want to express my profound gratitude to my father, Wen Li, my mother, Yong Lin, my brother, Jiayi Li, and my grandmother, Wenshu Peng, for their unconditional love and trust. Thank you for inspiring me since an early age and encouraging me to pursue my own path. I am deeply thankful to my partner, Jingyi Qiu, for her compassion, support, and companionship, especially during difficult times. Thank you for bringing me joy, exploring the world with me, and growing together. This thesis is dedicated to you all. Thank you! v ABSTRACT Scientific computing, which aims to accurately simulate complex physical phenomena, often requires substantial computational resources. By viewing data as continuous functions, we leverage the smoothness structures of function spaces to enable efficient large-scale simulations. We introduce the neural operator, a universal machine learning framework designed to approximate solution operators in infinite-dimensional spaces, achieving scalable physical simulations. The thesis begins with the introduction and definition of neural operators. Chapters 2-4 discuss architecture designs of neural operators including graph neural operator, multipole neural operator, and Fourier neural operator. Chapters 5-7 discuss physics-based learning techniques such as dissipative loss, physics-informed loss, and scale consistency loss. Chapters 8-10 discuss geometric neural operators with various boundary shapes, including latent space embedding, learned deformation, and optimal transport. Chapters 11-12 discuss further applications of neural operator in weather forecast and carbon capture & storage. vi PUBLISHED CONTENT AND CONTRIBUTIONS Li, Zongyi, Daniel Zhengyu Huang, et al. (2023). “Fourier neural operator with learned deformations for pdes on general geometries”. In: Journal of Machine Learning Research 24.388, pp. 1–26. Z.L. participated in the conception of the project, formulated and implemented the method, conducted experiments and analyzed results, and wrote the majority of the manuscript. Li, Zongyi, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, et al. (2020). “Neural operator: Graph kernel network for partial differential equations”. In: url: https://arxiv.org/abs/2003.03485. Z.L. participated in the conception of the project, formulated and implemented the method, conducted experiments and analyzed results, and wrote the majority of the manuscript. – (2021). “Fourier neural operator for parametric partial differential equations”. In: International Conference on Learning Representations. Z.L. participated in the conception of the project, formulated and implemented the method, conducted experiments and analyzed results, and wrote the majority of the manuscript. Li, Zongyi, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Andrew Stuart, et al. (2020). “Multipole graph neural operator for parametric partial differential equations”. In: Advances in Neural Information Processing Systems (Neurips) 33, pp. 6755–6766. Z.L. participated in the conception of the project, formulated and implemented the method, conducted experiments and analyzed results, and wrote the majority of the manuscript. Li, Zongyi, Nikola Kovachki, Chris Choy, et al. (2023). “Geometry-informed neural operator for large-scale 3D PDEs”. In: Advances in Neural Information Processing Systems (Neurips) 36. Z.L. participated in the conception of the project, formulated and implemented the method, conducted experiments and analyzed results, and wrote parts of manuscript. Li, Zongyi, Miguel Liu-Schiaffini, et al. (2022). “Learning Dissipative Dynamics in Chaotic Systems”. In: Advances in Neural Information Processing Systems (Neurips). Z.L. participated in the conception of the project, formulated and implemented the method, conducted experiments and analyzed results, and wrote the majority of the manuscript. Li, Zongyi, Hongkai Zheng, et al. (2024). “Physics-informed neural operator for learning partial differential equations”. In: ACM/JMS Journal of Data Science 1.3, pp. 1–27. vii Z.L. participated in the conception of the project, formulated and implemented the method, conducted experiments and analyzed results, and wrote the majority of the manuscript. Pathak, Jaideep et al. (2023). “Fourcastnet: A global data-driven high-resolution weather model using adaptive fourier neural operators”. In: PASC ’23: Proceedings of the Platform for Advanced Scientific Computing Conference. Z.L. participated in designing the model and methodology. Wen, Gege et al. (2023). “Real-time high-resolution CO 2 geological storage prediction using nested Fourier neural operators”. In: Energy & Environmental Science 16.4, pp. 1732–1741. Z.L. participated in designing the model and methodology. viii TABLE OF CONTENTS Acknowledgements . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . iii Abstract . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . v Published Content and Contributions . . . . . . . . . . . . . . . . . . . . . . vi Table of Contents . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . vii List of Illustrations . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . x List of Tables . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . xvi Chapter I: Introduction . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 1 1.1 Machine Learning for Scientific Computing . . . . . . . . . . . . . . 1 1.2 Neural Operator: Mapping Between Function Spaces . . . . . . . . . 4 1.3 Applications in Science and Engineering . . . . . . . . . . . . . . . 7 1.4 Foundation Models . . . . . . . . . . . . . . . . . . . . . . . . . . . 10 1.5 Outline of the Thesis. . . . . . . . . . . . . . . . . . . . . . . . . . 13 Chapter II: Model: Graph Neural Operator . . . . . . . . . . . . . . . . . . . 16 2.1 Introduction . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 16 2.2 Problem Setting . . . . . . . . . . . . . . . . . . . . . . . . . . . . 21 2.3 Graph Neural Operator . . . . . . . . . . . . . . . . . . . . . . . . . 23 2.4 Experiments . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 29 2.5 Discussion and Conclusion . . . . . . . . . . . . . . . . . . . . . . 35 Chapter III: Model: Multipole Neural Operator . . . . . . . . . . . . . . . . 40 3.1 Introduction . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 40 3.2 Related Works . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 43 3.3 Preliminaries . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 44 3.4 Multipole Graph Kernel Network . . . . . . . . . . . . . . . . . . . 46 3.5 Experiments . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 49 3.6 Discussion and Conclusion . . . . . . . . . . . . . . . . . . . . . . 52 Chapter IV: Model: Fourier Neural Operator . . . . . . . . . . . . . . . . . . 56 4.1 Introduction . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 56 4.2 Neural Operator . . . . . . . . . . . . . . . . . . . . . . . . . . . . 60 4.3 Fourier Neural Operator . . . . . . . . . . . . . . . . . . . . . . . . 61 4.4 Numerical Experiments . . . . . . . . . . . . . . . . . . . . . . . . 63 4.5 Discussion and Conclusion . . . . . . . . . . . . . . . . . . . . . . 69 Chapter V: Learning: Dissipative Loss for Operator Learning . . . . . . . . . 73 5.1 Introduction . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 73 5.2 Problem Setting . . . . . . . . . . . . . . . . . . . . . . . . . . . . 76 5.3 Learning the Markov neural operator in Chaotic Dynamics . . . . . . 79 5.4 Experiments . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 82 5.5 Discussion and Conclusion . . . . . . . . . . . . . . . . . . . . . . 91 Chapter VI: Learning: Physics Informed Learning for Neural Operator . . . . 97 6.1 Introduction . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 97 ix 6.2 Problem Settings . . . . . . . . . . . . . . . . . . . . . . . . . . . . 103 6.3 Physics-Informed Neural Operator . . . . . . . . . . . . . . . . . . . 108 6.4 Experiments . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 114 6.5 Discussion and Conclusion . . . . . . . . . . . . . . . . . . . . . . 126 Chapter VII: Learning: Scale Consistency Learning for Neural Operator . . . 131 7.1 Introduction . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 131 7.2 Related Work . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 135 7.3 Scale-Consistency . . . . . . . . . . . . . . . . . . . . . . . . . . . 136 7.4 Scale-Informed Neural Operator . . . . . . . . . . . . . . . . . . . . 141 7.5 Experiments . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 144 7.6 Discussion and Conclusion . . . . . . . . . . . . . . . . . . . . . . 147 Chapter VIII: Geometry: Neural Operator with Latent Space . . . . . . . . . 151 8.1 Introduction . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 151 8.2 Related Work . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 155 8.3 Problem setting . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 157 8.4 Geometric-Informed Neural Operator . . . . . . . . . . . . . . . . . 160 8.5 Experiments . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 163 8.6 Discussion and Conclusion . . . . . . . . . . . . . . . . . . . . . . 169 Chapter IX: Geometry: Neural Operator with Learned Deformation . . . . . 175 9.1 Introduction . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 175 9.2 Problem Settings and Preliminaries . . . . . . . . . . . . . . . . . . 179 9.3 Geometry-Aware Fourier Neural Operator . . . . . . . . . . . . . . 182 9.4 Numerical Examples . . . . . . . . . . . . . . . . . . . . . . . . . . 187 9.5 Discussion and Conclusion . . . . . . . . . . . . . . . . . . . . . . 199 Chapter X: Geometry: Neural Operator with Optimal Transport . . . . . . . . 204 10.1 Introduction . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 204 10.2 Problem Setting: Geometric Operator Learning . . . . . . . . . . . . 210 10.3 Geometry Embedding as Optimal Transport . . . . . . . . . . . . . 212 10.4 Optimal Transport Neural Operator (OTNO) . . . . . . . . . . . . . 216 10.5 Experiments . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 222 10.6 Ablation Studies . . . . . . . . . . . . . . . . . . . . . . . . . . . . 225 10.7 Discussion and Conclusion . . . . . . . . . . . . . . . . . . . . . . 232 Chapter XI: Application: Weather Forecast . . . . . . . . . . . . . . . . . . 237 11.1 Introduction . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 237 11.2 Training Methods . . . . . . . . . . . . . . . . . . . . . . . . . . . 242 11.3 Results . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 250 11.4 Computational Cost of FourCastNet . . . . . . . . . . . . . . . . . . 263 11.5 Comparison Against State-of-the-art DL Weather Prediction . . . . . 265 11.6 Implications, Discussion, and Future Work . . . . . . . . . . . . . . 266 Chapter XII: Application: Carbon Capture and Storage . . . . . . . . . . . . 273 12.1 Introduction . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 273 12.2 Methods . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 276 12.3 Results . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 282 12.4 Discussion . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 285 x LIST OF ILLUSTRATIONS Number Page 1.1 Neural operator: mapping between function spaces . . . . . . . . . . 3 1.2 Applications of neural operators: (a) FourCastNet on the ERA5 weather prediction, (b) Multi-level FNO on 3D-time multi-phase flow, (c) Geometry-informed FNO on vehicle aerodynamics, (d) FNO on camera data of MAST reactor. . . . . . . . . . . . . . . . . . . . . 7 2.1 Graph Neural Operator trained on 16 × 16 and test on 241 × 241 . . . 20 2.2 Kernel for one-dimensional Green’s function . . . . . . . . . . . . . 25 2.3 Comparison between the performance of graph kernel network and the dense neural network on the approximation of (2.3). Plots of relative 𝑙2 test approximation errors versus the number of training samples for the Graph kernel network and dense neural networks when approximating the problem (2.3). . . . . . . . . . . . . . . . . 26 3.1 V-Cycle Algorithm on Graph . . . . . . . . . . . . . . . . . . . . . 43 3.2 Hierarchical matrix decomposition . . . . . . . . . . . . . . . . . . 47 3.3 Properties of multipole graph kernel network (MGKN) on Darcy flow 50 3.4 Comparsion with benchmarks on Burgers equation, with examples of inputs and outputs. . . . . . . . . . . . . . . . . . . . . . . . . . . . 52 4.1 Zero-shot super-resolution on Navier-Stokes. . . . . . . . . . . . . . 58 4.2 The architecture of the neural operators and Fourier layer . . . . . . . 60 4.3 Benchmark on Burger’s equation, Darcy Flow, and Navier-Stokes . . 65 4.4 Results of the Bayesian inverse problem for the Navier-Stokes equation. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 68 5.1 Enforcing dissipativity on the Lorenz 63 system. . . . . . . . . . . . 75 5.2 Markov neural operator (MNO): learn global dynamics from local data. (a) Learn the MNO from the local time-evolution data. (b) Compose MNO to evaluate the long-time dynamic. (c) Architecture of the MNO. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 77 5.3 Predicted flow field of dissipative network on Lorenz-63. The red points are training points on the attractor. Dissipativity is enforced uniformly within the blue shell. . . . . . . . . . . . . . . . . . . . . 84 5.4 Trajectory and error on the KS equation . . . . . . . . . . . . . . . . 85 xi 5.5 Invariant statistics for the KS equation . . . . . . . . . . . . . . . . . 86 5.6 Choice of time step . . . . . . . . . . . . . . . . . . . . . . . . . . . 87 5.7 The learned attractor of the Kolmogorov flow. . . . . . . . . . . . . . 88 5.8 Visualization of the NS equation, Re40 . . . . . . . . . . . . . . . . 89 5.9 Invariant statistics for the NS equation . . . . . . . . . . . . . . . . . 90 5.10 Error on velocity with respect to the order of derivative. . . . . . . . 92 6.1 Physics-Informed Neural Operator (PINO) extrapolates to unseen higher frequencies. . . . . . . . . . . . . . . . . . . . . . . . . . . . 98 6.2 Architecture of Physics-Informed Neural Operator . . . . . . . . . . 100 6.3 Fourier Continuation by padding zeros. The x-axis is the spatial dimension; the y-axis is the temporal dimension. FNO extends the output smoothly on the padded domain. . . . . . . . . . . . . . . . . 114 6.4 PINO on Kolmogorov flow (left) and Lid-cavity flow (right) . . . . . 116 6.5 The accuracy-complexity trade-off on PINO, PINN, and the GPUbased pseudo-spectral solver. . . . . . . . . . . . . . . . . . . . . . . 116 6.6 Plot of test relative 𝐿 2 error versus runtime step for the Kolmogorov flow with Re500, T=0.5s. Left: resolution 64 × 64 × 65; right: resolution 128 × 128 × 129. Averaged over 20 instances. LAAFPINN: PINN with locally adaptive activation functions. SA-PINN: self-adaptive PINN. . . . . . . . . . . . . . . . . . . . . . . . . . . 117 6.7 Plot of relative 𝐿 2 error versus update step for the Kolmogorov flow with Reynolds number 500, 𝑇 = 1. . . . . . . . . . . . . . . . . . . . 122 6.8 Inverse Problem with Physics Informed Neural Operator . . . . . . . 123 6.9 Darcy inverse problem: comparing PINO forward, inverse models with numerical solver with MCMC. . . . . . . . . . . . . . . . . . . 124 7.1 Multi-scale PDE dataset: Continuum mechanics at different scales (kilometer- or millimeter-scale) can be formulated to a unit-scaled domain with corresponding scale parameters. Row 1: Darcy Flows, Row 2: Helmholtz Equation, Row 3: Navier Stokes equation. In this work, we aim to design a learning framework to capture the consistency across the scales. . . . . . . . . . . . . . . . . . . . . . 132 xii 7.2 7.3 7.4 8.1 8.2 8.3 8.4 8.5 8.6 9.1 9.2 Scale-consistency loss via sub-domain sampling and re-scaling: given data instance of input coefficient, boundary, scale parameter, and solution we restrict them to a sub-domain to obtain new data instance, which is rescaled to unit length according to the ratio of subdomain size. The scale-consistency loss is defined as the discrepancy between the global and sub-domain prediction. . . . . . . . . . . . . . . . . 133 The scale-informed neural operator has a U-shape structure on the Fourier space. The scale parameter (such as 𝑅𝑒) are embedded at each spectral layer. In the down block, the input tensors are truncated and lifted by complex layer 𝑅; in the up block, the tensors are projected and added to the inputs. Skip connections are added across the blocks. 𝑃 is the encoder and 𝑄 is the decoder. . . . . . . . . . . . . . . . . 142 Ablation Studies of Scale Consistency. left: Cost-Accuracy: we train and test each model at various sizes on Kolmogorov Flow with RE=5000. Our model (u-shape) converges faster than baseline models. Further, the model (shared) achieves comparative accuracy with 1/10 of the parameters. right: discretization convergence: the proposed model does not truncate to a fixed bandwidth. As the training resolution increases, the model’s error converges while the baseline FNO remains the same. . . . . . . . . . . . . . . . . . . . . . . . . 144 The architecture of GINO . . . . . . . . . . . . . . . . . . . . . . . 152 Comparison of GNN and GNO as the discretization becomes finer. GNN is discretization dependent, while GNO is discretization convergent. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 158 Visualization of a ground-truth pressure and corresponding prediction by GINO from the Shape-Net Car (top) and Ahmed-body (bottom) datasets, as well as the absolute error. . . . . . . . . . . . . . . 163 Cost-accuracy trade-off analysis for the drag coefficient. . . . . . . . 166 Discretization-convergence studies and zero-shot super-resolution. . . 167 Illustrations of the Ahmed-body dataset . . . . . . . . . . . . . . . . 170 Geometry-aware FNO . . . . . . . . . . . . . . . . . . . . . . . . . 178 Elasticity (a) and plasticity problems (b) introduced in section 9.4. The comparison is shown between the reference obtained using a traditional solver (left) and the Geo-FNO result (right). . . . . . . . . 187 xiii 9.3 9.4 9.5 9.6 9.7 9.8 9.9 10.1 10.2 10.3 10.4 10.5 The cost-accuracy trade-off for Geo-FNO/UNet and traditional numerical solver with implicit/explicit temporal integrators on the airfoil problem. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 188 Interpolation into different meshes for the elasticity problem . . . . . 192 Visualization of the deformation map . . . . . . . . . . . . . . . . . 193 Geo-FNO on Plasticity . . . . . . . . . . . . . . . . . . . . . . . . . 194 Advection equation on the unit sphere. It shows that 2D Geo-FNO with domain decomposition can simulate PDE on general topologies. 196 The airfoil flows (a) and pipe flows (b) introduced in section 9.4; The comparison is shown between the reference obtained using a traditional solver (top) and the Geo-FNO result (bottom). . . . . . . 196 The inverse design for the airfoil flow problem with end-to-end optimization. The optimal design using the simulation from the numerical solver matches the prediction from Geo-FNO. . . . . . . . . . . . 199 OT plans as bi-partite graphs. In the figure, the green nodes represent the input shape and the red nodes represent the latent grid. Compared to other graph mapping strategies such as ball connection and nearest neighbor connection, OT preserves the global measure, which is essential for computing integral operators. . . . . . . . . . . . . . . 207 Illustration of the optimal transport neural operator (OTNO). (a) OT Encoder: The surface mesh is encoded onto a latent computational mesh, and the OT coupling is visualized by representing the coordinates of points as RGB colors. (b) Fourier Neural Operator (FNO): Within the latent space, we apply S/FNO to calculate solutions on the latent mesh. (c) OT Decoder: We decode the solutions from the latent space back to the original surface mesh; here, the colors indicate solution values. . . . . . . . . . . . . . . . . . . . . . . . . 208 Convergence Plot: This figure presents a comparison of convergence rates among different models. We vary the physical mesh size and scale the latent mesh size along the physical size. The data indicates that our model, OTNO, exhibits the fastest exponential convergence rate of 1.85, surpassing both GINO at 1.37 and GeoFNO at 1.32. . . . 209 Illustration of OT encoder and OT decoder. The curve represents the boundary sampling points and the line denotes the latent computational grid . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 219 Results visualization for OTNO on Car Datasets . . . . . . . . . . . 224 xiv 10.6 10.7 10.8 11.1 11.2 11.3 11.4 11.5 11.6 11.7 11.8 11.9 11.10 12.1 12.2 12.3 Results Visualization for OTNO on FlowBench Dataset . . . . . . . . 227 Ablation Studies of expand factor for OTNO(Plan) . . . . . . . . . . 230 Ablation Studies of Sampling Mesh Size for OTNO(Map) and OTNO(Plan) on DrivAerNet Dataset . . . . . . . . . . . . . . . . . . . . . . . . 231 Illustrative example of a global near-surface wind forecast generated by FourCastNet over the entire globe at a resolution of 0.25° . . . . . 240 Architecture of FourCastNet . . . . . . . . . . . . . . . . . . . . . . 246 Illustration of a global Total Precipitation (TP) forecast using the FourCastNet model. . . . . . . . . . . . . . . . . . . . . . . . . . . 252 The FourCastNet model has excellent skill on forecasting fine-scale, rapidly changing variables relevant to a hurricane forecast. . . . . . . 255 Illustrative example of the utility of the FourCastNet model for forecasting atmospheric rivers. . . . . . . . . . . . . . . . . . . . . . . . 256 Latitude weighted ACC for the FourCastNet model forecasts (red line with markers) and the corresponding matched IFS forecasts (blue line with markers) . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 257 Illustration of the improvement in forecast skill of FourCastNet by utilizing large ensembles. . . . . . . . . . . . . . . . . . . . . . . . 260 The FourCastNet model shows excellent skill on forecasting overland wind speed, a challenging problem due to topographic features such as mountains and lakes. . . . . . . . . . . . . . . . . . . . . . . . . 261 Comparison of extreme percentiles between ERA5, FourCastNet, and IFS. The left panel shows the top percentiles of the 𝑇 𝑃 and 𝑈10 distribution at a forecast time of 24 hours, for a randomly sampled initial condition. The right panel shows the 𝑇 𝑃 and 𝑈10 relative quantile error (RQE, defined in the text) as a function of forecast time, averaged over 𝑁 𝑓 initial conditions in the calendar year 2018 (filled region spans the 1st and 3rd quartiles). On average, RQE trends slightly negative for both models as they under-predict the most extreme values for these variables, especially for 𝑇 𝑃. . . . . . . 262 Comparison of ACC and RMSE metrics between the (downsampled) FourCastNet predictions, (downsampled) IFS, and baseline state-ofthe-art DLWP model . . . . . . . . . . . . . . . . . . . . . . . . . . 267 Introduction to Nested-FNO. . . . . . . . . . . . . . . . . . . . . . . 275 Gas saturation prediction. . . . . . . . . . . . . . . . . . . . . . . . 277 Fine-tuning. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 280 xv 12.4 12.5 Pressure buildup prediction. . . . . . . . . . . . . . . . . . . . . . . 283 Probabilistic assessment. . . . . . . . . . . . . . . . . . . . . . . . . 284 xvi LIST OF TABLES Number Page 2.1 Comparing resolutions on full grids . . . . . . . . . . . . . . . . . . 30 2.2 Comparing number of training pairs . . . . . . . . . . . . . . . . . . 31 2.3 Generalization of resolutions on sampled Grids . . . . . . . . . . . . 32 2.4 Number of training pairs and sampling . . . . . . . . . . . . . . . . 32 2.5 Number of nodes in the training and testing . . . . . . . . . . . . . . 32 2.6 The radius and the number of nodes . . . . . . . . . . . . . . . . . . 33 2.7 Comparing the width and depth for the inner kernel network 𝜅 . . . . 34 2.8 Error of different methods . . . . . . . . . . . . . . . . . . . . . . . 35 4.1 Benchmarks on 1-d Burgers’ equation . . . . . . . . . . . . . . . . . 66 4.2 Benchmarks on 2-d Darcy Flow . . . . . . . . . . . . . . . . . . . . 66 5.1 Benchmark on vorticity for the Kolmogorov flow with Re= 40 . . . . 88 5.2 Vorticity 𝑤, velocity 𝑢, and stream function 𝑓 for the Kolmogorov flow with 𝑅𝑒 = 500 . . . . . . . . . . . . . . . . . . . . . . . . . . . 92 6.1 Operator-learning using fixed resolution data and PDE loss allows us to train operators with high accuracy on higher resolution unseen data. 118 6.2 Operator learning on Darcy Flow equation. Incorporating physics constraints in operator learning improves the performance of neural operators. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 118 6.3 Physics-informed neural operator learning on Kolmogorov flow 𝑅𝑒 = 500. PINO is effective and flexible in combining physics constraints and any amount of available data. The mean and standard error of the relative 𝐿 2 test error is reported over 200 instances and evaluated on resolution 256 × 256 × 65. . . . . . . . . . . . . . . . . . . . . . 118 6.4 Instance-wise fine-tuning on Kolmogorov flow 𝑅𝑒 = 500, 𝑇 = 0.5. Using the learned operator as the initial condition helps fine-tuning converge faster. . . . . . . . . . . . . . . . . . . . . . . . . . . . . 120 6.5 Reynolds number transfer learning. . . . . . . . . . . . . . . . . . . 121 6.6 relative L2 error of PINO (Finite-difference in time) on Kolmogorov flow with 𝑅𝑒 = 500 and 𝑇 = 0.125. PINO inherits the convergence rate of its differentiation method with no limitation of optimization. . 122 xvii 6.7 7.1 7.2 7.3 7.4 8.1 8.2 8.3 8.4 9.1 9.2 9.3 10.1 10.2 10.3 10.4 10.5 Each neural operator is trained with 400 or 4000 data points additionally sampled free initial conditions. The Reynolds number is 500. The reported generalization error is averaged over 300 instances. Training on additional initial conditions boosts the generalization ability of the operator. . . . . . . . . . . . . . . . . . . . . . . . . . 125 Comparison of FNO, UNet, UNO, and SINO with and without scaleconsistency. Models are trained at certain scale and zero-shot test across others. Overall, scale-consistency helps each model extrapolate to unseen scales. Errors are in relative-L2 (%). The Darcy Flow is scale-invariant so the SINO does not apply. . . . . . . . . . . . . 143 Comparison of scale-consistency with existing symmetries for data augmentation, in relative-L2 error (%). We train FNO on Darcy flow at 𝑠𝑐𝑎𝑙𝑒 = 4 and zero-shot test at other scales. . . . . . . . . . . . . . 146 Ablation for scale-informed neural operator on different equations in relative-L2 error (%). For Burgers’ equation, we train on viscosity 𝜈 = 1/400 and zero-shot test on other scales. For Helmholtz equation, we train on wavenumber 𝑘 = 5, 10, 25. . . . . . . . . . . . . . . . . . 147 Navier-Stokes equation trained on RE1000, zero-shot test on various RE (2+1 dimensional models). Values in percentage (%). . . . . . . . 147 Computational complexity of standard deep learning models . . . . 155 Ablation on the Ahmed-body with different sizes of the latent space . 164 Shape-Net Car dataset (3.7k mesh points). . . . . . . . . . . . . . . . 165 Ahmed-body dataset (100k mesh points). . . . . . . . . . . . . . . . 166 Benchmark on the elasticity problem. Inputs are point clouds. The table presents the mesh size as the number of input mesh points, the model size as the quantity of training parameters, the training time for each epoch, and the relative training and test errors. . . . . . . . 191 Benchmark on airfoils, pipe flows, and plasticity. Inputs are structured meshes. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 195 Advection equation on sphere . . . . . . . . . . . . . . . . . . . . . 195 Predict pressure field on ShapeNet Car Dataset . . . . . . . . . . . . 223 Predict drag coefficient (Cd) on DrivAerNet Car Dataset (single A100)224 Prediction under M1 Metric (global) on FlowBench Dataset (single A100) . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 226 Prediction under M2 Metric (boundary) on FlowBench Dataset . . . 226 Comparison of Max vs. Mean Strategy . . . . . . . . . . . . . . . . 228 xviii 10.6 10.7 10.8 10.9 Comparison of Multi vs. Single Encoder/Decoder Strategy . . . . . . 228 Comparison of Different Normal Features . . . . . . . . . . . . . . . 229 Comparison of Different Mesh Shapes in Latent Space . . . . . . . . 229 Results on the DrivAerNet Dataset for Different Voxel Sizes and Different Iteration Numbers of PPMM (MSE (e-5)) . . . . . . . . . . 230 10.10 Demonstrate topology independence of OTNO on elasticity . . . . . 233 11.1 Prognostic variables modeled by the DL model. Abbreviations are as follows. 𝑈10 (𝑉10 ): zonal (meridional) wind at 10 m; 𝑇2𝑚 : temperature at 2 m above ground; 𝑇, 𝑈, 𝑉, 𝑍, 𝑅𝐻: temperature, zonal velocity, meridional velocity, geopotential, and relative humidity (at the specified pressure level); 𝑇𝐶𝑊𝑉: total column water vapor. . . . 248 11.2 The FourCastNet model can compute a 100-member ensemble forecast on a single 4GPU A100 node. . . . . . . . . . . . . . . . . . . . 264 1 Chapter 1 INTRODUCTION 1.1 Machine Learning for Scientific Computing Advancements in machine learning (ML) are driving paradigm shifts in the field of scientific computing. Traditionally, many scientific applications rely on brute-force numerical methods using high-performance computing (HPC). Despite increasing computational power, these traditional approaches face significant speed and scalability limitations for complex systems. For instance, simulations of turbulent flows or climate systems can take weeks on supercomputers, hindering scientific progress. This thesis aims to address these challenges by leveraging ML to accelerate and improve complex scientific simulations. To address these limitations, researchers have been exploring data-driven surrogate models, particularly neural networks. However, standard neural networks, such as convolutional neural networks and transformers, were initially designed for machine learning tasks with image data and language tokens. The challenge lies in adapting these AI methods for scientific computing, where data typically consists of functions which are high-resolution, multi-scale, with complex geometries, and must adhere to physical principles. Direct use of standard deep learning models without careful adaptation or physical constraints may lead to ill-definedness and inconsistency (Kovachki et al., 2021). This raises key questions: How can AI augment or replace traditional brute-force calculations? What is the most general yet mathematically well-defined model architecture for physical systems? And can we eventually extrapolate beyond the training data to find unique designs and discover new science? Answering these questions has the potential to lead to paradigm shifts in scientific research. Continuum Mechanics. Physical phenomena can be classified by scale: quantum, microscopic, mesoscopic, and macroscopic. In recent years, massive progress has been made in microscale simulation, especially molecular dynamics for drug discovery (Jumper et al., 2021; Baek et al., 2021). This progress is largely enabled by machine learning models that can learn and predict complex molecular interactions or force fields with high accuracy, significantly accelerating the simulation process. While architectures like the transformer (Vaswani et al., 2017) have proven powerful 2 in related biomolecular tasks like protein structure prediction, the challenge of accurately and efficiently simulating the dynamics of large numbers of atoms at scale remains. However, larger scale physics and mechanics, such as fluid mechanics, solid mechanics, and climate science, present significant computational and modeling challenges. For macroscopic or mesoscopic scale problems, the domain is so large that explicitly modeling every discrete particle is computationally prohibitive. Instead, the continuum assumption is made, treating the material as continuous where nearby particles are similar. Therefore, instead of studying particle-based representations, function-based representations are used. For continuum mechanics, the physical interactions can often be described by partial differential equations (PDEs). Partial Differential Equations Consider a leading example: the Poisson-type second order elliptic equation (Darcy Flow) that describes steady state diffusion of heat or pressure: −∇ · (𝑎∇𝑢) = 𝑓 Ω (1.1) 𝑢 = 𝑔 𝜕Ω Here, a is the diffusion coefficient or heat conductivity; u is the pressure field or the temperature field; f is the forcing function; and g is the boundary condition. A standard forward problem is to find the corresponding solution function u, given the coefficient function a, forcing f, and boundary g. Conversely, the inverse problem is to find a or f given u. In the past fifty years, massive progress has been made on numerical solvers for partial differential equations, such as finite difference methods, finite element methods, and spectral methods. These methods discretize the domains or function spaces and then solve discrete linear or non-linear problems. Researchers have developed numerical methods with excellent convergence rates and stability guarantees. However, two significant challenges remain, particularly for complex, large-scale systems. First, standard numerical simulations rely on mathematical models in the form of partial differential equations, which inherently introduce aspects of simplification and approximation. For example, while PDEs like the Navier-Stokes equation are powerful models, they are not always a perfect representation of all aspects of realworld fluid motions under all conditions. Mathematically, fundamental questions remain, such as the existence of smooth solutions for the 3D incompressible Navier- 3 Figure 1.1: Neural operator: mapping between function spaces Stokes equations, with some evidence regarding related equations like the Euler equation suggesting potential blowups (Chen and Hou, 2022). Physically, viscous dissipation in the Navier-Stokes equation inherently involves energy transformation, often resulting in thermal effects. For accurate modeling of many complex real-world flows, particularly those with significant temperature variations or compressibility, a coupled set of equations, including an energy equation, is necessary to capture all relevant physics. Furthermore, standard numerical simulations can be computationally prohibitive for large-scale problems. Numerical methods require a sufficiently fine discretization that can resolve the differential operators and relevant physical scales. While methods relying on convergence (ℎ ↦→ 0) work very well for smaller scale and lower dimensional problems, Direct Numerical Simulation (DNS) of large-scale, highReynolds number flows is usually computationally unaffordable. To address this, researchers have studied sub-grid models such as the Reynolds Averaged Navier Stokes Equation (RANS) or Large Eddy Simulation (LES) with turbulence models. Such models usually introduce further approximations with empirical parameters often found in experiments. There has been ongoing debate on the universal validity and accuracy of these subgrid models for all flow regimes (Pope, 2004). The computational cost associated with fine discretization has been an obstacle for larger scale, high-resolution numerical simulations such as weather forecasting and 3D aerodynamics design. Machine learning models offer the hope to complement or go beyond traditional numerical simulations, potentially by learning directly from or being constrained by real-world observational data. 4 1.2 Neural Operator: Mapping Between Function Spaces This research focuses on developing machine learning approaches that effectively handle the complex, smooth structures inherent in scientific computing data. A central theme of this work is the development of operator learning methods, particularly neural operators, for solving partial differential equations (PDEs). For the example of Equation 1.1, the neural operator aim to approximation the solution operator mapping from the space of coefficient A = L ∞ to the space of solution U = H 1 , as shown in Figure 1.1. G : 𝑎 ↦→ 𝑢 Neural operators are a generalization of neural networks designed to learn mappings between infinite-dimensional function spaces, from input coefficients (or initial conditions) to output solutions. The neural operator is formulated as a composition of linear integral operators K with nonlinear functions 𝑓 (channel mixing), augmented with lifting (encoder) P and projection (decoder) Q. Neural operators are discretization-convergent, i.e., they share the same model parameters across different discretizations and converge to the target continuum operator as the discretization refines. G = Q ◦ ( 𝑓𝑙 ◦ K𝑙 ) ◦ . . . ◦ ( 𝑓1 ◦ K1 ) ◦ P Earlier works implemented the integral operators as message passing on the graph, leading to Graph Neural Operator (GNO) (Z. Li, Kovachki, Azizzadenesheli, Liu, Bhattacharya, et al., 2020) and its enhancement based on the Multi-pole method (MGNO) (Z. Li, Kovachki, Azizzadenesheli, Liu, A. Stuart, et al., 2020). These methods are highly flexible and they showed promising results on smaller-scale PDEs; however, the graph implementation is relatively slow due to its local connections. To efficiently capture long-range interactions, the Fourier Neural Operator (FNO), which has been a key contribution in this field, was proposed. ∫ (K𝑣)(𝑥) : 𝜅(𝑥, 𝑦)𝑣(𝑦)d𝑦 Fourier Neural Operator. Building on the concept of neural operators, the Fourier Neural Operator (FNO) was introduced (Z. Li, Kovachki, Azizzadenesheli, Liu, Bhattacharya, et al., 2021). FNO parameterizes the integral kernel directly in Fourier space while keeping nonlinear activations in the spatial domain. Utilizing the Fast Fourier Transform (𝐹), FNO scales quasi-linearly with the resolution, which offers higher efficiency compared to previous graph-based implementations. 5 (K𝑣)(𝑥) : 𝐹 −1 ( 𝜅ˆ · 𝐹𝑣) Since operations that are local in spectral or spatial space are global in the other space, this approach enables FNO to effectively perform global operations in both domains. FNO achieves state-of-the-art performance on several benchmark problems, including Burgers’ equation, Darcy Flow, and the Navier-Stokes equations. Notably, it is the first machine learning-based method to successfully model turbulent flows with zero-shot super-resolution capabilities. This approach demonstrates up to three orders of magnitude speedup compared to traditional PDE solvers while maintaining superior accuracy. Furthermore, FNO’s discretization-invariance allows it to generalize across different grid resolutions and geometries, which yields a substantial advantage over existing deep learning methods for PDEs. Several extensions based on FNO are further investigated such as the Physics-Informed Neural Operator (PINO) (Z. Li, Zheng, et al., 2024b), which employs equation loss, and the Adaptive Fourier Neural Operator (AFNO) (Guibas et al., 2021), which does not truncate to a specific Fourier mode but uses thresholding. These models have been successfully applied to various domains (Azizzadenesheli et al., 2024; Pathak et al., 2022; Liu et al., 2021). Physics Informed Learning. While the FNO demonstrated significant improvements in solving PDEs, the need to incorporate physical constraints and address limitations in data availability was recognized. This led to the development of Physics-Informed Neural Operators (PINO) (Z. Li, Zheng, et al., 2024a), a novel approach that bridges the gap between physics-informed optimization and datadriven neural operator learning for solving PDEs. PINO addresses key limitations of existing methods: (1) Physics-Informed Neural Networks (PINNs) often struggle with optimization challenges and fail to solve complex multi-scale dynamic systems. (2) Purely data-driven approaches like FNOs are limited by the availability and resolution of training data. To overcome these drawbacks, PINO combines training data with PDE constraints at different resolutions, allowing for high-fidelity reconstruction of the ground-truth operator. This hybrid approach enables PINO to learn accurate solution operators even with limited or low-resolution data, while also improving optimization stability. By incorporating PDE constraints at higher resolutions, PINO overcomes both the limitations of purely data-driven methods and the optimization challenges faced by PINNs. 6 Complex Geometry. One of the major challenges in scientific computing is handling complex 3D geometries, an area where standard FNO does not excel. While finite difference and pseudo-spectral methods work well on simple regular domains, meshing on complex geometries is extremely tedious, often requiring iterative mesh refinement. The goal was to develop a flexible method capable of handling arbitrary geometries through decomposition and deformation. To address this challenge, the Geometry-aware FNO (Geo-FNO) (Z. Li, Huang, et al., 2022), a novel framework for solving PDEs on arbitrary geometries, was developed. The core innovation lies in learning a deformation that maps the irregular physical domain to a uniform latent space where the FFT can be efficiently applied. The deformation induces generalized harmonic series on the manifold, which provides both efficiency and flexibility. For example, the cosine transform induces Chebyshev polynomials. This approach combines the computational efficiency of the FFT with the flexibility of learned deformations, resulting in a model that is as fast as FNO but more versatile and accurate for complex geometries. Geo-FNO’s effectiveness was demonstrated on a variety of PDEs, including elasticity, plasticity, Euler’s equation, and the Navier-Stokes equation, for both forward modeling and inverse design problems. Experiments yielded impressive results: Geo-FNO achieved up to 105 times acceleration compared to standard numerical solvers and halved the error rate relative to previous interpolation-based methods. Attention Mechanism and Transformer Operator Interdisciplinary efforts often foster meaningful progress, which connect diverse fields and opening new pathways to address complex challenges. Building on neural operator approaches, concepts from other areas of machine learning, particularly the attention mechanism widely used in natural language processing, were explored. In (Guibas et al., 2021), Adaptive FNO (AFNO) was introduced, a novel approach to efficient token mixing in vision transformers. AFNO leverages the geometric structure of images to frame token mixing as a continuous global convolution, which enable highresolution image processing without the quadratic complexity of self-attention. To adapt FNO for visual tasks, key modifications were made: a block-diagonal structure was imposed on channel mixing weights, adaptive weight sharing across tokens was implemented, and frequency modes were sparsified through soft-thresholding. AFNO achieves quasi-linear complexity in sequence length while maintaining high expressiveness. Experiments demonstrate AFNO’s effectiveness across various tasks, including few-shot segmentation, high-resolution image segmentation, and 7 Figure 1.2: Applications of neural operators: (a) FourCastNet on the ERA5 weather prediction, (b) Multi-level FNO on 3D-time multi-phase flow, (c) Geometryinformed FNO on vehicle aerodynamics, (d) FNO on camera data of MAST reactor. fluid dynamics simulations, where it outperforms self-attention in both efficiency and accuracy. as it outperforms self-attention in efficiency and accuracy for few-shot segmentation, high-resolution image segmentation, and fluid dynamics simulations. Universal Approximation Theorem. As neural operator architectures became more sophisticated and demonstrated impressive empirical results, it is important to understand their theoretical foundations. This led to an investigation into the universal approximation capabilities of these neural operators. In (Kovachki et al., 2021), techniques from linear approximation are used to show that neural operators, in general, can approximate any continuous solution operators, including nonlinear ones, between pairs of arbitrary Sobolev spaces. Further, in (Lanthaler, Z. Li, and A. M. Stuart, 2023) the minimum necessary components for universal approximation are identified. It is demonstrated that combining nonlocality and nonlinearity in a simple way is sufficient to approximate any arbitrary continuous operator. A key contribution is the introduction of the Averaging Neural Operator (ANO), which reduces FNO to its core essence. It is proven that even this minimal ANO architecture, which utilizes only a spatial average as its nonlocal component, exhibits universal approximation capabilities. This finding challenges the prevailing view that an increasing number of Fourier modes is necessary for improved accuracy. This analysis unifies and extends universal approximation results across a wide range of neural operator architectures, including FNO, DeepONet, and other common models. 1.3 Applications in Science and Engineering Research on neural operators has led to significant advancements in various scientific fields. By applying these techniques to real-world problems, the power and versatility of these approaches have been demonstrated: 8 Weather Forecast. Accurate and efficient weather forecasting is important for mitigating the impacts of extreme weather events, optimizing agricultural practices, and planning renewable energy resources. In collaboration with climate scientists and the Nvidia Earth-2 team, FNO’s ability to handle complex, multi-scale phenomena was leveraged to develop FourCastNet (Pathak et al., 2022). This global data-driven weather forecasting model provides accurate short to medium-range global predictions at 0.25° resolution. FourCastNet excels in forecasting highresolution, fast-timescale variables such as surface wind speed, precipitation, and atmospheric water vapor. Its performance matches or exceeds the standard ECMWF Integrated Forecasting System (IFS), particularly in small-scale features and precipitation. Remarkably, FourCastNet generates week-long forecasts in less than 2 seconds, approximately 45,000 times faster than traditional Numerical Weather Prediction (NWP) models. This speed facilitates the rapid creation of large-ensemble forecasts for improved probabilistic forecasting and uncertainty qualification. The model demonstrates exceptional performance in forecasting extreme events such as hurricanes and atmospheric rivers. It shows excellent capability in predicting near-surface wind speeds over land and coastal areas, which are critical for planning wind energy resource and managing extreme weather events. FourCastNet’s impact on weather prediction mirrors AlexNet’s revolutionary role in image classification. After validating this methodology, other leading tech companies like DeepMind and Microsoft built upon this work, helping establish a vibrant research field of weather forecast. Carbon Capture and Storage. As global efforts to combat climate change grow, carbon capture and storage (CCS) in underground reservoirs has emerged as a critical technology for reducing atmospheric CO2 levels. Accurate simulation of CO2 migration is essential for ensuring the safety and efficiency of these storage systems. Extending the neural operator approach to multi-phase flow problems, innovative machine learning frameworks for high-resolution modeling of CO2 geological storage were developed (Wen, Z. Li, Azizzadenesheli, et al., 2022). The U-FNO model accurately predicts gas saturation and pressure buildup in complex CO2-water multiphase flow scenarios, outperforming traditional CNN-based approaches in accuracy and data efficiency. Building on this, Nested FNO was introduced, which enables real-time, high-resolution 3D modeling at basin scale (Wen, Z. Li, Long, et al., 2023). Nested FNO produces forecasts nearly 700,000 times faster than existing methods, facilitating rapid probabilistic simulations that are crucial for scaling up 9 global CCS deployment. These models demonstrate superior performance in heterogeneous geological formations and critical applications such as determining gas saturation and pressure buildup fronts. By dramatically reducing computation time while maintaining high accuracy, these models provide invaluable tools for CCS project planning, risk assessment, and decision-making, enabling more comprehensive scenario analyses and risk evaluations crucial for accelerating global CCS deployment. Computational Fluid Dynamics. Aerodynamic simulation plays an important role in optimizing the performance and efficiency of vehicles, aircraft, and other engineering systems, but traditional methods often struggle with high computational costs and complexity. For aerodynamic simulation on complex shapes, the Geometry-Informed Neural Operator (GINO) (Z. Li, Kovachki, Choy, et al., 2023) is introduced. GINO combines the strengths of GNO and FNO to handle irregular grids and capture global interactions efficiently. GINO was validated on two large-scale 3D computational fluid dynamics datasets, including industry-standard Ahmed’s body geometries. Each dataset consists of 7M mesh points in the volume and 100k points on the surface. Results demonstrated that GINO achieves superior accuracy compared to existing methods, with error rates reduced by up to 25%. Moreover, GINO exhibited remarkable computational efficiency, accelerating drag coefficient calculations by a factor of 26,000 compared to the optimized GPU-based CFD simulators. This breakthrough in accuracy and speed positions GINO as a promising tool for rapid simulation and design optimization in complex engineering applications. Nuclear Fusion. Nuclear Fusion, as the ultimate energy source, has been studied for more than 50 years. Controlling nuclear fusion requires accurate and efficient plasma modeling to avoid blow-up events. Collaborating with scientists from the UK Atomic Energy Authority, FNO was applied for modeling plasma evolution in both simulations and experiments (Gopakumar et al., 2024). The FNO model accurately solved the Magneto-Hydrodynamic (MHD) equations that describe plasma dynamics, achieving speeds 6 orders of magnitude faster than traditional numerical solvers while maintaining high accuracy. This approach was further extended to model plasma evolution observed by cameras in the MAST spherical tokamak, successfully forecasting filament formation and heat deposits. The camera-based FNO model could predict the full length of a plasma shot in half the actual shot 10 time. This work represents a significant advancement in real-time plasma modeling and prediction, with potential applications in plasma control and fusion reactor optimization. 1.4 Foundation Models The long term goal of neural operator is to go beyond merely accelerating numerical solvers and focus on developing machine-learning models that can generalize outside their training data. While previous supervised machine learning approaches have shown success in scientific domains, generating the necessary training data often requires substantial computing resources for thousands of simulations, severely limiting its applicability. The diverse physical behaviors associated with different PDEs make it challenging to create a single universal model. Existing models focus on subsets of inputs, such as specific PDEs, but are often restricted to simple cases with one or two dimensions. This constraint leads to limited generalization, with models typically targeting only a single input variable while keeping other parameters fixed—unlike the flexibility of traditional numerical solvers. To address these limitations and amortize costs, the development of foundation models for scientific computing, trained on diverse datasets to capture a wide range of conditions, is proposed. Multi-Scale Learning. Capturing multi-scale phenomena in scientific and engineering simulations remains one of the biggest challenges in computational modeling, particularly for complex systems like turbulence and lithography. Large-scale simulations in turbulence or lithography can involve billions of mesh points, reaching resolution of 10, 0002 or 1, 0003 . Such multi-scale phenomena are difficult to capture accurately with single-scale models. To address this, the aim is to develop a multi-scale foundation model that can effectively handle a wide range of scales in PDEs. This approach is rooted in the principle of scale-consistency, a property inherent in many PDEs where solutions on larger domains, when restricted to subdomains, match solutions directly computed on those subdomains. To leverage this property, work is underway on a novel data augmentation scheme using sub-sampling and super-sampling techniques to bridge multi scales. This method creates virtual instances at various scales, potentially enhancing model generalization without requiring additional simulation data. This research has the potential to significantly improve the efficiency and generalizability of data-driven PDE solvers. Once generalization across different scales is achieved, the key question arises: Can 11 emergent behavior be found beyond the training scale? Specifically, if training data is available at the micro scale, can consistency be used to scale up and discover distinct behaviors at the macro scale such as hurricanes and blow-ups? By reducing the need for extensive training data and enabling more flexible, foundational models, this work could have far-reaching implications for a wide range of scientific and engineering applications, from fluid dynamics to electromagnetic simulations. Multi-Fidelity Data. In scientific computing, high-quality data are often scarce. How can learning occur from a wide variety of data including low- and high-fidelity physical simulations, sparse observational and experimental data, graphics data, and existing physics knowledge? Current computational fluid dynamics (CFD) datasets for 3D simulations are especially resource-intensive to generate. Consequently, existing CFD datasets are limited to the order of thousands of cases, which is way lower than standard ML setting. To overcome this limitation, leveraging the abundance and accessibility of large-scale graphics datasets such as ShapeNet and Objverse is proposed. The approach involves pretraining a foundation model on diverse complex geometries from these graphics datasets, incorporating various representations, including Signed Distance Functions (SDFs) and solutions to the Poisson equation with both Dirichlet and Neumann boundary conditions. This strategy aims to create a versatile foundation capable of handling a wide array of geometric challenges. This model will then be fine-tuned by integrating fluid dynamics data from coarse simulations like RANS, LES, progressing to DNS and real experimental data. The resulting foundation model has the potential to not only accelerate simulations but also enable more efficient shape design and optimization in fields requiring both intricate geometries and sophisticated physics simulations, effectively bridging the gap between graphics and CFD. AI-Aided Design. It is believed that machine learning methods will revolutionize the way design problems are approached. Design optimization problems require searching through a parameter space, which typically requires a large number of trials and runs. Due to excessive computational cost, traditional methods based on Bayesian methods and numerical solvers are usually limited to a low-dimensional parameter space. Especially in computational fluid dynamics, the standard industriallevel solvers based on RANS and LES omit all small-scale behaviors, designed in a highly empirical manner. In contrast, machine learning methods are fast, differentiable, and capable of handling high dimensional design space. The research 12 question is: Can AI discover novel designs and materials with unique geometry and structure? Initial works have shown neural operator can solve 2D inverse problem (Z. Li, Kovachki, Azizzadenesheli, Liu, Bhattacharya, et al., 2021), enhanced with physics constraint (Z. Li, Zheng, et al., 2024b) which show application in 2D airfoil (Z. Li, Kovachki, Choy, et al., 2023) and biomedical catheter design (Zhou et al., 2024). Combined with geometric methods, the aim is to further push design to 3D vehicle shapes and drone propellers. To go beyond numerical simulation used as training data, real-world experimental data can be used to augment the model (Renn et al., 2023). The goal is to discover new design shapes with less drag, and less noise, but higher resilience. In five to ten years, AI-based methods are expected to be as prevalent as standard methods such as RANS. More is Different. While classical science based on reductionism attempts to break systems down into simpler parts and identify basic rules such as PDEs, this approach aims to develop machine learning-based models that capture complex global behaviors. AI-driven models have achieved several order of speedups in forecasting weather and simulating carbon storage (Pathak et al., 2022; Wen, Z. Li, Long, et al., 2023; Ye et al., 2024; Z. Li, Kovachki, Choy, et al., 2023). In nearterm future, there has an ongoing progress towards a universal foundation model to handle a wide range of inputs such as initial conditions, coefficient conditions, domain sizes, geometries, and even multiple families of PDEs. As the model scales up to sufficient complexity, it will also reach unprecedented generalization. In long term, AI-driven methods are expected to find emergent behaviors in unseen scales and discover novel shapes and geometries for engineering design. Realizing such capabilities could lead to significant advancements in computational methods, greatly improving efficiency and accuracy in fields like rocket design and nuclear fusion. The realization of such a model would enable unprecedented efficiency and accuracy in computational methods, with far-reaching implications across multiple domains such as rocket design and nuclear fusion. Future research directions converge towards addressing the fundamental challenges in applying neural operators to scientific computing: large scales, complex geometries, and extrapolation beyond the training regime. By developing increasingly flexible and compositional models, the aim is to progress towards a universal foundation model capable of solving a wide range of simulations with different initial condition, boundary condition, coefficient function, forcing function, domain size, geometry, and even multiple families of PDEs. The realization of such a model 13 would enable unprecedented efficiency and accuracy in computational methods, with far-reaching implications across multiple domains. For instance, it could accelerate climate modeling to enhance our understanding and mitigation of climate change, optimize spacecraft design for interplanetary exploration, and improve control systems for nuclear fusion reactors. Ultimately, this research seeks to push the boundaries of scientific computing, creating tools capable of tackling some of humanity’s most challenging engineering problems and discovering new science. 1.5 Outline of the Thesis. Chapters 2-4 discuss architecture designs of neural operators including graph neural operator, multipole neural operator, and Fourier neural operator. Chapters 5-7 discuss physics-based learning techniques such as dissipative loss, physics-informed loss, and scale consistency loss. Chapters 8-10 discuss geometric neural operators with various boundary shapes, including latent space embedding, learned deformation, and optimal transport. Chapters 11-12 discuss further applications of neural operator in weather forecast, carbon capture & storage. References Azizzadenesheli, Kamyar et al. (2024). “Neural operators for accelerating scientific simulations and design”. In: Nature Reviews Physics, pp. 1–9. Baek, Minkyung et al. (2021). “Accurate prediction of protein structures and interactions using a three-track neural network”. In: Science 373.6557, pp. 871– 876. Chen, Jiajie and Thomas Y Hou (2022). “Stable nearly self-similar blowup of the 2D Boussinesq and 3D Euler equations with smooth data I: Analysis”. In: arXiv preprint arXiv:2210.07191. Gopakumar, Vignesh et al. (2024). “Fourier neural operator for plasma modelling”. In: Nuclear Fusion. Guibas, John et al. (2021). “Adaptive Fourier Neural Operators: Efficient Token Mixers for Transformers”. In: International Conference on Learning Representations, 2022. Jumper, John et al. (2021). “Highly accurate protein structure prediction with AlphaFold”. In: nature 596.7873, pp. 583–589. Kovachki, Nikola et al. (2021). “Neural operator: Learning maps between function spaces”. In: Journal of Machine Learning Research, pp. 89–1. Lanthaler, Samuel, Zongyi Li, and Andrew M Stuart (2023). “The nonlocal neural operator: Universal approximation”. In: arXiv preprint arXiv:2304.13221. 14 Li, Zongyi, Daniel Zhengyu Huang, et al. (2022). “Fourier Neural Operator with Learned Deformations for PDEs on General Geometries”. In: To appear at JMLR. Li, Zongyi, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, et al. (2020). “Neural operator: Graph kernel network for partial differential equations”. In: ICLR 2020 Workshop on Integration of Deep Neural Models and Differential Equations. – (2021). “Fourier neural operator for parametric partial differential equations”. In: International Conference on Learning Representations. Li, Zongyi, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Andrew Stuart, et al. (2020). “Multipole graph neural operator for parametric partial differential equations”. In: Advances in Neural Information Processing Systems (Neurips) 33, pp. 6755–6766. Li, Zongyi, Nikola Kovachki, Chris Choy, et al. (2023). “Geometry-informed neural operator for large-scale 3D PDEs”. In: Advances in Neural Information Processing Systems (Neurips) 36. Li, Zongyi, Hongkai Zheng, et al. (2024a). “Physics-informed neural operator for learning partial differential equations”. In: ACM-JDS. – (2024b). “Physics-informed neural operator for learning partial differential equations”. In: ACM/JMS Journal of Data Science 1.3, pp. 1–27. Liu, Burigede et al. (2021). “A learning-based multiscale method and its application to inelastic impact problems”. In: Journal of the Mechanics and Physics of Solids. Pathak, Jaideep et al. (2022). “Fourcastnet: A global data-driven high-resolution weather model using adaptive fourier neural operators”. In: arXiv preprint arXiv:2202.11214. Pope, Stephen B (2004). “Ten questions concerning the large-eddy simulation of turbulent flows”. In: New journal of Physics 6.1, p. 35. Renn, Peter I et al. (2023). “Forecasting subcritical cylinder wakes with Fourier Neural Operators”. In: arXiv preprint arXiv:2301.08290. Vaswani, Ashish et al. (2017). “Attention is All you Need”. In: Advances in Neural Information Processing Systems. Ed. by I. Guyon et al. Vol. 30. Curran Associates, Inc. Wen, Gege, Zongyi Li, Kamyar Azizzadenesheli, et al. (2022). “U-FNO—An enhanced Fourier neural operator-based deep-learning model for multiphase flow”. In: Advances in Water Resources 163, p. 104180. Wen, Gege, Zongyi Li, Qirui Long, et al. (2023). “Real-time high-resolution CO 2 geological storage prediction using nested Fourier neural operators”. In: Energy & Environmental Science 16.4, pp. 1732–1741. Ye, Teng et al. (2024). “Using Artificial Intelligence to Unlock Crowdfunding Success for Small Businesses”. In: arXiv preprint arXiv:2407.09480. 15 Zhou, Tingtao et al. (2024). “AI-aided geometric design of anti-infection catheters”. In: Science Advances 10.1, eadj1741. 16 Chapter 2 MODEL: GRAPH NEURAL OPERATOR The classical development of neural networks has been primarily for mappings between a finite-dimensional Euclidean space and a set of classes, or between two finite-dimensional Euclidean spaces. The purpose of this work is to generalize neural networks so that they can learn mappings between infinite-dimensional spaces (operators). The key innovation in our work is that a single set of network parameters, within a carefully designed network architecture, may be used to describe mappings between infinite-dimensional spaces and between different finite-dimensional approximations of those spaces. We formulate approximation of the infinite-dimensional mapping by composing nonlinear activation functions and a class of integral operators. The kernel integration is computed by message passing on graph networks. This approach has substantial practical consequences which we will illustrate in the context of mappings between input data to partial differential equations (PDEs) and their solutions. In this context, such learned networks can generalize among different approximation methods for the PDE (such as finite difference or finite element methods) and among approximations corresponding to different underlying levels of resolution and discretization. Experiments confirm that the proposed graph kernel network does have the desired properties and show competitive performance compared to the state of the art solvers. 2.1 Introduction There are numerous applications in which it is desirable to learn a mapping between Banach spaces. In particular, either the input or the output space, or both, may be infinite-dimensional. The possibility of learning such mappings opens up a new class of problems in the design of neural networks, with widespread potential applicability. New ideas are required to build on traditional neural networks which are mappings from finite-dimensional Euclidean spaces into classes, or into another finite-dimensional Euclidean space. We study the development of neural networks in the setting in which the input and output spaces comprise real-valued functions defined on subsets in R𝑑 . 17 Related Work We formulate a new class of neural networks, which are defined to map between spaces of functions on a bounded open set 𝐷 in R𝑑 . Such neural networks, once trained, have the important property that they are discretization invariant, sharing the same network parameters between different discretizations. In contrast, standard neural network architectures depend heavily on the discretization and have difficulty in generalizing between different grid representations. Our methodology has an underlying Nyström approximation formulation (Nyström, 1930) which links different grids to a single set of network parameters. We illustrate the new conceptual class of neural networks within the context of partial differential equations, and the mapping between input data (in the form of a function) and output data (the function which solves the PDE). Both supervised and semisupervised settings are considered. In PDE applications, the defining equations are often local, whilst the solution operator has non-local effects which, nonetheless, decay. Such non-local effects can be described by integral operators with graph approximations of Nyström type (Belongie et al., 2002) providing a consistent way of connecting different grid or data structures arising in computational methods. For this reason, graph networks hold great potential for the solution operators of PDEs, which is the point of departure for our work. Partial differential equations (PDEs). A wide range of important engineering and physical problems are governed by PDEs. Over the past few decades, significant progress has been made on formulating (Gurtin, 1982) and solving (Johnson, 2012) the governing PDEs in many scientific fields from micro-scale problems (e.g., quantum and molecular dynamics) to macro-scale applications (e.g., civil and marine engineering). Despite the success in the application of PDEs to solve real-life problems, two significant challenges remain. First, identifying/formulating the underlying PDEs appropriate for the modeling of a specific problem usually requires extensive prior knowledge in the corresponding field which is then combined with universal conservation laws to design a predictive model; for example, modeling the deformation and fracture of solid structures requires detailed knowledge on the relationship between stress and strain in the constituent material. For complicated systems such as living cells, acquiring such knowledge is often elusive and formulating the governing PDE for these systems remains prohibitive; the possibility of learning such knowledge from data may revolutionize such fields. Second, solving complicated non-linear PDE systems (such as those arising in turbulence and plas- 18 ticity) is computationally demanding; again the possibility of using instances of data from such computations to design fast approximate solvers holds great potential. In both these challenges, if neural networks are to play a role in exploiting the increasing volume of available data, then there is a need to formulate them so that they are well-adapted to mappings between function spaces. We first outline two major neural network based approaches for PDEs. We consider PDEs of the form (L𝑎 𝑢)(𝑥) = 𝑓 (𝑥), 𝑥∈𝐷 (2.1) 𝑢(𝑥) = 0, 𝑥 ∈ 𝜕𝐷, with solution 𝑢 : 𝐷 → R, and parameter 𝑎 : 𝐷 → R entering the definition of L𝑎 . The domain 𝐷 is discretized into 𝐾 points (see Section 2.2) and 𝑁 training pairs of coefficient functions and (approximate) solution functions {𝑎 𝑗 , 𝑢 𝑗 } 𝑁𝑗=1 are used to design a neural network. The first approach parametrizes the solution operator as a deep convolutional neural network between finite Euclidean space G : R𝐾 × Θ → R𝐾 (Guo, Li, and Iorio, 2016; Zhu and Zabaras, 2018; Adler and Oktem, 2017; Bhatnagar et al., 2019). Such an approach is, by definition, not mesh independent and will need modifications to the architecture for different resolution and discretization of 𝐾 in order to achieve consistent error (if at all possible). We demonstrate this issue in section 4.4 using the architecture of (Zhu and Zabaras, 2018) which was designed for the solution of (2.3) on a uniform 64 × 64 mesh. Furthermore, this approach is limited to the discretization size and geometry of the training data hence it is not possible to query solutions at new points in the domain. In contrast we show, for our method, both invariance of the error to grid resolution, and the ability to transfer the solution between meshes in section 4.4. Especially, (Ummenhofer et al., 2020) purposed continuous convolution network for fluid problems, where they sample points off grid and do linear interpolation on the grid, which may the convolution continuous compared to common CNN based method. However continuous convolution is still constrained by the underlying grid and preventing generalization of resolutions. The second approach directly parameterizes the solution 𝑢 as a neural network G : 𝐷 × Θ → R (E and Yu, 2018; Raissi, Perdikaris, and Karniadakis, 2019; Bar and Sochen, 2019). This approach is, of course, mesh independent since the solution is defined on the physical domain. However, the parametric dependence is accounted for in a mesh-dependent fashion. Indeed, for any given new equation with new coefficient function 𝑎, one would need to train a new neural network G𝑎 . Such an 19 approach closely resembles classical methods such as finite elements, replacing the linear span of a finite set of local basis functions with the space of neural networks. This approach suffers from the same computational issue as the classical methods: one needs to solve an optimization problem for every new parameter. Furthermore, the approach is limited to a setting in which the underlying PDE is known; purely data-driven learning of a map between spaces of functions is not possible. The methodology we introduce circumvents these issues. Our methodology most closely resembles the classical reduced basis method (R. A. DeVore, 2014) or the method of (Cohen and R. DeVore, 2015). Along with the contemporaneous work (Bhattacharya, Kovachki, and Stuart, 2020), our method, to the best of our knowledge, is the first practical deep learning method that is designed to learn maps between infinite-dimensional spaces. It remedies the mesh-dependent nature of the approach in (Guo, Li, and Iorio, 2016; Zhu and Zabaras, 2018; Adler and Oktem, 2017; Bhatnagar et al., 2019) by producing a quality of approximation that is invariant to the resolution of the function and it has the ability to transfer solutions between meshes. Moreover it need only be trained once on the equations set {𝑎 𝑗 , 𝑢 𝑗 } 𝑁𝑗=1 ; then, obtaining a solution for a new 𝑎 ∼ 𝜇, only requires a forward pass of the network, alleviating the major computational issues incurred in (E and Yu, 2018; Raissi, Perdikaris, and Karniadakis, 2019; Herrmann, Schwab, and Zech, 2020; Bar and Sochen, 2019). Lastly, our method requires no knowledge of the underlying PDE; the true map G † can be treated as a black-box, perhaps trained on experimental data or on the output of a costly computer simulation, not necessarily a PDE. Graph neural networks. Graph neural network (GNNs), a class of neural networks that apply on graph-structured data, have recently been developed and seen a variety of applications. Graph networks incorporate an array of techniques such as graph convolution, edge convolution, attention, and graph pooling (Kipf and Welling, 2016; Hamilton, Ying, and Leskovec, 2017; Gilmer et al., 2017; Veličković et al., 2017; Murphy et al., 2018). GNNs have also been applied to the modeling of physical phenomena such as molecules (Chen et al., 2019) and rigid body systems (Battaglia et al., 2018), as these problems exhibit a natural graph interpretation: the particles are the nodes and the interactions are the edges. The work (Alet et al., 2019) performed an initial study that employs graph networks on the problem of learning solutions to Poisson’s equation among other physical 20 Graph kernel network for the solution of (2.3). It can be trained on a small resolution and will generalize to a large one. The Error is point-wise absolute squared error. Figure 2.1: Graph Neural Operator trained on 16 × 16 and test on 241 × 241 applications. They propose an encoder-decoder setting, constructing graphs in the latent space and utilizing message passing between the encoder and decoder. However, their model uses a nearest neighbor structure that is unable to capture nonlocal dependencies as the mesh size is increased. In contrast, we directly construct a graph in which the nodes are located on the spatial domain of the output function. Through message passing, we are then able to directly learn the kernel of the network which approximates the PDE solution. When querying a new location, we simply add a new node to our spatial graph and connect it to the existing nodes, avoiding interpolation error by leveraging the power of the Nyström extension for integral operators. Continuous neural networks. The concept of defining neural networks in infinitedimensional spaces is a central problem that long been studied (Williams, 1996; Neal, 1996; Roux and Bengio, 2007; Globerson and Livni, 2016; Guss, 2016). The general idea is to take the infinite-width limit which yields a non-parametric method and has connections to Gaussian Process Regression (Neal, 1996; Matthews et al., 2018; Garriga-Alonso, Rasmussen, and Aitchison, 2018), leading to the introduction of deep Gaussian processes (aretha; Damianou and Lawrence, 2013). Thus far, such methods have not yielded efficient numerical algorithms that can parallel the success of convolutional or recurrent neural networks in finite dimensions in the setting of mappings between function spaces. Another idea is to simply define a sequence of compositions where each layer is a map between infinite-dimensional spaces with a finite-dimensional parametric dependence. This is the approach we take in this work, going a step further by sharing parameters between each layer. 21 Contributions We introduce the concept of Neural Operator and instantiate it through graph kernel networks, a novel deep neural network method to learn the mapping between infinitedimensional spaces of functions defined on bounded open subsets of R𝑑 . • Unlike existing methods, our approach is demonstrably able to learn the mapping between function spaces, and is invariant to different approximations and grids. It can do zero-shot super-resolution as demonstrated in Figure 2.1. • By composing nonlinear activation functions and nonlocal integral operators, the learn operators can approximate complex nonlinear nonlocal solution operator raised in various PDEs. • The standard integral operation requires quadratic time complexity in terms of the number of points, which is usually intractable. We study four efficient approximation schemes that have tractable complexity. • Numerical results show that the Neural Operator has state-of-the-art accuracy comparing to previous data-driven finite-dimensional operators in the fixeddiscretization setting. Meanwhile, it is also up to three order of magnitude faster than the conventional numerical solvers. These concepts are illustrated in the context of a family of elliptic PDEs prototypical of a number of problems arising throughout the sciences and engineering. 2.2 Problem Setting Our goal is to learn a mapping between two infinite dimensional spaces by using a finite collection of observations of input-output pairs from this mapping. Let A and U be separable Banach spaces and G † : A → U a (typically) non-linear map. Suppose we have observations {𝑎 𝑗 , 𝑢 𝑗 } 𝑁𝑗=1 where 𝑎 𝑗 ∼ 𝜇 is an i.i.d. sequence from the probability measure 𝜇 supported on A and 𝑢 𝑗 = G † (𝑎 𝑗 ) is possibly corrupted with noise. We aim to build an approximation of G † by constructing a parametric map G : A×Θ→U (2.2) for some finite-dimensional parameter space Θ and then choosing 𝜃 † ∈ Θ so that G(·, 𝜃 † ) ≈ G † . 22 This is a natural framework for learning in infinite-dimensions as one could define a cost functional 𝐶 : U × U → R and seek a minimizer of the problem min E𝑎∼𝜇 [𝐶 (G(𝑎, 𝜃), G † (𝑎))] 𝜃∈Θ which directly parallels the classical finite-dimensional setting (Vapnik, 1998). Showing the existence of minimizers, in the infinite-dimensional setting, remains a challenging open problem. We will approach this problem in the test-train setting in which empirical approximations to the cost are used. We conceptualize our methodology in the infinite-dimensional setting. This means that all finitedimensional approximations can share a common set of network parameters which are defined in the (approximation-free) infinite-dimensional setting. To be concrete we will consider infinite-dimensional spaces which are Banach spaces of real-valued functions defined on a bounded open set in R𝑑 . We then consider mappings G † which take input functions to a PDE and map them to solutions of the PDE, both input and solutions being real-valued functions on R𝑑 . A common instantiation of the preceding problem is the approximation of the second order elliptic PDE −∇ · (𝑎(𝑥)∇𝑢(𝑥)) = 𝑓 (𝑥), 𝑢(𝑥) = 0, 𝑥∈𝐷 𝑥 ∈ 𝜕𝐷 (2.3) for some bounded, open set 𝐷 ⊂ R𝑑 and a fixed function 𝑓 ∈ 𝐿 2 (𝐷; R). This equation is prototypical of PDEs arising in numerous applications including hydrology (Bear and Corapcioglu, 2012) and elasticity (Antman, 2005). For a given 𝑎 ∈ A = 𝐿 ∞ (𝐷; R+ ) ∩ 𝐿 2 (𝐷; R+ ), equation (2.3) has a unique weak solution 𝑢 ∈ U = 𝐻01 (𝐷; R) (Evans, 2010) and therefore we can define the solution operator G † as the map 𝑎 ↦→ 𝑢. Note that while the PDE (2.3) is linear, the solution operator G † is not. Since our data 𝑎 𝑗 and 𝑢 𝑗 are , in general, functions, to work with them numerically, we assume access only to point-wise evaluations. To illustrate this, we will continue with the example of the preceding paragraph. To this end let 𝑃𝐾 = {𝑥 1 , . . . , 𝑥 𝐾 } ⊂ 𝐷 be a 𝐾-point discretization of the domain 𝐷 and assume we have observations 𝑎 𝑗 | 𝑃𝐾 , 𝑢 𝑗 | 𝑃𝐾 ∈ R𝐾 , for a finite collection of input-output pairs indexed by 𝑗. In the next section, we propose a kernel inspired graph neural network architecture which, while trained on the discretized data, can produce an answer 𝑢(𝑥) for any 𝑥 ∈ 𝐷 given a new input 𝑎 ∼ 𝜇. That is to say that our approach is independent of the 23 discretization 𝑃𝐾 and therefore a true function space method; we verify this claim numerically by showing invariance of the error as 𝐾 → ∞. Such a property is highly desirable as it allows a transfer of solutions between different grid geometries and discretization sizes. We note that, while the application of our methodology is based on having point-wise evaluations of the function, it is not limited by it. One may, for example, represent a function numerically as a finite set of truncated basis coefficients. Invariance of the representation would then be with respect to the size of this set. Our methodology can, in principle, be modified to accommodate this scenario through a suitably chosen architecture. We do not pursue this direction in the current work. 2.3 Graph Neural Operator We propose the graph neural operator (graph kernel neural network) for the solution of the problem outlined in section 2.2. As a guiding principle for our architecture, we take the following example. Let L𝑎 be a differential operator depending on a parameter 𝑎 ∈ A and consider the PDE (L𝑎 𝑢)(𝑥) = 𝑓 (𝑥), 𝑥∈𝐷 𝑢(𝑥) = 0, 𝑥 ∈ 𝜕𝐷 (2.4) for a bounded, open set 𝐷 ⊂ R𝑑 and some fixed function 𝑓 living in an appropriate function space determined by the structure of L𝑎 . The elliptic operator L𝑎 · = −div(𝑎∇·) from equation (2.3) is an example. Under fairly general conditions on L𝑎 (Evans, 2010), we may define the Green’s function 𝐺 : 𝐷 × 𝐷 → R as the unique solution to the problem L𝑎 𝐺 (𝑥, ·) = 𝛿𝑥 where 𝛿𝑥 is the delta measure on R𝑑 centered at 𝑥. Note that 𝐺 will depend on the parameter 𝑎 thus we will henceforth denote it as 𝐺 𝑎 . The solution to (2.4) can then be represented as ∫ 𝐺 𝑎 (𝑥, 𝑦) 𝑓 (𝑦) 𝑑𝑦. 𝑢(𝑥) = 𝐷 This is easily seen via the formal computation ∫ (L𝑎 𝑢)(𝑥) = (L𝑎 𝐺 (𝑥, ·))(𝑦) 𝑓 (𝑦) 𝑑𝑦 𝐷 ∫ = 𝛿𝑥 (𝑦) 𝑓 (𝑦) 𝑑𝑦 𝐷 = 𝑓 (𝑥) (2.5) 24 Generally the Green’s function is continuous at points 𝑥 ≠ 𝑦, for example, when L𝑎 is uniformly elliptic (Gilbarg and Trudinger, 2015), hence it is natural to model it via a neural network. Guided by the representation (3.1), we propose the following iterative architecture for 𝑡 = 0, . . . , 𝑇 − 1.   ∫ 𝑣 𝑡+1 (𝑥) = 𝜎 𝑊𝑣 𝑡 (𝑥)+ 𝜅 𝜙 (𝑥, 𝑦, 𝑎(𝑥), 𝑎(𝑦))𝑣 𝑡 (𝑦) 𝜈𝑥 (𝑑𝑦) (2.6) 𝐷 where 𝜎 : R → R is a fixed function applied element-wise, 𝜈𝑥 is a fixed Borel measure for each 𝑥 ∈ 𝐷. The matrix 𝑊 ∈ R𝑛×𝑛 , together with the parameters 𝜙 entering kernel 𝜅 𝜙 : R2(𝑑+1) → R𝑛×𝑛 , are to be learned from data. We model 𝜅 𝜙 as a neural network mapping R2(𝑑+1) to R𝑛×𝑛 . Discretization of the continuum picture may be viewed as replacing Borel measure 𝜈𝑥 by an empirical approximation based on the 𝐾 grid points being used. In this setting we may view 𝜅 𝜙 as a 𝐾 × 𝐾 kernel block matrix, where each entry 𝜅 𝜙 (𝑥, 𝑦) is itself a 𝑛 × 𝑛 matrix. Each block shares the same set of network parameters. This is the key to making a method that shares common parameters independent of the discretization used. Finally, we observe that, although we have focussed on neural networks mapping 𝑎 to 𝑢, generalizations are possible, such as mapping 𝑓 to 𝑢, or having non-zero boundary data 𝑔 on 𝜕𝐷 and mapping 𝑔 to 𝑢. More generally one can consider the mapping from (𝑎, 𝑓 , 𝑔) into 𝑢 and use similar ideas. Indeed to illustrate ideas we will consider the mapping from 𝑓 to 𝑢 below (which is linear and for which an analytic solution is known) before moving on to study the (nonlinear) mapping from 𝑎 to 𝑢. Example 1: Poisson equation. We consider a simplification of the foregoing in which we study the map from 𝑓 to 𝑢. To this end we set 𝑣 0 (𝑥) = 𝑓 (𝑥), 𝑇 = 1, 𝑛 = 1, 𝜎(𝑥) = 𝑥, 𝑊 = 𝑤 = 0, and 𝜈𝑥 (𝑑𝑦) = 𝑑𝑦 (the Lebesgue measure) in (2.6). We then obtain the representation (3.1) with the Green’s function 𝐺 𝑎 parameterized by the neural network 𝜅 𝜙 with explicit dependence on 𝑎(𝑥), 𝑎(𝑦). Now consider the setting where 𝐷 = [0, 1] and 𝑎(𝑥) ≡ 1, so that (2.3) reduces to the 1-dimensional Poisson equation with explicitly computable Green’s function. Indeed, 𝐺 (𝑥, 𝑦) = 1 (𝑥 + 𝑦 − |𝑦 − 𝑥|) − 𝑥𝑦. 2 Note that although the map 𝑓 ↦→ 𝑢 is, in function space, linear, the Green’s function itself is not linear in either argument. Figure 2.2 shows 𝜅 𝜙 after training with 25 Proof of concept: graph kernel network on 1 dimensional Poisson equation; comparison of learned and truth kernel. Figure 2.2: Kernel for one-dimensional Green’s function 𝑁 = 2048 samples 𝑓 𝑗 ∼ 𝜇 = N (0, (−Δ + 𝐼) −1 ) with periodic boundary conditions on the operator −Δ + 𝐼. (samples from this measure can be easily implemented by means of a random Fourier series – Karhunen-Loeve – see (Lord, Powell, and Shardlow, 2014)). Example 2: 2D Poisson equation. We further demonstrate the power of the graph kernel network by extending the Poisson equation studied in example 1 to the two dimensional (2D) case, where we approximate the map 𝑓 (𝑥) ↦→ 𝑢(𝑥) where 𝑥 ∈ 𝐷 = [0, 1] × [0, 1]. We consider two approaches: (𝑖) graph kernel network to approximate the 2D Green’s function 𝐺 (𝑥, 𝑦) and (𝑖𝑖) dense neural network with 𝑓 (𝑥) as input and 𝑢(𝑥) as output such that the mapping 𝑓 (𝑥) ↦→ 𝑢(𝑥) is directly approximated. The two neural networks are trained with the same training sets of different sizes ranging from 1 to 100 samples and tested with 1000 test samples. We observe that the Graph kernel network approximates the map with a minimum (≈ 5) number of training samples while still possessing a smaller test error comparing to that of a dense neural network trained with 100 samples. Therefore, the Graph kernel network potentially significantly reduces the number of required training samples to approximate the mapping. This property is especially important in practice as obtaining a huge number of training data for certain engineering/physics problems is always prohibitive. The reason for the Graph kernel network to have such strong approximation power is because that it is able to learn the truth Green’s function for the Poisson equation, as already demonstrated in the 1D case. 26 Figure 2.3: Comparison between the performance of graph kernel network and the dense neural network on the approximation of (2.3). Plots of relative 𝑙 2 test approximation errors versus the number of training samples for the Graph kernel network and dense neural networks when approximating the problem (2.3). Algorithmic framework. The initialization 𝑣 0 (𝑥) to our network (2.6) can be viewed as the initial guess we make for the solution 𝑢(𝑥) as well as any other dependence we want to make explicit. A natural choice is to start with the coefficient 𝑎(𝑥) itself as well as the position in physical space 𝑥. This (𝑑 + 1)-dimensional vector field is then lifted to a 𝑛-dimensional vector field, an operation which we may view as the first layer of the overarching neural network. This is then used as an initialization to the kernel neural network, which is iterated 𝑇 times. In the final layer, we project back to the scalar field of interest with another neural network layer. Due to the smoothing effect of the inverse elliptic operator in (2.3) with respect to the input data 𝑎 (and indeed 𝑓 when we consider this as input), we augment the initialization (𝑥, 𝑎(𝑥)) with a Gaussian smoothed version of the coefficients 𝑎 𝜖 (𝑥), together with their gradient ∇𝑎 𝜖 (𝑥). Thus we initialize with a 2(𝑑 + 1)-dimensional vector field. Throughout this paper the Gaussian smoothing is performed with a centred isotropic Gaussian with variance 5. The Borel measure 𝜈𝑥 is chosen to be 27 the Lebesgue measure supported on a ball at 𝑥 of radius 𝑟. Thus we have 𝑣 0 (𝑥) = 𝑃(𝑥, 𝑎(𝑥), 𝑎 𝜖 (𝑥), ∇𝑎 𝜖 (𝑥)) + 𝑝 ∫  
𝑣 𝑡+1 (𝑥) = 𝜎 𝑊𝑣 𝑡 (𝑥) + 𝜅 𝜙 𝑥, 𝑦, 𝑎(𝑥), 𝑎(𝑦) 𝑣 𝑡 (𝑦) d𝑦 (2.7) (2.8) 𝐵(𝑥,𝑟) 𝑢(𝑥) = 𝑄𝑣𝑇 (𝑥) + 𝑞 (2.9) where 𝑃 ∈ R𝑛×2(𝑑+1) , 𝑝 ∈ R𝑛 , 𝑣 𝑡 (𝑥) ∈ R𝑛 and 𝑄 ∈ R1×𝑛 , 𝑞 ∈ R. The integration in (2.8) is approximated by a Monte Carlo sum via a message passing graph network with edge weights (𝑥, 𝑦, 𝑎(𝑥), 𝑎(𝑦)). The choice of measure 𝜈𝑥 (d𝑦) = 1𝐵(𝑥,𝑟) d𝑦 is two-fold: 1) it allows for more efficient computation and 2) it exploits the decay property of the Green’s function. Note that if more information is known about the true kernel, it can be added into this measure. For example, if we know the true kernel has a Gaussian structure, we can define 𝜈𝑥 (d𝑦) = 1𝐵(𝑥,𝑟) 𝜌𝑥 (𝑦)d𝑦 where 𝜌𝑥 (𝑦) is a Gaussian density. Then 𝜅 𝜙 will need to learn a much less complicated function. We however do not pursue this direction in the current line of work. Message passing graph networks. Message passing graph networks comprise a standard architecture employing edge features (Gilmer et al., 2017). If we properly construct the graph on the spatial domain 𝐷 of the PDE, the kernel integration can be viewed as an aggregation of messages. Given node features 𝑣 𝑡 (𝑥) ∈ R𝑛 , edge features 𝑒(𝑥, 𝑦) ∈ R𝑛𝑒 , and a graph 𝐺, the message passing neural network with averaging aggregation is  𝑣 𝑡+1 (𝑥) = 𝜎 𝑊𝑣 𝑡 (𝑥) +  ∑︁
1 𝜅 𝜙 𝑒(𝑥, 𝑦) 𝑣 𝑡 (𝑦) |𝑁 (𝑥)| (2.10) 𝑦∈𝑁 (𝑥)
where 𝑊∈ R𝑛×𝑛 , 𝑁 (𝑥) is the neighborhood of 𝑥 according to the graph, 𝜅 𝜙 𝑒(𝑥, 𝑦) is a neural network taking as input edge features and as output a matrix in R𝑛×𝑛 . Relating to (2.8), 𝑒(𝑥, 𝑦) = (𝑥, 𝑦, 𝑎(𝑥), 𝑎(𝑦)) ∈ R2(𝑑+1) . Graph construction. To use the message passing framework (3.4), we need to design a graph that connects the physical domain 𝐷 of the PDE. The nodes are chosen to be the 𝐾 discretized spatial locations. Here we work on a standard uniform mesh, but there are many other possibilities such as finite-element triangulations and random points at which data is acquired. The edge connectivity is then chosen according to the integration measure in (2.8), namely Lebesgue restricted to a ball. Each node 𝑥 ∈ R𝑑 is connected to nodes which lie within 𝐵(𝑥, 𝑟), defining the 28 neighborhood set 𝑁 (𝑥). Then for each neighbor 𝑦 ∈ 𝑁 (𝑥), we assign the edge weight 𝑒(𝑥, 𝑦) = (𝑥, 𝑦, 𝑎(𝑥), 𝑎(𝑦)). Equation (3.4) can then be viewed as a Monte Carlo approximation of (2.8). This local structure allows for more efficient computation while remaining invariant to mesh-refinement. Indeed, since the radius parameter 𝑟 is chosen in physical space, the size of the set 𝑁 (𝑥) grows as the discretization size 𝐾 grows. This is a key feature that makes our methodology mesh-independent. Nyström approximation of the kernel. While the aforementioned graph structure severely reduces the computational overhead of integrating over the entire domain 𝐷 (corresponding to a fully-connected graph), the number of edges still scales like O (𝐾 2 ). To overcome this, we employ a random Nyström-type approximation of the kernel. In particular, we uniformly sample 𝑚 ≪ 𝐾 nodes from the original graph, constructing a new random sub-graph. This process is repeated 𝑙 ∈ N times, yielding 𝑙 random sub-graphs each with 𝑚 nodes. This can be thought of as a way of reducing the variance in the estimator. We use these sub-graphs when evaluating (3.4) during training, leading to the more favorable scaling O (𝑙𝑚 2 ). Indeed, numerically we find that 𝑙 = 4 and 𝑚 = 200 is sufficient even when 𝐾 = 4212 = 177, 241. In the evaluation phase, when we want the solution on a particular mesh geometry, we simply partition the mesh into sub-graphs each with 𝑚 nodes and evaluate each separately. We will now highlight the quality of this kernel approximation in a RHKS setting. A real Reproducing Kernel Hilbert Space (RKHS) (H , ⟨·, ·⟩, ∥ · ∥) is a Hilbert space of functions 𝑓 : 𝐷 → R where point-wise evaluation is a continuous linear functional, i.e. | 𝑓 (𝑥)| ≤ 𝐶 ∥ 𝑓 ∥ for some constant 𝐶 ≥ 0, independent of 𝑥. For every RHKS, there exists a unique, symmetric, positive definite kernel 𝜅 : 𝐷 × 𝐷 → R, which gives the representation 𝑓 (𝑥) = ⟨ 𝑓 , 𝜅(·, 𝑥)⟩. Let 𝑇 : H → H be a linear operator on H acting via the kernel ∫ 𝑇𝑓 = 𝜅(·, 𝑦) 𝑓 (𝑦)𝜈(𝑑𝑦). 𝐵(·,𝑟) Let 𝑇𝑚 : H → H be its 𝑚-point empirical approximation ∫ 𝑇𝑚 = 𝜅(·, 𝑦) 𝑓 (𝑦)𝜈𝑚 (𝑑𝑦), 𝐵(·,𝑟) 𝑚 1 ∑︁ 𝛿 𝑦 (𝑑𝑦), 𝜈𝑚 (𝑑𝑦) = 𝑚 𝑘=1 𝑘 29 so that 𝑚 1 ∑︁ 𝑇𝑚 𝑓 = 𝜅(·, 𝑦 𝑘 ) 𝑓 (𝑦 𝑘 ). 𝑚 𝑘=1 The error of this approximation achieves the Monte Carlo rate 𝑂 (𝑚 −1/2 ): Proposition 2.3.1 Suppose E𝑦∼𝜈 [𝜅(·, 𝑦) 4 ] < ∞ then there exists a constant 𝐶 ≥ 0 such that 𝐶 E∥𝑇 − 𝑇𝑚 ∥ 𝐻𝑆 ≤ √ 𝑚 where ∥ · ∥ 𝐻𝑆 denotes the Hilbert-Schmidt norm on operators acting on H . With stricter assumptions similar results can also be proven with high probability (Rosasco, Belkin, and Vito, 2010). Lemma 2.3.1 (Rosasco, Belkin, and Vito, 2010) 𝑇 and 𝑇𝐾 are Hilbert-Schmidt. Furthermore, with probability greater than 1 − 2𝑒 𝜏 √ √ 2 2𝑘 𝜏 ∥𝑇 − 𝑇𝐾 ∥ 𝐻𝑆 ≤ √ 𝐾 where 𝑘 = sup𝑥∈𝐷 𝜅(𝑥, 𝑥). We note that, in our algorithm, 𝜅 : 𝐷 × 𝐷 → R𝑛×𝑛 whereas the preceding results are proven only in the setting 𝑛 = 1; nonetheless they provide useful intuition regarding the approximations used in our methodology. 2.4 Experiments In this section we illustrate the claimed properties of our methodology, and compare it to existing approaches in the literature. All experimental results concern the mapping 𝑎 ↦→ 𝑢 defined by (2.3) with 𝐷 = [0, 1] 2 . Coefficients are generated according to 𝑎 ∼ 𝜇 where 𝜇 = 𝜓# N (0, (−Δ + 9𝐼) −2 ) with a Neumann boundry condition on the operator −Δ + 9𝐼. The mapping 𝜓 : R → R takes the value 12 on the positive part of the real line and 3 on the negative. Hence the coefficients are piecewise constant with random geometry and a fixed contrast of 4. Such constructions are prototypical of physical properties such as permeability in subsurface flows and material microstructures in elasticity. Solutions 𝑢 are obtained by using a second-order finite difference scheme on a 241 × 241 grid. Different resolutions are downsampled from this dataset. 30 Resolutions 𝑠′ = 16 𝑠′ = 31 𝑠′ = 61 𝑠 = 16 0.0525 0.0591 0.0585 0.0787 0.0538 0.0588 𝑠 = 31 𝑟 = 0.10, 𝑁 = 100, relative 𝑙2 test error Table 2.1: Comparing resolutions on full grids To be concrete we set the dimension of representation 𝑛 (i.e. the width of graph network) to be 64, the number of iterations 𝑇 to be 6, 𝜎 to be the ReLU function, and the inner kernel network 𝜅 to be a 3−layer feed-forward network with widths (6, 512, 1024, 𝑛2 ) and ReLU activation. We use the Adam optimizer with a learning rate 1𝑒 − 4 and train for 200 epochs with respect to the absolute mean squared error on the normalized data unless otherwise stated. These chosen hyperparameters are not optimized and could be adapted to improve performance. We employ the message passing network from the standard Pytorch graph network library Torchgeometric (Fey and Lenssen, 2019). All the test errors are relative 𝐿 2 (𝐷) errors on the original data. The code and data can be found at https://github.com/ wumming/graph-pde. Supervised Setting First we consider the supervised scenario that we are given 𝑁 training pairs {𝑎 𝑗 , 𝑢 𝑗 }1𝑁 , where each 𝑎 𝑗 and 𝑢 𝑗 are provided on a 𝑠 × 𝑠 grid (𝐾 = 𝑠2 ). Generalization of resolutions on full grids. To examine the generalization property, we train the graph kernel network on resolution 𝑠 × 𝑠 and test on another resolution 𝑠′ × 𝑠′. We fix the radius to be 𝑟 = 0.10, train on 𝑁 = 100 equation pairs and test on 40 equation pairs. As shown in Table 2.1, for each row, the test errors at different resolutions remain on the same scale, illustrating the desired design feature that graph kernel networks can train on one resolution and generalize to another resolution. The test errors on the diagonal (𝑠 = 𝑠′ = 16 and 𝑠 = 𝑠′ = 31) are the smallest, which means the network has the best performance when the training grid and the test grid are the same. Interestingly, for the second row, when training on 𝑠 = 31, it is easier to general to 𝑠′ = 61 than to 𝑠′ = 16. This is because when generalizing to a larger grid, the support of the kernel becomes large which does not hurt the performance. But when generalizing to a smaller grid, part of the support of the kernel is lost, which causes the kernel to be inaccurate. 31 Training Size Training Error Test Error 𝑁 = 10 0.0089 0.0931 0.0183 0.0478 𝑁 = 100 𝑁 = 1000 0.0255 0.0345 2000, 500, 100 epochs, respectively. Table 2.2: Comparing number of training pairs Expressiveness and overfitting. We compare the training error and test error with a different number of training pairs 𝑁 to see if the kernel network can learn the kernel structure even with a small amount of data. We study the expressiveness of the kernel network, examining how it overfits. We fix 𝑟 = 0.10 on the 𝑠 = 𝑠′ = 31 grid and train with 𝑁 = 10, 100, 1000 whilst employing 2000, 500, 100 epochs respectively. We see from Table 2.2 that the kernel network already achieves a reasonable result when 𝑁 = 10, and the accuracy is competitive when 𝑁 = 100. In all three cases, the test error is larger than the training error suggesting that the kernel network has enough expressiveness to overfit the training set. This overfitting is not severe as the training error will not be pushed to zero even for 𝑁 = 10, after 2000 epochs. Semi-Supervised Setting In the semi-supervised setting, we are only given 𝑚 nodes sampled from a 𝑠 × 𝑠 grid for each training pair, and want to evaluate on 𝑚′ nodes sampled from an 𝑠′ × 𝑠′ grid for each test pair. To be concrete, we set the number of sampled nodes 𝑚 = 𝑚′ = 200. For each training pair, we sample twice 𝑙 = 2; for each test pair, we sample once 𝑙 ′ = 1. We train on 𝑁 = 100 pairs and test on 𝑁 ′ = 100 pairs. The radius for both training and testing is set to 𝑟 = 𝑟 ′ = 0.25. Generalization of resolutions on sampled grids. Similar to the first experiments, we train the graph kernel network with nodes sampled from the 𝑠 × 𝑠 resolution and test on nodes sampled from the 𝑠′ × 𝑠′ resolution. As shown in Table 2.3, for each row, the test errors on different resolutions are about the same, which means the graph kernel network can also generalize in the semi-supervised setting. Comparing the rows, large training resolutions 𝑠 tend to have a smaller error. When sampled from a finer grid, there are more edges because the support of the kernel is larger on the finer grid. Still, the performance is best when 𝑠 = 𝑠′. 32 Resolutions 𝑠 = 16 𝑠 = 31 𝑠 = 61 𝑠 = 121 𝑠 = 241 𝑠′ = 61 0.0717 0.0726 0.0687 0.0687 0.0649 𝑠′ = 121 0.0768 0.0710 0.0728 0.0664 0.0658 𝑠′ = 241 0.0724 0.0722 0.0723 0.0685 0.0649 𝑁 = 100, 𝑚 = 𝑚 ′ = 200, 𝑟 = 𝑟 ′ = 0.25, 𝑙 = 2 Table 2.3: Generalization of resolutions on sampled Grids The number of examples vs the times of sampling. Increasing the number of times we sample, 𝑙, reduces the error from the Nyström approximation. By comparing different 𝑙 we determine which value will be sufficient. When we sample 𝑙 times for each equation, we get a total of 𝑁𝑙 sampled training pairs, Table 2.4. 𝑁 = 10 𝑁 = 100 𝑁 = 1000 𝑙=1 𝑙=2 𝑙=4 𝑙=8 0.1259 0.1069 0.0967 0.1026 0.0786 0.0687 0.0690 0.0621 0.0604 0.0579 0.0540 0.0483 𝑠 = 241, 𝑚 = 𝑚 ′ = 200, 𝑟 = 𝑟 ′ = 0.25 Table 2.4: Number of training pairs and sampling Table 2.4 indicates that the larger 𝑙 the better, but 𝑙 = 2 already gives a reasonable performance. Moreover, the same order of sampled training pairs, (𝑁 = 100, 𝑙 = 8), and (𝑁 = 1000, 𝑙 = 1), result in a similar performance. It implies that in a low training data regime, increasing 𝑙 improves the performance. Different number of nodes in training and testing. To further examine the Nyström approximation, we compare different numbers of node samples 𝑚, 𝑚′ for both training and testing. 𝑚 = 100 𝑚 = 200 𝑚 = 400 𝑚 = 800 𝑚′ = 100 0.0871 0.0972 0.0991 0.1084 𝑚′ = 200 0.0716 0.0734 0.0699 0.0751 𝑚′ = 400 0.0662 0.0606 0.0560 0.0573 𝑚′ = 800 0.0609 0.0562 0.0506 0.0478 𝑠 = 121, 𝑟 = 𝑟 ′ = 0.15, 𝑙 = 5 Table 2.5: Number of nodes in the training and testing 33 As can be seen from Table 2.5, in general the larger 𝑚 and 𝑚′ the better. For each row, fixing 𝑚, the larger 𝑚′ the better. But for each column, when fixing 𝑚′, increasing 𝑚 may not lead to better performance. This is again due to the fact that when learning on a larger grid, the kernel network learns a kernel with larger support. When evaluating on a smaller grid, the learned kernel will be truncated to have small support, leading to an increased error. In general, 𝑚 = 𝑚′ will be the best choice. The radius and the number of nodes. The computation and storage of graph networks directly scale with the number of edges. In this experiment, we want to study the trade-off between the number of nodes 𝑚 and the radius 𝑟 when fixing the number of edges. 𝑟 = 0.05 𝑟 = 0.15 𝑟 = 0.40 𝑟 = 1.00 𝑚 = 100 0.110(176) 0.086(512) 0.064(1596) 0.059(9756) 𝑚 = 200 0.109(666) 0.070(2770) 0.051(9728) 0.048(38690) 𝑚 = 400 0.099(3354) 0.053(14086) 0.040(55919) − Error (Edges), 𝑠 = 121, 𝑙 = 5, 𝑚 ′ = 𝑚 Table 2.6: The radius and the number of nodes Table 2.6 shows the test error with the number of edges for different 𝑟 and 𝑚. In general, the more edges, the better. For a fixed number of edges, the performance depends more on the radius 𝑟 than on the number of nodes 𝑚. In other words, the error of truncating the kernel locally is larger than the error from the Nyström approximation. It would be better to use larger 𝑟 with smaller 𝑚. Inner Kernel Network 𝜅. To find the best network structure of 𝜅, we compare different combinations of width and depth. We consider three cases of 𝜅: 1. a 2−layer feed-forward network with widths (6, 𝑤𝑖𝑑𝑡ℎ, 𝑛2 ), 2. a 3−layer feed-forward network with widths (6, 𝑤𝑖𝑑𝑡ℎ/2, 𝑤𝑖𝑑𝑡ℎ, 𝑛2 ), and 3. a 5−layer feed-forward network with widths (6, 𝑤𝑖𝑑𝑡ℎ/4, 𝑤𝑖𝑑𝑡ℎ/2, 𝑤𝑖𝑑𝑡ℎ, 𝑤𝑖𝑑𝑡ℎ, 𝑛2 ), all with ReLU activation and learning rate 1𝑒 − 4. As shown in Table 2.7, have wider and deeper network increase the expensiveness of the kernel network. The diagonal combinations 𝑤𝑖𝑑𝑡ℎ = 256, 𝑑𝑒 𝑝𝑡ℎ = 2, 𝑤𝑖𝑑𝑡ℎ = 1024, 𝑑𝑒 𝑝𝑡ℎ = 3, and 𝑤𝑖𝑑𝑡ℎ = 4096, 𝑑𝑒 𝑝𝑡ℎ = 5 have better test error. Notice the wide but shallow network 𝑤𝑖𝑑𝑡ℎ = 4096, 𝑑𝑒 𝑝𝑡ℎ = 2 has very bad 34 width = 64 width = 128 width = 256 width = 1024 width = 4096 depth = 2 0.0685 0.0630 0.0617 0.0641 0.2934 depth = 3 0.0695 0.0633 0.0610 0.0591 0.0690 depth = 5 0.0770 0.0702 0.0688 0.0608 0.0638 𝑠 = 241, 𝑚 = 200, 𝑟 = 0.25 Table 2.7: Comparing the width and depth for the inner kernel network 𝜅 performance. Both its training and testing error once decreased to 0.09 around 10th epochs, but then blow off. The 1𝑒 − 4 learning rate is probably too high for this combination. In general, depth of 3 with with of 256 is a good combination of Dracy Equation dataset. Comparison with Different Benchmarks In the following section, we compare the Graph Kernel Network with different benchmarks on a larger dataset of 𝑁 = 1024 training pairs computed on a 421 × 421 grid. The network is trained and evaluated on the same full grid. The results are presented in Table 2.8. • NN is a simple point-wise feedforward neural network. It is mesh-free, but performs badly due to lack of neighbor information. • FCN is the state of the art neural network method based on Fully Convolution Network (Zhu and Zabaras, 2018). It has a dominating performance for small grids 𝑠 = 61. But fully convolution networks are mesh-dependent and therefore their error grows when moving to a larger grid. • PCA+NN is an instantiation of the methodology proposed in (Bhattacharya, Kovachki, and Stuart, 2020): using PCA as an autoencoder on both the input and output data and interpolating the latent spaces with a neural network. The method provably obtains mesh-independent error and can learn purely from data, however the solution can only be evaluated on the same mesh as the training data. • RBM is the classical Reduced Basis Method (using a PCA basis), which is widely used in applications and provably obtains mesh-independent error (R. A. DeVore, 2014). It has the best performance but the solutions can only 35 be evaluated on the same mesh as the training data and one needs knowledge of the PDE to employ it. • GKN stands for our graph kernel network with 𝑟 = 0.25 and 𝑚 = 300. It enjoys competitive performance against all other methods while being able to generalize to different mesh geometries. Networks NN FCN PCA+NN RBM GKN 𝑠 = 85 0.1716 0.0253 0.0299 0.0244 0.0346 𝑠 = 141 0.1716 0.0493 0.0298 0.0251 0.0332 𝑠 = 211 0.1716 0.0727 0.0298 0.0255 0.0342 𝑠 = 421 0.1716 0.1097 0.0299 0.0259 0.0369 Table 2.8: Error of different methods 2.5 Discussion and Conclusion We have introduced the concept of Neural Operator and instantiated it through graph kernel networks designed to approximate mappings between function spaces. They are constructed to be mesh-free and our numerical experiments demonstrate that they have the desired property of being able to train and generalize on different meshes. This is because the networks learn the mapping between infinite-dimensional function spaces, which can then be shared with approximations at various levels of discretization. A further advantage is that data may be incorporated on unstructured grids, using the Nyström approximation. We demonstrate that our method can achieve competitive performance with other mesh-free approaches developed in the numerical analysis community, and beats state-of-the-art neural network approaches on large grids, which are mesh-dependent. The methods developed in the numerical analysis community are less flexible than the approach we introduce here, relying heavily on the variational structure of divergence form elliptic PDEs. Our new mesh-free method has many applications. It has the potential to be a faster solver that learns from only sparse observations in physical space. It is the only method that can work in the semi-supervised scenario when we only have measurements on some parts of the grid. It is also the only method that can transfer between different geometries. For example, when computing the flow dynamic of many different airfoils, we can construct different graphs and train together. When learning from irregular grids and querying new locations, our method does not require any interpolation, avoid subsequently interpolation error. 36 Disadvantage. Graph kernel network’s runtime and storage scale with the number of edges 𝐸 = 𝑂 (𝐾 2 ). While other mesh-dependent methods such as PCA+NN and RBM require only 𝑂 (𝐾). This is somewhat inevitable, because to learn the continuous function or the kernel, we need to capture pairwise information between every two nodes, which is 𝑂 (𝐾 2 ); when the discretization is fixed, one just needs to capture the point-wise information, which is 𝑂 (𝐾). Therefore training and evaluating the whole grid is costly when the grid is large. On the other hand, subsampling to ameliorate cost loses information in the data, and causes errors which make our method less competitive than PCA+NN and RBM. Future works. To deal with the above problem, we propose that ideas such as multi-grid and fast multipole methods (Gholami et al., 2016) may be combined with our approach to reduce complexity. In particular, a multi-grid approach will construct multi-graphs corresponding to different resolutions so that, within each graph, nodes only connect to their nearest neighbors. The number of edges then scale as 𝑂 (𝐾) instead of 𝑂 (𝐾 2 ). The error terms from Nyström approximation and local truncation can be avoided. Another direction is to extend the framework for time-dependent PDEs. Since the graph kernel network is itself an iterative solver with the time step 𝑡, it is natural to frame it as an RNN that each time step corresponds to a time step of the PDEs. References Adler, Jonas and Ozan Oktem (Nov. 2017). “Solving ill-posed inverse problems using iterative deep neural networks”. In: Inverse Problems. doi: 10.1088/13616420/aa9581. url: https://doi.org/10.1088%2F1361-6420%2Faa9581. Alet, Ferran et al. (2019). “Graph Element Networks: adaptive, structured computation and memory”. In: 36th International Conference on Machine Learning. PMLR. url: http://proceedings.mlr.press/v97/alet19a.html. Antman, Stuart S (2005). Problems In Nonlinear Elasticity. Springer. Bar, Leah and Nir Sochen (2019). “Unsupervised deep learning algorithm for PDEbased forward and inverse problems”. In: arXiv preprint arXiv:1904.05417. Battaglia, Peter W et al. (2018). “Relational inductive biases, deep learning, and graph networks”. In: arXiv preprint arXiv:1806.01261. Bear, Jacob and M Yavuz Corapcioglu (2012). Fundamentals of transport phenomena in porous media. Springer Science & Business Media. Belongie, Serge et al. (2002). “Spectral partitioning with indefinite kernels using the Nyström extension”. In: European conference on computer vision. Springer. 37 Bhatnagar, Saakaar et al. (2019). “Prediction of aerodynamic flow fields using convolutional neural networks”. In: Computational Mechanics, pp. 1–21. Bhattacharya, Kaushik, Nikola B. Kovachki, and Andrew M. Stuart (2020). “Model Reduction and Neural Networks for Parametric PDE(s)”. In: arXiv preprint arXiv:2005.03180. Chen, Chi et al. (2019). “Graph networks as a universal machine learning framework for molecules and crystals”. In: Chemistry of Materials 31.9, pp. 3564–3572. Cohen, Albert and Ronald DeVore (2015). “Approximation of high-dimensional parametric PDEs”. In: Acta Numerica. doi: 10.1017/S0962492915000033. Damianou, Andreas and Neil Lawrence (2013). “Deep gaussian processes”. In: Artificial Intelligence and Statistics, pp. 207–215. DeVore, Ronald A. (2014). “Chapter 3: The Theoretical Foundation of Reduced Basis Methods”. In: Model Reduction and Approximation. doi: 10. 1137/1. 9781611974829 . ch3. eprint: https : / / epubs . siam . org / doi / pdf / 10 . 1137/1.9781611974829.ch3. url: https://epubs.siam.org/doi/abs/ 10.1137/1.9781611974829.ch3. E, Weinan and Bing Yu (Mar. 2018). “The Deep Ritz Method: A Deep LearningBased Numerical Algorithm for Solving Variational Problems”. English (US). In: Communications in Mathematics and Statistics. issn: 2194-6701. doi: 10.1007/ s40304-018-0127-z. Evans, Lawrence C (2010). Partial Differential Equations. Vol. 19. American Mathematical Soc. Fey, Matthias and Jan E. Lenssen (2019). “Fast Graph Representation Learning with PyTorch Geometric”. In: ICLR Workshop on Representation Learning. Garriga-Alonso, Adrià, Carl Edward Rasmussen, and Laurence Aitchison (Aug. 2018). “Deep Convolutional Networks as shallow Gaussian Processes”. In: arXiv e-prints, arXiv:1808.05587, arXiv:1808.05587. arXiv: 1808.05587 [stat.ML]. Gholami, Amir et al. (2016). “FFT, FMM, or multigrid? A comparative study of state-of-the-art Poisson solvers for uniform and nonuniform grids in the unit cube”. In: SIAM Journal on Scientific Computing. Gilbarg, David and Neil S Trudinger (2015). Elliptic partial differential equations of second order. springer. Gilmer, Justin et al. (2017). “Neural message passing for quantum chemistry”. In: Proceedings of the 34th International Conference on Machine Learning. Globerson, Amir and Roi Livni (2016). “Learning Infinite-Layer Networks: Beyond the Kernel Trick”. In: CoRR abs/1606.05316. arXiv: 1606.05316. url: http: //arxiv.org/abs/1606.05316. 38 Guo, Xiaoxiao, Wei Li, and Francesco Iorio (2016). “Convolutional neural networks for steady flow approximation”. In: Proceedings of the 22nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining. Gurtin, Morton E (1982). An introduction to continuum mechanics. Academic press. Guss, William H. (Dec. 2016). “Deep Function Machines: Generalized Neural Networks for Topological Layer Expression”. In: arXiv e-prints, arXiv:1612.04799, arXiv:1612.04799. arXiv: 1612.04799 [stat.ML]. Hamilton, Will, Zhitao Ying, and Jure Leskovec (2017). “Inductive representation learning on large graphs”. In: Advances in neural information processing systems, pp. 1024–1034. Herrmann, L, Ch Schwab, and J Zech (2020). “Deep ReLU Neural Network Expression Rates for Data-to-QoI Maps in Bayesian PDE Inversion”. In. Johnson, Claes (2012). Numerical solution of partial differential equations by the finite element method. Courier Corporation. Kipf, Thomas N and Max Welling (2016). “Semi-supervised classification with graph convolutional networks”. In: arXiv preprint arXiv:1609.02907. Lord, Gabriel J, Catherine E Powell, and Tony Shardlow (2014). An introduction to computational stochastic PDEs. Vol. 50. Cambridge University Press. Matthews, Alexander G. de G. et al. (Apr. 2018). “Gaussian Process Behaviour in Wide Deep Neural Networks”. In. Murphy, Ryan L et al. (2018). “Janossy pooling: Learning deep permutationinvariant functions for variable-size inputs”. In: arXiv preprint arXiv:1811.01900. Neal, Radford M. (1996). Bayesian Learning for Neural Networks. Springer-Verlag. isbn: 0387947248. Nyström, Evert J (1930). “Über die praktische Auflösung von Integralgleichungen mit Anwendungen auf Randwertaufgaben”. In: Acta Mathematica. Raissi, Maziar, Paris Perdikaris, and George E Karniadakis (2019). “Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations”. In: Journal of Computational Physics 378, pp. 686–707. Rosasco, Lorenzo, Mikhail Belkin, and Ernesto De Vito (Mar. 2010). “On Learning with Integral Operators”. In: J. Mach. Learn. Res. 11, pp. 905–934. issn: 15324435. Roux, Nicolas Le and Yoshua Bengio (2007). “Continuous Neural Networks”. In: Proceedings of the Eleventh International Conference on Artificial Intelligence and Statistics. Ed. by Marina Meila and Xiaotong Shen. Ummenhofer, Benjamin et al. (2020). “Lagrangian fluid simulation with continuous convolutions”. In: International Conference on Learning Representations. 39 Vapnik, Vladimir N. (1998). Statistical Learning Theory. Wiley-Interscience. Veličković, Petar et al. (2017). “Graph attention networks”. In. Williams, Christopher K. I. (1996). “Computing with Infinite Networks”. In: Proceedings of the 9th International Conference on Neural Information Processing Systems. Cambridge, MA, USA: MIT Press. Zhu, Yinhao and Nicholas Zabaras (2018). “Bayesian deep convolutional encoder–decoder networks for surrogate modeling and uncertainty quantification”. In: Journal of Computational Physics. issn: 0021-9991. doi: https://doi.org/10.1016/ j.jcp.2018.04.018. url: http://www.sciencedirect.com/science/ article/pii/S0021999118302341. 40 Chapter 3 MODEL: MULTIPOLE NEURAL OPERATOR One of the main challenges in using deep learning-based methods for simulating physical systems and solving partial differential equations (PDEs) is formulating physics-based data in the desired structure for neural networks. Graph neural networks (GNNs) have gained popularity in this area since graphs offer a natural way of modeling particle interactions and provide a clear way of discretizing the continuum models. However, the graphs constructed for approximating such tasks usually ignore long-range interactions due to unfavorable scaling of the computational complexity with respect to the number of nodes. The errors due to these approximations scale with the discretization of the system, thereby not allowing for generalization under mesh-refinement. Inspired by the classical multipole methods, we propose a novel multi-level graph neural network framework that captures interaction at all ranges with only linear complexity. Our multi-level formulation is equivalent to recursively adding inducing points to the kernel matrix, unifying GNNs with multi-resolution matrix factorization of the kernel. Experiments confirm our multigraph network learns discretization-invariant solution operators to PDEs and can be evaluated in linear time. 3.1 Introduction A wide class of important scientific applications involve numerical approximation of parametric PDEs. There has been immense research efforts in formulating and solving the governing PDEs for a variety of physical and biological phenomena ranging from the quantum to the cosmic scale. While this endeavor has been successful in producing solutions to real-life problems, major challenges remain. Solving complex PDE systems such as those arising in climate modeling, turbulent flow of fluids, and plastic deformation of solid materials requires considerable time, computational resources, and domain expertise. Producing accurate, efficient, and automated data-driven approximation schemes has the potential to significantly accelerate the rate of innovation in these fields. Machine learning based methods enable this since they are much faster to evaluate and require only observational data to train, in stark contrast to traditional Galerkin methods (Zienkiewicz et al., 1977) and classical reduced order models (Rozza, Huynh, and Patera, 2007). 41 While deep learning approaches such as convolutional neural networks can be fast and powerful, they are usually restricted to a specific format or discretization. On the other hand, many problems can be naturally formulated on graphs. An emerging class of neural network architectures designed to operate on graph-structured data, Graph neural networks (GNNs), have gained popularity in this area. GNNs have seen numerous applications on tasks in imaging, natural language modeling, and the simulation of physical systems (Wu et al., 2020). In the latter case, graphs are typically used to model particles systems (the nodes) and the their interactions (the edges). Recently, GNNs have been directly used to learn solutions to PDEs by constructing graphs on the physical domain (Alet et al., 2019), and it is was further shown that GNNs can learn mesh-invariant solution operators (Z. Li et al., 2020). Since GNNs offer great flexibility in accurately representing solutions on any unstructured mesh, finding efficient algorithms is an important open problem. The computational complexity of GNNs depends on the sparsity structure of the underlying graph, scaling with the number of edges which may grow quadratically with the number of nodes in fully connected regions (Wu et al., 2020). Therefore, to make computations feasible, GNNs make approximations using nearest neighbor connection graphs which ignore long-range correlations. Such approximations are not suitable in the context of approximating solution operators of parametric PDEs since they will not generalize under refinement of the discretization, as we demonstrate in Section 8.5. However, using fully connected graphs quickly becomes computationally infeasible. Indeed evaluation of the kernel matrices outlined in Section 3.3 is only possible for coarse discretizations due to both memory and computational constraints. Throughout this work, we aim to develop approximation techniques that help alleviate this issue. To efficiently capture long-range interaction, multi-scale methods such as the classical fast multipole methods (FMM) have been developed (Greengard and Rokhlin, 1997). Based on the insight that long-range interaction are smooth, FMM decomposes the kernel matrix into different ranges and hierarchically imposes low-rank structures to the long-range components (hierarchical matrices)(Börm, Grasedyck, and Hackbusch, 2003). This decomposition can be viewed as a specific form of the multi-resolution matrix factorization of the kernel (Kondor, Teneva, and Garg, 2014; Börm, Grasedyck, and Hackbusch, 2003). However, the classical FMM requires nested grids as well as the explicit form of the PDEs. We generalize this idea to arbitrary graphs in the data-driven setting, so that the corresponding graph neural 42 networks can learn discretization-invariant solution operators. Main contributions. Inspired by the fast multi-pole method (FMM), we propose a novel hierarchical, and multi-scale graph structure which, when deployed with GNNs, captures global properties of the PDE solution operator with a linear timecomplexity (Greengard and Rokhlin, 1997; L. Ying, Biros, and Zorin, 2004). As shown in Figure 3.1, starting with a nearest neighbor graph, instead of directly adding edges to connect every pair of nodes, we add inducing points which help facilitate long-range communication. The inducing points may be thought of as forming a new subgraph which models long-range correlations. By adding a small amount of inducing nodes to the original graph, we make computation more efficient. Repeating this process yields a hierarchy of new subgraphs, modeling correlations at different length scales. We show that message passing through the inducing points is equivalent to imposing a low-rank structure on the corresponding kernel matrix, and recursively adding inducing points leads to multi-resolution matrix factorization of the kernel (Kondor, Teneva, and Garg, 2014; Börm, Grasedyck, and Hackbusch, 2003). We propose the graph V-cycle algorithm (Figure 3.1) inspired by FMM, so that message passing through the V-cycle directly computes the multi-resolution matrix factorization. We show that the computational complexity of our construction is linear in the number of nodes, achieving the desired efficiency, and we demonstrate the linear complexity and competitive performance through experiments on Darcy flow (Tek et al., 1957), a linear second-order elliptic equation, and Burgers’ equation(Su and Gardner, 1969), which considered a stepping stone to Naiver-Stokes, is nonlinear, long-range correlated and more challenging. Our primary contributions are listed below. • We develop the multipole graph kernel neural network (MGKN) that can capture long-range correlations in graph-based data with a linear time complexity in the nodes of the graph. • We unify GNNs with multi-resolution matrix factorization through the V-cycle algorithm. • We verify, analytically and numerically, the linear time complexity of our proposed methodology. 43 Left: the multi-level graph. Right: one V-cycle iteration for the multipole graph kernel network. Figure 3.1: V-Cycle Algorithm on Graph • We demonstrate numerically our method’s ability to capture global information by learning mesh-invariant solution operators to the Darcy flow and Burgers’ equations. 3.2 Related Works Deep learning approaches for PDEs: There have been two primary approaches in the application of deep learning for the solution of PDEs. The first parametrizes the solution operator as a deep convolutional neural network (CNN) between finitedimensional Euclidean spaces G𝜃 : R𝑛 → R𝑛 (Guo, W. Li, and Iorio, 2016; Zhu and Zabaras, 2018; Adler and Oktem, 2017; Bhatnagar et al., 2019; Ummenhofer et al., 2020). As demonstrated in Figure 3.4, such approaches are tied to a discritezation and cannot generalize. The second approach directly parameterizes the solution 𝑢 as a neural network G𝜃 : 𝐷 → R (E and Yu, 2018; Raissi, Perdikaris, and George E Karniadakis, 2019; Bar and Sochen, 2019; Smith, Azizzadenesheli, and Ross, 2020; Jiang et al., 2020). This approach is close to classical Galerkin methods and therefore suffers from the same issues w.r.t. the parametric dependence in the PDE. For any new parameter, an optimization problem must be solved which requires backpropagating through the differential operator L𝑎 many times, making it too slow for many practical application. Only very few recent works have attempted to capture the infinite-dimensional solution operator of the PDE (Alet et al., 2019; Lu, Jin, and George Em Karniadakis, 2019; Z. Li et al., 2020; Bhattacharya, Kovachki, and Andrew M. Stuart, 2020; Nelsen and Andrew M Stuart, 2020). The current work advances this direction. GNN and non-sparse graphs: A multitude of techniques such as graph convolution, edge convolution, attention, and graph pooling, have been developed for 44 improving GNNs (Kipf and Welling, 2016; Hamilton, Z. Ying, and Leskovec, 2017; Gilmer et al., 2017; Veličković et al., 2017; Murphy et al., 2018). Most, however, have been designed for sparse graphs and become computationally infeasible as the number of edges grow. The work (Alfke and Stoll, 2019) proposes using a low-rank decomposition to address this issue. Works on multi-resolution graphs with a Unet like structure have also recently began to emerge and seen success in imaging, classification, and semi-supervised learning (Ronneberger, Fischer, and Brox, 2015; Abu-El-Haija, Kapoor, et al., 2018; Abu-El-Haija, Perozzi, et al., 2019; Gao and Ji, 2019; M. Li et al., 2020). Our work ties together many of these ideas and provides a principled way of designing multi-scale GNNs(Gao and Ji, 2019; Abu-El-Haija, Kapoor, et al., 2018; Abu-El-Haija, Perozzi, et al., 2019; M. Li et al., 2020). All of these works focus on build multi-scale structure on a given graph. Our method, on the other hand, studies how to construct randomized graphs on the spatial domain for physics and applied math problems. We carefully craft the multi-level graph that corresponds to multi-resolution decomposition of the kernel matrix. Multipole and multi-resolution methods: The works (Fan, Lin, et al., 2019; Fan, Feliu-Faba, et al., 2019) propose a similar multi-pole expansion for solving parametric PDEs on structured grids. Our work generalizes on this idea by allowing for arbitrary discretizations through the use of GNNs. Multi-resolution matrix factorizations have been proposed in (Kondor, Teneva, and Garg, 2014; Ithapu et al., 2017). We employ such ideas to build our approximation architecture. 3.3 Preliminaries Graph Kernel Network (GKN) Suppose that L𝑎 in (2.4) is uniformly elliptic then the Green’s representation formula implies ∫ 𝐺 𝑎 (𝑥, 𝑦) [ 𝑓 (𝑦) + (Γ𝑎 𝑢)(𝑦)] 𝑑𝑦. 𝑢(𝑥) = (3.1) 𝐷 where 𝐺 𝑎 is a Newtonian potential and Γ𝑎 is an operator defined by appropriate sums and compositions of the modified trace and co-normal derivative operators (Sauter and Schwab, 2010). We have turned the PDE (2.4) into the integral equation (3.1) which lends itself to an iterative approximation architecture. Kernel operator. Since 𝐺 𝑎 is continuous for all points 𝑥 ≠ 𝑦, it is sensible to model the action of the integral operator in (3.1) by a neural network 𝜅 𝜙 with parameters 𝜙. To that end, define the operator K𝑎 : U → U as the action of the 45 kernel 𝜅 𝜙 on 𝑢: ∫ (K𝑎 𝑢)(𝑥) = 𝜅 𝜙 (𝑎(𝑥), 𝑎(𝑦), 𝑥, 𝑦)𝑢(𝑦) 𝑑𝑦 (3.2) 𝐷 where the kernel neural network 𝜅 𝜙 takes as inputs spatial locations 𝑥, 𝑦 as well as the values of the parameter 𝑎(𝑥), 𝑎(𝑦). Since Γ𝑎 is itself an operator, its action cannot be fully accounted for by the kernel 𝜅 𝜙 , we therefore add local parameters 𝑊 = 𝑤𝐼 to K𝑎 and apply a non-linear activation 𝜎, defining the iterative architecture 𝑢 (𝑡) = 𝜎((𝑊 + K𝑎 )𝑢 (𝑡−1) ) for 𝑡 = 1, . . . 𝑇 with 𝑢 (0) = 𝑎. Since Γ𝑎 is local w.r.t 𝑢, we only need local parameters to capture its effect, while we expect its non-locality w.r.t. 𝑎 to manifest via the initial condition (Sauter and Schwab, 2010). To increase expressivity, we lift 𝑢(𝑥) ∈ R to a higher dimensional representation 𝑣(𝑥) ∈ R𝑑 𝑣 by a point-wise linear transformation 𝑣 0 = 𝑃𝑢 0 , and update the representation
𝑣 (𝑡) = 𝜎 (𝑊 + K𝑎 )𝑣 (𝑡−1) , 𝑡 = 1, . . . , 𝑇 (3.3) projecting it back 𝑢 (𝑇) = 𝑄𝑣 (𝑇) at the last step. Hence the kernel is a mapping 𝜅 𝜙 : R2(𝑑+1) → R𝑑 𝑣 ×𝑑 𝑣 and 𝑊 ∈ R𝑑 𝑣 ×𝑑 𝑣 . Note that, since our goal is to approximate the mapping 𝑎 ↦→ 𝑢 with 𝑓 in (2.4) fixed, we do not need explicit dependence on 𝑓 in our architecture as it will remain constant for any new 𝑎 ∈ A. For a specific discretization 𝐷 𝑗 , 𝑎 𝑗 | 𝐷 𝑗 , 𝑢 𝑗 | 𝐷 𝑗 ∈ R𝑛 are 𝑛-dimensional vectors, and the evaluation of the kernel network can be viewed as a 𝑛 × 𝑛 matrix 𝐾, with its 𝑥, 𝑦 entry (𝐾)𝑥𝑦 = 𝜅 𝜙 (𝑎(𝑥), 𝑎(𝑦), 𝑥, 𝑦). Then the action of K𝑎 becomes the matrixvector multiplication 𝐾𝑢. For the lifted representation (3.3), for each 𝑥 ∈ 𝐷 𝑗 , 𝑣 (𝑡) (𝑥) ∈ R𝑑 𝑣 and the output of the kernel (𝐾)𝑥𝑦 = 𝜅 𝜙 (𝑎(𝑥), 𝑎(𝑦), 𝑥, 𝑦) is a 𝑑𝑣 × 𝑑𝑣 matrix. Therefore 𝐾 becomes a fourth order tensor with shape 𝑛 × 𝑛 × 𝑑𝑣 × 𝑑𝑣 . Kernel convolution on graphs. Since we assume a non-uniform discretization of 𝐷 that can differ for each data pair, computing with (3.3) cannot be implemented in a standard way. Graph neural networks offer a natural solution since message passing on graphs can be viewed as the integration (3.2). Since the message passing is computed locally, it avoids storing the full kernel matrix 𝐾. Given a discretization 𝐷 𝑗 , we can adaptively construct graphs on the domain 𝐷. The structure of the graph’s adjacency matrix transfers to the kernel matrix 𝐾. We define the edge attributes 𝑒(𝑥, 𝑦) = (𝑎(𝑥), 𝑎(𝑦), 𝑥, 𝑦) ∈ R2(𝑑+1) and update the graph nodes following (3.3) 46 which mimics the message passing neural network (Gilmer et al., 2017):  𝑣 (𝑡+1) (𝑥) = (𝜎(𝑊 + K𝑎 )𝑣 (𝑡) )(𝑥) ≈ 𝜎 𝑊𝑣 (𝑡) +  ∑︁
1 𝜅 𝜙 𝑒(𝑥, 𝑦) 𝑣 (𝑡) (𝑦) |𝑁 (𝑥)| 𝑦∈𝑁 (𝑥) (3.4) where 𝑁 (𝑥) is the neighborhood of 𝑥, in this case, the entire discritized domain 𝐷 𝑗 . Domain of Integration. Construction of fully connected graphs is memory intensive and can become computationally infeasible for fine discretizations, i.e., when |𝐷 𝑗 | is large. To partially alleviate this, we can ignore the longest range kernel interactions as they have decayed the most and change the integration domain in (3.2) from 𝐷 to 𝐵(𝑥, 𝑟) for some fixed radius 𝑟 > 0. This is equivalent to imposing a sparse structure on the kernel matrix 𝐾 so that only entries around the diagonal are non-zero and results in the complexity 𝑂 (𝑛2𝑟 𝑑 ). Nyström approximation. To further relieve computational complexity, we use Nyström approximation or the inducing points method by uniformly sampling 𝑚 < 𝑛 nodes from the 𝑛 nodes discretization, which is to approximate the kernel matrix by a low-rank decomposition 𝐾𝑛𝑛 ≈ 𝐾𝑛𝑚 𝐾𝑚𝑚 𝐾𝑚𝑛 (3.5) where 𝐾𝑛𝑛 = 𝐾 is the original 𝑛×𝑛 kernel matrix and 𝐾𝑚𝑚 is the 𝑚 ×𝑚 kernel matrix corresponding to the 𝑚 inducing points. 𝐾𝑛𝑚 and 𝐾𝑚𝑛 are transition matrices which could include restriction, prolongation, and interpolation. Nyström approximation further reduces the complexity to 𝑂 (𝑚 2𝑟 𝑑 ). 3.4 Multipole Graph Kernel Network The fast multipole method (FMM) is a systematic approach of combining the aforementioned sparse and low-rank approximations while achieving linear complexity. The kernel matrix is decomposed into different ranges and a hierarchy of low-rank structures is imposed on the long-range components. We employ this idea to construct hierarchical, multi-scale graphs, without being constraint to particular forms of the kernel (L. Ying, Biros, and Zorin, 2004). We elucidate the workings of the FMM through matrix factorization. The key to the fast multipole method’s linear complexity lies in the subdivision of 47 the kernel matrix according to the range of interaction, as shown in Figure 3.2: 𝐾 = 𝐾1 + 𝐾2 + . . . + 𝐾 𝐿 (3.6) where 𝐾1 corresponds to the shortest-range interaction, and 𝐾 𝐿 corresponds to the longest-range interaction. While the uniform grids depicted in Figure 3.2 produce an orthogonal decomposition of the kernel, the decomposition may be generalized to arbitrary graphs by allowing slight overlap. The kernel matrix 𝐾 is decomposed respect to ranges. 𝐾1 corresponds to short-range interaction; it is sparse but high-rank. 𝐾3 corresponds to long-range interaction; it is dense but low-rank. Figure 3.2: Hierarchical matrix decomposition Multi-scale graphs We construct 𝐿 graph levels, where the finest graph corresponds to the shortest-range interaction 𝐾1 , and the coarsest graph corresponds to the longest-range interaction 𝐾 𝐿 . In what follows, we will drop the time dependence from (3.3) and use the subscript 𝑣 𝑙 to denote the representation at each level of the graph. Assuming the underlying graph is a uniform grid with resolution 𝑠 such that 𝑠 𝑑 = 𝑛, the 𝐿 multi-level graphs will be grids with resolution 𝑠𝑙 = 𝑠/2𝑙−1 , and consequentially 𝑛𝑙 = 𝑠𝑙𝑑 = (𝑠/2𝑙−1 ) 𝑑 for 𝑙 = 1, . . . , 𝐿. In general, the underlying discretization can be arbitrary and we produce a hierarchy of 𝐿 graphs with a decreasing number of nodes 𝑛1 , . . . , 𝑛 𝐿 . The coarse graph representation can be understood as recursively applying an inducing points approximation: starting from a graph with 𝑛1 = 𝑛 nodes, we impose inducing points of size 𝑛2 , 𝑛3 , . . . which all admit a low-rank kernel matrix decomposition of the form (3.5). The original 𝑛 × 𝑛 kernel matrix 𝐾𝑙 is represented by a much smaller 𝑛𝑙 × 𝑛𝑙 kernel matrix, denoted by 𝐾𝑙,𝑙 . As shown in Figure (3.2), 𝐾1 is full-rank but very sparse while 𝐾 𝐿 is dense but low-rank. Such structure can be achieved by applying equation (3.5) recursively to equation (3.6), leading to the 48 multi-resolution matrix factorization (Kondor, Teneva, and Garg, 2014): 𝐾 ≈ 𝐾1,1 + 𝐾1,2 𝐾2,2 𝐾2,1 + 𝐾1,2 𝐾2,3 𝐾3,3 𝐾3,2 𝐾2,1 + · · · (3.7) where 𝐾1,1 = 𝐾1 represents the shortest range, 𝐾1,2 𝐾2,2 𝐾2,1 ≈ 𝐾2 , represents the second shortest range, etc. The center matrix 𝐾𝑙,𝑙 is a 𝑛𝑙 × 𝑛𝑙 kernel matrix corresponding to the 𝑙-level of the graph described above. The long matrices 𝐾𝑙+1,𝑙 , 𝐾𝑙,𝑙+1 are 𝑛𝑙+1 × 𝑛𝑙 and 𝑛𝑙+1 × 𝑛𝑙 transition matrices. We define them as moving the representation 𝑣 𝑙 between different levels of the graph via an integral kernel that we learn. In general, 𝑣 (𝑡) (𝑥) ∈ R𝑑 𝑣 and the output of the kernel (𝐾𝑙,𝑙 ′ )𝑥𝑦 = 𝜅 𝜙 (𝑎(𝑥), 𝑎(𝑦), 𝑥, 𝑦) is itself a 𝑑𝑣 × 𝑑𝑣 matrix, so all these matrices are again fourthorder tensors. ∫ 𝐾𝑙,𝑙 : 𝑣 𝑙 ↦→ 𝑣 𝑙 = 𝜅 𝜙𝑙,𝑙 (𝑎(𝑥), 𝑎(𝑦), 𝑥, 𝑦)𝑣 𝑙 (𝑦) 𝑑𝑦 (3.8) 𝐵(𝑥,𝑟 𝑙,𝑙 ) ∫ 𝐾𝑙+1,𝑙 : 𝑣 𝑙 ↦→ 𝑣 𝑙+1 = 𝜅 𝜙𝑙+1,𝑙 (𝑎(𝑥), 𝑎(𝑦), 𝑥, 𝑦)𝑣 𝑙 (𝑦) 𝑑𝑦 (3.9) 𝐵(𝑥,𝑟 𝑙+1,𝑙 ) ∫ 𝐾𝑙,𝑙+1 : 𝑣 𝑙+1 ↦→ 𝑣 𝑙 = 𝜅 𝜙𝑙,𝑙+1 (𝑎(𝑥), 𝑎(𝑦), 𝑥, 𝑦)𝑣 𝑙+1 (𝑦) 𝑑𝑦 (3.10) 𝐵(𝑥,𝑟 𝑙,𝑙+1 ) Linear complexity. Each matrix in the decomposition (3.6) is represented by the kernel matrix 𝐾𝑙,𝑙 corresponding to the appropriate sub-graph. Since the number of non-zero entries of each row in these matrices is constant, we obtain that the comÍ putational complexity is 𝑙 𝑂 (𝑛𝑙 ). By designing the sub-graphs so that 𝑛𝑙 decays fast enough, we can obtain linear complexity. For example, choose 𝑛𝑙 = 𝑂 (𝑛/2𝑙 ) Í Í then 𝑙 𝑂 (𝑛𝑙 ) = 𝑙 𝑛/2𝑙 = 𝑂 (𝑛). Combined with a Nyström approximation, we obtain 𝑂 (𝑚) complexity. V-cycle Algorithm We present a V-cycle algorithm, see Figure 3.1, for efficiently computing (3.7). It consists of two steps: the downward pass and the upward pass. Denote the representation in downward pass and upward pass by 𝑣ˇ and 𝑣ˆ , respectively. In the downward step, the algorithm starts from the fine graph representation 𝑣ˇ 1 and updates it by applying a downward transition 𝑣ˇ 𝑙+1 = 𝐾𝑙+1,𝑙 𝑣ˇ 𝑙 . In the upward step, the algorithm starts from the coarse presentation 𝑣ˆ 𝐿 and updates it by applying an upward transition and the center kernel matrix 𝑣ˆ 𝑙 = 𝐾𝑙,𝑙−1 𝑣ˆ 𝑙−1 + 𝐾𝑙,𝑙 𝑣ˇ 𝑙 . Notice that the one level downward and upward exactly computes 𝐾1,1 + 𝐾1,2 𝐾2,2 𝐾2,1 , and a full 𝐿-level v-cycle leads to the multi-resolution decomposition (3.7). 49 Multipole graph kernel network Employing (3.8)-(3.10), we use 𝐿 neural networks 𝜅 𝜙1,1 , . . . , 𝜅 𝜙 𝐿,𝐿 to approximate the kernel 𝐾𝑙,𝑙 , and 2(𝐿 − 1) neural networks 𝜅 𝜙1,2 , 𝜅 𝜙2,1 , . . . to approximate the transitions 𝐾𝑙+1,𝑙 , 𝐾𝑙,𝑙+1 . Following the iterative architecture (3.3), we also introduce the linear operator 𝑊, denoting it by 𝑊𝑙 for each corresponding resolution. Since it acts on a fixed resolution, we employ it only along with the kernel 𝐾𝑙,𝑙 and not the transitions. At each time step 𝑡 = 0, . . . , 𝑇 − 1, we perform a full V-cycle: Downward Pass: For 𝑙 = 1, . . . , 𝐿 : (𝑡) (𝑡+1) + 𝐾𝑙+1,𝑙 𝑣ˇ 𝑙(𝑡+1) ) = 𝜎( 𝑣ˆ 𝑙+1 𝑣ˇ 𝑙+1 (3.11) Upward Pass: For 𝑙 = 𝐿, . . . , 1 : (𝑡+1) 𝑣ˆ 𝑙(𝑡+1) = 𝜎((𝑊𝑙 + 𝐾𝑙,𝑙 ) 𝑣ˇ 𝑙(𝑡+1) + 𝐾𝑙,𝑙−1 𝑣ˆ 𝑙−1 ). (3.12) We initialize as 𝑣 1(0) = 𝑃𝑢 (0) = 𝑃𝑎 and output 𝑢 (𝑇) = 𝑄𝑣 (𝑇) = 𝑄 𝑣ˆ 1(𝑇) . The algorithm unifies multi-resolution matrix decomposition with iterative graph kernel networks. Combined with a Nyström approximation it leads to 𝑂 (𝑚) computational complexity that can be implemented with message passing neural networks. Notice GKN is a specific case of V-cycle when 𝐿 = 1. 3.5 Experiments Properties of the multipole graph kernel network In this section, we show that MGKN has linear complexity and learns discretization invariant solutions by solving the steady-state of Darcy flow. In particular, we consider the 2-d PDE −∇ · (𝑎(𝑥)∇𝑢(𝑥)) = 𝑓 (𝑥) 𝑢(𝑥) = 0 𝑥 ∈ (0, 1) 2 𝑥 ∈ 𝜕 (0, 1) 2 (3.13) and approximate the mapping 𝑎 ↦→ 𝑢 which is non-linear despite the fact that (3.13) is a linear PDE. We model the coefficients 𝑎 as random piece-wise constant functions and generate data by solving (3.13) using a second-order finite difference scheme on a fine grid. Data of coarser resolutions are sub-sampled. See the supplements for further details. The code depends on Pytorch Geometric(Fey and Lenssen, 2019), also included in the supplements. 50 Left: compared to GKN whose complexity scales quadratically with the number of nodes, MGKN has a linear complexity; Mid: Adding more levels reduces test error; Right: MGKN can be trained on a coarse resolution and perform well when tested on a fine resolution, showing invariance to discretization. Figure 3.3: Properties of multipole graph kernel network (MGKN) on Darcy flow We use Nyström approximation by sampling 𝑚 1 , . . . , 𝑚 𝐿 nodes for each level. When changing the number of levels, we fix coarsest level 𝑚 𝐿 = 25, 𝑟 𝐿,𝐿 = 2−1 , and let 𝑚 𝑙 = 25 · 4 𝐿−𝑙 , 𝑟 𝑙,𝑙 = 2−(𝐿−𝑙) , and 𝑟 𝑙,𝑙+1 = 𝑟 𝑙+1,𝑙 = 2−(𝐿−𝑙)+1/2 . This set-up is one example that can obtain linear complexity. In general, any choice satisfying Í 2 2 𝑙 𝑚 𝑙 𝑟 𝑙,𝑙 = 𝑂 (𝑚 1 ) also works. We set width 𝑑 𝑣 = 64, iteration 𝑇 = 5 and kernel network 𝜅 𝜙𝑙,𝑙 as a three-layer neural network with width 256/2 (𝑙−1) , so that coarser grids will have smaller kernel networks. 1. Linear complexity: The left most plot in Figure 3.3 shows that MGKN (blue line) achieves linear time complexity (the time to evaluate one equation) w.r.t. the number of nodes, while GKN (red line) has quadratic complexity (the solid line is interpolated; the dash line is extrapolated). Since the GPU memory used for backpropagation also scales with the number of edges, GKN is limited to 𝑚 ≤ 800 on a single 11G-GPU while MGKN can scale to much higher resolutions . In other words, MGKN can be applicable for larger settings where GKN cannot. 2. Comparing with single-graph: As shown in Figure 3.3 (mid), adding multileveled graphs helps decrease the error. The MGKN depicted in blue bars starts from a fine sampling 𝐿 = 1; 𝑚 = [1600], and adding subgraphs, 𝐿 = 2; 𝑚 = [400, 1600], 𝐿 = 3; 𝑚 = [100, 400, 1600], up to, 𝐿 = 4; 𝑚 = [25, 100, 400, 1600]. When 𝐿 = 1, MGKN and GKN are equivalent. This experiment shows using multi-level graphs helps improve accuracy without increasing much of time-complexity. 51 3. Generalization to resolution: The MGKN is discretization invariant, and therefore capable of super-resolution. We train with nodes sampled from a 𝑠 × 𝑠 resolution mesh and test on nodes sampled from a 𝑠′×𝑠′ resolution mesh. As shown in on the right of Figure 3.3, MGKN achieves similar testing error on 𝑠′ = 61, 121, 241, independently of the training discretization. Comparison with benchmarks We compare the accuracy of our methodology with other deep learning methods as well as reduced order modeling techniques that are commonly used in practice. As a test bed, we use Darcy flow (3.13) and the 1-d viscous Burger’s equations: 𝜕𝑡 𝑢(𝑥, 𝑡) + 𝜕𝑥 (𝑢 2 (𝑥, 𝑡)/2) = 𝜈𝜕𝑥𝑥 𝑢(𝑥, 𝑡), 𝑢(𝑥, 0) = 𝑢 0 (𝑥), 𝑥 ∈ (0, 2𝜋), 𝑡 ∈ (0, 1] 𝑥 ∈ (0, 2𝜋) (3.14) with periodic boundary conditions. We consider mapping the initial condition to the solution at time one 𝑢 0 ↦→ 𝑢(·, 1). Burger’s equation re-arranges low to mid range energies resulting in steep discontinuities that are dampened proportionately to the size of the viscosity 𝜈. It acts as a simplified model for the Navier-Stokes equation. We sample initial conditions as Gaussian random fields and solve (6.18) via a split-step method on a fine mesh, sub-sampling other data as needed; two examples of 𝑢 0 and 𝑢(·, 1) are shown in the middle and right of Figure 3.4. Figure 3.4 shows the relative test errors for a variety of methods on Burger’s (6.18) (left) as a function of the grid resolution. First notice that MGKN achieves a constant steady state test error, demonstrating that it has learned the true infinitedimensional solution operator irrespective of the discretization. This is in contrast to the state-of-the-art fully convolution network ( FCN) proposed in (Zhu and Zabaras, 2018) which has the property that what it learns is tied to a specific discretization. Indeed, we see that, in both cases, the error increases with the resolution since standard convolution layer are parametrized locally and therefore cannot capture the long-range correlations of the solution operator. Using linear spaces, the ( PCA+NN) method proposed in (Bhattacharya, Kovachki, and Andrew M. Stuart, 2020) utilizes deep learning to produce a fast, fully data-driven reduced order model. The graph convolution network ( GCN) method follows (Alet et al., 2019)’s architecture, with naive nearest neighbor connection. It shows simple nearest-neighbor graph structures are insufficient. The graph kernel network ( GKN) employs an architecture similar to (3.3) but, without the multi-level graph extension, it can be slow due the quadratic time complexity. For Burger’s equation, when 52 1-d Burgers equation with viscosity 𝜈 = 0.1. Left: performance of different methods. MGKN has competitive performance. Mid: input functions (𝑢 0 ) of two examples. Right: corresponding outputs from MGKN of the two examples, and their ground truth (dash line). The error is minimal on both examples. Figure 3.4: Comparsion with benchmarks on Burgers equation, with examples of inputs and outputs. linear spaces are no longer near-optimal, MGKN is the best performing method. This is a very encouraging result since many challenging applied problem are not well approximated by linear spaces and can therefore greatly benefit from non-linear approximation methods such as MGKN. For details on the other methods, see the supplements. 3.6 Discussion and Conclusion We introduced the multipole graph kernel network (MGKN), a graph-based algorithm able to capture correlations in data at any length scale with a linear time complexity. Our work ties together ideas from graph neural networks, multi-scale modeling, and matrix factorization. Using a kernel integration architecture, we validate our methodology by showing that it can learn mesh-invariant solutions operators to parametric PDEs. Ideas in this work are not tied to our particular applications and can be used provide significant speed-up for processing densely connected graph data. References Abu-El-Haija, Sami, Amol Kapoor, et al. (2018). “N-gcn: Multi-scale graph convolution for semi-supervised node classification”. In: arXiv preprint arXiv:1802.08888. Abu-El-Haija, Sami, Bryan Perozzi, et al. (2019). “Mixhop: Higher-order graph convolutional architectures via sparsified neighborhood mixing”. In: arXiv preprint arXiv:1905.00067. 53 Adler, Jonas and Ozan Oktem (Nov. 2017). “Solving ill-posed inverse problems using iterative deep neural networks”. In: Inverse Problems. doi: 10.1088/13616420/aa9581. url: https://doi.org/10.1088%2F1361-6420%2Faa9581. Alet, Ferran et al. (2019). “Graph Element Networks: adaptive, structured computation and memory”. In: 36th International Conference on Machine Learning. PMLR. url: http://proceedings.mlr.press/v97/alet19a.html. Alfke, Dominik and Martin Stoll (2019). “Semi-Supervised Classification on NonSparse Graphs Using Low-Rank Graph Convolutional Networks”. In: CoRR abs/1905.10224. arXiv: 1905.10224. url: http://arxiv.org/abs/1905. 10224. Bar, Leah and Nir Sochen (2019). “Unsupervised deep learning algorithm for PDEbased forward and inverse problems”. In: arXiv preprint arXiv:1904.05417. Bhatnagar, Saakaar et al. (2019). “Prediction of aerodynamic flow fields using convolutional neural networks”. In: Computational Mechanics, pp. 1–21. Bhattacharya, Kaushik, Nikola B. Kovachki, and Andrew M. Stuart (2020). “Model Reduction and Neural Networks for Parametric PDE(s)”. In: arXiv preprint arXiv:2005.03180. Börm, Steffen, Lars Grasedyck, and Wolfgang Hackbusch (2003). “Hierarchical matrices”. In: Lecture notes 21, p. 2003. E, Weinan and Bing Yu (Mar. 2018). “The Deep Ritz Method: A Deep LearningBased Numerical Algorithm for Solving Variational Problems”. English (US). In: Communications in Mathematics and Statistics. issn: 2194-6701. doi: 10.1007/ s40304-018-0127-z. Fan, Yuwei, Jordi Feliu-Faba, et al. (2019). “A multiscale neural network based on hierarchical nested bases”. In: Research in the Mathematical Sciences 6.2, p. 21. Fan, Yuwei, Lin Lin, et al. (2019). “A multiscale neural network based on hierarchical matrices”. In: Multiscale Modeling & Simulation 17.4, pp. 1189–1213. Fey, Matthias and Jan E. Lenssen (2019). “Fast Graph Representation Learning with PyTorch Geometric”. In: ICLR Workshop on Representation Learning. Gao, Hongyang and Shuiwang Ji (2019). “Graph U-Nets”. In: arXiv preprint arXiv:1905.05178. Gilmer, Justin et al. (2017). “Neural message passing for quantum chemistry”. In: Proceedings of the 34th International Conference on Machine Learning. Greengard, Leslie and Vladimir Rokhlin (1997). “A new version of the fast multipole method for the Laplace equation in three dimensions”. In: Acta numerica 6, pp. 229–269. Guo, Xiaoxiao, Wei Li, and Francesco Iorio (2016). “Convolutional neural networks for steady flow approximation”. In: Proceedings of the 22nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining. 54 Hamilton, Will, Zhitao Ying, and Jure Leskovec (2017). “Inductive representation learning on large graphs”. In: Advances in neural information processing systems, pp. 1024–1034. Ithapu, Vamsi K et al. (2017). “The incremental multiresolution matrix factorization algorithm”. In: Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 2951–2960. Jiang, Chiyu Max et al. (2020). “MeshfreeFlowNet: A Physics-Constrained Deep Continuous Space-Time Super-Resolution Framework”. In: arXiv preprint arXiv:2005.01463. Kipf, Thomas N and Max Welling (2016). “Semi-supervised classification with graph convolutional networks”. In: arXiv preprint arXiv:1609.02907. Kondor, Risi, Nedelina Teneva, and Vikas Garg (2014). “Multiresolution matrix factorization”. In: International Conference on Machine Learning, pp. 1620– 1628. Li, Maosen et al. (2020). “Dynamic Multiscale Graph Neural Networks for 3D Skeleton-Based Human Motion Prediction”. In: arXiv preprint arXiv:2003.08802. Li, Zongyi et al. (2020). “Neural operator: Graph kernel network for partial differential equations”. In: arXiv preprint arXiv:2003.03485. Lu, Lu, Pengzhan Jin, and George Em Karniadakis (2019). “DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators”. In: arXiv preprint arXiv:1910.03193. Murphy, Ryan L et al. (2018). “Janossy pooling: Learning deep permutationinvariant functions for variable-size inputs”. In: arXiv preprint arXiv:1811.01900. Nelsen, Nicholas H and Andrew M Stuart (2020). “The Random Feature Model for Input-Output Maps between Banach Spaces”. In: arXiv preprint arXiv:2005.10224. Raissi, Maziar, Paris Perdikaris, and George E Karniadakis (2019). “Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations”. In: Journal of Computational Physics 378, pp. 686–707. Ronneberger, Olaf, Philipp Fischer, and Thomas Brox (2015). “U-net: Convolutional networks for biomedical image segmentation”. In: International Conference on Medical image computing and computer-assisted intervention. Springer, pp. 234– 241. Rozza, Gianluigi, Dinh Bao Phuong Huynh, and Anthony T Patera (2007). “Reduced basis approximation and a posteriori error estimation for affinely parametrized elliptic coercive partial differential equations”. In: Archives of Computational Methods in Engineering 15.3, p. 1. Sauter, Stefan A. and Christoph Schwab (2010). Boundary Element Methods. Springer Series in Computational Mathematics. 55 Smith, Jonathan D, Kamyar Azizzadenesheli, and Zachary E Ross (2020). “Eikonet: Solving the eikonal equation with deep neural networks”. In: arXiv preprint arXiv:2004.00361. Su, Chau Hsing and Clifford S Gardner (1969). “Korteweg-de Vries equation and generalizations. III. Derivation of the Korteweg-de Vries equation and Burgers equation”. In: Journal of Mathematical Physics 10.3, pp. 536–539. Tek, MR et al. (1957). “Development of a generalized Darcy equation”. In: Journal of Petroleum Technology 9.06, pp. 45–47. Ummenhofer, Benjamin et al. (2020). “Lagrangian fluid simulation with continuous convolutions”. In: International Conference on Learning Representations. Veličković, Petar et al. (2017). “Graph attention networks”. In. Wu, Zonghan et al. (2020). “A comprehensive survey on graph neural networks”. In: IEEE Transactions on Neural Networks and Learning Systems. Ying, Lexing, George Biros, and Denis Zorin (2004). “A kernel-independent adaptive fast multipole algorithm in two and three dimensions”. In: Journal of Computational Physics 196.2, pp. 591–626. Zhu, Yinhao and Nicholas Zabaras (2018). “Bayesian deep convolutional encoder–decoder networks for surrogate modeling and uncertainty quantification”. In: Journal of Computational Physics. issn: 0021-9991. doi: https://doi.org/10.1016/ j.jcp.2018.04.018. url: http://www.sciencedirect.com/science/ article/pii/S0021999118302341. Zienkiewicz, Olgierd Cecil et al. (1977). The finite element method. Vol. 3. McGrawhill London. 56 Chapter 4 MODEL: FOURIER NEURAL OPERATOR The classical development of neural networks has primarily focused on learning mappings between finite-dimensional Euclidean spaces. Recently, this has been generalized to neural operators that learn mappings between function spaces. For partial differential equations (PDEs), neural operators directly learn the mapping from any functional parametric dependence to the solution. Thus, they learn an entire family of PDEs, in contrast to classical methods which solve one instance of the equation. In this work, we formulate a new neural operator by parameterizing the integral kernel directly in Fourier space, allowing for an expressive and efficient architecture. We perform experiments on Burgers’ equation, Darcy flow, and NavierStokes equation. Experiments shows that Fourier Neural Operator achieves superior accuracy compared to previous learning-based solvers under fixed resolution. 4.1 Introduction Many problems in science and engineering involve solving complex partial differential equation (PDE) systems repeatedly for different values of some parameters. Examples arise in molecular dynamics, micro-mechanics, and turbulent flows. Often such systems require fine discretization in order to capture the phenomenon being modeled. As a consequence, traditional numerical solvers are slow and sometimes inefficient. For example, when designing materials such as airfoils, one needs to solve the associated inverse problem where thousands of evaluations of the forward model are needed. A fast method can make such problems feasible. Conventional solvers vs. Data-driven methods. Traditional solvers such as finite element methods (FEM) and finite difference methods (FDM) solve the equation by discretizing the space. Therefore, they impose a trade-off on the resolution: coarse grids are fast but less accurate; fine grids are accurate but slow. Complex PDE systems, as described above, usually require a very fine discretization, and therefore very challenging and time-consuming for traditional solvers. On the other hand, data-driven methods can directly learn the trajectory of the family of equations from the data. As a result, the learning-based method can be orders of magnitude faster than the conventional solvers. 57 Machine learning methods may hold the key to revolutionizing scientific disciplines by providing fast solvers that approximate or enhance traditional ones (Raissi, Perdikaris, and George E Karniadakis, 2019; C. M. Jiang et al., 2020; Greenfeld et al., 2019; Kochkov et al., 2021). However, classical neural networks map between finite-dimensional spaces and can therefore only learn solutions tied to a specific discretization. This is often a limitation for practical applications and therefore the development of mesh-invariant neural networks is required. We first outline two mainstream neural network-based approaches for PDEs: the finite-dimensional operators and physics-informed neural network. Finite-dimensional operators. These approaches parameterize the solution operator as a deep convolutional neural network between finite-dimensional Euclidean spaces Guo, W. Li, and Iorio, 2016; Zhu and Zabaras, 2018; Adler and Oktem, 2017; Bhatnagar et al., 2019; Khoo, J. Lu, and Ying, 2017. Such approaches are, by definition, mesh-dependent and will need modifications and tuning for different resolutions and discretizations in order to achieve consistent error (if at all possible). Furthermore, these approaches are limited to the discretization size and geometry of the training data and hence, it is not possible to query solutions at new points in the domain. In contrast, we show, for our method, both invariance of the error to grid resolution, and the ability to transfer the solution between meshes. Physics-Informed Neural Network. The second approach directly parameterizes the solution function as a neural network (E and Yu, 2018; Raissi, Perdikaris, and George E Karniadakis, 2019; Bar and Sochen, 2019; Smith, Azizzadenesheli, and Ross, 2020; Pan and Duraisamy, 2020). This approach is designed to model one specific instance of the PDE, not the solution operator. It is mesh-independent and accurate, but for any given new instance of the functional parameter/coefficient, it requires training a new neural network. The approach closely resembles classical methods such as finite elements, replacing the linear span of a finite set of local basis functions with the space of neural networks. The Neural-FEM approach suffers from the same computational issue as classical methods: the optimization problem needs to be solved for every new instance. Furthermore, the approach is limited to a setting in which the underlying PDE is known. Neural Operators. Recently, a new line of work proposed learning mesh-free, infinite-dimensional operators with neural networks (L. Lu, Jin, and George Em 58 Zero-shot super-resolution: Navier-Stokes Equation with viscosity 𝜈 = 1e−4; Ground truth on top and prediction on bottom; trained on 64 × 64 × 20 dataset; evaluated on 256 × 256 × 80 (see Section 4.4). Figure 4.1: Zero-shot super-resolution on Navier-Stokes. Karniadakis, 2019; Bhattacharya, Kovachki, and Andrew M. Stuart, 2020; Nelsen and A. Stuart, 2020; Z. Li et al., 2020b; Z. Li et al., 2020a; Patel et al., 2021). The neural operator remedies the mesh-dependent nature of the finite-dimensional operator methods discussed above by producing a single set of network parameters that may be used with different discretizations. It has the ability to transfer solutions between meshes. Furthermore, the neural operator needs to be trained only once. Obtaining a solution for a new instance of the parameter requires only a forward pass of the network, alleviating the major computational issues incurred in Neural-FEM methods. Lastly, the neural operator requires no knowledge of the underlying PDE, only data. Thus far, neural operators have not yielded efficient numerical algorithms that can parallel the success of convolutional or recurrent neural networks in the finite-dimensional setting due to the cost of evaluating integral operators. Through the fast Fourier transform, our work alleviates this issue. Fourier Transform. The Fourier transform is frequently used in spectral methods for solving differential equations, since differentiation is equivalent to multiplication in the Fourier domain. Fourier transforms have also played an important role in the development of deep learning. In theory, they appear in the proof of the universal approximation theorem (Hornik, Stinchcombe, White, et al., 1989) and, empirically, they have been used to speed up convolutional neural networks (Mathieu, Henaff, and LeCun, 2013). Neural network architectures involving the Fourier transform or the use of sinusoidal activation functions have also been proposed and studied (mingo2004Fourier; Bengio, LeCun, et al., 2007; Sitzmann et al., 2020). Recently, 59 some spectral methods for PDEs have been extended to neural networks (Fan, Bohorquez, and Ying, 2019; Fan, Lin, et al., 2019; Kashinath, Marcus, et al., 2020). We build on these works by proposing a neural operator architecture defined directly in Fourier space with quasi-linear time complexity and state-of-the-art approximation capabilities. Our Contributions. We introduce the Fourier neural operator, a novel deep learning architecture able to learn mappings between infinite-dimensional spaces of functions; the integral operator is restricted to a convolution, and instantiated through a linear transformation in the Fourier domain. • The Fourier neural operator learns the resolution-invariant solution operator for the family of Navier-Stokes equation in the turbulent regime, where previous graph-based neural operators do not converge. • By construction, the method shares the same learned network parameters irrespective of the discretization used on the input and output spaces. It can do zero-shot super-resolution: trained on a lower resolution directly evaluated on a higher resolution, as shown in Figure 4.1. • The proposed method consistently outperforms all existing deep learning methods even when fixing the resolution to be 64 × 64. It achieves error rates that are 30% lower on Burgers’ Equation, 60% lower on Darcy Flow, and 30% lower on Navier Stokes (turbulent regime with viscosity 𝜈 = 1e−4). When learning the mapping for the entire time series, the method achieves < 1% error with viscosity 𝜈 = 1e−3 and 8% error with viscosity 𝜈 = 1e−4. • On a 256 × 256 grid, the Fourier neural operator has an inference time of only 0.005s compared to the 2.2𝑠 of the pseudo-spectral method used to solve Navier-Stokes. Despite its tremendous speed advantage, the method does not suffer from accuracy degradation when used in downstream applications such as solving the Bayesian inverse problem, as shown in Figure 4.4. We observed that the proposed framework can approximate complex operators raising in PDEs that are highly non-linear, with high frequency modes and slow energy decay. The power of neural operators comes from combining linear, global integral operators (via the Fourier transform) and non-linear, local activation functions. Similar to the way standard neural networks approximate highly non-linear functions by 60 (a) The full architecture of neural operator: start from input 𝑎. 1. Lift to a higher dimension channel space by a neural network 𝑃. 2. Apply four layers of integral operators and activation functions. 3. Project back to the target dimension by a neural network 𝑄. Output 𝑢. (b) Fourier layers: Start from input 𝑣. On top: apply the Fourier transform G; a linear transform 𝑅 on the lower Fourier modes and filters out the higher modes; then apply the inverse Fourier transform G −1 . On the bottom: apply a local linear transform 𝑊. Figure 4.2: The architecture of the neural operators and Fourier layer combining linear multiplications with non-linear activations, the proposed neural operators can approximate highly non-linear operators. 4.2 Neural Operator The neural operator, proposed in (Z. Li et al., 2020b), is formulated as an iterative architecture 𝑣 0 ↦→ 𝑣 1 ↦→ . . . ↦→ 𝑣𝑇 where 𝑣 𝑗 for 𝑗 = 0, 1, . . . , 𝑇 − 1 is a sequence of functions each taking values in R𝑑 𝑣 . As shown in Figure 4.2 (a), the input 𝑎 ∈ A is first lifted to a higher dimensional representation 𝑣 0 (𝑥) = 𝑃(𝑎(𝑥)) by the local transformation 𝑃 which is usually parameterized by a shallow fully-connected neural network. Then we apply several iterations of updates 𝑣 𝑡 ↦→ 𝑣 𝑡+1 (defined below). The output 𝑢(𝑥) = 𝑄(𝑣𝑇 (𝑥)) is the projection of 𝑣𝑇 by the local transformation 𝑄 : R𝑑 𝑣 → R𝑑𝑢 . In each iteration, the update 𝑣 𝑡 ↦→ 𝑣 𝑡+1 is defined as the composition of a non-local integral operator K and a local, nonlinear activation function 𝜎. Definition 4.2.1 (Iterative updates) Define the update to the representation 𝑣 𝑡 ↦→ 𝑣 𝑡+1 by 
 𝑣 𝑡+1 (𝑥) := 𝜎 𝑊𝑣 𝑡 (𝑥) + K (𝑎; 𝜙)𝑣 𝑡 (𝑥) , ∀𝑥 ∈ 𝐷 (4.1) where K : A × ΘK → L (U (𝐷; R𝑑 𝑣 ), U (𝐷; R𝑑 𝑣 )) maps to bounded linear operators on U (𝐷; R𝑑 𝑣 ) and is parameterized by 𝜙 ∈ ΘK , 𝑊 : R𝑑 𝑣 → R𝑑 𝑣 is a linear transformation, and 𝜎 : R → R is a non-linear activation function whose action is defined component-wise. 61 We choose K (𝑎; 𝜙) to be a kernel integral transformation parameterized by a neural network. Definition 4.2.2 (Kernel integral operator K) Define the kernel integral operator mapping in (4.1) by ∫

K (𝑎; 𝜙)𝑣 𝑡 (𝑥) := 𝜅 𝑥, 𝑦, 𝑎(𝑥), 𝑎(𝑦); 𝜙 𝑣 𝑡 (𝑦)d𝑦, ∀𝑥 ∈ 𝐷 (4.2) 𝐷 where 𝜅 𝜙 : R2(𝑑+𝑑 𝑎 ) → R𝑑 𝑣 ×𝑑 𝑣 is a neural network parameterized by 𝜙 ∈ ΘK . Here 𝜅 𝜙 plays the role of a kernel function which we learn from data. Together definitions 1 and 2 constitute a generalization of neural networks to infinite-dimensional spaces as first proposed in Z. Li et al., 2020b. Notice even the integral operator is linear, the neural operator can learn highly non-linear operators by composing linear integral operators with non-linear activation functions, analogous to standard neural networks. If we remove the dependence on the function 𝑎 and impose 𝜅 𝜙 (𝑥, 𝑦) = 𝜅 𝜙 (𝑥 − 𝑦), we obtain that (4.2) is a convolution operator, which is a natural choice from the perspective of fundamental solutions. We exploit this fact in the following section by parameterizing 𝜅 𝜙 directly in Fourier space and using the Fast Fourier Transform (FFT) to efficiently compute (4.2). This leads to a fast architecture that obtains state-of-the-art results for PDE problems. 4.3 Fourier Neural Operator We propose replacing the kernel integral operator in (4.2), by a convolution operator defined in Fourier space. Let G denote the Fourier transform of a function 𝑓 : 𝐷 → R𝑑 𝑣 and G −1 its inverse then ∫ ∫ −2𝑖𝜋⟨𝑥,𝑘⟩ −1 (G 𝑓 ) 𝑗 (𝑘) = 𝑓 𝑗 (𝑥)𝑒 d𝑥, (G 𝑓 ) 𝑗 (𝑥) = 𝑓 𝑗 (𝑘)𝑒 2𝑖𝜋⟨𝑥,𝑘⟩ d𝑘 𝐷 𝐷 √ for 𝑗 = 1, . . . , 𝑑 𝑣 where 𝑖 = −1 is the imaginary unit. By letting 𝜅 𝜙 (𝑥, 𝑦, 𝑎(𝑥), 𝑎(𝑦)) = 𝜅 𝜙 (𝑥 − 𝑦) in (4.2) and applying the convolution theorem, we find that

K (𝑎; 𝜙)𝑣 𝑡 (𝑥) = G −1 G(𝜅 𝜙 ) · G(𝑣 𝑡 ) (𝑥), ∀𝑥 ∈ 𝐷. We, therefore, propose to directly parameterize 𝜅 𝜙 in Fourier space. 62 Definition 4.3.1 (Fourier integral operator K) Define the Fourier integral operator  
−1 K (𝜙)𝑣 𝑡 (𝑥) = G 𝑅𝜙 · (G𝑣 𝑡 ) (𝑥) ∀𝑥 ∈ 𝐷 (4.3) where 𝑅𝜙 is the Fourier transform of a periodic function 𝜅 : 𝐷¯ → R𝑑 𝑣 ×𝑑 𝑣 parameterized by 𝜙 ∈ ΘK . An illustration is given in Figure 4.2 (b). For frequency mode 𝑘 ∈ 𝐷, we have (G𝑣 𝑡 )(𝑘) ∈ C 𝑑 𝑣 and 𝑅𝜙 (𝑘) ∈ C 𝑑 𝑣 ×𝑑 𝑣 . Notice that since we assume 𝜅 is periodic, it admits a Fourier series expansion, so we may work with the discrete modes 𝑘 ∈ Z𝑑 . We pick a finite-dimensional parameterization by truncating the Fourier series at a maximal number of modes 𝑘 max = |𝑍 𝑘 max | = |{𝑘 ∈ Z𝑑 : |𝑘 𝑗 | ≤ 𝑘 max, 𝑗 , for 𝑗 = 1, . . . , 𝑑}|. We thus parameterize 𝑅𝜙 directly as complex-valued (𝑘 max × 𝑑𝑣 × 𝑑𝑣 )-tensor comprising a collection of truncated Fourier modes and therefore drop 𝜙 from our notation. Since 𝜅 is real-valued, we impose conjugate symmetry. We note that the set 𝑍 𝑘 max is not the canonical choice for the low frequency modes of 𝑣 𝑡 . Indeed, the low frequency modes are usually defined by placing an upper-bound on the ℓ1 -norm of 𝑘 ∈ Z𝑑 . We choose 𝑍 𝑘 max as above since it allows for an efficient implementation. The discrete case and the FFT. Assuming the domain 𝐷 is discretized with 𝑛 ∈ N points, we have that 𝑣 𝑡 ∈ R𝑛×𝑑 𝑣 and G(𝑣 𝑡 ) ∈ C 𝑛×𝑑 𝑣 . Since we convolve 𝑣 𝑡 with a function which only has 𝑘 max Fourier modes, we may simply truncate the higher modes to obtain G(𝑣 𝑡 ) ∈ C 𝑘 max ×𝑑 𝑣 . Multiplication by the weight tensor 𝑅 ∈ C 𝑘 max ×𝑑 𝑣 ×𝑑 𝑣 is then
𝑅 · (G𝑣 𝑡 ) 𝑘,𝑙 = 𝑑𝑣 ∑︁ 𝑅 𝑘,𝑙, 𝑗 (G𝑣 𝑡 ) 𝑘, 𝑗 , 𝑘 = 1, . . . , 𝑘 max , 𝑗 = 1, . . . , 𝑑 𝑣 . (4.4) 𝑗=1 When the discretization is uniform with resolution 𝑠1 ×· · ·×𝑠 𝑑 = 𝑛, G can be replaced by the Fast Fourier Transform. For 𝑓 ∈ R𝑛×𝑑 𝑣 , 𝑘 = (𝑘 1 , . . . , 𝑘 𝑑 ) ∈ Z𝑠1 × · · · × Z𝑠 𝑑 , and 𝑥 = (𝑥 1 , . . . , 𝑥 𝑑 ) ∈ 𝐷, the FFT Ĝ and its inverse Ĝ −1 are defined as ( Ĝ 𝑓 )𝑙 (𝑘) = 𝑠∑︁ 1 −1 ··· 𝑥 1 =0 ( Ĝ −1 𝑓 )𝑙 (𝑥) = 𝑠∑︁ 1 −1 𝑘 1 =0 𝑠∑︁ 𝑑 −1 𝑓𝑙 (𝑥 1 , . . . , 𝑥 𝑑 )𝑒 −2𝑖𝜋 Í𝑑 𝑗=1 𝑥𝑗𝑘𝑗 𝑠𝑗 , 𝑥 𝑑 =0 ··· 𝑠∑︁ 𝑑 −1 𝑓𝑙 (𝑘 1 , . . . , 𝑘 𝑑 )𝑒 2𝑖𝜋 Í𝑑 𝑗=1 𝑥𝑗𝑘𝑗 𝑠𝑗 𝑘 𝑑 =0 for 𝑙 = 1, . . . , 𝑑 𝑣 . In this case, the set of truncated modes becomes 𝑍 𝑘 max = {(𝑘 1 , . . . , 𝑘 𝑑 ) ∈ Z𝑠1 ×· · ·×Z𝑠 𝑑 | 𝑘 𝑗 ≤ 𝑘 max, 𝑗 or 𝑠 𝑗 −𝑘 𝑗 ≤ 𝑘 max, 𝑗 , for 𝑗 = 1, . . . , 𝑑}. 63 When implemented, 𝑅 is treated as a (𝑠1 × · · · × 𝑠 𝑑 × 𝑑𝑣 × 𝑑𝑣 )-tensor and the above definition of 𝑍 𝑘 max corresponds to the “corners” of 𝑅, which allows for a straight-forward parallel implementation of (4.4) via matrix-vector multiplication. In practice, we have found that choosing 𝑘 max, 𝑗 = 12 which yields 𝑘 max = 12𝑑 parameters per channel to be sufficient for all the tasks that we consider. Parameterizations of 𝑅. In general, 𝑅 can be defined to depend on (G𝑎) to parallel (4.2). Indeed, we can define 𝑅𝜙 : Z𝑑 × R𝑑 𝑣 → R𝑑 𝑣 ×𝑑 𝑣 as a parametric function that maps 𝑘, (G𝑎)(𝑘)) to the values of the appropriate Fourier modes. We have experimented with linear as well as neural network parameterizations of 𝑅𝜙 . We find that the linear parameterization has a similar performance to the previously described direct parameterization, while neural networks have worse performance. This is likely due to the discrete structure of the space Z𝑑 . Our experiments in this work focus on the direct parameterization presented above. Invariance to discretization. The Fourier layers are discretization-invariant because they can learn from and evaluate functions which are discretized in an arbitrary way. Since parameters are learned directly in Fourier space, resolving the functions in physical space simply amounts to projecting on the basis 𝑒 2𝜋𝑖⟨𝑥,𝑘⟩ which are welldefined everywhere on R𝑑 . This allows us to achieve zero-shot super-resolution as shown in Section 4.4. Furthermore, our architecture has a consistent error at any resolution of the inputs and outputs. On the other hand, notice that, in Figure 4.3, the standard CNN methods we compare against have an error that grows with the resolution. Quasi-linear complexity. The weight tensor 𝑅 contains 𝑘 max < 𝑛 modes, so the inner multiplication has complexity 𝑂 (𝑘 max ). Therefore, the majority of the computational cost lies in computing the Fourier transform G(𝑣 𝑡 ) and its inverse. General Fourier transforms have complexity 𝑂 (𝑛2 ), however, since we truncate the series the complexity is in fact 𝑂 (𝑛𝑘 max ), while the FFT has complexity 𝑂 (𝑛 log 𝑛). Generally, we have found using FFTs to be very efficient. However a uniform discretization is required. 4.4 Numerical Experiments In this section, we compare the proposed Fourier neural operator with multiple finite-dimensional architectures as well as operator-based approximation methods 64 on the 1-d Burgers’ equation, the 2-d Darcy Flow problem, and 2-d Navier-Stokes equation. We do not compare against traditional solvers (FEM/FDM) or neuralFEM type methods since our goal is to produce an efficient operator approximation that can be used for downstream applications. We demonstrate one such application to the Bayesian inverse problem in Section 4.4. We construct our Fourier neural operator by stacking four Fourier integral operator layers as specified in (4.1) and (9.5) with the ReLU activation as well as batch normalization. Unless otherwise specified, we use 𝑁 = 1000 training instances and 200 testing instances. We use Adam optimizer to train for 500 epochs with an initial learning rate of 0.001 that is halved every 100 epochs. We set 𝑘 max, 𝑗 = 16, 𝑑𝑣 = 64 for the 1-d problem and 𝑘 max, 𝑗 = 12, 𝑑𝑣 = 32 for the 2-d problems. Lower resolution data are downsampled from higher resolution. All the computation is carried on a single Nvidia V100 GPU with 16GB memory. Remark on Resolution. Traditional PDE solvers such as FEM and FDM approximate a single function and therefore their error to the continuum decreases as the resolution is increased. On the other hand, operator approximation is independent of the ways its data is discretized as long as all relevant information is resolved. Resolution-invariant operators have consistent error rates among different resolutions as shown in Figure 4.3. Further, resolution-invariant operators can do zero-shot super-resolution, as shown in Section 4.4. Benchmarks for time-independent problems (Burgers and Darcy): NN: a simple point-wise feedforward neural network. RBM: the classical Reduced Basis Method (using a POD basis) (DeVore, 2014). FCN: a the-state-of-the-art neural network architecture based on Fully Convolution Networks (Zhu and Zabaras, 2018). PCANN: an operator method using PCA as an autoencoder on both the input and output data and interpolating the latent spaces with a neural network (Bhattacharya, Kovachki, and Andrew M. Stuart, 2020). GNO: the original graph neural operator (Z. Li et al., 2020b). MGNO: the multipole graph neural operator (Z. Li et al., 2020a). LNO: a neural operator method based on the low-rank Í decomposition of the kernel 𝜅(𝑥, 𝑦) := 𝑟𝑗=1 𝜙 𝑗 (𝑥)𝜓 𝑗 (𝑦), similar to the unstacked DeepONet proposed in (L. Lu, Jin, and George Em Karniadakis, 2019). FNO: the newly purposed Fourier neural operator. 65 Left: benchmarks on Burgers equation; Mid: benchmarks on Darcy Flow for different resolutions; Right: the learning curves on Navier-Stokes 𝜈 = 1e−3 with different benchmarks. Train and test on the same resolution. Figure 4.3: Benchmark on Burger’s equation, Darcy Flow, and Navier-Stokes Benchmarks for time-dependent problems (Navier-Stokes): ResNet: 18 layers of 2-d convolution with residual connections (He et al., 2016). U-Net: A popular choice for image-to-image regression tasks consisting of four blocks with 2-d convolutions and deconvolutions (Ronneberger, Fischer, and Brox, 2015). TFNet: A network designed for learning turbulent flows based on a combination of spatial and temporal convolutions (Wang et al., 2020). FNO-2d: 2-d Fourier neural operator with a RNN structure in time. FNO-3d: 3-d Fourier neural operator that directly convolves in space-time. Burgers’ Equation The 1-d Burgers’ equation is a non-linear PDE with various applications including modeling the one dimensional flow of a viscous fluid. It takes the form 𝜕𝑡 𝑢(𝑥, 𝑡) + 𝜕𝑥 (𝑢 2 (𝑥, 𝑡)/2) = 𝜈𝜕𝑥𝑥 𝑢(𝑥, 𝑡), 𝑢(𝑥, 0) = 𝑢 0 (𝑥), 𝑥 ∈ (0, 1), 𝑡 ∈ (0, 1] 𝑥 ∈ (0, 1) (4.5) 2 ((0, 1); R) is the initial condition with periodic boundary conditions where 𝑢 0 ∈ 𝐿 per and 𝜈 ∈ R+ is the viscosity coefficient. We aim to learn the operator mapping the 2 ((0, 1); R) → 𝐻 𝑟 ((0, 1); R) initial condition to the solution at time one, G † : 𝐿 per per defined by 𝑢 0 ↦→ 𝑢(·, 1) for any 𝑟 > 0. The results of our experiments are shown in Figure 4.3 (a) and Table 4.1. Our proposed method obtains the lowest relative error compared to any of the benchmarks. Further, the error is invariant with the resolution, while the error of convolution neural network based methods (FCN) grows with the resolution. Compared to other neural operator methods such as GNO and MGNO that use Nyström sam- 66 pling in physical space, the Fourier neural operator is both more accurate and more computationally efficient. Networks NN GCN FCN PCANN GNO LNO MGNO FNO 𝑠 = 256 0.4714 0.3999 0.0958 0.0398 0.0555 0.0212 0.0243 0.0149 𝑠 = 512 0.4561 0.4138 0.1407 0.0395 0.0594 0.0221 0.0355 0.0158 𝑠 = 1024 0.4803 0.4176 0.1877 0.0391 0.0651 0.0217 0.0374 0.0160 𝑠 = 2048 0.4645 0.4157 0.2313 0.0383 0.0663 0.0219 0.0360 0.0146 𝑠 = 4096 0.4779 0.4191 0.2855 0.0392 0.0666 0.0200 0.0364 0.0142 𝑠 = 8192 0.4452 0.4198 0.3238 0.0393 0.0699 0.0189 0.0364 0.0139 Table 4.1: Benchmarks on 1-d Burgers’ equation Networks NN FCN PCANN RBM GNO LNO MGNO FNO 𝑠 = 85 0.1716 0.0253 0.0299 0.0244 0.0346 0.0520 0.0416 0.0108 𝑠 = 141 0.1716 0.0493 0.0298 0.0251 0.0332 0.0461 0.0428 0.0109 𝑠 = 211 0.1716 0.0727 0.0298 0.0255 0.0342 0.0445 0.0428 0.0109 𝑠 = 421 0.1716 0.1097 0.0299 0.0259 0.0369 − 0.0420 0.0098 Table 4.2: Benchmarks on 2-d Darcy Flow Darcy Flow We consider the steady-state of the 2-d Darcy Flow equation on the unit box which is the second order, linear, elliptic PDE −∇ · (𝑎(𝑥)∇𝑢(𝑥)) = 𝑓 (𝑥) 𝑢(𝑥) = 0 𝑥 ∈ (0, 1) 2 𝑥 ∈ 𝜕 (0, 1) 2 (4.6) with a Dirichlet boundary where 𝑎 ∈ 𝐿 ∞ ((0, 1) 2 ; R+ ) is the diffusion coefficient and 𝑓 ∈ 𝐿 2 ((0, 1) 2 ; R) is the forcing function. This PDE has numerous applications including modeling the pressure of subsurface flow, the deformation of linearly elastic materials, and the electric potential in conductive materials. We are interested in learning the operator mapping the diffusion coefficient to the solution, G † : 𝐿 ∞ ((0, 1) 2 ; R+ ) → 𝐻01 ((0, 1) 2 ; R+ ) defined by 𝑎 ↦→ 𝑢. Note that although the PDE is linear, the operator G † is not. 67 The results of our experiments are shown in Figure 4.3 (b) and Table 4.2. The proposed Fourier neural operator obtains nearly one order of magnitude lower relative error compared to any benchmarks. We again observe the invariance of the error with respect to the resolution. 2D and 3D Convolutions. FNO-2D, U-Net, TF-Net, and ResNet all do 2Dconvolution in the spatial domain and recurrently propagate in the time domain (2D+RNN). The operator maps the solution at the previous 10 time steps to the next time step (2D functions to 2D functions). On the other hand, FNO-3D performs convolution in space-time. It maps the initial time steps directly to the full trajectory (3D functions to 3D functions). The 2D+RNN structure can propagate the solution to any arbitrary time 𝑇 in increments of a fixed interval length Δ𝑡, while the Conv3D structure is fixed to the interval [0, 𝑇] but can transfer the solution to an arbitrary time-discretization. We find the 3-d method to be more expressive and easier to train compared to its RNN-structured counterpart. Zero-shot super-resolution. The neural operator is mesh-invariant, so it can be trained on a lower resolution and evaluated at a higher resolution, without seeing any higher resolution data (zero-shot super-resolution). Figure 4.1 shows an example where we train the FNO-3D model on 64 × 64 × 20 resolution data in the setting above with (𝜈 = 1e−4, 𝑁 = 10000) and transfer to 256 × 256 × 80 resolution, demonstrating super-resolution in space-time. Fourier neural operator is the only model among the benchmarks (FNO-2D, U-Net, TF-Net, and ResNet) that can do zero-shot super-resolution. And surprisingly, it can do super-resolution not only in the spatial domain but also in the temporal domain. Bayesian Inverse Problem In this experiment, we use a function space Markov chain Monte Carlo (MCMC) method (Cotter et al., 2013) to draw samples from the posterior distribution of the initial vorticity in Navier-Stokes given sparse, noisy observations at time 𝑇 = 50. We compare the Fourier neural operator acting as a surrogate model with the traditional solvers used to generate our train-test data (both run on GPU). We generate 25,000 samples from the posterior (with a 5,000 sample burn-in period), requiring 30,000 evaluations of the forward operator. As shown in Figure 4.4, FNO and the traditional solver recover almost the same posterior mean which, when pushed forward, recovers well the late-time dynamic 68 The top left panel shows the true initial vorticity while bottom left panel shows the true observed vorticity at 𝑇 = 50 with black dots indicating the locations of the observation points placed on a 7 × 7 grid. The top middle panel shows the posterior mean of the initial vorticity given the noisy observations estimated with MCMC using the traditional solver, while the top right panel shows the same thing but using FNO as a surrogate model. The bottom middle and right panels show the vorticity at 𝑇 = 50 when the respective approximate posterior means are used as initial conditions. Figure 4.4: Results of the Bayesian inverse problem for the Navier-Stokes equation. of Navier Stokes. In sharp contrast, FNO takes 0.005𝑠 to evaluate a single instance while the traditional solver, after being optimized to use the largest possible internal time-step which does not lead to blow-up, takes 2.2𝑠. This amounts to 2.5 minutes for the MCMC using FNO and over 18 hours for the traditional solver. Even if we account for data generation and training time (offline steps) which take 12 hours, using FNO is still faster! Once trained, FNO can be used to quickly perform multiple MCMC runs for different initial conditions and observations, while the traditional solver will take 18 hours for every instance. Furthermore, since FNO is differentiable, it can easily be applied to PDE-constrained optimization problems without the need for the adjoint method. Spectral analysis. Due to the way we parameterize 𝑅𝜙 , the function output by (9.5) has at most 𝑘 max, 𝑗 Fourier modes per channel. This, however, does not mean that the Fourier neural operator can only approximate functions up to 𝑘 max, 𝑗 modes. Indeed, the activation functions which occur between integral operators and the final decoder network 𝑄 recover the high frequency modes. As an example, consider a solution to the Navier-Stokes equation with viscosity 𝜈 = 1e−3. Truncating this 69 function at 20 Fourier modes yields an error around 2% while our Fourier neural operator learns the parametric dependence and produces approximations to an error of ≤ 1% with only 𝑘 max, 𝑗 = 12 parameterized modes. Non-periodic boundary condition. Traditional Fourier methods work only with periodic boundary conditions. However, the Fourier neural operator does not have this limitation. This is due to the linear transform 𝑊 (the bias term) which keeps the track of non-periodic boundary. As an example, the Darcy Flow and the time domain of Navier-Stokes have non-periodic boundary conditions, and the Fourier neural operator still learns the solution operator with excellent accuracy. 4.5 Discussion and Conclusion Requirements on Data. Data-driven methods rely on the quality and quantity of data. To learn Navier-Stokes equation with viscosity 𝜈 = 1e−4, we need to generate 𝑁 = 10000 training pairs {𝑎 𝑗 , 𝑢 𝑗 } with the numerical solver. However, for more challenging PDEs, generating a few training samples can be already very expensive. A future direction is to combine neural operators with numerical solvers to levitate the requirements on data. Recurrent structure. The neural operator has an iterative structure that can naturally be formulated as a recurrent network where all layers share the same parameters without sacrificing performance. (We did not impose this restriction in the experiments.) Computer vision. Operator learning is not restricted to PDEs. Images can naturally be viewed as real-valued functions on 2-d domains and videos simply add a temporal structure. Our approach is therefore a natural choice for problems in computer vision where invariance to discretization crucial is important (Chi, B. Jiang, and Mu, 2020). References Adler, Jonas and Ozan Oktem (Nov. 2017). “Solving ill-posed inverse problems using iterative deep neural networks”. In: Inverse Problems. doi: 10.1088/13616420/aa9581. url: https://doi.org/10.1088%2F1361-6420%2Faa9581. Bar, Leah and Nir Sochen (2019). “Unsupervised deep learning algorithm for PDEbased forward and inverse problems”. In: arXiv preprint arXiv:1904.05417. Bengio, Yoshua, Yann LeCun, et al. (2007). “Scaling learning algorithms towards AI”. In: Large-scale kernel machines 34.5, pp. 1–41. Bhatnagar, Saakaar et al. (2019). “Prediction of aerodynamic flow fields using convolutional neural networks”. In: Computational Mechanics, pp. 1–21. 70 Bhattacharya, Kaushik, Nikola B. Kovachki, and Andrew M. Stuart (2020). “Model Reduction and Neural Networks for Parametric PDE(s)”. In: preprint. Chi, Lu, Borui Jiang, and Yadong Mu (2020). “Fast Fourier Convolution”. In: Advances in Neural Information Processing Systems 33. Cotter, S. L. et al. (Aug. 2013). “MCMC Methods for Functions: Modifying Old Algorithms to Make Them Faster”. In: Statistical Science 28.3, pp. 424–446. issn: 0883-4237. doi: 10.1214/13-sts421. url: http://dx.doi.org/10.1214/ 13-STS421. DeVore, Ronald A. (2014). “Chapter 3: The Theoretical Foundation of Reduced Basis Methods”. In: Model Reduction and Approximation. doi: 10. 1137/1. 9781611974829 . ch3. eprint: https : / / epubs . siam . org / doi / pdf / 10 . 1137/1.9781611974829.ch3. url: https://epubs.siam.org/doi/abs/ 10.1137/1.9781611974829.ch3. E, Weinan and Bing Yu (Mar. 2018). “The Deep Ritz Method: A Deep LearningBased Numerical Algorithm for Solving Variational Problems”. English (US). In: Communications in Mathematics and Statistics. issn: 2194-6701. doi: 10.1007/ s40304-018-0127-z. Fan, Yuwei, Cindy Orozco Bohorquez, and Lexing Ying (2019). “BCR-Net: A neural network based on the nonstandard wavelet form”. In: Journal of Computational Physics 384, pp. 1–15. Fan, Yuwei, Lin Lin, et al. (2019). “A multiscale neural network based on hierarchical matrices”. In: Multiscale Modeling & Simulation 17.4, pp. 1189–1213. Greenfeld, Daniel et al. (2019). “Learning to optimize multigrid PDE solvers”. In: International Conference on Machine Learning. PMLR, pp. 2415–2423. Guo, Xiaoxiao, Wei Li, and Francesco Iorio (2016). “Convolutional neural networks for steady flow approximation”. In: Proceedings of the 22nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining. He, Kaiming et al. (2016). “Deep residual learning for image recognition”. In: Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770–778. Hornik, Kurt, Maxwell Stinchcombe, Halbert White, et al. (1989). “Multilayer feedforward networks are universal approximators.” In: Neural networks 2.5, pp. 359–366. Jiang, Chiyu Max et al. (2020). “MeshfreeFlowNet: A Physics-Constrained Deep Continuous Space-Time Super-Resolution Framework”. In: arXiv preprint arXiv:2005.01463. Kashinath, Karthik, Philip Marcus, et al. (2020). “Enforcing Physical Constraints in CNNs through Differentiable PDE Layer”. In: ICLR 2020 Workshop on Integration of Deep Neural Models and Differential Equations. 71 Khoo, Yuehaw, Jianfeng Lu, and Lexing Ying (2017). “Solving parametric PDE problems with artificial neural networks”. In: arXiv preprint arXiv:1707.03351. Kochkov, Dmitrii et al. (2021). “Machine learning accelerated computational fluid dynamics”. In: arXiv preprint arXiv:2102.01010. Li, Zongyi et al. (2020a). Multipole Graph Neural Operator for Parametric Partial Differential Equations. arXiv: 2006.09535 [cs.LG]. – (2020b). “Neural operator: Graph kernel network for partial differential equations”. In: arXiv preprint arXiv:2003.03485. Lu, Lu, Pengzhan Jin, and George Em Karniadakis (2019). “DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators”. In: arXiv preprint arXiv:1910.03193. Mathieu, Michael, Mikael Henaff, and Yann LeCun (2013). Fast Training of Convolutional Networks through FFTs. arXiv: 1312.5851 [cs.CV]. Nelsen, NH and AM Stuart (2020). “The Random Feature Model for Input-Output Maps between Banach Spaces”. In: arXiv preprint arXiv:2005.10224. Pan, Shaowu and Karthik Duraisamy (2020). “Physics-informed probabilistic learning of linear embeddings of nonlinear dynamics with guaranteed stability”. In: SIAM Journal on Applied Dynamical Systems 19.1, pp. 480–509. Patel, Ravi G et al. (2021). “A physics-informed operator regression framework for extracting data-driven continuum models”. In: Computer Methods in Applied Mechanics and Engineering 373, p. 113500. Raissi, Maziar, Paris Perdikaris, and George E Karniadakis (2019). “Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations”. In: Journal of Computational Physics 378, pp. 686–707. Ronneberger, Olaf, Philipp Fischer, and Thomas Brox (2015). “U-net: Convolutional networks for biomedical image segmentation”. In: International Conference on Medical image computing and computer-assisted intervention. Springer, pp. 234– 241. Sitzmann, Vincent et al. (2020). “Implicit neural representations with periodic activation functions”. In: arXiv preprint arXiv:2006.09661. Smith, Jonathan D, Kamyar Azizzadenesheli, and Zachary E Ross (2020). “Eikonet: Solving the eikonal equation with deep neural networks”. In: arXiv preprint arXiv:2004.00361. Wang, Rui et al. (2020). “Towards physics-informed deep learning for turbulent flow prediction”. In: Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pp. 1457–1466. 72 Zhu, Yinhao and Nicholas Zabaras (2018). “Bayesian deep convolutional encoder–decoder networks for surrogate modeling and uncertainty quantification”. In: Journal of Computational Physics. issn: 0021-9991. doi: https://doi.org/10.1016/ j.jcp.2018.04.018. url: http://www.sciencedirect.com/science/ article/pii/S0021999118302341. 73 Chapter 5 LEARNING: DISSIPATIVE LOSS FOR OPERATOR LEARNING Chaotic systems are notoriously challenging to predict because of their sensitivity to perturbations. Small changes in the initialization or errors during time-stepping accumulate and lead to vastly diverging trajectories. Despite this unpredictable behavior, in dissipative chaotic systems, the long term trajectories converge to a low dimensional set, known as the global attractor. Furthermore this set supports an invariant measure which determines the statistical properties of long-term trajectories. In this work, we propose a machine learning framework for dissipative chaotic systems guided by the goal of accurate statistical prediction. The global attractor, and its invariant measure, are determined by the Markov operator that maps the evolution of the system during infinitesimal time steps. By learning the Markov operator, we can predict the statistical behavior of dissipative chaotic systems, even if we cannot predict exact trajectories. We learn the operator with only the local one-step evolution information and then compose the learned operator to obtain the global attractor and its invariant measure. We further develop a data augmentation technique to impose dissipativity on the learned model, enforcing the desirable property that the composition does not lead to blow ups. Experiments show improved prediction of statistics associated with chaotic dynamical systems, in comparison with previous methods; our experiments are conducted on the Lorenz-63 system, the Kuramoto-Sivashinsky equation and the Navier-Stokes equation with Reynolds number up to 5000. 5.1 Introduction Machine learning methods for chaotic systems. Chaotic systems are characterized by strong instabilities. A tiny perturbation in the initial condition leads to a completely different trajectory. Such instability makes chaotic systems challenging, both for mathematical analysis and numerical simulation. Because of its intrinsic instability, it is infeasible for any method to capture the exact trajectory of a chaotic system for long periods. Therefore, prior works either fit RNNs on extremely short trajectories or only learn a step-wise projection from a randomly generated evolution using reservoir computing (RC) (Vlachas, Byeon, et al., 2018; Pathak et al., 2018; Vlachas, Pathak, et al., 2020; Fan et al., 2020). These previous attempts are able to 74 push the limits of faithful prediction to moderate periods on low dimension systems such as ordinary differential equations (ODEs), e.g. the Lorenz ’63 system or onedimensional partial differential equations (PDEs), e.g. the Kuramoto-Sivashinsky (KS) equation. However, they are less effective at modeling more complicated turbulent systems such as the Navier-Stokes equation (NS). Indeed, predicting long trajectories of such chaotic systems is an ill-posed problem and we cannot expect such attempts to be successful. Instead, we take a new perspective: we aim to capture statistical properties of long trajectories, even if we cannot precisely predict them. Invariants in chaos. Despite the instability, many chaotic systems exhibit certain reproducible statistical properties, such as the auto-correlation and, for PDEs, the energy spectrum. Such properties remain the same in different realizations of the initial condition (Gleick, 2011). This is provably the case for the Lorenz ’63 model Tucker, 1999; Holland and Melbourne, 2007 and empirically holds for numerous dissipative PDEs, such as the KS equation and the two-dimensional NS equation (Kolmogorov flows) (Temam, 2012). Mathematically, this behavior is characterized by the existence of invariant measures. Furthermore there exists an attractor which is a set towards which the system tends to evolve. For dissipative systems there is a compact set which any given trajectory will enter in a finite time, depending on the initial condition, and thereafter remain inside. The attractor is defined by mapping all initial conditions from this compact set, forward in time (technically it is the 𝜔−limit set of the compact set A. Stuart and A.R. Humphries, 1998). While learning infinite time-horizon trajectories is intractable, it is possible to approximate the attractor, and statistics of trajectories on it, using the Markov operator. The dissipativity property helps to make this problem tractable (A. Stuart and A.R. Humphries, 1998; Humphries and Stuart, 1994). For Markovian systems, samples from this measure can be obtained through access to a Markov operator, a memoryless deterministic operator that captures the evolution of the dynamics along an infinitesimal time step. In this case, the Markov operator forms a discrete semigroup, defined by the compositions of Markov operators (A. Stuart and A.R. Humphries, 1998). By learning the Markov operator, we are able to generate the attractor and estimate the invariant measure for a variety of chaotic systems that are of interest to the physics and applied mathematics communities (Fukami, Fukagata, and Taira, 2021; Cardesa, Vela-Martín, and Jiménez, 2017; Chandler and Kerswell, 2013; Koltai and Weiss, 2020; Bramburger, Kutz, and Brunton, 2021; Page, Brenner, and Kerswell, 75 The flow map extrapolated on a larger domain without (left) and with (right) the enforced dissipativity. The red points are training points on the attractor. The dissipativity is imposed by augmenting the data on the blue shell. Figure 5.1: Enforcing dissipativity on the Lorenz 63 system. 2021). Neural operators. To learn the Markov operators for PDEs, we need to model the time-evolution of functions in infinite-dimensional function spaces. This is especially challenging when we need to generate long trajectories. Since even a small error will accumulate over multiple compositions of the learned operator, potentially causing exponential build up or collapse. Because we study the evolution of functions time, we propose to use the recently developed operator learning method known as the neural operator (Li et al., 2020b; Kovachki et al., n.d.). The neural operator remedies the mesh-dependent nature of finite-dimensional operator methods such as RNNs, CNNs, and RC. Neural operators are guaranteed to universally approximate any operator in a mesh independent manner, and hence, can capture the Markov operator of chaotic systems. This local guarantee plus the absorption of trajectories by the global attractor makes it possible to accurately follow the it over long time horizons, allowing us to capture the invariant measure of chaotic systems. Our contributions. In this work, we formulate a machine learning framework for chaotic systems exploiting their ergodicity and Markovian property. We learn a Markov neural operator (MNO) given only one-step evolution data and compose it over a long horizons to obtain the global attractor of the chaotic system, as shown in Figure 5.2. The MNO forms a discrete semigroup defined by the composition of operators, allowing us to reach any state further in the future (Sviridyuk, 1994). The goal is to predict the global attractor which does not collapse or blow up over a long or infinite time horizon, and shares the same distribution and statistics as the true function space trajectories. 76 • We formulate a machine learning framework to learn the Markov neural operator with only local time evolution training data. At test time, we compose it over a long horizon to obtain the global attractor. • We study the ability to learn statistical invariant measures, since predicting exact trajectories is generally intractable. • We impose dissipativity by augmenting the data on an outer shell to enforce that the dynamic evolution stays close to the attractor. The resulting system is stable against large perturbations. • We study the choice of time steps d𝑡 used for training the MNO and investigate various Sobolev losses in operator learning. Experiments show the Sobolev norm significantly outperforms standard loss functions in preserving invariant statistics such as the spectrum. Invariant statistics. We study the invariant statistics proposed in the literature such as the Fourier spectrum, the proper orthogonal decomposition (POD), the point-wise distribution, the auto-correlation, and other domain-specific statistics such as turbulence kinetic energy and dissipation rate. These invariant statistics are usually a combination of different orders of derivatives and moments. By using a standard mean square error (MSE) for training, the models often fail to capture the higher frequency information induced from the derivatives. Therefore, we use the Sobolev norm to address higher-order derivatives and moments (Czarnecki et al., 2017; Beatson et al., 2020). This is similar in spirit to pre-multiplying the spectrum by the wavenumber, an approach commonly used in fluid analysis (Smits, McKeon, and Marusic, 2011). When used for training, the Sobolev norm yields a significant improvement in capturing high frequency details compared to the standard MSE or 𝐿 2 loss, and therefore we obtain an approximation of the Markov operator that is able to better characterize the invariant measure. 5.2 Problem Setting We consider potentially infinitely dimensional dynamical systems where the phase space U is a Banach space and, in particular, a function space on a Lipschitz domain 𝐷 ⊂ R𝑑 (for finite dimensional systems, U will be a Euclidian space). We 77 Figure 5.2: Markov neural operator (MNO): learn global dynamics from local data. (a) Learn the MNO from the local time-evolution data. (b) Compose MNO to evaluate the long-time dynamic. (c) Architecture of the MNO. are interested in the initial-value problem 𝑑𝑢 (𝑡) = 𝐹 (𝑢(𝑡)), 𝑑𝑡 𝑢(0) = 𝑢 0 , 𝑡 ∈ (0, ∞), (5.1) for initial conditions 𝑢 0 ∈ U where 𝐹 is usually a non-linear operator. We will assume, given some appropriate boundary conditions on 𝜕𝐷 when applicable, that the solution 𝑢(𝑡) ∈ U exists and is unique for all times 𝑡 ∈ (0, ∞). When making the spatial dependence explicit, if it is present, we will write 𝑢(𝑥, 𝑡) to indicate the evaluation 𝑢(𝑡)| 𝑥 for any 𝑥 ∈ 𝐷. We adopt the viewpoint of casting time-dependent PDEs into function space ODEs (5.1), as this leads to the semigroup approach to evolutionary PDEs which underlies our learning methodology. Dissipativity. Systems for which there exists some bounded, positively-invariant set 𝐸 such that for any bounded 𝐵 ⊂ U, there is some time 𝑡 ∗ = 𝑡 ∗ (𝐸, 𝐵) beyond which the dynamics of any trajectory starting in 𝐵 enters and remains in 𝐸 are known as dissipative systems (A. Stuart and A.R. Humphries, 1998). The set 𝐸 is known as the absorbing set of the system. For such systems, the global attractor 𝐴, defined subsequently, is characterized as the 𝜔-limit set of 𝐸. In particular, for any initial condition 𝑢 0 ∈ U, the trajectory 𝑢(𝑡) approaches 𝐴 as 𝑡 → ∞. In this work, we consider dissipative dynamical systems where there exist some 𝛼 ≥ 0 and 𝛽 > 0 such that 1𝑑 ||𝑢|| 2 ≤ 𝛼 − 𝛽||𝑢|| 2 2 𝑑𝑡 (5.2) for all 𝑢 ∈ U. It can be shown that systems which satisfy this inequality are 78 dissipative (A. Stuart and A.R. Humphries, 1998) with the absorbing set  √︂  𝛼 𝐸 = 𝐵 0, +𝜀 , 𝛽 (5.3) for any 𝜀 > 0, where 𝐵(𝑢, 𝑟) is the open ball centered at 𝑢 ∈ U with radius 𝑟. There are several well-known examples of dynamical systems that satisfy the above inequality. In this paper we consider the finite-dimensional Lorenz-63 system and the infinite-dimensional cases of the Kuramoto-Sivashinsky and 2D incompressible Navier-Stokes equations (Temam, 2012). Global Attractors. The long time behavior of the solution to (5.1) is characterized by the set 𝑈 = 𝑈 (𝑢 0 ) ⊂ U which is invariant under the dynamic i.e. 𝑆𝑡 (𝑈) = 𝑈 for all 𝑡 ≥ 0, and the orbit 𝑢(𝑡) converges inf ∥𝑢(𝑡) − 𝑣∥ U → 0 as 𝑡 → ∞. 𝑣∈𝑈 When it exists, 𝑈 is often identified as the 𝜔-limit set of 𝑢 0 . The chaotic nature of certain dynamical systems arises due to the complex structure of this set because 𝑢(𝑡) follows 𝑈 and 𝑈 can be, for example, a fractal set. A compact, invariant set 𝐴 is called a global attractor if, for any bounded set 𝐵 ⊂ U and any 𝜖 > 0 there exists a time 𝑡 ∗ = 𝑡 ∗ (𝜖, 𝐴, 𝐵) such that 𝑆𝑡 (𝐵) is contained within an 𝜖-neighborhood of 𝐴 for all 𝑡 ≥ 𝑡 ∗ . Many PDEs arising in physics such as reaction-diffusion equations describing chemical dynamics or the Navier-Stokes equation describing the flow of fluids are dissipative and possess a global attractor which is often times finitedimensional (Temam, 2012). Therefore, numerically characterizing the attractor is an important problem in scientific computing with many potential applications. Data distribution. For many applications, an exact form for the possible initial conditions to (5.1) is not available; it is therefore convenient to use a stochastic model to describe the initial states. To that end, let 𝜇0 be a probability measure on U and assume that all possible initial conditions to (5.1) come as samples from 𝜇0 i.e. 𝑢 0 ∼ 𝜇0 . Then any possible state of the dynamic (5.1) after some time 𝑡 > 0 can be thought of as being distributed according to the pushforward measure 𝜇𝑡 := 𝑆𝑡♯ 𝜇0 i.e. 𝑢(𝑡) ∼ 𝜇𝑡 . Therefore as the dynamic evolves, so does the type of likely functions that result. This further complicates the problem of long time predictions since training data may only be obtained up to finite time horizons hence the model will need the ability to predict not only on data that is out-of-sample but also out-of-distribution. 79 Ergodic systems. To alleviate some of the previously presented challenges, we consider ergodic systems. Roughly speaking, a system is ergodic if there exists an invariant measure 𝜇 such that after some time 𝑡 ∗ > 0, we have 𝜇𝑡 ≈ 𝜇 for any 𝑡 ≥ 𝑡 ∗ (in fact, 𝜇 can be defined without any reference to 𝜇0 or its pushforwards, see (Pavliotis and Andrew Stuart, 2008) for details). That is, after some large enough time, the distribution of possible states that the system can be in is fixed for any time further into the future. Indeed, 𝜇 charges the global attractor 𝐴. Notice that ergodicity is a much more general property than having stationary states which means that the system has a fixed period in time, or having steady states which means the system is unchanged in time. Ergodicity mitigates learning a model that is able to predict out-of-distribution since both the input and the output of 𝑆ˆ ℎ will approximately be distributed according to 𝜇. Furthermore, we may use 𝑆ˆ ℎ to learn about 𝜇 since sampling it simply corresponds to running the dynamic forward. Indeed, we need only generate data on a finite time horizon in order to learn 𝑆ˆ ℎ , and, once learned, we may use it to sample 𝜇 indefinitely by simply repeatedly composing 𝑆ˆ ℎ with itself. Having samples of 𝜇 then allows us to compute statistics which characterize the long term behavior of the system and therefore the global attractor 𝐴. Note further that this strategy avoids the issue of accumulating errors in long term trajectory predictions since we are only interested in the property that 𝑆ˆ ℎ (𝑢(𝑡)) ∼ 𝜇. Notably, the existence of a global attractor does not imply the existence of an invariant measure. Indeed, the only deterministic and chaotic systems that are proven to possess an invariant measure are certain ODEs such as the Lorenz-63 system (Holland and Melbourne, 2007). On the other hand, proving the existence of an invariant measure for deterministic and chaotic PDEs such as the KS or NS equations are still open problems, despite ergodic behavior being observed empirically. 5.3 Learning the Markov neural operator in Chaotic Dynamics We propose the Markov neural operator, a method for learning the underlying Markov operators of chaotic dynamical systems. In particular, we approximate the operator mapping the solution from the current step to the next step 𝑆ˆ ℎ : 𝑢(𝑡) ↦→ 𝑢(𝑡 + ℎ). The architecture of the Markov neural operator is detailed in Figure 5.2. 80 Markov operators and the semigroup Since the system (5.1) is autonomous, that is, 𝐹 does not explicitly depend on time, under the well-posedness assumption, we may define, for any 𝑡 ∈ [0, ∞), a Markov operator 𝑆𝑡 : U → U such that 𝑢(𝑡) = 𝑆𝑡 𝑢(0). This map satisfies the properties 1. 𝑆0 = 𝐼, 2. 𝑆𝑡 (𝑢 0 ) = 𝑢(𝑡), 3. 𝑆𝑡 (𝑆 𝑠 (𝑢 0 )) = 𝑢(𝑡 + 𝑠), for any 𝑠, 𝑡 ∈ [0, ∞) and any 𝑢 0 ∈ U where 𝐼 denotes the identity operator on U. In particular, the family {𝑆𝑡 : 𝑡 ∈ [0, ∞)} defines a semigroup of operators acting on U. Our goal is to approximate a particular element of this semigroup associated to some fixed time step ℎ > 0 given observations of the trajectory from (5.1). We build an approximation 𝑆ˆ ℎ : U → U such that 𝑆ˆ ℎ ≈ 𝑆 ℎ . (5.4) Long-term predictions. Note that having access to the map 𝑆ˆ ℎ from (5.4) allows for approximating long time trajectories of (5.1) by repeatedly composing 𝑆ˆ ℎ with its own output. Indeed, by (5.4) and the semigroup property, we can approximate the evolution by composing the local Markov operators 𝑢(𝑛 · ℎ) ≈ 𝑆ˆ𝑛ℎ (𝑢 0 ) := ( 𝑆ˆ ℎ ◦ · · · ◦ 𝑆ˆ ℎ ) (𝑢 0 ) | {z } (5.5) 𝑛 times for any 𝑛 ∈ N . This fact is formalized in Theorem 5.3.1 which is stated subsequently. Invariant statistics. A useful application of the Markov operators is to estimate statistics of the invariant measure of a chaotic system. Assume the target system is ergodic and there exists an invariant measure 𝜇 such that 𝑢(𝑡) ∼ 𝜇 for any 𝑡 as discussed in Section 5.2. An invariant statistic is defined as ∫ ∫ 1 𝑇 G(𝑢(𝑡)) d𝑡 (5.6) 𝑇𝐺 := G(𝑢) d𝜇(𝑢) = lim 𝑇→∞ 𝑇 0 U for any functional G : U → R𝑑 . Examples include the 𝐿 2 norm, any spectral coefficients, and the spatial correlation, as well as problem-specific statistics such as the turbulence kinetic energy and dissipation rates in fluid flow problems. Given 81 the property (5.5) and using the ergodicity from (5.6), the approximate model 𝑆ˆ ℎ can be used to estimate any invariant statistic simply by computing 𝑛 1 ∑︁ 𝐺 ( 𝑆ˆ ℎ𝑘 𝑢 0 ) 𝑇𝐺 ≈ 𝑛 𝑘=1 for some 𝑛 = 𝑇/ℎ and 𝑇 > 0 large enough. Examples of fast approximation 𝑆ˆ ℎ which accurately predict invariant statics are given in Section 8.5. Enforcing dissipativity. In practice, training data for dynamical systems is typically drawn from trajectories that lie close to the global attractor of the system, so a priori there is no guarantee of a learned model’s behavior far from the attractor. Thus, if we seek to learn the global attractor and invariant statistics of a dynamical system, it is crucial that we place contraints on the model to enforce that it learns a dissipative dynamical system. In particular, given some Markov operator mapping between time-steps 𝑆ˆ ℎ : 𝑢(𝑡) ↦→ 𝑢(𝑡 + ℎ) and cost functional 𝐶𝐷 : U × U → R, we supplement the loss function with the additional term h  i ∫ 2 𝑆ˆ ℎ (𝑢) − 𝜆𝑢 d𝜈(𝑢), (5.7) E𝑢∼𝜈 𝐶𝐷 𝑆ˆ ℎ (𝑢), 𝑢 = U up to some scaling constant with respect to the other terms in the loss function, where 0 < 𝜆 < 1 is some constant factor for scaling down (i.e., enforcing dissipativity) inputs 𝑢 drawn from probability measure 𝜈. We choose 𝜈 to be a uniform probability distribution supported on some shell with a fixed inner and outer radii from the origin in U. Our choice of cost functional 𝐶𝐷 as given in eq. 5.7 scales down 𝑢 by some constant factor 𝜆, but in principle alternative dissipative cost functionals can be used. The norm | · | used in our 𝐶𝐷 is task-specific (e.g., Euclidean norm in the case of the Lorenz-63 system). We find that enforcing this dissipativity constraint on a shell at a sufficiently large radius encourages the learned Markov operator to produce dissipative predictions arbitrarily far away from the shell. In Section 8.5 we demonstrate that the constraint in eq. 5.7 prevents blow-up of a Markov operator trained on the turbulent Kolmogorov flow system. Backbone model for 𝑆ˆ ℎ The above semigroup formulation can be applied with various choices of the backbone model 𝑆ˆ ℎ . In general, we prefer models that can be evaluated quickly and have 82 approximation guarantees so the per-step error can be controlled. Therefore, we choose the standard feed-forward neural network Park and Sandberg, 1991 for ODE systems, and the Fourier neural operator Li et al., 2020a for infinite dimensional PDE systems. For the the neural operator parametric class, we prove the following theorem regarding the Markov neural operator. The result states that our construction can approximate trajectories of infinite-dimensional dynamical systems arbitrary well. The proof is given in the Appendix. Theorem 5.3.1 Let 𝐾 ⊂ U be a compact set and assume that, for some ℎ > 0, the Markov operator 𝑆 ℎ : U → U associated to the dynamic (5.1) is locally Lipschitz. Then, for any 𝑛 ∈ N and 𝜖 > 0 there exists a neural operator 𝑆ˆ ℎ : U → U such that sup sup ∥𝑢(𝑘 ℎ) − 𝑆ˆ ℎ𝑘 (𝑢 0 )∥ U < 𝜖 . 𝑢 0 ∈𝐾 𝑘∈{1,...,𝑛} Theorem 5.3.1 indicates that of choice of backbone model is rich enough to approximate many chaotic dynamical systems for a arbitrarily long period. For finitedimensional systems, the same theorem holds with feed-forward neural networks instead of neural operators. We note that standard neural network such as RNNs and CNNs do not possess such approximation theorems in the infinite-dimensional setting. Furthermore we expect that it is possible to show that the global attractor can be approximated arbitrarily well by adapting ideas from the standard theory of numerical integrators such as (A. R. Humphries and A. M. Stuart, 1994). 5.4 Experiments In this section, we motivate our proposed approach for enforcing dissipativity described in Section 5.3 by applying it to learn a dissipative neural network model for the Lorenz-63 system. We also present the experiment for the KuramotoSivashinsky equation and the Kolmogorov flow raising from the 2-d Navier-Stokes equation. Both PDEs exhibit chaotic behavior. Lorenz-63 system To motivate and justify our framework for learning chaotic systems in the infinitedimensional setting (e.g., Navier-Stokes equations), we first apply our framework (e.g., enforcing model dissipativity and evaluating invariant statistics) on the relatively simple Lorenz-63 ODE system. 83 Lorenz-63 (Lorenz, 1963) is a system of coupled ordinary differential equations, and it is typically considered to be the simplest dynamical system that exhibits chaotic properties. The Lorenz-63 system is given by 𝑢¤ 𝑥 = 𝛼(𝑢 𝑦 − 𝑢 𝑥 ), 𝑢¤ 𝑦 = −𝛼𝑢 𝑥 − 𝑢 𝑦 − 𝑢 𝑥 𝑢 𝑧 , (5.8) 𝑢¤ 𝑧 = 𝑢 𝑥 𝑢 𝑦 − 𝑏𝑢 𝑧 − 𝑏(𝑟 + 𝛼) In this paper, we use the canonical parameters of (𝛼, 𝑏, 𝑟) = (10, 8/3, 28). Learning the Markov operator. Since the Markov operator of the Lorenz-63 system is finite-dimensional, we learn the Markov operator by training a simple feedforward neural network on a single trajectory on the Lorenz attractor. Specifically, we use a neural network with 6 hidden layers and 150 neurons per layer. We discretize the training trajectory into time-steps of 0.05 seconds. We train two networks: one with dissipativity enforced and one without. Enforcing dissipativity. We enforce dissipativity during training with the criterion described in eq. 5.7, with 𝜆 = 0.5 and 𝜈 being a uniform probability distribution supported on a shell around the origin with inner radius 90 and outer radius 130. As shown in figure 5.1, enforcing dissipativity produces predictions that isotropically point towards the attractor, implying that the attractive properties of the Lorenz attractor are learned in the process. Observe that the dissipative network is also dissipative outside the shell in which dissipativity was enforced during training. We believe this property of ReLU networks prevents model blow-up even in more difficult learning problems . Results. We find that enforcing dissipativity does not reduce the relative L2 error compared to the baseline neural network. Further, the invariant statistics of the dissipative model aligns well with the ground-truth data distribution. Our results suggest that dissipativity can be enforced without significantly affecting the model’s step-wise error and the learned statistsical properties of the attractor. In the setting of learning Markovian operator, we consider the following onedimensional Kuramoto-Sivashinsky equation, 𝜕𝑢 𝜕 2 𝑢 𝜕 4 𝑢 𝜕𝑢 = −𝑢 − − , 𝜕𝑡 𝜕𝑥 𝜕𝑥 2 𝜕𝑥 4 𝑢(·, 0) = 𝑢 0 , on [0, 𝐿] × (0, ∞) on [0, 𝐿] (5.9) 84 (a) 𝑥𝑦-plane at 𝑧 = 0 (b) 𝑥𝑧-plane at 𝑦 = 0 (c) 𝑦𝑧-plane at 𝑥 = 0 (d) 𝑥𝑦-plane at 𝑧 = 0 (e) 𝑥𝑧-plane at 𝑦 = 0 (f) 𝑦𝑧-plane at 𝑥 = 0 Figure 5.3: Predicted flow field of dissipative network on Lorenz-63. The red points are training points on the attractor. Dissipativity is enforced uniformly within the blue shell. where 𝐿 = 32𝜋 or 64𝜋 and the spatial domain [0, 𝐿] is equipped with periodic 2 ([0, 𝐿]; R), where boundary conditions. We assume the initial condition 𝑢 0 ∈ 𝐿¤ per 2 ([0, 𝐿]; R) is the space of all mean zero 𝐿 2 -functions that are periodic on [0, 𝐿]. 𝐿¤ per 2 ([0, 𝐿]; R) → 𝐿 2 ([0, 𝐿]; R) is established in ¤ per Existence of the semigroup 𝑆𝑡 : 𝐿¤ per Temam, 2012, Theorem 3.1. Data is obtained by solving the equation using the exponential time-differencing fourth-order Runge-Kutta method from (Kassam and Trefethen, 2005). Random initial conditions are generated according to a mean zero 1 Gaussian measure with covariance 𝐿 −2/𝛼 𝜏 2 (2𝛼−1) (−Δ + (𝜏 2 /𝐿 2 )𝐼) −𝛼 where 𝛼 = 2, 𝜏 = 7, and periodic boundary conditions on [0, 𝐿]; for details see (Lord, Powell, and Shardlow, 2014). Invariant statistics for the KS equations As shown in Figure 5.5, we present enormous invariant statistics for the KS equation. We use 150000 snapshots to train the MNO, LSTM, and GRU to model the evolution operator of the KS equation with ℎ = 1𝑠. We compose each model for 𝑇 = 1000 time steps to obtain a long trajectory (attractor), and estimate various invariant statistics from them. 85 The x-axis is the spatial domain; the y-axis is the temporal domain. The figure shows that LSTM and GRU start to diverge at 𝑡 = 20𝑠 while MNO is able to keep up with the exact trajectory until 𝑡 = 50𝑠. Figure 5.4: Trajectory and error on the KS equation Choice of time discretization. We further study the choice of time steps ℎ. As shown in Figure 5.6, when the time steps are too large, the correlation is chaotic and hard to capture. But counter-intuitively, when the time steps are too small, the evolution is also hard to capture. In this case, the input and output of the network will be very close, and the identity map will be a local minimum. An easy fix is to use MNO to learn the time-derivative or residual. This is shown in the figure, where the residual model (blue line) has a better per-step error and accumulated error at smaller ℎ. When the time step is large, there is no difference in modeling the residual. This idea can generalize to other integrators as an extension of Neural ODEs to PDEs (Chen et al., 2018). Navier-Stokes equation We also consider the two-dimensional Navier-Stokes equation for a viscous, incompressible fluid, 1 𝜕𝑢 = −𝑢 · ∇𝑢 − ∇𝑝 + Δ𝑢 + sin(𝑛𝑦) 𝑥, ˆ 𝜕𝑡 𝑅𝑒 ∇·𝑢 =0 𝑢(·, 0) = 𝑢 0 on [0, 2𝜋] 2 × (0, ∞) on [0, 2𝜋] 2 × [0, ∞) (5.10) on [0, 2𝜋] 2 where 𝑢 denotes the velocity, 𝑝 the pressure, and 𝑅𝑒 > 0 is the Reynolds number. The domain [0, 2𝜋] 2 is equipped with periodic boundary conditions. The specific choice of forcing sin(𝑛𝑦) 𝑥ˆ constitutes a Kolmogorov flow; we choose 𝑛 = 4 in all experiments. We define U to be the closed subspace of 𝐿 2 ([0, 2𝜋] 2 ; R2 ), 86 Figure 5.5: Invariant statistics for the KS equation 87 Figure 5.6: Choice of time step  2 ([0, 2𝜋] 2 ; R2 ) : ∇ · 𝑢 = 0 and assume 𝑢 ∈ U. We define the U = 𝑢 ∈ 𝐿¤ per 0 vorticity 𝑤 = (∇ × 𝑢) 𝑧ˆ and the stream function 𝑓 as the solution to the Poisson equation −Δ 𝑓 = 𝑤. Existence of the semigroup 𝑆𝑡 : U → U is established in Temam, 2012, Theorem 2.1. We denote turbulence kinetic energy (TKE) ⟨(𝑢 − 𝑢) ¯ 2 ⟩, and dissipation 𝜖 = ⟨𝑤 2 ⟩/𝑅𝑒. Data is obtained by solving the equation in vorticity form using the pseudo-spectral split step method from (Chandler and Kerswell, 2013). Random initial conditions are generated according to a mean zero Gaussian measure with covariance 73/2 (−Δ + 49𝐼) −2.5 with periodic boundary conditions on [0, 2𝜋] 2 . Benchmarks for 2d Navier-Stokes. We compare MNO with common standard two-dimensional dynamic models including U-Net(Ronneberger, Fischer, and Brox, 2015) and LSTM-CNN(Shi et al., 2015) on modeling the vorticity 𝑤. We choose the discretization ℎ = 1𝑠. The training dataset consists of 180 realizations of trajectories on time interval 𝑡 ∈ [100, 500] (the first 100 seconds are discarded) which adds up to 180 × 400 = 72, 000 snapshots in total. Another 20 realizations are generated for testing. Each single snapshot has resolution 64 × 64. We use the Adam optimizer to minimize the relative 𝐿 2 loss with learning rate = 0.0005, and step learning rate scheduler that decays by half every 10 epochs for 50 epochs in total. U-Net: we use five layers of convolution and deconvolution with width from 64 to 1024. LSTM-CNN: we use one layer of LSTM with width = 64. MNO: we parameterize the 2-d Fourier neural operator consists of four Fourier layers with 20 frequencies 88 The 10000 time steps trajectory generated by MNO projected onto the first two components of PCA. Each point corresponds to an snapshot on the attractor. Two points are selected for further visualization of vorticity field. Figure 5.7: The learned attractor of the Kolmogorov flow. Model MNO U-Net ConvLSTM training 2 𝐿 loss 𝐻 1 loss 𝐻 2 loss 𝐿 2 loss 𝐻 1 loss 𝐻 2 loss 𝐿 2 loss 𝐻 1 loss 𝐻 2 loss loss 0.0166 0.0184 0.0202 0.0269 0.0377 0.0565 0.2436 0.2855 0.3367 L2 error H1 error H2 error 0.0187 0.0474 0.1729 0.0151 0.0264 0.0656 0.0143 0.0206 0.0361 0.0549 0.1023 0.3055 0.0570 0.0901 0.2164 0.0676 0.0936 0.1749 0.2537 0.3854 1.0130 0.2577 0.3308 0.5336 0.2727 0.3311 0.4352 TKE error 𝜖 error 0.0136 0.0303 0.0256 0.0017 0.0226 0.0193 0.0958 0.0934 0.1688 0.1127 0.0482 0.0841 0.0140 24.1119 0.6977 6.9167 0.8594 4.0976 Table 5.1: Benchmark on vorticity for the Kolmogorov flow with Re= 40 per channel and width = 64. Enforcing dissipativity. We enforce dissipativity during training with the criterion described in eq. 5.7, with 𝜆 = 0.5 and 𝜈 being a uniform probability distribution supported on a shell around the origin. Accuracy with respect to various norms. MNO shows near one order of magnitude better accuracy compared to U-Net and LSTM-CNN. As shown in Table 5.1, we train each model using the balanced 𝐿 2 (= 𝐻 0 ), 𝐻 1 , and 𝐻 2 losses, defined as the sum of the relative 𝐿 2 loss grouped by each order of derivative. And we measure the error with respect to the standard (unbalanced) norms. The MNO with 𝐻 2 loss consistently achieves the smallest error on vorticity on all of the 𝐿 2 , 𝐻 1 , and 89 Figure 5.8: Visualization of the NS equation, Re40 𝐻 2 norms. However, 𝐿 2 loss achieves the smallest error on the turbulence kinetic energy (TKE); 𝐻 1 loss achieves the smallest error on the dissipation 𝜖. The NS equation with Re40 is shown in Figure 5.8. (a) show the ground truth data of vorticity field, each column represents a snapshot at 𝑡 = 100𝑠 with a different initial condition. (b), (c), (d) show the predicted trajectory of MNO on vorticity, using 𝐿 2 , 𝐻 1 , and 𝐻 2 losses respectively. We are able to generate a long trajectory with the MNO model. The five columns represent 𝑡 = 1000𝑠, 2000𝑠, 3000𝑠, 4000𝑠, 5000𝑠 respectively. As shown in the figure, the predicted trajectories (b) (c) (d) share the same behaviors as in the ground truth (a). It indicates the MNO model is stable. Estimate the attractor and invariant statistics. We compose MNO 10000 times to obtain the global attractor, and we compute the PCA (POD) basis of these 10000 snapshots and project them onto the first two components. As shown in Figure (a), we obtain a cycle-shaped attractor. The true attractor has a degree of freedom 90 Figure 5.9: Invariant statistics for the NS equation 91 around 𝑂 (100) (Temam, 2012). If the attractor is a high-dimensional sphere, then most of the mass concentrates around its equator. Therefore, when projected to low-dimension, the attractor will have the shape of a ring. Most of the points are located on the ring, while a few other points are located in the center. The points in the center have high dissipation, implying they are intermittent states. In Figure (b) we add the time axis. While the trajectory jumps around the cycle, we observe there is a rough period of 2000s. We perform the same PCA analysis on the training data, which shows the same behavior. Similarly, we present enormous invariant statistics for the NS equation (Re40), as shown in Figure 5.9,. We use 72000 snapshots to train the MNO, UNet, and ConvLSTM to model the evolution operator of the KS equation with ℎ = 1𝑠. We compose each model for 𝑇 = 10000 time steps to obtain a long trajectory (attractor), and estimate various invariant statistics from them. Order of derivatives. Roughly speaking, vorticity is the derivative of velocity; velocity is the derivative of the stream function. Therefore we can denote the order of derivative of vorticity, velocity, and stream function as 2, 1, and 0 respectively. Combining vorticity, velocity, and stream function, with 𝐿 2 , 𝐻 1 , and 𝐻 2 loss, we have in total the order of derivatives ranging from 0 to 4. We observe, in general, it is best practice to keep the order of derivatives in the model at a number slightly higher than that of the target quantity. For example, as shown in Figure 5.10, when querying the velocity (first-order quantity), it is best to use second-order (modeling velocity plus 𝐻 1 loss or modeling vorticity plus 𝐿 2 loss). This is further illustrated in Table 5.2. In general, using a higher order of derivatives as the loss will increase the power of the model and capture the invariant statistics more accurately. However, a higher-order of derivative means higher irregularity. It in turn requires a higher resolution for the model to resolve and for computing the discrete Fourier transform. This trade-off again suggests it is best to pick a Sobolev norm not too low or too high. 5.5 Discussion and Conclusion In this work, we propose a machine learning framework that trains from local data and predict the global attractor and invariant statistic of chaotic systems. The Markov operator forms a discrete semigroup defined by the composition of operators that empirically does not collapse or blow up over a long or infinite time horizon. Experiments also show MNO predicts the attractor that shares the same distribution and statistics as the true function space trajectories. The Markov operator and 92 Figure 5.10: Error on velocity with respect to the order of derivative. Model 𝑓 𝑓 𝑓 𝑢 𝑢 𝑢 𝑤 𝑤 𝑤 training 𝐿 2 loss 𝐻 1 loss 𝐻 2 loss 𝐿 2 loss 𝐻 1 loss 𝐻 2 loss 𝐿 2 loss 𝐻 1 loss 𝐻 2 loss loss (order) 𝑡ℎ 0.0379 (0 order) 0.0512 (1𝑠𝑡 order) 0.0973 (2𝑛𝑑 order) 0.0688 (1𝑠𝑡 order) 0.1246 (2𝑛𝑑 order) 0.2662 (3𝑟 𝑑 order) 0.1710 (2𝑛𝑑 order) 0.3383 (3𝑟 𝑑 order) 0.4590 (4𝑡ℎ order) error on 𝑓 0.0383 0.0268 0.0198 0.0217 0.0170 0.0178 0.0219 0.0268 0.0312 error on 𝑢 0.2057 0.0769 0.0522 0.0691 0.0467 0.0482 0.0415 0.0463 0.0536 error on 𝑤 2.0154 0.3656 0.2227 0.3217 0.1972 0.1852 0.1736 0.1694 0.1854 Table 5.2: Vorticity 𝑤, velocity 𝑢, and stream function 𝑓 for the Kolmogorov flow with 𝑅𝑒 = 500 chaotic systems are mathematically interesting and physically relevant. MNO shows potential for application studying the bifurcation of complicated system, for example, in the modeling of chemical reactions and the flow of turbulent fluids. This work provides a method for fast computation in many scientific computing problems. These methods have two main long-term impacts beyond the immediate interests of scientific computing communities. Since our methods are orders of magnitude faster than traditional solvers that are prominently used in supercomputers, edge devices, and servers, the deployment of our methods significantly reduces the carbon footprints caused by scientific studies. Furthermore, the proposed methods are extremely flexible. The off the shelf usage of our methods allows scientists from a variety of disciplines, from chemistry, biology, ecology, epidemiology, to physics and applied mathematics to deploy our methods to their complex PDE of interest (or any other problem that involves learning maps between function spaces). 93 MNO has two major limitations. First, it assumes the target system is approximately Markovian. If the system is heavily path-dependent, then the MNO framework does not directly apply. Second, although we develop an approximation theorem for finite period, it does not hold for infinite time horizon. As discussed previously, it is infeasible to track the exact trajectory of chaotic systems on an infinite time horizon. Even very small errors will accumulate in each step, and eventually cause the simulation to diverge from the true trajectory. However, it is possible to track the attractor of the system. An attractor is absorbing. If the simulated trajectory only makes a small error, the attractor will absorb it back, so that the simulated trajectory will never diverge from the true attractor. Therefore, it is possible to have the simulated trajectory capture the true attractor. To obtain an infinite-time approximation error bound is non-trivial. Previously, (A. Stuart and A.R. Humphries, 1998; Humphries and Stuart, 1994) (cf. Theorem 3.12) show a result for (finite-dimensional) ODE systems. If the system is Lipschitz then there exists a numerical simulation that forms a dissipative dynamical system that does not blow up or collapse. And the the simulated attractor A ℎ approximates the true attractor A with the time step ℎ 𝑑𝑖𝑠𝑡 (A ℎ , A) → 0, as ℎ → 0 To generalize such theorem to Markov neural operator (MNO), we need to overcome two difficulties (1) generalize the formulation from (finite-dimensional) ODE systems to (infinite-dimensional) PDE systems, and (2) show MNO can obtain a sufficient error rate with respect to the time step ℎ. The first aspect requires extending the theory from finite dimension to infinite dimension, which is non-trivial since the operator 𝐹 in (5.1) is not compact or bounded. This makes it hard to bound the error with respect to the attractor (Andrew Stuart, 1995). The second aspect requires to formulate MNO slightly differently. In the current formulation, the evolution operator is chosen for a fixed time step ℎ. To achieve 𝑂 (ℎ) error we need to formulate the evolution operator continuously for infinitesimal ℎ. Especially, for a semi-linear PDE system 𝑑𝑢 + 𝐴𝑢 = 𝐹 (𝑢) 𝑑𝑡 where 𝐴 is a linear, self-adjoint operator and 𝐹 is a continuous but nonlinear operator (This formulation includes the KS and NS equations). The evolution can be written 94 as 𝑢(𝑡 + ℎ) = 𝑒 −𝐴ℎ ∫ ℎ 𝑢(𝑡) + 𝑒 −𝐴(ℎ−𝑠) 𝐹 (𝑢(𝑡 + 𝑠))𝑑𝑠 0 ∫ℎ Where Φ(𝑢(𝑡), 𝐴, 𝑡) := 0 𝑒 −𝐴(ℎ−𝑠) 𝐹 (𝑢(𝑠))𝑑𝑠 is bounded despite 𝐹 is not. If one can approximate Φ(𝑢(𝑡), 𝐴, ℎ) by a neural operator, then MNO can potentially achieve the needed error rate. This shows hope to obtain an approximation error bound for infinite time zero. We leave this as a promising future direction. References Beatson, Alex et al. (2020). “Learning composable energy surrogates for PDE order reduction”. In: Advances in Neural Information Processing Systems 33. Bramburger, Jason J, J Nathan Kutz, and Steven L Brunton (2021). “Data-Driven Stabilization of Periodic Orbits”. In: IEEE Access 9, pp. 43504–43521. Cardesa, José I, Alberto Vela-Martín, and Javier Jiménez (2017). “The turbulent cascade in five dimensions”. In: Science 357.6353, pp. 782–784. Chandler, Gary J and Rich R Kerswell (2013). “Invariant recurrent solutions embedded in a turbulent two-dimensional Kolmogorov flow”. In: Journal of Fluid Mechanics 722, pp. 554–595. Chen, Tian Qi et al. (2018). “Neural ordinary differential equations”. In: Advances in neural information processing systems, pp. 6571–6583. Czarnecki, Wojciech Marian et al. (2017). “Sobolev training for neural networks”. In: arXiv preprint arXiv:1706.04859. Fan, Huawei et al. (2020). “Long-term prediction of chaotic systems with machine learning”. In: Physical Review Research 2.1, p. 012080. Fukami, Kai, Koji Fukagata, and Kunihiko Taira (2021). “Machine-learning-based spatio-temporal super resolution reconstruction of turbulent flows”. In: Journal of Fluid Mechanics 909. Gleick, James (2011). Chaos: Making a new science. Open Road Media. Holland, Mark and Ian Melbourne (Oct. 2007). “Central limit theorems and invariance principles for Lorenz attractors”. In: Journal of the London Mathematical Society 76.2, pp. 345–364. Humphries, A. R. and A. M. Stuart (1994). “Runge–Kutta Methods for Dissipative and Gradient Dynamical Systems”. In: SIAM Journal on Numerical Analysis 31.5, pp. 1452–1485. – (1994). “Runge–Kutta methods for dissipative and gradient dynamical systems”. In: SIAM journal on numerical analysis 31.5, pp. 1452–1485. Kassam, Aly-Khan and Lloyd N. Trefethen (2005). “Fourth-Order Time-Stepping for Stiff PDEs”. In: SIAM Journal on Scientific Computing 26.4, pp. 1214–1233. 95 Koltai, Péter and Stephan Weiss (2020). “Diffusion maps embedding and transition matrix analysis of the large-scale flow structure in turbulent Rayleigh–Bénard convection”. In: Nonlinearity 33.4, p. 1723. Kovachki, Nikola B. et al. (n.d.). “Neural Operator: Learning Maps Between Function Spaces”. In: In Preparation (). Li, Zongyi et al. (2020a). “Fourier neural operator for parametric partial differential equations”. In: arXiv preprint arXiv:2010.08895. – (2020b). “Neural operator: Graph kernel network for partial differential equations”. In: arXiv preprint arXiv:2003.03485. Lord, Gabriel J, Catherine E Powell, and Tony Shardlow (2014). An introduction to computational stochastic PDEs. Vol. 50. Cambridge University Press. Lorenz, Edward N (1963). “Deterministic nonperiodic flow”. In: Journal of atmospheric sciences 20.2, pp. 130–141. Page, Jacob, Michael P Brenner, and Rich R Kerswell (2021). “Revealing the state space of turbulence using machine learning”. In: Physical Review Fluids 6.3, p. 034402. Park, Jooyoung and Irwin W Sandberg (1991). “Universal approximation using radial-basis-function networks”. In: Neural computation 3.2, pp. 246–257. Pathak, Jaideep et al. (2018). “Model-free prediction of large spatiotemporally chaotic systems from data: A reservoir computing approach”. In: Physical review letters 120.2, p. 024102. Pavliotis, Grigorios and Andrew Stuart (Jan. 2008). Multiscale Methods: Averaging and Homogenization. Vol. 53. isbn: 978-0-387-73828-4. doi: 10.1007/978-0387-73829-1. Ronneberger, Olaf, Philipp Fischer, and Thomas Brox (2015). “U-net: Convolutional networks for biomedical image segmentation”. In: International Conference on Medical image computing and computer-assisted intervention. Springer, pp. 234– 241. Shi, Xingjian et al. (2015). “Convolutional LSTM network: A machine learning approach for precipitation nowcasting”. In: arXiv preprint arXiv:1506.04214. Smits, Alexander J, Beverley J McKeon, and Ivan Marusic (2011). “High–Reynolds number wall turbulence”. In: Annual Review of Fluid Mechanics 43, pp. 353–375. Stuart, A. and A.R. Humphries (1998). Dynamical Systems and Numerical Analysis. Cambridge Monographs on Applie. Cambridge University Press. isbn: 9780521645638. Stuart, Andrew (1995). “Perturbation theory for infinite dimensional dynamical systems”. In. 96 Sviridyuk, Georgy Anatolevich (1994). “On the general theory of operator semigroups”. In: Russian Mathematical Surveys 49.4, pp. 45–74. Temam, Roger (2012). Infinite-dimensional dynamical systems in mechanics and physics. Vol. 68. Springer Science & Business Media. Tucker, Warwick (1999). “The Lorenz attractor exists”. In: Comptes Rendus de l’Academie des Sciences - Series I - Mathematics 328.12, pp. 1197–1202. issn: 0764-4442. Vlachas, Pantelis R, Wonmin Byeon, et al. (2018). “Data-driven forecasting of high-dimensional chaotic systems with long short-term memory networks”. In: Proceedings of the Royal Society A: Mathematical, Physical and Engineering Sciences 474.2213, p. 20170844. Vlachas, Pantelis R, Jaideep Pathak, et al. (2020). “Backpropagation algorithms and reservoir computing in recurrent neural networks for the forecasting of complex spatiotemporal dynamics”. In: Neural Networks. 97 Chapter 6 LEARNING: PHYSICS INFORMED LEARNING FOR NEURAL OPERATOR In this paper, we propose physics-informed neural operators (PINO) that combine training data and physics constraints to learn the solution operator of a given family of parametric Partial Differential Equations (PDE). PINO is the first hybrid approach incorporating data and PDE constraints at different resolutions to learn the operator. Specifically, in PINO, we combine coarse-resolution training data with PDE constraints imposed at a higher resolution. The resulting PINO model can accurately approximate the ground-truth solution operator for many popular PDE families and shows no degradation in accuracy even under zero-shot super-resolution, i.e., being able to predict beyond the resolution of training data. PINO uses the Fourier neural operator (FNO) framework that is guaranteed to be a universal approximator for any continuous operator and discretization convergent in the limit of mesh refinement. By adding PDE constraints to FNO at a higher resolution, we obtain a high-fidelity reconstruction of the ground-truth operator. Moreover, PINO succeeds in settings where no training data is available and only PDE constraints are imposed, while previous approaches, such as the Physics-Informed Neural Network (PINN), fail due to optimization challenges, e.g., in multi-scale dynamic systems such as Kolmogorov flows. 6.1 Introduction Machine learning methods have recently shown promise in solving partial differential equations (PDEs) (Kovachki et al., n.d.; Z. Li, Kovachki, et al., 2021a; Z. Li, Kovachki, et al., 2021b; L. Lu, P. Jin, and George Em Karniadakis, 2021; Brunton, Noack, and Koumoutsakos, 2020). A recent breakthrough is the paradigm of operator learning for solving PDEs. Unlike standard neural networks that learn using inputs and outputs of fixed dimensions, neural operators learn operators, which are mappings between function spaces (Kovachki et al., n.d.; Z. Li, Kovachki, et al., 2021a; Z. Li, Kovachki, et al., 2021b). The class of neural operators is guaranteed to be a universal approximator for any continuous operator (Kovachki et al., n.d.) and hence, has the capacity to approximate any operator including any solution operator of a given family of parametric PDEs. Note that the solution operator is 98 PINO uses both training data and PDE loss function and perfectly extrapolates to unseen frequencies in Kolmogorov Flows. FNO uses only training data and does not have further information on higher frequencies, but still follows the general trend of the ground-truth spectrum. On the other hand, using a trained UNet with trilinear interpolation (NN+Interpolation) has severe distortions at higher frequencies. Figure 6.1: Physics-Informed Neural Operator (PINO) extrapolates to unseen higher frequencies. the mapping from the input function (initial and boundary conditions) to the output solution function. Previous works show that neural operators can capture complex multi-scale dynamic processes and are significantly faster than numerical solvers (B. Liu et al., 2022; H. Yang et al., 2022; Pathak, Subramanian, et al., 2023; Z. Li, D. Z. Huang, et al., n.d.; Wen et al., 2023b; Bonev et al., 2023; Wen et al., 2023a). Neural operators are proven to be discretization convergent in the limit of mesh refinement (Kovachki et al., n.d.), meaning they converge to a continuum operator in the limit as the discretization is refined. Consequently, they can be evaluated at any data discretization or resolution at inference time without the need for retraining. For example, neural operators such as Fourier neural operator (FNO) can extrapolate to frequencies that are not seen during training in Kolmogorov Flows, as shown in Figure 6.1, while standard approaches such as training a UNet and adding trilinear interpolation leads to significantly worse results at higher resolutions. Even though FNO follows the general trend of the Kolmogorov flow in Figure 6.1, it cannot perfectly match it in the regime of super-resolution, i.e., beyond the frequencies seen during training. More generally, neural operators cannot perfectly approximate the ground-truth operator when only coarse-resolution training data is 99 provided. This is a fundamental limitation of data-driven operator learning methods which depend on the availability of training data, which can come either from existing numerical solvers or direct observations of the physical phenomena. In many scenarios, such data can be expensive to generate, unavailable, or available only as low-resolution observations (Hersbach et al., 2020). This limits the ability of neural operators to learn high-fidelity models. Moreover, the generalization of the learned neural operators to unseen scenarios and conditions that are different from training data is challenging. Our Approach and Contributions In this paper, we remedy the above shortcomings of data-driven operator learning methods through the framework of physics-informed neural operators (PINO). Here, we take a hybrid approach of combining training data (when available) with a PDE loss function at a higher resolution. This allows us to approximate the solution operator of many PDE families nearly perfectly. While there have been many physics-informed approaches proposed previously (discussed in 6.1), ours is the first to incorporate PDE constraints at a higher resolution as a remedy for low resolution training data. We show that this results in high-fidelity solution operator approximations. As shown in Figure 6.1, PINO extrapolates to unseen frequencies in Kolmogorov Flows. Thus, we show that the PINO model learned from such multiresolution hybrid loss functions has almost no degradation in accuracy even on highresolution test instances when only low-resolution training data is available. Further, our PINO approach also overcomes the optimization challenges in approaches such as Physics-Informed Neural Network (PINN) (Raissi, Perdikaris, and George E Karniadakis, 2019) that are purely based on PDE loss functions and do not use training data, and thus, PINO can solve more challenging problems such as timedependent PDEs. PINO utilizes both the data and equation loss functions (whichever are available) for operator learning. To further improve accuracy at test time, we fine-tune the learned operator on the given PDE instance using only the equation loss function. This allows us to provide a nearly-zero error for the given PDE instance at all resolutions. A schematic of PINO is shown in Figure 8.1, where the neural operator architecture is based on Fourier neural operator (FNO) (Z. Li, Kovachki, et al., 2021a). The derivatives needed for the equation loss in PINO are computed explicitly through the operator layers in function spaces. In particular, we efficiently compute the explicit gradients on function space through Fourier-space computations. In 100 PINO trains neural operator with both training data and PDE loss function. The figure shows the neural operator architecture with the lifting point-wise operator that receives input function 𝑎 and outputs function 𝑣 0 with a larger co-dimension. This operation is followed by 𝐿 blocks that compute linear integral operators followed by non-linearity, and the last layer of which outputs the function 𝑣 𝐿 . The pointwise projection operator projects 𝑣 𝐿 to output function 𝑢. Both 𝑣 𝐿 and 𝑢 are functions and all their derivatives (D𝑣 L , D𝑢) can be computed in their exact forms at any query points 𝑥. Figure 6.2: Architecture of Physics-Informed Neural Operator contrast, previous auto-differentiation methods must compute the derivatives at sampling locations. The PDE loss function added to PINO vastly improves generalization and physical validity in operator learning compared to purely data-driven methods. PINO requires fewer to no training data and generalizes better compared to the data-driven FNO (Z. Li, Kovachki, et al., n.d.), especially on high-resolution test instances. On average, the relative error is 7% lower on both transient and Kolmogorov flows, while matching the speedup of data-trained FNO architecture (400x) compared to the GPU-based pseudo-spectral solver (He and W. Sun, 2007). Further, the PINO model on the Navier-Stokes equation can be easily transferred to different Reynolds numbers ranging from 100 to 500 using instance-wise fine-tuning. We also use PINO for solving inverse problems by either: (1) learning the forward solution operator and using gradient-based optimization to obtain the inverse solution, or (2) learning the inverse solution operator directly. Imposing the PDE loss guarantees the inverse solution is physically valid in both approaches. We find that of these two approaches, the latter is more accurate for recovering the coefficient function in Darcy flow. We show this approach is 3000x faster than the conventional solvers using accelerated Markov Chain Monte-Carlo (MCMC) (Cotter et al., 2013). 101 Related Work Learning solution to PDEs has been proposed under two paradigms: (i) data-driven learning and (ii) physics-informed optimization. The former utilizes data from existing solvers or experiments, while the latter is purely based on PDE constraints. An example of data-driven methods is the class of neural operators for learning the solution operator of a given family of parametric PDEs. An example of the physicsbased approach is the Physics-Informed Neural Network (PINN) for optimizing the PDE constraints to obtain the solution function of a given PDE instance. Both these approaches have shortcomings. Neural operators require data, and when that is limited or not available, they are unable to learn the solution operator faithfully. PINN, on the other hand, does not require data but is prone to failure, especially on multi-scale dynamic systems due to optimization challenges. Previous work by (Yingzhou Li, J. Lu, and A. Mao, 2020) has shown promise in learning discretized solution map with variational loss. In this work we generalize it to operator learning. Neural operators learn the solution operator of a family of PDEs, defined by the map from the input–initial conditions and boundary conditions, to the output–solution functions. In this case, usually, a dataset of input-output pairs from an existing solver or real-world observation is given. There are two main aspects to consider (a) models: the design of models for learning highly complicated PDE solution operators, and (b) data: minimizing data requirements and improving generalization. Recent advances in operator learning replace traditional convolutional neural networks and U-Nets from computer vision with operator-based models tailored to PDEs with greatly improved model expressiveness (Z. Li, Kovachki, et al., n.d.; L. Lu, P. Jin, and George Em Karniadakis, 2021; Patel et al., 2021; Rui Wang et al., 2020; Duvall, Duraisamy, and Pan, 2021). Specifically, the neural operator generalizes the neural network to the operator setting where the input and output spaces are infinite-dimensional. The framework has successfully approximated solution operators for highly non-linear problems such as turbulence flow (Z. Li, Kovachki, et al., 2021b; Z. Li, Kovachki, et al., 2021a). However, the data challenges remain. In particular, (1) training data from an existing solver or an experimental setup is costly to obtain, (2) models struggle in generalizing away from the training distribution, and (3) constructing the most efficient approximation for given data remains elusive. Moreover, it is also evident that in many real-world applications, observational data often is available at only low resolutions (Hersbach et al., 2020), limiting the model’s ability to generalize. 102 Alternatives to data-driven approaches for solving PDEs are physics-based approaches that require no training data. A popular framework known as PhysicsInformed Neural Network (PINN) (Raissi, Perdikaris, and George E Karniadakis, 2019) uses optimization to find the solution function of a given PDE instance. PINN uses a neural network as the ansatz of the solution function and optimizes a loss function to minimize the violation of the given equation by taking advantage of auto-differentiation to compute the exact derivatives. PINN overcomes the need to choose a discretization grid that most numerical solvers require, e.g., finite difference methods (FDM) and finite element methods (FEM). It has shown promise in solving PDEs for a wide range of applications, including higher dimensional problems. (L. Lu, Meng, Z. Mao, et al., 2021; J. Han, Jentzen, and Weinan, 2018; Hennigh et al., 2021; Kashinath et al., 2021). Recently, researchers have developed many variations of PINN with promising results for solving inverse problems and partially observed tasks (D. Lu et al., 2021; Zhu et al., 2019; Smith et al., 2021). However, PINN fails in many multi-scale dynamic PDE systems (S. Wang, X. Yu, and Perdikaris, 2022; Fuks and Tchelepi, 2020; Raissi, Yazdani, and George Em Karniadakis, 2020) due to two main reasons, viz., (1) the challenging optimization landscape of the PDE constraints (S. Wang, Teng, and Perdikaris, 2021) and its sensitivity to hyper-parameter selection (L. Sun et al., 2020), and (2) the difficulty in propagating information from the initial or boundary conditions to unseen parts of the interior or to future times (Dwivedi and Srinivasan, 2020). Moreover, PINN only learns the solution function of a single PDE instance and cannot generalize to other instances without re-optimization. Previous work on physics-informed DeepONet that imposes PDE losses on DeepONet (S. Wang, H. Wang, and Perdikaris, 2021) overcomes this limitation and can learn across multiple PDE instances. While the PDE loss is computed at any query points, the input is limited to a fixed grid in standard DeepONet (Kovachki et al., n.d.), and its architecture is a linear method of approximation (Lanthaler et al., 2023). Our work overcomes all previously mentioned limitations. Further, a unique feature that PINO enjoys over other hybrid learning methods (Zhu et al., 2019; Zhang et al., 2022; X. Huang et al., 2022) is its ability to incorporate data and PDE loss functions at different resolutions. This has not been attempted before, and none of the previous works focus on extrapolation to higher resolutions, beyond what is seen in training data. 103 6.2 Problem Settings In this section, we first define the stationary and dynamic PDE systems that we consider. We give an overview of the physics-informed setting and operator-learning setting. In the end, we define the Fourier neural operator as a specific model for operator learning. Problem settings We consider two natural classes of PDEs. In the first, we consider the stationary system P (𝑢, 𝑎) = 0, in 𝐷 ⊂ R𝑑 (6.1) 𝑢 = 𝑔, in 𝜕𝐷 where 𝐷 is a bounded domain, 𝑎 ∈ A ⊆ V is a PDE coefficient/parameter, 𝑢 ∈ U is the unknown, and P : U × A → F is a possibly non-linear partial differential operator with (U, V, F ) a triplet of Banach spaces. Usually, the function 𝑔 is a fixed boundary condition but can also potentially enter as a parameter. This formulation gives rise to the solution operator G † : A → U defined by 𝑎 ↦→ 𝑢. A prototypical example is the second-order elliptic equation P (𝑢, 𝑎) = −∇ · (𝑎∇𝑢) + 𝑓 . In the second setting, we consider the dynamical system 𝑑𝑢 = R (𝑢), 𝑑𝑡 𝑢 = 𝑔, in 𝐷 × (0, ∞) 𝑢=𝑎 in 𝐷¯ × {0} in 𝜕𝐷 × (0, ∞) (6.2) where 𝑎 = 𝑢(0) ∈ A ⊆ V is the initial condition, 𝑢(𝑡) ∈ U for 𝑡 > 0 is the unknown, and R is a possibly non-linear partial differential operator with U, and V Banach spaces. As before, we take 𝑔 to be a known boundary condition. We assume that 𝑢 exists and is bounded for all time and for every 𝑢 0 ∈ U. This formulation gives rise
to the solution operator G † : A → 𝐶 (0, 𝑇]; U defined by 𝑎 ↦→ 𝑢. Prototypical examples include the Burgers’ equation and the Navier-Stokes equation. Solving equation using the physics-informed neural networks Given an instance 𝑎 and a solution operator G † defined by equations (6.1) or (9.1) , we denote by 𝑢 † = G † (𝑎) the unique ground truth. The equation-solving task is to approximate 𝑢 † . This setting consists of the ML-enhanced conventional solvers such as learned finite element, finite difference, and multigrid solvers (Kochkov et al., 2021; Pathak, Mustafa Mustafa, et al., 2021; Greenfeld et al., 2019), as well as 104 purely neural network-based solvers such as the Physics-Informed Neural Networks (PINNs), Deep Galerkin Method, and Deep Ritz Method (Raissi, Perdikaris, and George E Karniadakis, 2019; Sirignano and Spiliopoulos, 2018; Weinan and B. Yu, 2018). Especially, these PINN-type methods use a neural network 𝑢 𝜃 with parameters 𝜃 as the ansatz to approximate the solution function 𝑢 † . The parameters 𝜃 are found by minimizing the physics-informed loss with exact derivatives computed using automatic differentiation (autograd). In the stationary case, the physicsinformed loss is defined by minimizing the l.h.s. of equation (6.1) in the squared norm of F . A typical choice is F = 𝐿 2 (𝐷), giving the loss function 2 2 Lpde (𝑎, 𝑢 𝜃 ) = P (𝑎, 𝑢 𝜃 ) 2 + 𝛼 𝑢 𝜃 | 𝜕𝐷 − 𝑔 2 𝐿 (𝐷) ∫ ∫ 𝐿 (𝜕𝐷) |P (𝑢 𝜃 (𝑥), 𝑎(𝑥))| 2 d𝑥 + 𝛼 |𝑢 𝜃 (𝑥) − 𝑔(𝑥)| 2 d𝑥 = 𝐷 (6.3) 𝜕𝐷 In the case of a dynamical system, one minimizes the residual of equation (9.1) in some natural norm up to a fixed final time 𝑇 > 0. A typical choice is the
𝐿 2 (0, 𝑇]; 𝐿 2 (𝐷) norm, yielding 2 2 2 𝑑𝑢 𝜃 − R (𝑢 𝜃 ) 2 + 𝛼 𝑢 𝜃 | 𝜕𝐷 − 𝑔 2 + 𝛽 𝑢 𝜃 | 𝑡=0 − 𝑎 2 𝐿 (𝑇;𝐷) 𝐿 (𝑇;𝜕𝐷) 𝐿 (𝐷) 𝑑𝑡 ∫ 𝑇∫ 𝑑𝑢 𝜃 (𝑡, 𝑥) − R (𝑢 𝜃 )(𝑡, 𝑥)| 2 d𝑥d𝑡 = | 𝑑𝑡 0 𝐷 ∫ 𝑇∫ +𝛼 |𝑢 𝜃 (𝑡, 𝑥) − 𝑔(𝑡, 𝑥)| 2 d𝑥d𝑡 ∫0 𝜕𝐷 +𝛽 |𝑢 𝜃 (0, 𝑥) − 𝑎(𝑥)| 2 d𝑥 Lpde (𝑎, 𝑢 𝜃 ) = 𝐷 (6.4) The PDE loss consists of the physics loss in the interior and the data loss on the boundary and initial conditions, with hyper-parameters 𝛼, 𝛽 > 0. Alternatively, the optimization can be formulated via the variational form (Weinan and B. Yu, 2018). Challenges of PINN PINNs take advantage of the universal approximability of neural networks, but, in return, suffer from the low-frequency induced bias. Em- 105 pirically, PINNs often fail to solve challenging PDEs when the solution exhibits high-frequency or multi-scale structure (S. Wang, Teng, and Perdikaris, 2021; S. Wang, X. Yu, and Perdikaris, 2022; Fuks and Tchelepi, 2020; Raissi, Yazdani, and George Em Karniadakis, 2020). Further, as an iterative solver, PINNs have difficulty propagating information from the initial condition or boundary condition to unseen parts of the interior or to future times (Dwivedi and Srinivasan, 2020). For example, in challenging problems such as turbulence, PINNs are only able to solve the PDE on a relatively small domain (X. Jin et al., 2021), or otherwise, require extra observational data which is not always available in practice (Raissi, Yazdani, and George Em Karniadakis, 2020; Cai et al., 2021). In this work, we propose to overcome the challenges posed by optimization by integrating operator learning with PINNs. Learning the solution operator via neural operator An alternative setting is to learn the solution operator G. Given a PDE as defined in (6.1) or (9.1) and the corresponding solution operator G † , one can use a neural operator G𝜃 with parameters 𝜃 as a surrogate model to approximate G † . Usually we assume a dataset {𝑎 𝑗 , 𝑢 𝑗 } 𝑁𝑗=1 is available, where G † (𝑎 𝑗 ) = 𝑢 𝑗 and 𝑎 𝑗 ∼ 𝜇 are i.i.d. samples from some distribution 𝜇 supported on A. In this case, one can optimize the solution operator by minimizing the empirical data loss on a given data pair ∫ 2 Ldata (𝑢, G𝜃 (𝑎)) = ∥𝑢 − G𝜃 (𝑎)∥ U = |𝑢(𝑥) − G𝜃 (𝑎)(𝑥)| 2 d𝑥 (6.5) 𝐷 where we assume the setting of (6.1) for simplicity of the exposition. The operator data loss is defined as the average error across all possible inputs Jdata (G𝜃 ) = ∥G † − G𝜃 ∥ 2𝐿 2 (A;U) = E𝑎∼𝜇 [Ldata (G † (𝑎), G𝜃 (𝑎))] 𝜇 𝑁 ∫ 1 ∑︁ |𝑢 𝑗 (𝑥) − G𝜃 (𝑎 𝑗 )(𝑥)| 2 d𝑥. ≈ 𝑁 𝑗=1 𝐷 (6.6) Similarly, one can define the operator PDE loss as Jpde (G𝜃 ) = E𝑎∼𝜇 [Lpde (𝑎, G𝜃 (𝑎))]. (6.7) In general, it is non-trivial to compute the derivatives 𝑑G𝜃 (𝑎)/𝑑𝑥 and 𝑑G𝜃 (𝑎)/𝑑𝑡 for model G𝜃 . In the following section, we will discuss how to compute these derivatives for Fourier neural operator. 106 Neural operators In this work, we will focus on the neural operator model designed for the operator learning problem (Z. Li, Kovachki, et al., n.d.). The neural operator is formulated as a generalization of standard deep neural networks to the operator setting. Neural operator composes linear integral operator K with pointwise non-linear activation function 𝜎 to approximate highly non-linear operators. Definition 6.2.1 (Neural operator G𝜃 ) Define the neural operator G𝜃 := Q ◦ (W𝐿 + K𝐿 ) ◦ · · · ◦ 𝜎(W1 + K1 ) ◦ P (6.8) where P and Q are pointwise operators, parameterized with neural networks 𝑃 : R𝑑 𝑎 → R𝑑1 and 𝑄 : R𝑑 𝐿 → R𝑑𝑢 , where 𝑑 𝑎 is the co-dimension of an input function 𝑎 ∈ A and 𝑑𝑢 is the co-dimension of the output function 𝑢. P operator lifts the lower dimension function into higher dimensional space and Q operator projects the higher dimension function back to the lower dimensional space. The model stacks 𝐿 layers of 𝜎(W𝑙 + K𝑙 ) where W are pointwise linear operators parameterized as matrices 𝑊𝑙 ∈ R𝑑𝑙+1 ×𝑑𝑙 , K𝑙 : {𝐷 → R𝑑𝑙 } → {𝐷 → R𝑑𝑙+1 } are integral kernel operators, and 𝜎 are fixed activation functions. The parameters 𝜃 consists of all the parameters in P, Q, W𝑙 , K𝑙 . Definition 6.2.2 (Kernel Integral Operators) We define the kernel integral operator K used in (9.18). Let 𝜅 (𝑙) ∈ 𝐶 (𝐷 × 𝐷; R𝑑𝑙+1 ×𝑑𝑙 ) and let 𝜈 be a Borel measure on 𝐷. Then we define K by ∫ (K𝑣 𝑙 )(𝑥) = 𝜅 (𝑙) (𝑥, 𝑦)𝑣 𝑙 (𝑦) d𝜈(𝑦) ∀𝑥 ∈ 𝐷. (6.9) 𝐷 The kernel integral operator can be discretized and implemented with graph neural networks as in graph neural operators (Z. Li, Kovachki, et al., n.d.). ∑︁ (K𝑣 𝑙 )(𝑥) = 𝜅 (𝑙) (𝑥, 𝑦)𝑣 𝑙 (𝑦) ∀𝑥 ∈ 𝐷. (6.10) 𝐵(𝑥) where 𝐵(𝑥) is a ball of center at 𝑥. As a generalization, the kernel function can also Í take the form of (K𝑣 𝑙 )(𝑥) = 𝐵(𝑥) 𝜅 (𝑙) (𝑥, 𝑦, 𝑣 𝑙 (𝑦)). Recently, (Z. Li, Kovachki, et al., 2021a) proposes the Fourier neural operator (FNO) that restricts the integral operator K to convolution. In this case, it can apply the Fast Fourier Transform (FFT) to efficiently compute K. This leads to a fast architecture that obtains state-of-the-art results for PDE problems. 107 Definition 6.2.3 (Fourier convolution operator) One specific form of the kernel integral operator is the Fourier convolution operator  
K𝑣 𝑙 (𝑥) = F −1 𝑅 · (F 𝑣 𝑙 ) (𝑥) ∀𝑥 ∈ 𝐷 (6.11) where F , F −1 are the Fast Fourier transform and its inverse; 𝑅 is part of the parameter 𝜃 to be learn. One can build a neural operator with mixed kernel integral layers and Fourier convolution layers. If the input and output query points are sampled from nonuniform mesh, we can use the graph kernel operator as the first and last layer for continuous evaluation, while using the Fourier layers in the middle for efficient computation, similar to (Z. Li, D. Z. Huang, et al., n.d.). Challenges of operator learning. Operator learning is similar to supervised learning in computer vision and language where data play a very important role. One needs to assume the training points and testing points follow the same problem setting and the same distribution. Especially, the previous FNO model trained on one coefficient (e.g. Reynolds number) or one geometry cannot be easily generalized to another. Moreover, for more challenging PDEs where the solver is very slow or the solver is even not existent, it is hard to gather a representative dataset. On the other hand, since prior training methods for FNO do not use any knowledge of the equation, the trained models cannot get arbitrarily close to the ground truth by using the higher resolution as in conventional solvers, leaving a gap of generalization error. These challenges limit the applications of the prior works beyond accelerating the solver and modeling real-world experiments. In the following section, we will introduce the PINO framework to overcome these problems by using the equation constraints. Discretization convergent. Resolution and discretization convergence is defined as obtaining a unique continuum operator in the limit of mesh refinement (Kovachki et al., n.d.). The work (Bartolucci et al., 2023) recently introduced a new concept of representation equivalence, which requires zero aliasing error at each layer, which PINO does not fulfill. When all the Fourier modes in FNO are active, an aliasing error is inevitably present. However, in many practical applications, this is typically not an issue, and degraded performance due to aliasing is rarely observed, since the maximum number of modes in an FNO is typically far fewer 108 than the grid size. In fact, the non-linear transformations allow for re-capturing the truncated high-frequency modes which allows for generalization beyond the see training data. Requiring representation equivalence leads to linear methods of approximation which are know to be sub-optimal in their representation capacity (Lanthaler et al., 2023). Related work. Many machine learning models have been explored for operator learning (Bhattacharya et al., 2021; Nelsen and Stuart, 2021; Patel et al., 2021). Besides the above line of work, the deep operator network (DeepONet) is one of the most famous operator models that have shown enormous promise in applications (L. Lu, P. Jin, and George Em Karniadakis, 2021). The work from (Kontolati et al., 2023) compares the polynomial chaos expansion (PCE), DeepONet, and FNO, and shows that DeepONet has a higher approximation accuracy over PCE. According to Figure 5 in (Kontolati et al., 2023), standard DeepONet and FNO share a similar convergence rate. A similar comparison study is reported by de Hoop et. al. (Hoop et al., 2023) where FNO seems to converge faster. We choose FNO as our base model for its scalability to large problems. 6.3 Physics-Informed Neural Operator We propose the Physics-Informed Neural Operator (PINO) framework that uses one neural operator model G𝜃 for solving both operator learning problems and equation solving problems. It consists of two phases • Operator learning: learn a neural operator G𝜃 to approximate the target solution operator G † using either/both the data loss J𝑑𝑎𝑡𝑎 or/and the PDE loss J𝑝𝑑𝑒 . • Instance-wise fine-tuning: use G𝜃 (𝑎) as the ansatz to approximate 𝑢 † with the pde loss L 𝑝𝑑𝑒 and/or an additional operator loss L 𝑜 𝑝 obtained from the operator learning phase. Physics-informed operator learning For operator learning, we use the physics constraints J𝑝𝑑𝑒 and supervision from data to train the neural operator. Especially one can sample an unlimited amount of virtual PDE instances by drawing additional initial conditions or coefficient conditions 𝑎 𝑗 ∼ 𝜇 for training. In this sense, we have access to the unlimited dataset by sampling new input 𝑎 𝑗 in each training iteration. This advantage of using 109 PDE constraints removes the requirement on the dataset and makes the supervised problem into a semi-supervised one. While PINO can be trained with physics constraints J𝑝𝑑𝑒 only, the J𝑑𝑎𝑡𝑎 can provide stronger supervision than physics constraints and thus make the optimization much easier. PINO leverages the supervision from any available data to combine with physics constraints for better optimization landscape and thus learning accurate neural operators. A special case is to train a neural operator on low-resolution data instances with high-resolution PDE constraint, which will be studied in section 6.4. Instance-wise fine-tuning of trained operator ansatz Given a learned operator G𝜃 , we use G𝜃 (𝑎) as the ansatz to solve for 𝑢 † . The optimization procedure is similar to PINNs where it computes the PDE loss L 𝑝𝑑𝑒 on 𝑎, except that we propose to use a neural operator instead of a neural network. Since the PDE loss is a soft constraint and challenging to optimize, we also add an optional operator loss L 𝑜 𝑝 (anchor loss) to bound the further fine-tuned model from the learned operator model
L 𝑜 𝑝 G𝜃 𝑖 (𝑎), G𝜃 0 (𝑎) := ∥G𝜃 𝑖 (𝑎) − G𝜃 0 (𝑎)∥ 2U where G𝜃 𝑖 (𝑎) is the model at 𝑖 𝑡ℎ training epoch. We update the operator G𝜃 using the loss L 𝑝𝑑𝑒 + 𝛼L 𝑜 𝑝 . It is possible to further apply optimization techniques to fine-tune the last fewer layers of the neural operator and progressive training that gradually increase the grid resolution and use finer resolution in test time. Optimization landscape. Using the operator as the ansatz has two major advantages: (1) PINN does point-wise optimization, while PINO does optimization in the space of functions. In the linear integral operation K, the operator parameterizes the solution function as a sum of the basis function. Optimization of the set of coefficients and basis is easier than just optimizing a single function as in PINNs. (2) we can learn these basis functions in the operator learning phase which makes the later instance-wise fine-tuning even easier. In PINO, we do not need to propagate the information from the initial condition and boundary condition to the interior. It just requires fine-tuning the solution function parameterized by the solution operator. Trade-off. (1) complexity and accuracy: instance-wise fine-tuning is an option to spend more computation in exchange for better accuracy. The learned operator is extremely fast since it is performing inference on the neural operator. On the 110 other hand, instance-wise fine-tuning can improve accuracy while incurring more computational costs. (2) resolution effects on optimization landscape and truncation error (i.e. the error of the numerical differentiation): using a higher resolution and finer grid will reduce the truncation error. However, it has a higher computational complexity and memory consumption. A higher resolution may also potentially make the optimization unstable. Using hard constraints such as the anchor loss L 𝑜 𝑝 relieves such a problem. Derivatives of neural operators In order to use the equation loss L 𝑝𝑑𝑒 , one of the major technical challenges is to efficiently compute the derivatives 𝐷 (G𝜃 𝑎) = 𝑑 (G𝜃 𝑎)/𝑑𝑥 for neural operators. In this section, we discuss three efficient methods to compute the derivatives of the neural operator G𝜃 as defined in (9.18). Numerical differentiation. A simple but efficient approach is to use conventional numerical derivatives such as finite difference and Fourier differentiation (Zhu et al., 2019; Gao, L. Sun, and J.-X. Wang, 2021). These numerical differentiation methods are fast and memory-efficient: given a 𝑛-points grid, finite difference requires 𝑂 (𝑛), and the Fourier method requires 𝑂 (𝑛 log 𝑛). These differentiation methods are agnostic to the underlying neural network architecture. It can be applied to the neural operator with Graph layer 6.2.2 or Fourier layer 6.2.3 or neural networks such as UNet. However, the numerical differentiation methods face the same challenges as the corresponding numerical solvers: finite difference methods require a fine-resolution uniform grid; spectral methods require smoothness and uniform grids. Especially. These numerical errors on the derivatives will be amplified on the output solution. Pointwise differentiation with autograd. Similar to PINN (Raissi, Perdikaris, and George E Karniadakis, 2019), the most general way to compute the exact derivatives is to use the auto-differentiation library of neural networks (autograd). To apply autograd, one needs to use a neural network to parameterize the solution function 𝑢 : 𝑥 ↦→ 𝑢(𝑥). However, it is not straightforward to write out the solution function in the neural operator which directly outputs the numerical solution 𝑢 = G𝜃 (𝑎) on a grid, especially for FNO which uses FFT. To apply autograd, we design a query function 𝑢 that input 𝑥 and output 𝑢(𝑥). Recall G𝜃 := Q ◦ (W𝐿 + K𝐿 ) ◦ · · · ◦ 𝜎(W1 + K1 ) ◦ P and 𝑢 = G𝜃 𝑎 = Q𝑣 𝐿 = Q (W𝐿 + K𝐿 )𝑣 𝐿−1 . . .. Since Q is 111 pointwise, 𝑢(𝑥) = Q (𝑣 𝐿 )(𝑥) = 𝑄(𝑣 𝐿 (𝑥)) = 𝑄 (W𝐿 𝑣 𝐿−1 )(𝑥) + K𝐿 𝑣 𝐿−1 (𝑥)
(6.12) For both the kernel operator and Fourier operator, we either remove the pointwise residual term of the last layer (W𝐿 𝑣 𝐿−1 )(𝑥) or define the query function as an interpolation function on W𝐿 . For kernel integral operator 6.2.2, the kernel function can directly take the query points are input. So the query function ∑︁
𝜅 (𝑙) (𝑥, 𝑦, 𝑣 𝐿−1 (𝑦)) 𝑢(𝑥) = 𝑄 𝐵(𝑥) where we omit the derivative of the support 𝐵(𝑥). We can apply auto-differentiation to compute the derivatives
∑︁ (𝑙) ′ 𝑢′ (𝑥) = 𝑄 ′ 𝑣 𝐿 (𝑥) · 𝜅 (𝑥, 𝑦, 𝑣 𝐿−1 (𝑦)) (6.13) 𝐵(𝑥) Similarly, for the Fourier convolution operator, we need to evaluate the Fourier convolution K𝐿 𝑣 𝐿−1 (𝑥) on the query points 𝑥. It can be done by writing out the output function as Fourier series composing with 𝑄 : 𝑢(𝑥) = Q ◦ F −1  𝑚𝑎𝑥  1 𝑘∑︁
𝑖2𝜋𝑘  𝑅 · (F 𝑣 𝐿−1 )) (𝑥) = 𝑄 𝑅 𝑘 (F 𝑣 𝐿−1 ) 𝑘 exp (𝑥) 𝑘 𝑚𝑎𝑥 𝑘=0 𝐷  where F is the discrete Fourier transform. The inverse discrete Fourier transform
is the sum of 𝑘 𝑚𝑎𝑥 Fourier series with the coefficients 𝑅 𝑘 (F 𝑣 𝐿−1 ) 𝑘 coming from the previous layer.
𝑢′ (𝑥) = 𝑄 ′ 𝑣 𝐿 (𝑥) · 1 𝑘∑︁ 𝑚𝑎𝑥 𝑘 𝑚𝑎𝑥 𝑘=0
𝑖2𝜋𝑘 𝑅 𝑘 (F 𝑣 𝐿−1 ) 𝑘 exp′ (𝑥) 𝐷 (6.14) 𝑖2𝜋𝑘 𝑖2𝜋𝑘 Notice exp′ 𝑖2𝜋𝑘 𝐷 (𝑥) = 𝐷 exp 𝐷 (𝑥), just as the numerical Fourier method. If the query points 𝑥 are a uniform grid, the derivative can be efficiently computed with the Fast Fourier transform. The autograd method is general and exact, however, it is less efficient. Since the number of parameters |𝜃| is usually much greater than the grid size 𝑛, the numerical methods are indeed significantly faster. Empirically, the autograd method is usually slower and memory-consuming. 112 Function-wise differentiation. While it is expensive to apply the auto differentiation per query point, the derivative can be batched and computed in a function-wise manner. We develop an efficient and exact derivatives method based on the architecture of the neural operator that can compute the full gradient field. The idea is similar to the autograd, but we explicitly write out the derivatives on the Fourier space and apply the chain rule. Given the explicit form (6.14), 𝑢′ can be directly computed on the Fourier space. ′ ′
𝑢 = Q 𝑣𝐿 · F −1  𝑖2𝜋 𝐾 · (F 𝑣 𝐿 ))  (6.15) 𝐷 Therefore, to exactly compute the derivative of the Fourier neural operator, one just needs to run the numerical Fourier differentiation. Especially the derivative and be efficiently computed with the Fast Fourier transform when the query is uniform. Similarly, if the kernel function 𝜅 (𝑙) in (6.13) has a structured form, we can also write out its gradient field explicitly. Next, we show how to compute higher orders derivatives in their exact form, without evoking the autograd method. To this end, we can directly apply the chain rule for the higher-order derivatives without calling autograd. For example, the first order derivatives is 𝑢′ = (𝑄 ◦ 𝑣 𝐿 ) ′ = Q ′ (𝑣 𝐿 ) · 𝑣 ′𝐿 and the 2nd-order derivatives is ′′ ′ ′′ 𝑢′′ = (Q𝑣 𝐿 ) ′′ = 𝑣 ′2 𝐿 · Q (𝑣 𝐿 ) + Q (𝑣 𝐿 ) · 𝑣 𝐿 . Higher-order derivatives can be similarly computed using the chain rule. Furthermore, derivative-based quantities on 𝑣 𝐿 , e.g., 𝑣 ′𝐿 can be computed in its exact form in the Fourier domain. Similarly, we can write out the higher-order derivatives of 𝑄 using the chain rule. Usually 𝑄 is parameterized as a two layer neural networks 𝑄(𝑥) = ( 𝐴2 𝜎( 𝐴1 𝑥 + 𝑏 1 ) + 𝑏 2 ). So 𝑄 ′ (𝑥) = 𝐴2 𝜎′ ( 𝐴1 𝑥 + 𝑏 1 ) 𝐴1 . In this manner, we have got the explicit formula of the derivatives for all neural operators. Fourier continuation. The Fourier method has its best performance when applied to periodic problems. If the target function is non-periodic or non-smooth, the Fourier differentiation is not accurate. To deal with this issue, we apply the Fourier continuation method that embeds the problem domain into a larger and periodic space. The extension can be simply done by padding zeros in the input. The loss is computed at the original space during training. The Fourier neural operator will automatically generate a smooth extension. The details are given in Appendix 6.3. 113 Inverse problem The physics-informed method can be used to solve the inverse problem, where given the output function 𝑢, the goal is to recover (a distribution of) the input function 𝑎. By imposing the constraint P (𝑢, 𝑎) = 0, we can restrict 𝑎 to a physically valid manifold. We propose two formulations to do the optimization-based inverse problem with PINO: the forward operator model and the inverse operator model. • Forward model: learn the forward operator G𝜃 : 𝑎 ↦→ 𝑢 with data. Initialize 𝑎ˆ to approximate 𝑎 † . Optimize 𝑎ˆ using J𝑓 𝑜𝑟𝑤𝑎𝑟 𝑑 := L 𝑝𝑑𝑒 ( 𝑎, ˆ 𝑢 † ) + L 𝑑𝑎𝑡𝑎 (G𝜃 ( 𝑎)) ˆ + 𝑅( 𝑎). ˆ (6.16) • inverse model: learn the inverse operator F𝜃 : 𝑢 ↦→ 𝑎 with data. Use F𝜃 (𝑢 † ) to approximate 𝑎 † . Optimize F𝜃 using J𝑏𝑎𝑐𝑘𝑤𝑎𝑟 𝑑 := L 𝑝𝑑𝑒 (F𝜃 (𝑢 † ), 𝑢 † ) + L 𝑜 𝑝 (F𝜃 (𝑢 † ), F𝜃 0 (𝑢 † )) + 𝑅(F𝜃 (𝑢 † ) (6.17) where L 𝑝𝑑𝑒 is the PDE loss; L 𝑜 𝑝 is the operator loss from the learned operator; 𝑅(𝑎) is the regularization term. We use the PDE loss L 𝑝𝑑𝑒 to deal with the small error in G𝜃 and the ill-defining issue of F𝜃 . We provide a numerical study in section 6.4. Fourier continuation The Fourier neural operator can be applied to arbitrary geometry via Fourier continuations. Given any compact manifold M, we can always embed it into a periodic cube (torus), 𝑖 : M → T𝑛 where we can do the regular FFT. Conventionally, people would define the embedding 𝑖 as a continuous extension by fitting polynomials (Bruno, Y. Han, and Pohlman, 2007). However, in Fourier neural operator, it can be simply done by padding zeros in the input. The loss is computed at the original space during training. The Fourier neural operator will automatically generate a smooth extension to do a padded domain in the output, as shown in Figure 6.3. This technique is first used in the original Fourier neural operator paper (Z. Li, Kovachki, et al., 2021a) to deal with the time dimension in the Navier-Stokes equation. Similarly, (L. Lu, Meng, Cai, et al., 2022) apply FNO with extension 114 Figure 6.3: Fourier Continuation by padding zeros. The x-axis is the spatial dimension; the y-axis is the temporal dimension. FNO extends the output smoothly on the padded domain. and interpolation on diverse geometries on the Darcy equation. In the work, we use Fourier continuation widely for non-periodic boundary conditions (Darcy, time dimension). We also added an example of lid-cavity to demonstrate that PINO can work with non-periodic boundary conditions. Furthermore, this Fourier continuation technique helps to take the derivatives of the Fourier neural operator. Since the output of FNO is always on a periodic domain, the numerical Fourier gradient is usually efficient and accurate, except if there is shock (in this case, we will use the exact gradient method). 6.4 Experiments In this section, we conduct empirical experiments to examine the efficacy of the proposed PINO. We present the PDE settings, their domains, and function spaces. In 6.4, we show using PDE constraint in operator learning results in neural operators that (1) generalize to significantly high-resolution unseen data. (2) achieve smaller generalization errors with fewer to no data. Then in 6.4, we investigate how PINO uses the operator ansatz to solve harder equations with improved speed and accuracy. We study three concrete cases of PDEs on Burgers’ Equation, Darcy’s Equation, and Navier-Stokes equation. In 6.4 we study the inverse problems. Burgers’ Equation. The 1-d Burgers’ equation is a non-linear PDE with periodic 2 ((0, 1); R) is the initial condition and 𝜈 = 0.01 boundary conditions where 𝑢 0 ∈ 𝐿 per is the viscosity coefficient. We aim to learn the operator mapping the initial condition to the solution, G † : 𝑢 0 ↦→ 𝑢| [0,1] . 𝜕𝑡 𝑢(𝑥, 𝑡) + 𝜕𝑥 (𝑢 2 (𝑥, 𝑡)/2) = 𝜈𝜕𝑥𝑥 𝑢(𝑥, 𝑡), 𝑢(𝑥, 0) = 𝑢 0 (𝑥), 𝑥 ∈ (0, 1), 𝑡 ∈ (0, 1] 𝑥 ∈ (0, 1) (6.18) 115 Darcy Flow. The 2-d steady-state Darcy Flow equation on the unit box which is the second order linear elliptic PDE with a Dirichlet boundary where 𝑎 ∈ 𝐿 ∞ ((0, 1) 2 ; R+ ) is a piecewise constant diffusion coefficient and 𝑓 = 1 is a fixed forcing function. We are interested in learning the operator mapping the diffusion coefficient to the solution, G † : 𝑎 ↦→ 𝑢. Note that although the PDE is linear, the operator G † is not. −∇ · (𝑎(𝑥)∇𝑢(𝑥)) = 𝑓 (𝑥) 𝑢(𝑥) = 0 𝑥 ∈ (0, 1) 2 𝑥 ∈ 𝜕 (0, 1) 2 (6.19) Since 𝑎 is in 𝐿 𝑖𝑛 𝑓 , we considered both the strong form L 𝑝𝑑𝑒 (𝑢) = ∇ · (𝑎∇𝑢) − 𝑓 and the weak form minimization loss L 𝑝𝑑𝑒 (𝑢) = − 21 (𝑎∇𝑢, ∇𝑢) − (𝑢, 𝑓 ), 𝑢 ∈ 𝐻1 . Experiments show the strong form has a better performance. Navier-Stokes Equation. We consider the 2-d Navier-Stokes equation for a vis𝑟 ((0, 𝑙) 2 ; R2 )) cous, incompressible fluid in vorticity form on the unit torus, where 𝑢 ∈ 𝐶 ([0, 𝑇]; 𝐻per 2 ((0, 𝑙) 2 ; R) for any 𝑟 > 0 is the velocity field, 𝑤 = ∇ × 𝑢 is the vorticity, 𝑤 0 ∈ 𝐿 per 2 ((0, 𝑙) 2 ; R) is is the initial vorticity, 𝜈 ∈ R+ is the viscosity coefficient, and 𝑓 ∈ 𝐿 per the forcing function. We want to learn the operator mapping the vorticity from the initial condition to the full solution G † : 𝑤 0 ↦→ 𝑤| [0,𝑇] . 𝜕𝑡 𝑤(𝑥, 𝑡) + 𝑢(𝑥, 𝑡) · ∇𝑤(𝑥, 𝑡) = 𝜈Δ𝑤(𝑥, 𝑡) + 𝑓 (𝑥), 𝑥 ∈ (0, 𝑙) 2 , 𝑡 ∈ (0, 𝑇] ∇ · 𝑢(𝑥, 𝑡) = 0, 𝑥 ∈ (0, 𝑙) 2 , 𝑡 ∈ [0, 𝑇] 𝑤(𝑥, 0) = 𝑤 0 (𝑥), 𝑥 ∈ (0, 𝑙) 2 (6.20) Specially, we consider two problem settings: • Long temporal transient flow: we study the build-up of the flow from the initial condition 𝑢 0 near-zero velocity to 𝑢𝑇 that reaches the ergodic state. We choose 𝑡 ∈ [0, 50], 𝑙 = 1, 𝑅𝑒 = 20 as in (Z. Li, Kovachki, et al., 2021a). The main challenge is to predict the long time interval. • Chaotic Kolmogorov flow: In this case 𝑢 lies in the attractor where arbitrary starting time 𝑡0 . We choose 𝑡 ∈ [𝑡 0 , 𝑡0 + 0.5] or [𝑡0 , 𝑡0 + 1], 𝑙 = 2𝜋, 𝑅𝑒 = 500 similar to (Z. Li, Kovachki, et al., 2023). The main challenge is to capture the small details that evolve chaotically. 116 Figure 6.4: PINO on Kolmogorov flow (left) and Lid-cavity flow (right) (a) The long-temporal transient flow with 𝑅𝑒 ∼ 20, 𝑇 = 50. PINO outputs the full trajectory in one step, which leads to a 400x speedup compared to the GPU solver. PINN cannot converge to a reasonable error rate due to the long time window. (b) The chaotic Kolmogorov flow with 𝑅𝑒 =∼ 500, 𝑇 = 0.5. PINO converges faster compared to PINN, but their convergence rates with gradient descent are less effective compared to using higher resolutions in the GPU solver. Figure 6.5: The accuracy-complexity trade-off on PINO, PINN, and the GPU-based pseudo-spectral solver. • Lid cavity flow: In this case, we assume the no-slip boundary condition where 𝑢(𝑥, 𝑡) = (0, 0) at left, bottom, and right walls and 𝑢(𝑥, 𝑡) = (1, 0) on top, similar to (Bruneau and Saad, 2006). We choose 𝑡 ∈ [5, 10], 𝑙 = 1, 𝑅𝑒 = 500. The main challenge is to address the boundary using the velocity-pressure formulation. 117 5HODWLYHL2HUURU    3,12RSHUDWRUOHDUQLQJ 3,12LQVWDQFHZLVHILQHWXQLQJ 3,11 /$$)3,11 6$3,11  5HODWLYHL2HUURU      3,12RSHUDWRUOHDUQLQJ 3,12LQVWDQFHZLVHILQHWXQLQJ 3,11 /$$)3,11 6$3,11               5XQWLPHV           5XQWLPHV Figure 6.6: Plot of test relative 𝐿 2 error versus runtime step for the Kolmogorov flow with Re500, T=0.5s. Left: resolution 64×64×65; right: resolution 128×128×129. Averaged over 20 instances. LAAF-PINN: PINN with locally adaptive activation functions. SA-PINN: self-adaptive PINN. Operator learning with physics constraints We show that we can utilize the equation constraints to improve neural operator training. For this purpose, we train neural operators on fixed-resolution data in the presence of physics loss, J𝑝𝑑𝑒 , and test the performance of the trained operators on high-resolution data. In particular, we test the performance of the trained model on data with the same resolution of the training data, 1x, 2x, and 4x, of the training data resolution Table 6.1. We observe that incorporating the J𝑝𝑑𝑒 in the training results in operators that, with high accuracy, generalize across data resolution. In this experiment, the training data for Burgers equation setting is in 32 × 25 (spatio-temporal), and the J𝑝𝑑𝑒 is imposed in 128 × 100 resolution. We use 800 low-resolution data and the same 800 PDE instances. The mean relative 𝐿 2 error and its standard deviation are reported over 200 test instances at resolution 32 × 25, 64 × 50, and 128 × 100. Accordingly, the training data for the Darcy equation setting is at the spatial resolution of 11 × 11 and the J𝑝𝑑𝑒 is imposed in 61 × 61 resolution. We use 1000 low-resolution data and the same 1000 PDE instances. The mean and standard error are reported over 500 test instances at resolution 11 × 11, 61 × 61, and 211 × 211. Darcy flow is unresolved at the 11 × 11 resolution, training on such a coarse grid causes higher errors. However, adding higher resolution PDE loss helps the operator to resolve. The training data for Kolmogorov flow is in 64 × 64 × 33 and the J𝑝𝑑𝑒 is imposed in 256×256×65 resolution for the time interval [0, 0.125]. We use 800 low-resolution data and 2200 PDE instances. The mean and std of the relative 𝐿 2 error are reported over 200 test instances, at resolution 64×64×33, 128×128×33, and 256×256×65. 118 PDE Training setting Data Data & PDE loss Data Darcy Data & PDE loss Data Kolm. flow Data & PDE loss Burgers Error at low data resolution 0.32±0.01% 0.17±0.01% 5.41±0.12% 5.23±0.12% 8.28%±0.15% 6.04%±0.12% Error at 2× data resolution 3.32±0.02% 0.28±0.01% 9.01±0.07% 1.56±0.05% 8.27%±0.15% 6.02%±0.12% Error at 4× data resolution 3.76±0.02% 0.38±0.01% 9.46±0.07% 1.58±0.06% 8.30%±0.15% 6.01%±0.12% Table 6.1: Operator-learning using fixed resolution data and PDE loss allows us to train operators with high accuracy on higher resolution unseen data. Burgers equation and Darcy equation. PINO can learn the solution operator without any data on simpler problems such as Burgers and Darcy. Compared to other PDE-constrained operators, PINO is more expressive and thereby achieves better accuracy. On Burgers (6.18), PI-DeepONet achieves 1.38% (S. Wang, H. Wang, and Perdikaris, 2021); PINO achieves 0.38%. Similarly, on Darcy flow (7.2), PINO outperforms FNO by utilizing physics constraints, as shown in Table 6.2. For these simpler equations, instance-wise fine-tuning may not be needed. Method DeepONet with data PINO with data PINO w/o data Solution error 6.97 ± 0.09% 1.22 ± 0.03% 1.50 ± 0.03% Table 6.2: Operator learning on Darcy Flow equation. Incorporating physics constraints in operator learning improves the performance of neural operators. # data samples 0 800 2200 # PDE instances 2200 2200 2200 Solution error 6.22%±0.11% 6.01%±0.12% 5.04%±0.11% Table 6.3: Physics-informed neural operator learning on Kolmogorov flow 𝑅𝑒 = 500. PINO is effective and flexible in combining physics constraints and any amount of available data. The mean and standard error of the relative 𝐿 2 test error is reported over 200 instances and evaluated on resolution 256 × 256 × 65. Chaotic Kolmogorov flow. We conduct an empirical study on how PINO can improve the generalization of neural operators by enforcing more physics. In the first experiment, we consider the Kolmogorov flow with 𝑇 = 0.125. We train PINO 119 with 2200 initial conditions and different amounts of low-resolution data. As shown in Table 6.3, PINO achieves 6.22% error even without any data. We also observe that adding more low-resolution data to training makes the optimization easier and consistently improves the accuracy of the learned operator, showing that PINO is effective and flexible in combining physics constraints and any amount of available data. The second experiment considers the Kolmogorov flow with 𝑇 = 0.5. The training set consists of 4000 data points of the initial condition and corresponding solution. For operator learning, we sample high-resolution initial conditions from a Gaussian random field. Table 6.7 compare the generalization error of neural operators trained by different schemes and different amounts of simulated data. The result shows that training neural operator with additional PDE instances consistently improves the generalization error on all three resolutions we are evaluating. Note that the relative 𝐿 2 error in this setting is much higher than the first one because the time horizon is 4 times longer. Next, we show how to solve for specific instances by finetuning the learned operator. Solve equation using operator ansatz We solve specific equation instances by fine-tuning the learned operator ansatz. Long temporal transient flow. It is extremely challenging to propagate the information from the initial condition to future time steps over such a long interval 𝑇 = [0, 50] just using the soft physics constraint. Neither the PINN nor PINO (from scratch) can handle this case (error > 50%), no matter solving the full interval at once or solving per smaller steps. However, when the data is available for PINO, we can use the learned neural operator ansatz and the anchor loss L 𝑜 𝑝 . The anchor loss is a hard constraint that makes the optimization much easier. Providing 𝑁 = 4800 training data, the PINO without instance-wise fine-tuning achieves 2.87% error, lower than FNO 3.04% and it retains a 400x speedup compared to the GPU-based pseudo-spectral solver (He and W. Sun, 2007), matching FNO. Further doing test time optimization with the anchor loss and PDE loss, PINO reduces the error to 1.84%. Chaotic Kolmogorov flow. Based on the solution operators learned in Section 6.4, the second operator-learning setting, we continue to do instance-wise fine-tuning. We compare our method against other physics-informed learning methods including 120 PINN (Raissi, Perdikaris, and George E Karniadakis, 2019), LAAF-PINN (Jagtap, Kawaguchi, and Em Karniadakis, 2020), and SA-PINN (McClenny and Braga-Neto, 2023), as shown in Figure 6.6 and Table 6.4. Overall, PINO outperforms PINN and its improved variants by 20x smaller error and 25x speedup. Using a learned operator model makes PINO converge faster. Method PINNs PINO PINO PINO # data samples 0 0.4k 0.4k # PDE instances 0 0 160k Solution error (𝑤) 18.7% 0.9% 0.9% 0.9% Time cost 4577s 608s 536s 473s Table 6.4: Instance-wise fine-tuning on Kolmogorov flow 𝑅𝑒 = 500, 𝑇 = 0.5. Using the learned operator as the initial condition helps fine-tuning converge faster. Zero-shot super-resolution. The neural operator models are discretization-convergent, meaning they can take the training dataset of variant resolutions and be evaluated at higher resolution. As shown in Figure 6.1, we train the FNO, PINO, and UNet model with 64 × 64 × 32 Kolmogorov Flows dataset and evaluate them at 256 × 256 × 65 resolution. Any frequencies higher than 64 are unseen during the training time. Conventional models such as UNet are not capable of direct super-resolution. For compassion, we equip it with tri-linear interpolation. For PINO, we also do test-time optimization. As shown in Figure 6.1, the spectrums of the predictions are averaged over 50 instances. PINO with the test-time optimization achieves a very high accuracy rate, and its spectrum overlaps with the ground truth spectrum. However, Conventional models such as UNet+Interpolation have noising prediction with oscillating high frequencies. On the other hand, with the help of test-time optimization, PINO can extrapolate to unseen frequencies with high accuracy. Transfer Reynolds numbers. The extrapolation of different parameters and conditions is one of the biggest challenges for ML-based methods. It poses a domain shift problem. In this experiment, we train the source operator model on one Reynolds number and then fine-tune the model to another Reynolds number, on the Kolmogorov flow with 𝑇 = 1. As shown in Table 6.5 by doing instance-wise finetuning, PINO can be easily transferred to different Reynolds numbers ranging from 100 to 500. This transferability shows PINO learned the dynamics shared across different Reynolds numbers. Such property envisions broad applications including transferring the learned operator to different boundary conditions or geometries. 121 Transfer learning across Reynolds numbers We study the instance-wise fine-tuning with different Reynolds numbers on the 𝑇 = 1 Kolmogorov flow. For the higher Reynolds number problem 𝑅𝑒 = 500, 400, finetuning the source operator shows better convergence accuracy than learning from scratch. In all cases, the fine-tuning of the source operator shows better convergence speed as demonstrated in Figure 6.7. The results are shown in Table 6.5 where the error is averaged over 40 instances. Each row is a testing case, and each column is a source operator. Testing Re 500 400 350 300 250 200 100 From scratch 4.93 2.96 1.92 1.68 1.51 0.921 0.234 100 3.83 2.43 2.10 1.61 1.50 0.913 0.235 200 3.93 2.45 2.11 1.64 1.53 0.921 0.236 250 3.15 2.44 2.13 1.51 1.51 0.915 0.235 300 4.77 3.00 2.33 1.77 1.60 0.985 0.239 350 4.46 2.71 2.22 1.73 1.56 0.945 0.239 400 4.34 2.73 2.22 1.70 1.60 0.923 0.237 500 4.36 2.40 2.12 1.60 1.51 0.892 0.237 Each row is a test set of PDEs with corresponding Reynolds number. Each column represents the operator ansatz we use as the starting point of instance-wise fine-tuning. For example, column header “100” means the operator ansatz is trained over a set of PDEs with Reynolds number 100. The relative 𝐿 2 errors is averaged over 40 instances of the corresponding test set. Table 6.5: Reynolds number transfer learning. Lid cavity flow. We demonstrate an additional example using PINO to solve for lid-cavity flow on 𝑇 = [5, 10] with 𝑅𝑒 = 500. In this case, we do not have the operator-learning phase and directly solve the equation (instance-wise fine-tuning). We use PINO with the velocity-pressure formulation and resolution 65 × 65 × 50 plus the Fourier numerical gradient. It takes 2 minutes to achieve a relative error of 14.52%. Figure 6.4 shows the ground truth and prediction of the velocity field at 𝑡 = 10 where the PINO accurately predicts the ground truth. The experiment shows that PINO can address non-periodic boundary conditions and multiple output fields. Convergence of accuracy with respect to resolution. We study the convergence rate of PINO in the instance-wise optimization setting, where we minimize the PDE loss under different resolutions without any data. For PINO, using a higher resolution is more effective compared to running gradient descent for longer iterations. We test PINO on the Kolmogorov flow with Re = 500 and T = 0.125. We use the Fourier method in the spatial dimension and the finite difference method in the 122 The test error is averaged over 40 instances. We observe that all the operator ansatzs trained over PDE instances with different Reynolds numbers can boost the instance-wise fine-tuning accuracy and convergence speed compared to training from scratch. Figure 6.7: Plot of relative 𝐿 2 error versus update step for the Kolmogorov flow with Reynolds number 500, 𝑇 = 1. temporal dimension. As shown in Table 6.6, PINO shares the same convergence rate of its differentiation methods with no obvious limitation from optimization. It has an exponential convergence rate in space and a first-order convergence rate in time 𝑒𝑟𝑟 = 𝑂 (𝑒𝑥 𝑝(𝑑𝑥)) + 𝑂 (𝑑𝑡) . It implies the PDE constraint can achieve high accuracy given a reasonable computational cost, and the virtual instances are almost as good as the data instances generated by the solver. Since the PDE loss can be computed on an unlimited amount of virtual instances in the operator learning setting, it is possible to reduce the generalization error going to zero by sampling virtual instances. dx 2−4 2−5 2−6 2−7 dt 2−6 2−7 2−8 2−9 2−10 0.4081 0.1819 0.0730 0.0582 0.3150 0.1817 0.0436 0.0234 0.3149 0.1780 0.0398 0.0122 0.3179 0.1773 0.0386 0.0066 0.3196 0.1757 0.0382 0.0034 Table 6.6: relative L2 error of PINO (Finite-difference in time) on Kolmogorov flow with 𝑅𝑒 = 500 and 𝑇 = 0.125. PINO inherits the convergence rate of its differentiation method with no limitation of optimization. 123 0 0 0 10 10 10 20 20 20 30 30 30 40 40 40 50 50 50 60 0 10 20 30 40 50 60 60 0 10 20 30 40 50 60 60 0 10 20 30 40 50 60 (a) Ground truth input 𝑎 (b) Inversion using only (c) Inversion using data and data constraint PDE constraints 0 0 0 10 10 10 20 20 20 30 30 30 40 40 40 50 50 50 60 0 10 20 30 40 50 60 60 0 10 20 30 40 50 60 60 0 10 20 30 40 50 60 (d) Observed output func- (e) Output function of in- (f) Output function of intion version using only data con- version using data and PDE straint constraints In the above figures, (6.8a) represents the ground truth input function 𝑎 † , and (6.8d) demonstrates the corresponding solution 𝑢 † , i.e., the output function. Given the output 𝑢 † , we aim to recover what 𝑎 could have generated the output function 𝑢 † . Using only data constraint, (6.8b) shows that our method can find an 𝑎 that results in an output function very close to the ground truth 𝑢 † (6.8e). However, the recovered 𝑎 is far from satisfying the PDE equation. Using both data and PDE constraints, (6.8c) shows that our physics-informed method can find an 𝑎 that not only results in an output function very close to the ground truth 𝑢 † (6.8f), but also the recovered 𝑎 satisfies the PDE constraint and is close to the underlying 𝑎 † . Figure 6.8: Inverse Problem with Physics Informed Neural Operator Inverse problem One of the major advantages of the physics-informed method is to solve the inverse problem. In the experiment, we investigate PINO on the inverse problem of the Darcy equation to recover the coefficient function 𝑎 † from the given solution function 𝑢 † . We assume a dataset {𝑎 𝑗 , 𝑢 𝑗 } is available to learn the operator. The coefficient function 𝑎 is a piecewise constant (representing two types of media), so the inverse problem can be viewed as a classification problem. We define 𝑅(𝑎) as the total variance. The PDE loss strongly improves the prediction of the inverse problem. The plain 124 Figure 6.9: Darcy inverse problem: comparing PINO forward, inverse models with numerical solver with MCMC. neural operator model, while accurate in the forward problem, is vulnerable under perturbation and shift of the input 𝑎, which is a common behavior of deep-learning models. This domain-shift problem is critical for optimization-based inverse problems. During the optimization, 𝑎 is likely to go out of the training distribution, which makes the neural operator model inaccurate. As shown in Figure 6.8 (b), the prediction of 𝑎 is less accurate, while the model believes its output 6.8 (e) is the same as the target. This issue is mitigated by adding the PDE constraints, which restrict the prediction 𝑎 to the physically-valid manifold where P (𝑎, 𝑢) = 0. As shown in Figure 6.8 (c), the initial condition recovered with PDE loss is very close to the ground truth. Comparing the PINO forward model with the inverse model, the inverse model F𝜃 : 𝑢 ↦→ 𝑎 (6.17) has the best performance. As Shown in Figure 6.9, the inverse model has 2.29% relative l2 error on the output 𝑢 and 97.10% classification accuracy on the input 𝑎; the forward model has 6.43% error on the output and 95.38% accuracy on the input. Both models converge with 200 iterations. The major advantage of the PINO inverse model compared to the PINO forward model is that it uses a neural operator F𝜃 (𝑢 † ) as the ansatz for the coefficient function, which is used as regularization L 𝑜 𝑝 . Similar to the forward problem, the operator ansatz has an easier optimization landscape while being expressive. As a reference, we compare the PINO inverse frameworks with PINN and the conventional solvers using the accelerated MCMC method with 500,000 steps (Cotter et al., 2013). The posterior mean of the MCMC has a 4.52% error and 90.30% respectively ( Notice the Bayesian method outputs the posterior distribution, which is 125 beyond obtaining a maximum a posteriori estimation). Meanwhile, PINO methods are 3000x faster compared to MCMC PINN does not converge in this case. Besides the speedup with respect to the online cost, the offline training of PINO only takes around 1 hour on a single GPU on the Darcy problem. Once trained, the model can be used without any further training cost. As a comparison, it takes considerably longer to deploy a finite element solver and an MCMC solver compared to a machine learning model. Generally speaking, numerical solvers usually have a more complicated codebase and it is non-trivial to specify boundary conditions, time schemes, and meshes. In the end, it can be easier to prepare a machine learning model than a standard numerical solver. Flexibility and accessibility are some of the major advantages of these machine learning methods. # Data # PDE 400 0 400 40k 400 160k 4k 0 4k 100k 4k 400k 0 100k Resolution 128 × 128 × 65 64 × 64 × 65 32 × 32 × 33 128 × 128 × 65 64 × 64 × 65 32 × 32 × 33 128 × 128 × 65 64 × 64 × 65 32 × 32 × 33 128 × 128 × 65 64 × 64 × 65 32 × 32 × 33 128 × 128 × 65 64 × 64 × 65 32 × 32 × 33 128 × 128 × 65 64 × 64 × 65 32 × 32 × 33 128 × 128 × 65 64 × 64 × 65 32 × 32 × 33 Solution error 33.32% 33.31% 30.61% 31.74% 31.72% 29.60% 31.32% 31.29% 29.28% 25.15% 25.16% 21.41% 24.15% 24.11% 20.85% 24.22% 23.95% 20.10% 74.36% 74.38% 74.14% Equation error 1.8779 1.8830 1.8421 1.8179 1.8227 1.8296 1.7840 1.7864 1.8524 1.8223 1.8257 1.8468 1.6112 1.6159 1.8251 1.4596 1.4656 1.9146 0.3741 0.3899 0.5226 Table 6.7: Each neural operator is trained with 400 or 4000 data points additionally sampled free initial conditions. The Reynolds number is 500. The reported generalization error is averaged over 300 instances. Training on additional initial conditions boosts the generalization ability of the operator. 126 6.5 Discussion and Conclusion In this work, we develop the physics-informed neural operator (PINO) that bridges the gap between physics-informed optimization and data-driven neural operator learning. We introduce operator-learning and instance-wise fine-tuning schemes for PINO to utilize both the data and physics constraints. In the operator learning phase, PINO learns an operator ansatz over multiple instances of a parametric PDE family. The instance-wise fine-tuning scheme allows us to take advantage of the learned neural operator ansatz and solve for the solution function on the querying instance faster and more accurately. While PINO shows many promising applications, it also shares some limitations as in the previous work. For example, since PINO is currently implemented with the FNO backbone with the Fast-Fourier transform, it is hard to extend to higher dimensional problems. Besides, as shown in Figure 6.5, finetuning PINO using gradient descent methods does not converge as fast as using a finer grid as in Table 6.6. Further optimization techniques are to be developed. There are many exciting future directions. Most of the techniques and analyses of PINN can be transferred to PINO. It is also interesting to ask how to overcome the hard trade-off of accuracy and complexity, and how the PINO model transfers across different geometries. Furthermore, it is promising to develop a software library of pre-trained models. PINO’s excellent extrapolation property allows it to be applied on a broad set of conditions, as shown in Transfer Reynold’s number experiments. References Bartolucci, Francesca et al. (2023). “Representation Equivalent Neural Operators”. In: Conference on Neural Information Processing Systems. Bhattacharya, Kaushik et al. (2021). “Model reduction and neural networks for parametric PDE(s)”. In: The SMAI Journal of computational mathematics, Volume 7, pp. 121-157. Bonev, Boris et al. (2023). “Spherical Fourier Neural Operators: Learning Stable Dynamics on the Sphere”. In: arXiv preprint arXiv:2306.03838. Bruneau, Charles-Henri and Mazen Saad (2006). “The 2D lid-driven cavity problem revisited”. In: Computers & fluids 35.3, pp. 326–348. Bruno, Oscar P, Youngae Han, and Matthew M Pohlman (2007). “Accurate, highorder representation of complex three-dimensional surfaces via Fourier continuation analysis”. In: Journal of computational Physics 227.2, pp. 1094–1125. 127 Brunton, Steven L, Bernd R Noack, and Petros Koumoutsakos (2020). “Machine learning for fluid mechanics”. In: Annual Review of Fluid Mechanics 52, pp. 477– 508. Cai, Shengze et al. (2021). “Physics-informed neural networks (PINNs) for fluid mechanics: A review”. In: Acta Mechanica Sinica. Cotter, Simon L et al. (2013). “MCMC methods for functions: modifying old algorithms to make them faster”. In: Statistical Science 28.3, pp. 424–446. Duvall, James, Karthik Duraisamy, and Shaowu Pan (2021). Non-linear Independent Dual System (NIDS) for Discretization-independent Surrogate Modeling over Complex Geometries. arXiv: 2109.07018 [physics.comp-ph]. Dwivedi, Vikas and Balaji Srinivasan (2020). “Physics Informed Extreme Learning Machine (PIELM)–A rapid method for the numerical solution of partial differential equations”. In: Neurocomputing 391, pp. 96–118. Fuks, Olga and Hamdi A Tchelepi (2020). “Limitations of physics informed machine learning for nonlinear two-phase transport in porous media”. In: Journal of Machine Learning for Modeling and Computing 1.1. Gao, Han, Luning Sun, and Jian-Xun Wang (2021). “PhyGeoNet: physics-informed geometry-adaptive convolutional neural networks for solving parameterized steadystate PDEs on irregular domain”. In: Journal of Computational Physics 428, p. 110079. Greenfeld, Daniel et al. (2019). “Learning to optimize multigrid PDE solvers”. In: International Conference on Machine Learning. PMLR, pp. 2415–2423. Han, Jiequn, Arnulf Jentzen, and E Weinan (2018). “Solving high-dimensional partial differential equations using deep learning”. In: Proceedings of the National Academy of Sciences 115.34, pp. 8505–8510. He, Yinnian and Weiwei Sun (2007). “Stability and convergence of the Crank– Nicolson/Adams–Bashforth scheme for the time-dependent Navier–Stokes equations”. In: SIAM Journal on Numerical Analysis 45.2, pp. 837–869. Hennigh, Oliver et al. (2021). “NVIDIA SimNet™: An AI-accelerated multi-physics simulation framework”. In: International Conference on Computational Science. Springer, pp. 447–461. Hersbach, Hans et al. (2020). “The ERA5 global reanalysis”. In: Quarterly Journal of the Royal Meteorological Society 146.730, pp. 1999–2049. Hoop, Maarten V de et al. (2023). “The cost-accuracy trade-off in operator learning with neural networks”. In: Journal of Machine Learning. Huang, Xiang et al. (2022). “Meta-Auto-Decoder for Solving Parametric Partial Differential Equations”. In: Conference on Neural Information Processing Systems. 128 Jagtap, Ameya D, Kenji Kawaguchi, and George Em Karniadakis (2020). “Locally adaptive activation functions with slope recovery for deep and physics-informed neural networks”. In: Proceedings of the Royal Society A 476.2239, p. 20200334. Jin, Xiaowei et al. (2021). “NSFnets (Navier-Stokes flow nets): Physics-informed neural networks for the incompressible Navier-Stokes equations”. In: Journal of Computational Physics 426, p. 109951. Kashinath, K et al. (2021). “Physics-informed machine learning: case studies for weather and climate modelling”. In: Philosophical Transactions of the Royal Society A 379.2194, p. 20200093. Kochkov, Dmitrii et al. (2021). “Machine learning accelerated computational fluid dynamics”. In: Proceedings of the National Academy of Sciences (PNAS). Kontolati, Katiana et al. (2023). “On the influence of over-parameterization in manifold based surrogates and deep neural operators”. In: Journal of Computational Physics 479, p. 112008. Kovachki, Nikola et al. (n.d.). “Neural operator: Learning maps between function spaces”. In: Journal of Machine Learning Research. volume 24, number 89, pages 1-97. (2023) (). Lanthaler, Samuel et al. (2023). “Nonlinear Reconstruction for Operator Learning of PDEs with Discontinuities”. In: International Conference on Learning Representations. Li, Yingzhou, Jianfeng Lu, and Anqi Mao (2020). “Variational training of neural network approximations of solution maps for physical models”. In: Journal of Computational Physics 409, p. 109338. Li, Zongyi, Daniel Zhengyu Huang, et al. (n.d.). “Fourier neural operator with learned deformations for pdes on general geometries”. In: Journal of Machine Learning Research (388):1-26.(2023) (). Li, Zongyi, Nikola Kovachki, et al. (2021a). “Fourier Neural Operator for Parametric Partial Differential Equations”. In: International Conference on Learning Representations. – (2021b). “Multipole Graph Neural Operator for Parametric Partial Differential Equations”. In: Conference on Neural Information Processing Systems. – (2023). “Markov Neural Operators for Learning Chaotic Systems”. In: Conference on Neural Information Processing Systems. – (n.d.). “Neural operator: Graph kernel network for partial differential equations”. In: Journal of Machine Learning Research. volume 24, number 89, pages 1-97. (2023) (). Liu, Burigede et al. (2022). “A learning-based multiscale method and its application to inelastic impact problems”. In: Journal of the Mechanics and Physics of Solids 158, p. 104668. 129 Lu, Denghui et al. (2021). “86 PFLOPS Deep Potential Molecular Dynamics simulation of 100 million atoms with ab initio accuracy”. In: Computer Physics Communications 259, p. 107624. Lu, Lu, Pengzhan Jin, and George Em Karniadakis (2021). “DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators”. In: Nat Mach Intell 3, 218–229. Lu, Lu, Xuhui Meng, Shengze Cai, et al. (2022). “A comprehensive and fair comparison of two neural operators (with practical extensions) based on FAIR data”. In: Computer Methods in Applied Mechanics and Engineering. Lu, Lu, Xuhui Meng, Zhiping Mao, et al. (2021). “DeepXDE: A deep learning library for solving differential equations”. In: SIAM Review 63.1, pp. 208–228. doi: 10.1137/19M1274067. McClenny, Levi and Ulisses Braga-Neto (2023). “Self-adaptive physics-informed neural networks using a soft attention mechanism”. In: Journal of Computational Physics, 2023. Nelsen, NH and AM Stuart (2021). “The Random Feature Model for Input-Output Maps between Banach Spaces”. In: SIAM Journal on Scientific Computing. Patel, Ravi G et al. (2021). “A physics-informed operator regression framework for extracting data-driven continuum models”. In: Computer Methods in Applied Mechanics and Engineering 373, p. 113500. Pathak, Jaideep, Mustafa Mustafa, et al. (2021). “ML-PDE: A Framework for a Machine Learning Enhanced PDE Solver”. In: Bulletin of the American Physical Society. Pathak, Jaideep, Shashank Subramanian, et al. (2023). “Fourcastnet: A global datadriven high-resolution weather model using adaptive fourier neural operators”. In. Raissi, Maziar, Paris Perdikaris, and George E Karniadakis (2019). “Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations”. In: Journal of Computational Physics 378, pp. 686–707. Raissi, Maziar, Alireza Yazdani, and George Em Karniadakis (2020). “Hidden fluid mechanics: Learning velocity and pressure fields from flow visualizations”. In: Science 367.6481, pp. 1026–1030. Sirignano, Justin and Konstantinos Spiliopoulos (2018). “DGM: A deep learning algorithm for solving partial differential equations”. In: Journal of computational physics 375, pp. 1339–1364. Smith, Jonathan D et al. (2021). “HypoSVI: Hypocenter inversion with Stein variational inference and Physics Informed Neural Networks”. In: AGU Fall Meeting 2021. 130 Sun, Luning et al. (2020). “Surrogate modeling for fluid flows based on physicsconstrained deep learning without simulation data”. In: Computer Methods in Applied Mechanics and Engineering 361, p. 112732. Wang, Rui et al. (2020). “Towards physics-informed deep learning for turbulent flow prediction”. In: Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pp. 1457–1466. Wang, Sifan, Yujun Teng, and Paris Perdikaris (2021). “Understanding and Mitigating Gradient Flow Pathologies in Physics-Informed Neural Networks”. In: SIAM Journal on Scientific Computing 43.5, A3055–A3081. Wang, Sifan, Hanwen Wang, and Paris Perdikaris (2021). “Learning the solution operator of parametric partial differential equations with physics-informed DeepOnets”. In: Science advances. Wang, Sifan, Xinling Yu, and Paris Perdikaris (2022). “When and why pinns fail to train: A neural tangent kernel perspective”. In: Journal of Computational Physics, Volume 449. Weinan, E and Bing Yu (2018). “The deep Ritz method: a deep learning-based numerical algorithm for solving variational problems”. In: Communications in Mathematics and Statistics 6.1, pp. 1–12. Wen, Gege et al. (2023a). “Real-time high-resolution CO 2 geological storage prediction using nested Fourier neural operators”. In: Energy & Environmental Science 16.4, pp. 1732–1741. – (2023b). “Real-time high-resolution CO2 geological storage prediction using nested Fourier neural operators”. In: Energy Environ. Sci., 16, 1732-1741. Yang, Haoyu et al. (2022). “Generic Lithography Modeling with Dual-band OpticsInspired Neural Networks”. In: DAC ’22: Proceedings of the 59th ACM/IEEE Design Automation Conference. Zhang, Lulu et al. (2022). “MOD-Net: A Machine Learning Approach via ModelOperator-Data Network for Solving PDEs”. In: Commun. Comput. Phys. Zhu, Yinhao et al. (2019). “Physics-constrained deep learning for high-dimensional surrogate modeling and uncertainty quantification without labeled data”. In: Journal of Computational Physics 394, pp. 56–81. 131 Chapter 7 LEARNING: SCALE CONSISTENCY LEARNING FOR NEURAL OPERATOR Machine learning (ML) models have emerged as a promising approach for solving partial differential equations (PDEs) in science and engineering. Previous ML models typically cannot generalize outside the training data; for example, a trained ML model for the Navier-Stokes equations only works for a fixed Reynolds number (𝑅𝑒) on a pre-defined domain. To overcome these limitations, we propose a data augmentation scheme based on scale-consistency properties of PDEs and design a scale-informed neural operator that can model a wide range of scales. Our formulation leverages the facts: (i) PDEs can be rescaled, or more concretely, a given domain can be re-scaled to unit size, and the parameters and the boundary conditions of the PDE can be appropriately adjusted to represent the original solution, and (ii) the solution operators on a given domain are consistent on the sub-domains. We leverage these facts to create a scale-consistency loss that encourages matching the solutions evaluated on a given domain and the solution obtained on its subdomain from the rescaled PDE. Since neural operators can fit to multiple scales and resolutions, they are the natural choice for incorporating scale-consistency loss during training of neural PDE solvers. We experiment with scale-consistency loss and the scale-informed neural operator model on the Burgers’ equation, Darcy Flow, Helmholtz equation, and Navier-Stokes equations. With scale-consistency, the model trained on 𝑅𝑒 of 1000 can generalize to 𝑅𝑒 ranging from 250 to 10000, and reduces the error by 38% on average of all datasets compared to baselines. 7.1 Introduction ML for PDEs. Data-driven methods have become increasingly popular in learning Partial Differential Equations (PDEs) for scientific computing (Azizzadenesheli et al., 2024), showing various applications ranging from weather forecasting (Pathak et al., 2022) to nuclear fusion (Gopakumar et al., 2023). While conventional models are typically parameterized for a fixed resolution at a predefined scale, neural operators have recently been proposed to generalize across discretization by parameterizing the model on function spaces (Li, N. Kovachki, et al., 2020b; Nikola B Kovachki et al., 2023; Lu et al., 2021; Kissas et al., 2022; Raonic et al., 2023). Among these, 132 Figure 7.1: Multi-scale PDE dataset: Continuum mechanics at different scales (kilometer- or millimeter-scale) can be formulated to a unit-scaled domain with corresponding scale parameters. Row 1: Darcy Flows, Row 2: Helmholtz Equation, Row 3: Navier Stokes equation. In this work, we aim to design a learning framework to capture the consistency across the scales. the Fourier Neural Operator (FNO) (Li, N. Kovachki, et al., 2020a) stands out as one of the most efficient models. It learns dynamics on the frequency domain, which can be viewed as an efficient, resolution-invariant tokenization. Recent advances further improve the model with shared kernel (Guibas et al., 2021) and U-shape architectures (Rahman, Ross, and Azizzadenesheli, 2022). Given promising results, one of the major challenges of scientific machine learning has been the lack of high-quality training data. Self-supervised learning for PDEs. To overcome the limitation of data, many self-supervised learning techniques have been studied. Especially, the AI for Science community has investigated building-in physics knowledge to the models via equation loss (Li, Zheng, et al., 2021) and symmetry augmentation (R. Wang, Walters, and Yu, 2020; Brandstetter, Welling, and Worrall, 2022). For two-dimensional PDEs, the symmetry groups include translation, rotation, Galilean boost, and scaling. Among them, scale symmetry has been the least effective in improving performance (Brandstetter, Welling, and Worrall, 2022; Mialon et al., 2023). Our hypothesis is that previous formulations of scale-symmetry are defined as positional 133 original coefficient and boundary original wave number 𝑘 = 20 original solution model scale-consistency loss: ∥original solution in sub-domain − sub-domain prediction ∥ sub-domain wave number 𝑘 = 10 sub-domain coefficient and boundary sub-domain prediction Figure 7.2: Scale-consistency loss via sub-domain sampling and re-scaling: given data instance of input coefficient, boundary, scale parameter, and solution we restrict them to a sub-domain to obtain new data instance, which is rescaled to unit length according to the ratio of subdomain size. The scale-consistency loss is defined as the discrepancy between the global and sub-domain prediction. encoding, which does not incorporate scaling parameters and boundary conditions. Multi-scale behavior in physics. Many natural phenomena exhibit multiscale behavior, i.e., interact across a wide range of scales. This is especially the case with solutions of partial differential equations (PDEs), which model various phenomena in science and engineering. For instance, the Navier-Stokes equation, a classical model describing fluid motion, applies to kilometer-scale problems such as weather forecasting (Pathak et al., 2022), meter-scale problems such as airfoils (Li, Nikola Borislavov Kovachki, et al., 2023), and millimeter-scale problems such as catheters (Zhou et al., 2023). PDEs can be rescaled. While the physics at the kilometer and millimeter scales exhibit very different behaviors and frequency ranges, continuum mechanics can be universally reformulated in PDEs using scale parameters, such as the Reynolds number in the Navier-Stokes equation, as illustrated in figure 7.1. Fact 1 (Rescaling of PDEs) In general, a PDE 𝑅 with coefficient function 𝑎 and solution function 𝑢 on domain Ω 𝑅 (𝑎(𝑥), 𝑢(𝑥)) = 0, (𝑥 ∈ Ω) (7.1) 134 can be rescaled to a new domain size Ω𝜆 with scaling 𝜆, 𝑅𝜆 (𝑎(𝜆𝑥), 𝑢(𝜆𝑥)) = 0, (𝑥 ∈ Ω𝜆 ) For example, in the Darcy flow, 𝑅(𝑎, 𝑢) = ∇(𝑎∇𝑢). Further, the solution operators on a given domain are consistent on the sub-domains. Thus, given a domain, the values of the PDE solution in a subdomain can be equivalently obtained by rescaling the subdomain to unit size but by choosing a different set of appropriate parameters (known as scale parameters) and boundary conditions in the PDE. Our approach. Based on the above observation, we define the scale-consistency loss as the overall difference between the original solution of the PDE limited to the subdomain, and the one obtained from the modified PDE upon rescaling the subdomain to unit size, as shown in Figure 7.2. The ground-truth solution has a zero scale-consistency loss, or in other words, solution operators of PDEs are scale consistent. We apply scale-consistency loss as a data augmentation procedure during the training of neural operators for solving PDEs. We address the challenging task of modeling PDEs that exhibit dramatically different behaviors across scales. Our new dataset consists of four PDEs at different scales: the Darcy flow with varying coefficients, the Burgers’ equation with viscosity ranging from 1/100 to 1/1000, the Helmholtz equation with wavenumbers spanning 1 to 100, and the Navier-Stokes equation with Reynolds numbers from 250 to 10000, as shown in Figure 7.1. To evaluate generalization capabilities, we train models at specific scales and test their performance across different scales. In particularly challenging cases, such as the Helmholtz equation, different wavenumbers result in entirely distinct frequency ranges, causing all baseline models to fail at generalization. However, by incorporating our scale-consistency loss, the model successfully achieves zero-shot extrapolation to previously unseen scales, i.e., PDEs with scale parameters not available during training. Our main contributions are as follows. • We propose a data augmentation scheme based on scale-consistency loss that creates data instances with various scales via sub- and super-sampling. For time-dependent problems, we sample in the space-time domain. 135 • We show a theorem (7.3.1) for elliptic PDEs that, under mild assumptions, low scale-consistency loss guarantees recovery of the underlying solution operator. • We design a scale-informed neural operator that takes the scale parameter as input with weight-sharing parameterization and adaptive U-shape architecture to capture a wide range of scales. • We propose a challenging multiscale dataset including the Burgers’ equation, Darcy Flow, Helmholtz equation, and Navier-Stokes equation. The results show that the scale-consistency loss helps the scale-informed neural operator extrapolate to wider scales with a 38% error reduction on average compared to baseline FNO models at the cost of double runtime. To capture a wide range of scales, we propose a new architecture named the scaleinformed neural operator, as shown in Figure 7.3. We use the Fourier neural operator (FNO) (Li, N. Kovachki, et al., 2020a) as the backbone, as FNO naturally handles varying resolution by mapping inputs to the Fourier basis of unit domain size. We incorporate the scale parameter as an additional input and embed the scale features in the frequency space, helping the model capture different frequencies corresponding to different scale parameters. Inspired by (Guibas et al., 2021), we use a weight-sharing parameterization, where a single weight network is shared across all frequency modes. Additionally, it employs a multi-band U-shaped architecture similar to (Rahman, Ross, and Azizzadenesheli, 2022) that optimizes channel dimensions, using larger dimensions for lower frequency bands and smaller dimensions for higher frequency bands. 7.2 Related Work Neural operator and foundation models. Data-driven models have become a common methodology to complement or augment numerical solvers for physical simulation (H. Wang et al., 2023). However, existing data-driven models are typically targeted to a single input variable, such as the coefficient function or initial condition, while other parameters remain fixed, including the domain size, boundary condition, and forcing term (Takamoto et al., 2022). Recently, foundation models have been proposed to capture various datasets under a wide range of conditions, or even multiple families of PDEs (Subramanian et al., 2024; McCabe et al., 2023; Hao et al., 2024; Shen, Marwah, and Talwalkar, 2024; Rahman, George, et al., 136 2024). However, they do not explicitly capture relationships across a wide range of scales seen in physical systems. It is challenging for standard neural networks to capture different scales. In general, separate neural network models are trained for capturing each scale, making it cumbersome to couple them together and impose constraints across scales. Symmetry. Scaling symmetry has been explored as a data augmentation technique in several works (R. Wang, Walters, and Yu, 2020; Brandstetter, Welling, and Worrall, 2022; Mialon et al., 2023). In dynamical systems, this symmetry represents a fundamental relationship between spatial coordinates, time evolution, and field magnitudes. However, both (Brandstetter, Welling, and Worrall, 2022) and (Mialon et al., 2023) reported limited effectiveness of this approach. This limitation may stem from two key challenges: first, continuous scaling symmetry becomes ill-defined on periodic domains without boundaries (Brandstetter, Welling, and Worrall, 2022), and second, scaling velocity magnitudes disrupts the natural range of the input space. This is particularly problematic in applications like weather forecasting, where velocity fields typically maintain consistent magnitude ranges. To address these limitations, we propose a generalized scaling consistenct framework that explicitly incorporates scaling parameters and boundary condition. 7.3 Scale-Consistency Many PDEs possess symmetries, which are reflected by the fact that the equations remain invariant under transformations such as translation, rotation, or re-scaling. An example is the Darcy flow problem on a domain Ω ∈ R𝑑 . ( −∇ · (𝑎(𝑥)∇𝑢(𝑥)) = 0, 𝑢(𝑥) = 𝑔(𝑥), (𝑥 ∈ Ω), (7.2a) (𝑥 ∈ 𝜕Ω). (7.2b) The associated solution operator G is defined as a mapping (𝑎(𝑥), 𝑔(𝑥)) ↦→ G(𝑎, 𝑔)(𝑥) := 𝑢(𝑥). Scale symmetry and scale consistency Re-scale symmetry. Let T𝜆 be the re-scaling operator with 𝜆 ∈ R+ defined by (T𝜆 𝑎)(𝑥) := 𝑎(𝜆𝑥) (or more generally with translation (T𝜆 𝑎)(𝑥) = 𝑎(𝜆𝑥 + 𝑏) with 𝑏 ∈ R𝑑 ). In the absence of boundary conditions, the scale symmetry implies an equivariance property of G: G(T𝜆 𝑎, . . .) = T𝜆 G(𝑎, . . .). 137 The boundary condition (or simply the fact that the PDE is defined on a bounded domain Ω) breaks the scale symmetry; if 𝑢 : Ω → R is defined on the domain Ω, then T𝜆 𝑢 is defined on the rescaled domain Ω𝜆 = {𝜆−1 𝑥|𝑥 ∈ Ω}, and we are generally lacking information about the boundary condition of the re-scaled domain 𝜕Ω𝜆 . Thus, the presence of boundaries in most problems of practical interest makes it difficult to leverage the underlying symmetry properties of the equations in a straightforward way. Nevertheless, under some conditions on the domain Ω (e.g. Ω = [0, 1] 𝑑 is a cube), the formal scale symmetry of the solution operator of (7.2) implies that if 𝑢(𝑥) solves (7.2) with coefficient field 𝑎(𝑥) and with boundary condition 𝑔(𝑥), then the rescaled function 𝑢𝜆 (𝑥) = T𝜆 𝑢(𝑥) = 𝑢(𝜆𝑥), solves ( −∇ · (𝑎𝜆 (𝑥)∇𝑢𝜆 (𝑥)) = 0, (𝑥 ∈ Ω𝜆 ), 𝑢𝜆 (𝑥) = T𝜆 𝑢(𝑥), (𝑥 ∈ 𝜕Ω𝜆 ). i.e. 𝑢𝜆 (𝑥) is a solution of the Darcy flow problem on domain Ω𝜆 , with coefficient field 𝑎𝜆 = T𝜆 𝑎, and boundary condition (T𝜆 𝑢)| 𝜕Ω𝜆 . Another operation we can perform is the restriction from Ω𝜆 to Ω when 𝜆 ≤ 1. Intuitively, this condition expresses the fact that the solution operator of (7.2) is scale-consistent: The solution on a smaller subdomain Ω ⊂ Ω𝜆 can either be obtained leftmirgin=0.5cm by solving the PDE over the entire domain Ω𝜆 and then restricting the solution 𝑢 to the smaller domain 𝑢| Ω . leftmiirgiin=0.5cm by solving the PDE directly on the subdomain Ω, and imposing consistent boundary condition 𝑢| 𝜕Ω . Combining the scale symmetry with restriction, we obtain a new equation (7.2) corresponding to the sub-domain of the original equation (7.3). ( −∇ · (𝑎𝜆 (𝑥)∇𝑢𝜆 (𝑥)) = 0, (𝑥 ∈ Ω), 𝑢𝜆 (𝑥) = T𝜆 𝑢(𝑥), (𝑥 ∈ 𝜕Ω). (7.3a) (7.3b) By uniqueness of the equation, the solution of (7.3) must be consistent with the original solution in (7.2). Lemma 7.3.1 (Scale-consistency (solution function)) If a function 𝑢 satisfies equation (7.2), then 𝑢𝜆 = T𝜆 𝑢 is the unique solution of equation (7.3). 138 Therefore, we obtain the following identity in terms of the solution operator G: let 𝜆≤1 [T𝜆 G(𝑎, 𝑔)]| Ω = G([T𝜆 𝑎]| Ω , [T𝜆 𝑢]| 𝜕Ω ) (7.4) ≡ G([T𝜆 𝑎]| Ω , [T𝜆 G(𝑎, 𝑔)]| 𝜕Ω ). For the solution operator, this identity holds for arbitrary inputs 𝑎(𝑥) and 𝑔(𝑥). The scale-consistency (7.4) can be used as a loss to train solution operators. Informally, if an operator satisfies (7.4), then it must be the target solution operator. Theorem 7.3.1 (Scale-consistency (solution operator)) If an operator Ψ satisfies the scale-consistency (7.4) and it matches the ground truth solution operator G on nearly constant coefficient functions, then Ψ ≡ G. Scale-consistency loss. The first way to impose such a constraint is by introducing a loss of the form 𝐿 (𝑎, 𝑔) = ∥T𝜆 Ψ(𝑎, 𝑔) − Ψ(T𝜆 𝑎, T𝜆 Ψ(𝑎, 𝑔)| 𝜕Ω )∥. (7.5) Note that this is an self-supervised loss term that doesn’t require access to labeled data 𝑢 = G(𝑎, 𝑔). It only requires producing input function samples (𝑎, 𝑔). When solution data 𝑢 is available, the scale-consistency loss simplifies to 𝐿(𝑎, 𝑔) = ∥T𝜆 𝑢 − Ψ(T𝜆 𝑎, T𝜆 𝑢| 𝜕Ω ))∥. (7.6) Infinitesimal scale-consistency. Another way to impose this constraint is by taking the 𝜆-derivative of (7.4), leading to: 𝜕𝜆 T𝜆 G(𝑎, 𝑔) = 𝜕𝜆 [G(T𝜆 𝑎, T𝜆 G(𝑎, 𝑔)| 𝜕Ω )] . We note that if 𝑎(𝑥) is a function, then the derivative 𝜕𝜆 T𝜆 𝑎 evaluated at 𝜆 = 1, is given by 𝜕𝜆 T𝜆 𝑎| 𝜆=1 = [𝜕𝜆 𝑎(𝜆𝑥)] 𝜆=1 = 𝑥 · ∇𝑎(𝑥), i.e., a radial spatial derivative of 𝑎. Substitution of this identity, and noting that T𝜆=1 𝑎 = 𝑎 and T𝜆=1 G(𝑎, 𝑔)| 𝜕Ω = 𝑔, implies that 𝑥 · ∇𝑥 [G(𝑎, 𝑔)] (𝑥)     𝛿G(𝑎, 𝑔) 𝛿G(𝑎, 𝑔) , 𝑥 · ∇𝑥 𝑎 + , 𝑥 · ∇𝑥 [G(𝑎, 𝑔)] = 𝛿𝑎 𝛿𝑔 We observe that while (7.4) is highly non-linear, the infinitesimal constraint is quadratic in G. 139 Scale-dependent problem: extension beyond scale symmetry The scale-consistency constraint can be written in greater generality, even if the underlying PDE has no scale symmetry. In this case, the domain could be an input to the operator, and the relevant scale-consistency would be G(𝑎, 𝑔; Ω)| Ω′ = G(𝑎| Ω′ , G(𝑎, 𝑔, Ω)| 𝜕Ω′ ; Ω′), (Ω′ ⊂ Ω). In some cases, this is equivalent to scaling certain parameters in the PDE, as explained below. Helmholtz equation. An example not satisfying scale symmetry is the Helmholtz equation, −∇ · (𝑎(𝑥)∇𝑢(𝑥)) + 𝑘 2 𝑢(𝑥) = 𝑓 (𝑥). (7.7) In this case, a rescaling of the spatial variable corresponds to a rescaling of the frequency 𝑘 2 , i.e. 𝑢𝜆 (𝑥) = 𝑢(𝜆𝑥) solves −∇ · (𝑎𝜆 (𝑥)∇𝑢𝜆 (𝑥)) + 𝜆−2 𝑘 2 𝑢𝜆 (𝑥) = 𝜆−2 𝑓 (𝜆𝑥), or −∇ · (𝑎𝜆 (𝑥)∇𝑢𝜆 (𝑥)) + 𝑘 𝜆2 𝑢𝜆 (𝑥) = 𝑓𝜆 (𝑥), with 𝑘 𝜆 := 𝜆−1 𝑘, 𝑓𝜆 (𝑥) := 𝜆−2 𝑓 (𝜆𝑥). Thus, the scale-consistency constraint involves the whole family of PDEs, Δ𝑢 + 𝑘 2 𝑢 = 𝑓 , for 𝑘 > 0, with the transform on parameter T𝜆 (𝑘) = 𝜆𝑘. Time-dependent problem: rescale in space-time domain For time-dependent problems, in general, we could view the time dimension as another spatial dimension, and rescale both the spatial and temporal dimensions. Navier-Stokes equation. Another example is the two-dimensional incompressible Navier-Stokes equation. In the velocity form, without forcing, 𝜕𝑡 𝑢(𝑥, 𝑡) + 𝑢(𝑥, 𝑡) · ∇𝑢(𝑥, 𝑡) = 1 ∇𝑝(𝑥, 𝑡) + 𝜈Δ𝑢(𝑥, 𝑡), 𝜌 The scaling is by 𝑢𝜆 (𝑥, 𝑡) = 𝑢(𝜆𝑥, 𝜆𝑡), 𝑝 𝜆 (𝑥, 𝑡) = 𝑝(𝜆𝑥, 𝜆𝑡), and 𝜈𝜆 := 𝜆−1 𝜈. In the vorticity formulation where 𝜔 = curl(𝑢), we do not need to rescale the time. 𝜕𝑡 𝜔(𝑥, 𝑡) + 𝑢(𝑥, 𝑡) · ∇𝜔(𝑥, 𝑡) = 𝜈Δ𝜔(𝑥, 𝑡), (7.8) Rescaling the spatial variable 𝑥 corresponds to rescaling the viscosity 𝜈; 𝜔𝜆 (𝑥, 𝑡) = 𝜔(𝜆𝑥, 𝑡) and 𝑢𝜆 (𝑥, 𝑡) = 𝜆−1 𝑢(𝜆𝑥, 𝑡) solves 𝜕𝑡 𝜔𝜆 (𝑥, 𝑡) + 𝑢𝜆 (𝑥, 𝑡) · ∇𝜔𝜆 (𝑥, 𝑡) = 𝜈𝜆 Δ𝜔𝜆 (𝑥, 𝑡), 140 where 𝜈𝜆 := 𝜆−2 𝜈, here the coefficient 𝜆 in front of the term (𝑢𝜆 (𝑥, 𝑡) · ∇𝜔𝜆 (𝑥, 𝑡)) is absorbed by 𝑢𝜆 . Main Algorithms Remark: neural operator automatically rescales input to unit length. For standard neural networks such as convolution neural networks, re-scaling T needs to be implemented as interpolation. However, in the design of neural operators such as FNO, the domain size is implicitly re-scaled to unit length, where the Fourier basis is defined with length [0, 1]. Given T𝜆 𝑓 defined on domain [0, 𝜆], Fourier neural operator Ψ automatically rescales it to unit length, Ψ(T𝜆 𝑓 , . . .) := Ψ(T1/𝜆 T𝜆 𝑓 , . . .) = Ψ( 𝑓 , . . .). where 𝑓 is defined on unit size [0, 1]. Therefore, the re-scaling T is omitted in the algorithm. Sub-domain sampling. The sub-domain sampling algorithm is based on equation (7.6), where we use sub-sampling (i.e., restrict to sub-domain) to obtain instance with smaller scale 𝜆𝑘 < 𝑘. Given the input and output data {(𝑎, 𝑔, 𝑘), 𝑢} defined on domain Ω, we truncate the domain into a smaller sub-domain Ω̂. The input and output restriction to the sub-domain, along with the re-scaled parameter, become a ˆ 𝑢}. new data instance {( 𝑎, ˆ 𝑔, ˆ 𝑘), ˆ We compute the consistency loss as the difference ˆ and the restricted output between the model evaluated on restricted input Ψ( 𝑎, ˆ 𝑔, ˆ 𝑘) 𝑢. ˆ Algorithm 1 Sub-domain sampling 1: Input: data tuple of coefficient, boundary, scale parameter, and solution {(𝑎, 𝑔, 𝑘), 𝑢} on domain Ω = [0, 1] 2 , model Ψ, and sampling rate 𝜆 < 1. 2: Sample the sub-domain Ω̂ = [𝑤, 𝑤 +𝜆] × [ℎ, ℎ +𝜆], where 𝑤, ℎ ∼ Unif[0, 1−𝜆]. 3: Define new instance ( 𝑎ˆ = 𝑎| Ω̂ , 𝑔ˆ = 𝑢| 𝜕 Ω̂ , 𝑘ˆ = 𝜆𝑘), 𝑢ˆ = 𝑢| Ω̂ . ˆ − 𝑢∥. 4: Output: scale-consistency loss ∥Ψ( 𝑎, ˆ 𝑔, ˆ 𝑘) ˆ Super-domain sampling. The super-domain sampling algorithm is based on equation (7.5), where we sample new instances corresponding to larger scale 𝜆𝑘 > 𝑘. Given the distributions 𝜇 for 𝑎 and 𝜈 for 𝑔, we can sample new instance 𝑎, 𝑔 with larger scale 𝜆𝑘 and apply Algorithm 1. Different from 1, we do not have the ground truth output 𝑢 on the larger scale. Instead, we estimate using the model 𝑢 = Ψ(𝑎, 𝑔, 𝜆𝑘). 141 Algorithm 2 Super-domain sampling 1: Input: distributions of inputs coefficient and boundary 𝜇, 𝜈, model Ψ, scale parameter 𝑘, and sampling rate 𝜆 > 1. 2: Sample new instances 𝑎 ∼ 𝜇, 𝑔 ∼ 𝜈. 3: Define new scale as 𝜆𝑘. 4: Estimate the solution of new domain 𝑢 = Ψ(𝑎, 𝑔, 𝜆𝑘). 5: Call Algorithm 1 with input {(𝑎, 𝑔, 𝜆𝑘), 𝑢} and scale 1/𝜆. 6: Output: scale-consistency loss ∥Ψ(𝑎| Ω̂ , Ψ(𝑎, 𝑔, 𝜆𝑘)| 𝜕 Ω̂ , 𝑘) − Ψ(𝑎, 𝑔, 𝜆𝑘)| Ω̂ ∥. 7.4 Scale-Informed Neural Operator The scale-informed neural operator is based on the FNO Li, N. Kovachki, et al., 2020a, where convolution is implemented as a pointwise multiplication in the Fourier space. Since FNO automatically rescales its input to unit length, we design a scale embedding in the Fourier space to inform the model of the scale parameter 𝑘. Furthermore, we design a U-shaped architecture to optimize the channel dimension. Embed scale parameters in Fourier Space In the previous FNO, the weight tensor 𝑅 is defined as a (𝑀1 × · · · × 𝑀𝑑 ×𝐶𝑖𝑛 ×𝐶𝑜𝑢𝑡 )tensor, which is sufficient for lower-dimensional problems with fewer total modes 𝑀. For larger-scale problems, such as highly turbulent flows, the weight tensor 𝑅 becomes prohibitively large. Therefore, we propose an implicit representation of the weight tensor similar to AFNO Guibas et al., 2021, where the complex weight 𝑅 with the shape (𝐶𝑖𝑛 × 𝐶𝑜𝑢𝑡 ) is shared across all modes (𝑀1 × · · · × 𝑀𝑑 ). Different from AFNO, we further define the features of scale 𝑘 and mode index 𝜉 as input, so that the transform 𝑅 can behave correspondingly with respect to different scales 𝑘 and modes 𝜉. Let 𝐶 be the embedding channel dimension; we define scale features as ℎ(𝑘)𝑖 = 𝑘 𝑖/(𝐶−1) for 𝑖 = 0, 1, . . . , 𝐶 − 1, which covers a wide range from 𝑘 0/(𝐶−1) = 1 to 𝑘 (𝐶−1)/(𝐶−1) = 𝑘. The input 𝑓𝑡 (𝜉) ∈ C𝐶𝑖𝑛 is first element-wise multiplied with the features of the scale parameter and wavenumber ℎ(𝑘, 𝜉), and then multiplied with 𝑅, followed by a group normalization and a complex activation 𝜎. The transform K can be viewed as a kernel function defined on the Fourier space:
(K 𝑓𝑡+1 )(𝜉) = 𝜎 𝑅( 𝑓𝑡 (𝜉) ⊙ ℎ(𝑘, 𝜉)) . (7.9) 142 a(x) P Re PRe Spatial R1 Spatial R2 W W 𝐹 𝐹 -1 𝐹 Spectral R1 Q u(x) W 𝐹 -1 𝐹 … Spectral R2 skips 2 skips 1 Fourier Domain Spatial Rt … 𝐹 -1 Spectral Rt skips t skip_in1 R11 skip_out1 Re, k features skip_in2 R12 R21 R22 skip_out2 skip_in3 R23 R32 skip_out3 Figure 7.3: The scale-informed neural operator has a U-shape structure on the Fourier space. The scale parameter (such as 𝑅𝑒) are embedded at each spectral layer. In the down block, the input tensors are truncated and lifted by complex layer 𝑅; in the up block, the tensors are projected and added to the inputs. Skip connections are added across the blocks. 𝑃 is the encoder and 𝑄 is the decoder. Multi-band Architecture The Fourier signal usually follows an ordered structure, where the energy decays as the wavenumber increases. Therefore, previous methods such as FNO Li, N. Kovachki, et al., 2020a and SNO Fanaskov and Oseledets, 2022 choose to truncate to a fixed number of frequencies by omitting higher frequencies. Similar to previous works such as UNet Ronneberger, Fischer, and Brox, 2015 and UNO, we design a multi-band structure to gradually shrink the frequency bands, as shown in Figure 7.3. Different from UNO, which applies spectral convolutions at each down and up block, in this work, we define the U-shaped structure fully in the Fourier space. Given the initial channel dimension 𝐶, maximum input modes 𝑀, and a predefined number of levels 𝐿, we define 𝐶𝑙 and 𝑀𝑙 as 𝐶𝑙 = 2𝑙 𝐶 and 𝑀𝑙 = 2−𝑙 𝑀, where each block has shape 𝐶𝑙2 𝑀𝑙𝑑 . For 𝑑 = 2, 𝐶𝑙2 𝑀𝑙2 = 𝐶 2 𝑀 2 , so each level has the same size. We define the first level using the weight-sharing formulation, where 𝑅1 has the shape (𝐶𝑖𝑛 × 𝐶𝑜𝑢𝑡 ), and higher levels in tensor formulation with (𝑀𝑙𝑑 × 𝐶𝑖𝑛 × 𝐶𝑜𝑢𝑡 ). 143 Boundary condition For boundary value problems, we take the boundary as an additional input. For a 1dimensional boundary on a 2-dimensional square domain, we extend the boundary to 2D by repeating along the other dimension. For Dirichlet-type boundaries, it is known that the boundary is the restriction of the solution, and their magnitudes should be similar. Therefore, we define a normalization at the end of the model that multiplies the output by the magnitude of the boundary. Darcy Flow (Scale) 2 3 4 (training) 8 16 FNO FNO+scale UNet UNet+scale UNO UNO+scale 3.921 1.990 6.638 5.130 5.534 3.009 3.842 1.932 6.981 4.945 5.336 2.753 3.737 1.990 6.011 5.869 4.725 3.087 3.323 2.130 5.527 5.645 4.495 4.602 3.214 2.300 6.361 6.094 4.366 4.600 — — — — — — Burgers’ Equation (𝜈) 1/100 1/200 1/400 (training) 1/1000 FNO FNO+scale SINO SINO+scale UNet UNet+scale UNO UNO+scale 28.602 27.799 24.914 5.926 32.897 30.137 28.581 28.716 11.005 10.008 10.027 1.720 22.463 22.815 10.963 11.009 1.230 1.908 1.174 0.957 20.119 25.138 1.235 1.387 8.709 9.442 8.363 4.575 26.481 30.747 8.624 8.720 — — — — — — — — — — — — — — — — 1 2 5 (training) 10 (training) 25 (training) 50 136.847 44.625 69.437 8.960 164.945 51.441 120.742 125.742 131.200 36.026 63.283 6.960 156.775 64.313 101.478 91.541 4.285 3.186 3.666 3.081 48.341 63.827 9.350 10.821 12.575 11.924 12.503 12.490 51.028 52.731 16.172 19.605 21.060 19.744 19.728 19.001 21.189 53.541 32.280 36.017 107.186 108.916 102.980 112.940 112.914 104.477 118.570 117.776 Navier-Stokes, auto-reg. (Re) 250 500 1000 (training) 2000 4000 10000 FNO FNO+scale SINO-U SINO-U + scale UNet UNet+scale UNO UNO+scale 0.447 0.302 0.695 0.357 4.156 1.086 4.228 4.005 0.750 0.531 0.782 0.514 2.706 1.753 3.021 2.661 1.015 0.743 0.976 0.953 0.809 13.802 4.147 3.458 3.108 2.446 2.466 2.186 2.096 15.427 8.316 6.941 7.374 6.137 4.793 4.289 10.027 16.442 16.728 14.785 18.295 17.127 13.772 11.483 22.284 28.297 33.221 30.663 Helmholtz Equation (k) FNO FNO+scale SINO-U SINO-U+scale UNet UNet+scale UNO UNO+scale Table 7.1: Comparison of FNO, UNet, UNO, and SINO with and without scaleconsistency. Models are trained at certain scale and zero-shot test across others. Overall, scale-consistency helps each model extrapolate to unseen scales. Errors are in relative-L2 (%). The Darcy Flow is scale-invariant so the SINO does not apply. 144 Kolmogorov Flow with RE=10000 0.10 FNO Ours relative l2 error 0.08 0.06 0.069 0.063 0.062 0.060 0.061 0.046 0.04 0.036 0.030 0.02 0.00 S=128 S=256 S=512 resolution (training & testing) S=1024 Figure 7.4: Ablation Studies of Scale Consistency. left: Cost-Accuracy: we train and test each model at various sizes on Kolmogorov Flow with RE=5000. Our model (u-shape) converges faster than baseline models. Further, the model (shared) achieves comparative accuracy with 1/10 of the parameters. right: discretization convergence: the proposed model does not truncate to a fixed bandwidth. As the training resolution increases, the model’s error converges while the baseline FNO remains the same. 7.5 Experiments We generated datasets for the Darcy Flow, Helmholtz equation, and Navier-Stokes equation, each spanning a wide range of scales. For each test case, we trained the models on a narrow range of scales and compared the performance with and without self-consistency augmentation. All experiments were run on Nvidia A100 (80GB, 40GB) or P100 (16GB) GPUs. The error metric is relative L2 error. The results show that self-consistency augmentation helps the model generalize better to unseen scales. Self-consistency loss In the first part, we compare FNO, UNet, and our models, with and without the selfconsistency loss. For Darcy and Helmholtz equations, where the input distribution is given as a Gaussian random field, we apply both sub-sampling 1 and super-sampling 2. For the Navier-Stokes equation, the input distribution is unknown, so we only apply sub-sampling. Darcy Flow. We considered the Darcy Flow (7.2) with a non-zero Dirichlet boundary. The input coefficient is sampled at different scale. The resolutions were 𝑠 = 64, 96, 128, 256, 512, respectively. We train 1024 instances for training and 128 for testing. We used 𝜎 = 1 for training. Since Darcy has no scale parameters, we used FNO with and without scale-consistency. As shown in Table 7.1, FNO with 145 scale-consistency reduced the error by half compared to the baseline. Helmholtz Equation. We considered the Helmholtz equation (7.7) with a non-zero Dirichlet boundary. The input coefficients 𝑎, 𝑔 were sampled from a fixed Gaussian random field, with varying wavenumbers 𝑘 = 1, 2, 5, 10, 25, 50, 100. The resolutions were 64, 64, 64, 128, 256, 512, 1024, respectively. We train 1024 instances for training and 128 for testing. We used 𝑘 = 5, 10, 25 for training. The scale-informed neural operator with scale-consistency reduced the error by half compared to the baseline FNO on smaller wavenumbers 𝑘 = 1, 2, but neither model captured larger scales 𝑘 = 50, 100, since Helmholtz equation has very different behaviors on larger scales. Burgers’ Equation. We considered the Burgers’ equation. Given the initial condition and time-dependent boundary condition as input, the model predicts the solution over the next time interval. We train the FNO model and the multi-scale neural operator model (ours) with and without scale-consistency loss. The scale-consistent loss is across both the spatial and temporal domain. The models are trained on viscosity = 1/400 and tested on viscosities 𝜈 = 1/100, 1/200, 1/400, 1/1000. The multiscale neural operator with scale-consistency reduced the error up to 5x compared to the baseline FNO on unseen viscosity. Navier-Stokes Equation (autoregressive). We considered the Navier-Stokes equation (7.8) defined on sub-domain similar to applications in climate. The input is the vorticity field of the previous ten time frames 𝜔0 . We considered Reynolds numbers ranging from 𝑅𝑒 = 250, 500, 1000, 2000, 4000, 10000. The resolutions were 32, 64, 128, 128, 256, 512, respectively. We train 50 trajectories for training and 5 (per each 𝑅𝑒) for testing, where each trajectory consists of 300 time steps, with 𝑑𝑡 = 0.1. We used 𝑅𝑒 = 1000 for training. The multi-scale multi-band neural operator with scale-consistency reduced the error by 1/4 compared to the baseline UNet on unseen 𝑅𝑒 = 250, 500, 4000, 10000. Navier-Stokes Equation (space-time, 2+1 dimensional). We also considered the spatiotemporal modeling for the Navier Stokes equation, velocity formulation. Similar to the autoregressive setting above, we considered Reynolds numbers ranging from 𝑅𝑒 = 250, 500, 1000, 2000, 4000, 10000. For continuous-time modeling, we use 𝑑𝑡 = 1/256 and are given input of the history, consisting of 24 frames, to predict the next 24 frames. We observe significant improvements with scale embedding and spatiotemporal cropping for out-of-distribution Reynolds numbers. Improvements are highlighted in Table 7.1. 146 Scale 2 3 4 8 16 No aug. Rot. Ref. Rot.+Ref. 4.143 3.101 2.821 2.713 4.193 2.944 2.701 2.469 4.036 2.953 2.597 2.450 3.552 2.797 2.616 2.461 3.352 2.872 2.684 2.582 Scale (ours) All (ours) 1.918 1.903 2.075 1.816 2.035 1.910 2.159 2.237 2.095 2.309 Table 7.2: Comparison of scale-consistency with existing symmetries for data augmentation, in relative-L2 error (%). We train FNO on Darcy flow at 𝑠𝑐𝑎𝑙𝑒 = 4 and zero-shot test at other scales. Comparison with symmetry-based augmentation. On the Darcy flow, we compare the scale-consistency augmentation with existing symmetry-based augmentation as used in (R. Wang, Walters, and Yu, 2020; Brandstetter, Welling, and Worrall, 2022). As shown in Table 7.2, scale-consistency augmentation leads to better generalization compared to rotation plus reflection. Furthermore, scale-consistency works seamlessly with rotation and reflection. The best result is achieved by combining the three augmentation methods together. Ablation studies on the model architecture Embed scale parameter in the frequency space. We conducted several ablation studies on the proposed model architecture along with scale-consistency loss. We test the scale embedding and positional embedding on the frequency space (7.9) with Burgers’ equation, Helmholtz equation, and Navier-Stokes equation. As shown in Table 7.3, the embedding in general improves the performance. We sometimes find the scale parameter is unnecessary in the Navier-Stokes equation when it can be inferred from the history of trajectory. U-shape structure and shared kernel. We further conducted ablation studies on the U-shape structure model in the standard supervised learning setting on periodic Navier-Stokes equation with fixed scales 𝑅𝑒 = 5000 (with forcing) and 𝑅𝑒 = 10000 (zero forcing) as in (Li, Liu-Schiaffini, et al., 2022). For baselines, we consider FNO (Li, N. Kovachki, et al., 2020a), UNet (Ronneberger, Fischer, and Brox, 2015), FNO-UNet (Gupta and Brandstetter, 2022), and UNO (Rahman, Ross, and Azizzadenesheli, 2022). The results show that our model achieves a smaller error rate with one-tenth of the parameters compared to the previous FNO at the cost of longer runtime, as shown in Figure 7.4 (left). Since the model does not truncate 147 Burgers’ equation Scale Freq. Emb. 𝜈 = 1/100 𝜈 = 1/200 𝜈 = 1/400 𝜈 = 1/800 No No Yes Yes No Yes No Yes 28.531 28.660 10.731 6.334 10.756 10.832 2.540 1.887 1.087 0.916 1.055 1.042 8.889 8.725 5.477 4.636 — — — — — — — — k = 50 Helmholtz equation Scale Freq. Emb. k=1 k=2 k=5 k = 10 k = 25 No No Yes Yes No Yes No Yes 17.914 15.642 16.741 9.438 5.963 6.384 4.822 4.980 3.537 3.441 2.914 2.921 11.042 10.294 10.631 9.874 16.338 106.597 14.056 102.015 12.989 103.151 11.574 93.938 Table 7.3: Ablation for scale-informed neural operator on different equations in relative-L2 error (%). For Burgers’ equation, we train on viscosity 𝜈 = 1/400 and zero-shot test on other scales. For Helmholtz equation, we train on wavenumber 𝑘 = 5, 10, 25. Model Scale Informed Freq. Aug. Emb. size min Re=250 256 Re=500 256 Re=1000 512 Re=2000 512 Re=4000 1024 Re=10000 1024 2+1 dim FNO No No Yes Yes No Yes No Yes OFF OFF OFF OFF N/A N/A N/A N/A 2.040 1.727 1.937 1.756 2.901 2.632 2.779 2.551 4.460 4.051 4.194 4.003 8.573 8.158 8.319 8.029 12.081 11.603 11.889 11.274 19.554 18.847 19.588 18.551 2+1 dim FNO Yes Yes Yes Yes Yes Yes Yes Yes ON ON ON ON 24 32 40 48 1.352 1.342 1.468 1.756 2.082 2.016 2.031 2.515 3.547 3.382 3.348 3.869 7.420 7.285 7.083 7.732 11.074 10.876 10.444 11.194 18.526 18.469 17.692 18.408 2+1 dim SINO No No Yes Yes No Yes No Yes OFF OFF OFF OFF N/A N/A N/A N/A 3.945 3.586 4.032 2.125 5.430 4.348 5.580 2.661 6.768 2.827 6.803 2.701 11.215 6.307 11.358 6.164 14.987 15.211 16.708 11.917 21.862 23.422 23.437 19.405 2+1 dim SINO No No Yes Yes No Yes No Yes ON ON ON ON 32 32 32 32 4.083 1.419 6.584 1.750 5.681 2.157 6.874 2.457 6.516 2.917 6.802 2.863 10.959 6.323 13.861 6.271 15.138 9.880 22.819 13.394 22.287 17.095 35.919 23.217 Table 7.4: Navier-Stokes equation trained on RE1000, zero-shot test on various RE (2+1 dimensional models). Values in percentage (%). the maximum Fourier frequency, its accuracy improves as the resolution refines, as shown in Figure 7.4 (bottom right). 7.6 Discussion and Conclusion In this paper, we consider the scale consistency for learning solution operators on PDEs across various scales. By leveraging the scale-consistency properties of PDEs and designing a scale-informed neural operator, we demonstrated the ability to model 148 a wide range of scales. Experimental results showed significant improvements in generalization to unseen scales, with better generalization errors compared to baseline models. This approach holds promise for improving the efficiency and generalizability of data-driven PDE solvers, reducing the need for extensive training data, and enabling the development of more flexible and foundational models for scientific and engineering applications. While sub-sampling (Algorithm 1) is generally helpful, super-sampling (Algorithm 2) requires input distribution known to sample new instances. While the supersampling works well for Darcy and Burgers, it is challenging to subsample from the attractor for the Navier-Stokes equation. As a potential future direction, it could be an interesting direction to combine with generative models (Lim et al., 2023) to sample virtual inputs. References Azizzadenesheli, Kamyar et al. (2024). “Neural operators for accelerating scientific simulations and design”. In: Nature Reviews Physics, pp. 1–9. Brandstetter, Johannes, Max Welling, and Daniel E Worrall (2022). “Lie point symmetry data augmentation for neural pde solvers”. In: International Conference on Machine Learning. PMLR, pp. 2241–2256. Fanaskov, Vladimir and Ivan Oseledets (2022). “Spectral neural operators”. In: arXiv preprint arXiv:2205.10573. Gopakumar, Vignesh et al. (2023). “Fourier Neural Operator for Plasma Modelling”. In: arXiv preprint arXiv:2302.06542. Guibas, John et al. (2021). “Adaptive fourier neural operators: Efficient token mixers for transformers”. In: arXiv preprint arXiv:2111.13587. Gupta, Jayesh K and Johannes Brandstetter (2022). “Towards multi-spatiotemporalscale generalized pde modeling”. In: arXiv preprint arXiv:2209.15616. Hao, Zhongkai et al. (2024). “DPOT: Auto-Regressive Denoising Operator Transformer for Large-Scale PDE Pre-Training”. In: arXiv preprint arXiv:2403.03542. Kissas, Georgios et al. (2022). “Learning operators with coupled attention”. In: Journal of Machine Learning Research 23.215, pp. 1–63. Kovachki, Nikola B et al. (2023). “Neural Operator: Learning Maps Between Function Spaces With Applications to PDEs.” In: J. Mach. Learn. Res. 24.89, pp. 1– 97. Li, Zongyi, Nikola Kovachki, et al. (2020a). “Fourier Neural Operator for Parametric Partial Differential Equations”. In: arXiv preprint arXiv:2010.08895. 149 Li, Zongyi, Nikola Kovachki, et al. (2020b). “Neural operator: Graph kernel network for partial differential equations”. In: arXiv preprint arXiv:2003.03485. Li, Zongyi, Nikola Borislavov Kovachki, et al. (2023). “Geometry-Informed Neural Operator for Large-Scale 3D PDEs”. In: arXiv preprint arXiv:2309.00583. Li, Zongyi, Miguel Liu-Schiaffini, et al. (2022). “Learning chaotic dynamics in dissipative systems”. In: Advances in Neural Information Processing Systems 35, pp. 16768–16781. Li, Zongyi, Hongkai Zheng, et al. (2021). “Physics-informed neural operator for learning partial differential equations”. In: ACM/JMS Journal of Data Science. Lim, Jae Hyun et al. (2023). “Score-based diffusion models in function space”. In: arXiv preprint arXiv:2302.07400. Lu, Lu et al. (Mar. 2021). “Learning nonlinear operators via DeepONet based on the universal approximation theorem of operators”. In: Nature Machine Intelligence 3.3, pp. 218–229. McCabe, Michael et al. (2023). “Multiple physics pretraining for physical surrogate models”. In: arXiv preprint arXiv:2310.02994. Mialon, Grégoire et al. (2023). “Self-supervised learning with lie symmetries for partial differential equations”. In: Advances in Neural Information Processing Systems 36, pp. 28973–29004. Pathak, Jaideep et al. (2022). “Fourcastnet: A global data-driven high-resolution weather model using adaptive fourier neural operators”. In: arXiv preprint arXiv:2202.11214. Rahman, Md Ashiqur, Robert Joseph George, et al. (2024). “Pretraining Codomain Attention Neural Operators for Solving Multiphysics PDEs”. In: arXiv preprint arXiv:2403.12553. Rahman, Md Ashiqur, Zachary E Ross, and Kamyar Azizzadenesheli (2022). “U-no: U-shaped neural operators”. In: arXiv preprint arXiv:2204.11127. Raonic, Bogdan et al. (2023). “Convolutional neural operators”. In: ICLR 2023 Workshop on Physics for Machine Learning. Ronneberger, Olaf, Philipp Fischer, and Thomas Brox (2015). “U-net: Convolutional networks for biomedical image segmentation”. In: Medical Image Computing and Computer-Assisted Intervention–MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18. Springer, pp. 234– 241. Shen, Junhong, Tanya Marwah, and Ameet Talwalkar (2024). “Ups: Towards foundation models for pde solving via cross-modal adaptation”. In: arXiv preprint arXiv:2403.07187. Subramanian, Shashank et al. (2024). “Towards foundation models for scientific machine learning: Characterizing scaling and transfer behavior”. In: Advances in Neural Information Processing Systems 36. 150 Takamoto, Makoto et al. (2022). “PDEBench: An extensive benchmark for scientific machine learning”. In: Advances in Neural Information Processing Systems 35, pp. 1596–1611. Wang, Hanchen et al. (2023). “Scientific discovery in the age of artificial intelligence”. In: Nature 620.7972, pp. 47–60. Wang, Rui, Robin Walters, and Rose Yu (2020). “Incorporating symmetry into deep dynamics models for improved generalization”. In: arXiv preprint arXiv:2002.03061. Zhou, Tingtao et al. (2023). “AI-aided Geometric Design of Anti-infection Catheters”. In: arXiv preprint arXiv:2304.14554. 151 Chapter 8 GEOMETRY: NEURAL OPERATOR WITH LATENT SPACE We propose the geometry-informed neural operator (GINO), a highly efficient approach for learning the solution operator of large-scale partial differential equations with varying geometries. GINO uses a signed distance function (SDF) and pointcloud representations of the input shape and neural operators based on graph and Fourier architectures to learn the solution operator. The graph neural operator handles irregular grids and transforms them into and from regular latent grids on which Fourier neural operator can be efficiently applied. GINO is discretizationconvergent, meaning the trained model can be applied to arbitrary discretizations of the continuous domain and it converges to the continuum operator as the discretization is refined. To empirically validate the performance of our method on large-scale simulation, we generate the industry-standard aerodynamics dataset of 3D vehicle geometries with Reynolds numbers as high as five million. For this large-scale 3D fluid simulation, numerical methods are expensive to compute surface pressure. We successfully trained GINO to predict the pressure on car surfaces using only five hundred data points. The cost-accuracy experiments show a 26, 000× speed-up compared to optimized GPU-based computational fluid dynamics (CFD) simulators on computing the drag coefficient. When tested on new combinations of geometries and boundary conditions (inlet velocities), GINO obtains a one-fourth reduction in error rate compared to deep neural network approaches. 8.1 Introduction Computational sciences aim to understand natural phenomena and develop computational models to study the physical world around us. Many natural phenomena follow the first principles of physics and are often described as evolution on function spaces, governed by partial differential equations (PDE). Various numerical methods, including finite difference and finite element methods, have been developed as computational approaches for solving PDEs. However, these methods need to be run at very high resolutions to capture detailed physics, which are time-consuming and expensive, and often beyond the available computation capacity. For instance, in computational fluid dynamics (CFD), given a shape design, the goal is to solve the Navier-Stokes equation and estimate physical properties such as pressure and 152 Input Samples discretize GNO FNO GNO Operates on input geometry Operates on (latent) regular geometry Projects to output geometry query Prediction FNO layers -134. -48.3 37.2 123. The input geometries are irregular and change for each sample. These are discretized into point clouds and passed on to a GNO layer, which maps from the given geometry to a latent regular grid. The output of this GNO layer is concatenated with the SDF features and passed into an FNO model. The output from the FNO model is projected back onto the domain of the input geometry for each query point using another GNO layer. This is used to predict the target function (e.g., pressure), which is used to compute the loss that is optimized end-to-end for training. Figure 8.1: The architecture of GINO velocity. Finding the optimal shape design often requires solving thousands of trial shapes, each of which can take more than ten hours even with GPUs (Korzun et al., 2022). To overcome these computational challenges, recent works propose deep learningbased methods, particularly neural operators (Zongyi Li, Kovachki, Azizzadenesheli, B. Liu, Bhattacharya, et al., 2020b), to speed up the simulation and inverse design. Neural operators generalize neural networks and learn operators, which are mappings between infinite-dimensional function spaces (Zongyi Li, Kovachki, Azizzadenesheli, B. Liu, Bhattacharya, et al., 2020b). Neural operators are discretization convergent and can approximate general operators (Kovachki et al., 2021). The input function to neural operators can be presented at any discretization, grid, resolution, or mesh, and the output function can be evaluated at any arbitrary point. Neural operators have shown promise in learning solution operators in partial differential equations (PDE) (Kovachki et al., 2021) with numerous applications in scientific computing, including weather forecasting (Pathak et al., 2022), carbon dioxide storage and reservoir engineering (Wen et al., 2023), with a tremendous speedup over traditional methods. Prior works on neural operators developed a series of principled neural operator architectures to tackle a variety of scientific computing applications. Among the neural operators, graph neural operators (GNO) (Zongyi Li, Kovachki, Azizzadenesheli, B. Liu, Bhattacharya, et al., 2020b), and Fourier neural operators (FNO) (Zongyi Li, Kovachki, Azizzadenesheli, B. Liu, Bhattacharya, et al., 2020a) have been popular in various applications. 208. 153 GNO implements kernel integration with graph structures and is applicable to complex geometries and irregular grids. The kernel integration in GNO shares similarities with the message-passing implementation of graph neural networks (GNN) (P. Battaglia et al., 2016), which is also used in scientific computing (Sanchez-Gonzalez et al., 2020; Pfaff et al., 2020; Allen et al., 2022). However, the main difference is that GNO defines the graph connection in a ball defined on the physical space, while GNN typically assumes a fixed set of neighbors, e.g., k-nearest neighbors, see Figure 8.2. Such nearest-neighbor connectivity in GNN violates discretization convergence, and it degenerates into a pointwise operator at high resolutions, leading to a poor approximation of the ground-truth operator using GNN. In contrast, GNO adapts the graph based on points within a physical space, allowing for universal approximation of operators. However, one limitation of graph-based methods is the computational complexity when applied to problems with long-range global interactions. To overcome this, prior works propose using multi-pole methods or multi-level graphs (Zongyi Li, Kovachki, Azizzadenesheli, B. Liu, Stuart, et al., 2020; Lam et al., 2022) to help with global connectivity. However, they do not fully alleviate the problem since they require many such levels to capture global dependence, which still makes them expensive. While GNO performs kernel integration in the physical space using graph operations, FNO leverages Fourier transform to represent the kernel integration in the spectral domain using Fourier modes. This architecture is applicable to general geometries and domains since the (continuous) Fourier transform can be defined on any domain. However, it becomes computationally efficient when applied to regular input grids since the continuous Fourier transform can then be efficiently approximated using discrete Fast Fourier transform (FFT) (Cooley and Tukey, 1965), giving FNO a significant quasi-linear computational complexity. However, FFT limits FNO to regular grids and cannot directly deal with complex geometries and irregular grids. A recent model, termed GeoFNO, learns a deformation from a given geometry to a latent regular grid (Zongyi Li, D. Z. Huang, et al., 2022) so that the FFT can be applied in the latent space. In order to transform the latent regular grid back to the irregular physical domain, discrete Fourier transform (DFT) on irregular grids is employed. However, DFT on irregular grids is more expensive than FFT, quadratic vs. quasi-linear, and does not approximate the Fourier transform in a discretization convergent manner. This is because, unlike in the regular setting, the points are not sampled at regular intervals, and therefore the integral does not take into account the underlying measure. Other attempts share a similar computational barrier as 154 shown in Table 8.1, which we discussed in Section 8.2. In this paper, we consider learning the solution operator for large-scale PDEs, in particular, 3D CFD simulations. We propose the geometry-informed neural operator (GINO), a neural operator architecture for arbitrary geometries and mesh discretizations. It uses a signed distance function (SDF) to represent the geometry and composes GNO and FNO architectures together in a principled manner to exploit the strengths of both frameworks. The GNO by itself can handle irregular grids through graphs but is able to operate only locally under a limited computational budget, while the FNO can capture global interactions, but requires a regular grid. By using GNO to transform the irregular grid into a regular one for the FNO block, we can get the best of both worlds, i.e., computational efficiency and accuracy of the approximation. Thus, this architecture tackles the issue of expensive global integration operations that were unaddressed in prior works, while maintaining discretization convergence. Specifically, GINO has three main components, (𝑖) Geometry encoder: multiple local kernel integration layers through GNO with graph operations, (𝑖𝑖) Global model: a sequence of FNO layers for global kernel integration, and (𝑖𝑖𝑖) Geometry decoder: the final kernel integral layers, as shown in Figure 8.1. The input to the GINO is the input surface (as a point cloud) along with the SDF, representing the distance of each 3D point to the surface. GINO is trained end-to-end to predict output (e.g., car surface pressure in our experiments), a function defined on the geometry surfaces. Geometry encoder: the first component in the GINO architecture uses the surface (i.e., point cloud) and SDF features as inputs. The irregular grid representation of the surface is encoded through local kernel integration layers implemented with GNOs, consisting of local graphs that can handle different geometries and irregular grids. The encoded function is evaluated on a regular grid, which is concatenated with the SDF input evaluated on the same grid. Global model: the output of the first component is encoded on a regular grid, enabling efficient learning with an FNO using FFT. Our second component consists of multiple FNO layers for efficient global integration. In practice, we find that this step can be performed at a lower resolution without significantly impacting accuracy, giving a further computational advantage. Geometry decoder: the final component is composed of local GNObased layers with graph operations, that decode the output of the FNO and project it back onto the desired geometry, making it possible to efficiently query the output 155 Model Range Complexity Irreg. grid Disc. conv. GNN local 𝑂 (𝑁degree) Yes No CNN local 𝑂 (𝑁) No No UNet global 𝑂 (𝑁) No No 2 Transformer global 𝑂 (𝑁 ) Yes Yes GNO (kernel) radius 𝑟 𝑂 (𝑁degree) Yes Yes FNO (FFT) global 𝑂 (𝑁 log 𝑁) No Yes GINO [Ours] global 𝑂 (𝑁 log 𝑁 + 𝑁degree) Yes Yes 𝑁 is the number of mesh points; 𝑑 is the dimension of the domain and degree is the maximum degree of the graph. Even though GNO and transformer both work on irregular grids and are discretization convergent, they become too expensive on large-scale problems. Table 8.1: Computational complexity of standard deep learning models function on irregular meshes. The GNO layers in our framework are accelerated using our GPU-based hash-table implementation of neighborhood search for graph connectivity of meshes. We validate our findings on two large-scale 3D CFD datasets. We generate our own large-scale industry-standard Ahmed’s body geometries using GPU-based OpenFOAM (Jasak, Jemcov, Tukovic, et al., 2007), composed of 500+ car geometries with 𝑂 (105 ) mesh points on the surface and 𝑂 (107 ) mesh points in space. Each simulation takes 7-19 hours on 16 CPU cores and 2 Nvidia V100 GPUs. Further, we also study a lower resolution dataset with more realistic car shapes, viz., Shape-Net car geometries generated by Umetani and Bickel, 2018. GINO takes the point clouds and SDF features as the input and predicts the pressure fields on the surfaces of the vehicles. We perform a full cost-accuracy trade-off analysis. The result shows GINO is 26, 000× faster at computing the drag coefficients over the GPU-based OpenFOAM solver, while achieving 8.31% (Ahmed-body) and 7.29% (Shape-Net car) error rates on the full pressure field. Further, GINO is capable of zero-shot super-resolution, training with only one-eighth of the mesh points, and having a good accuracy when evaluated on the full mesh that is not seen during training. 8.2 Related Work The study of neural operators and their extended applications in learning solution operators in PDE has been gaining momentum (Kovachki et al., 2021; Bhattacharya et al., 2021; Kissas et al., 2022; Rahman, Ross, and Azizzadenesheli, 2022). A method that stands out is FNO, which uses Fourier transform (Zongyi Li, Kovachki, 156 Azizzadenesheli, B. Liu, Bhattacharya, et al., 2020a). The FNO and its variations have proven to highly accelerate the simulations for large-scale flow problems, including weather forecasting (Pathak et al., 2022), seismology (Yang et al., 2021; Sun et al., 2022) and multi-phase flow (Wen et al., 2023). However, a challenge with the FNO is that its computation superiority is gained when applied on a regular grid, where the Fourier transform is approximated using FFT. Therefore, its reliance on FFT limits its use with irregular grids or complex geometries. There have been attempts to modify the FNO to work with these irregular structures, but scalability to large-scale 3D PDEs remains an issue. One such attempt is GeoFNO, which learns a coordinate transformation to map irregular inputs to a regular latent space (Zongyi Li, D. Z. Huang, et al., 2022). This method, while innovative, requires a geometric discrete Fourier transform, which is computationally demanding and lacks discretization insurance. To circumvent this, GINO limits the Fourier transform to a local GNO to improve efficiency. The locality is defined assuming the metrics of the physical space. Additionally, the Non-Equispaced Fourier neural solvers (NFS) merge the FNO with non-equispaced interpolation layers, a method similar to global GNO (Lin et al., 2022). However, at the architecture level, their method replaces the integration of GNO with the summation of the nearest neighbor points on the graph. This step transitions this method to a neural network, failing to deliver a discretization convergent approach. The Domain-Agnostic Fourier Neural Operators (DAFNO) represents another attempt at improvement, applying an FNO to inputs where the geometry is represented as an indicator function (N. Liu, Jafarzadeh, and Yu, 2023). However, this method lacks a strategy for handling irregular point clouds. Simultaneously, researchers are exploring the combination of FNO with the attention mechanisms (Kovachki et al., 2021) for irregular meshes. This includes the Operator Transformer (OFormer) (Zijie Li, Meidani, and Farimani, 2022), Mesh-Independent Neural Operator (MINO) (Lee, n.d.), and the General Neural Operator Transformer (GNOT) (Hao et al., 2023). Besides, the Clifford neural layers (Brandstetter et al., 2022) use the Clifford algebra to compute multivectors, which provides CliffordFNO implementations as an extension of FNO. The work (Wandel, Weinmann, and Klein, 2021) uses a physics-informed loss with U-Net for 3D channel flow simulation with several shapes. The work (Balu et al., 2021) innovatively proposes the use of multigrid training for neural networks that improves the convergence. Although these methods incorporate attention layers, which are special types of kernel integration Kovachki et al., 2021 with quadratic complexity, they face challenges when 157 scaling up for large-scale problems. GNNs are incorporated in the prior attempts in physical simulations involving complex geometry, primarily due to the inherent flexibility of graph structures. Early research (P. Battaglia et al., 2016; Kipf and Welling, 2016; Hamilton, Ying, and Leskovec, 2017; P. W. Battaglia et al., 2018) laid the foundation for GNNs, demonstrating that physical entities, when represented as graph nodes, and their interactions, as edges, could predict the dynamics of various systems. The introduction of graph element networks (Alet et al., 2019) marked a significant development, being the first to apply GNNs to PDEs by discretizing the domain into elements. The similar idea has been explored in term of continuous 3D point cloud convolution (Y. Li et al., 2018; Hermosilla et al., 2018; W. Wu, Qi, and Fuxin, 2019). Another line of work, mesh graph networks (Sanchez-Gonzalez et al., 2020; Pfaff et al., 2020; Allen et al., 2022), further explored PDEs in the context of fluid and solid mechanics. Remelli et al., 2020; Durasov et al., 2021 train a Graph convolutional neural works on the ShapeNet car dataset for inverse design. However, GNN architectures’ limitations hinder their use in operator learning for PDEs. GNNs connect each node to its nearest neighbors according to the graph’s metrics, not the metrics of the physical domain. As the input function’s discretization becomes finer, each node’s nearest neighbors eventually converge to the same node, contradicting the expectation of improved model performance with finer discretization. Furthermore, GNNs’ model behavior at the continuous function limit lacks a unique definition, failing the discretization convergence criterion. Consequently, as pointwise operators in function spaces at the continuous limit, GNNs struggle to approximate general operators between function spaces, Figure 8.2. 8.3 Problem setting We are interested in learning the map from the geometry of a PDE to its solution. We will first give a general framework and then discuss the Navier-Stokes equation in CFD as an example. Let 𝐷 ⊂ R𝑑 be a Lipschitz domain and A a Banach space of real-valued functions on 𝐷. We consider the set of distance functions T ⊂ A so that, for each function 𝑇 ∈ T , its zero set 𝑆𝑇 = {𝑥 ∈ 𝐷 : 𝑇 (𝑥) = 0} defines a (𝑑 − 1)-dimensional sub-manifold. We assume 𝑆𝑇 is simply connected, closed, smooth, and that there exists 𝜖 > 0 such that 𝐵𝜖 (𝑥) ∩ 𝜕𝐷 = ∅ for every 𝑥 ∈ 𝑆𝑇 and 𝑇 ∈ T . We denote by 𝑄𝑇 ⊂ 𝐷, the open volume enclosed by the sub-manifold 𝑆𝑇 and assume that 𝑄𝑇 is a Lipschitz domain with 𝜕𝑄𝑇 = 𝑆𝑇 . We define the Lipschitz domain Ω𝑇 := 𝐷 \ 𝑄¯ 𝑇 so that, 𝜕Ω𝑇 = 𝜕𝐷 ∪ 𝑆𝑇 . Let L denote a partial differential 158 Input geometries Discretized geometries at various fine to coarse levels of discretizations Coarser discretization discretization … (a) An input geometry (continuous function) is first discretized into a series of points by subsampling it. Note that in practice, the discretization can be highly irregular. A key challenge with several scientific computing applications is that we want a method that can work on arbitrary geometries, but also that is discretization convergent, meaning that the method converges to a desired solution operator as we make the discretization finer. Finer discretization Finer discretization Target Geometry Target Geometry (b) GNN connects each point in the latent subspace (red) to its nearest neighbors in the original space (top). This is very discretization dependent, and as we increase the resolution (sample points more densely), the method becomes increasingly local and fails to capture context. In addition, the operator at the discretization limit is non-unique and depends on how the discretization is done. (c) GNO instead connects each point in the latent subspace (red) to all its neighbors within an epsilon ball in the original space (top). This induces convergence to a continuum solution operator as we increase the resolution (sample points more densely). This means GNO converges to a unique operator as the discretization becomes finer and scales to large problems. Figure 8.2: Comparison of GNN and GNO as the discretization becomes finer. GNN is discretization dependent, while GNO is discretization convergent. operator and consider the problem L (𝑢) = 𝑓 , in Ω𝑇 , 𝑢 = 𝑔, in 𝜕Ω𝑇 , (8.1) for some 𝑓 ∈ F , 𝑔 ∈ B where B, F denote Banach spaces of functions on R𝑑 with the assumption that the evaluation functional is continuous in B. We assume that L is such that, for any triplet (𝑇, 𝑓 , 𝑔), the PDE (8.1) has a unique solution 𝑢 ∈ U𝑇 where U𝑇 denotes a Banach space of functions on Ω𝑇 . Let U denote a Banach space of functions on 𝐷 and let {𝐸𝑇 : U𝑇 → U : 𝑇 ∈ T } be a family of extension operators which are linear and bounded. We define the mapping from the distance 159 function to the solution function Ψ:T ×F ×B →U (8.2) by (𝑇, 𝑓 , 𝑔) ↦→ 𝐸𝑇 (𝑢) which is our operator of interest. Navier-Stokes Equation. We illustrate the above abstract formulation with the ¯ We take following example. Let 𝐷 = (0, 1) 𝑑 be the unit cube and let A = 𝐶 ( 𝐷). T ⊂ A to be some subset such that the zero level set of every element defines a (𝑑 − 1)-dimensional closed surface which can be realized as the graph of a Lipschitz function and that there exists 𝜖 > 0 such that each surface is at least distance 𝜖 away from the boundary of 𝐷. We now consider the steady Naiver-Stokes equations, −𝜈Δ𝑣 + (𝑣 · ∇)𝑣 + ∇𝑝 = 𝑓 , in Ω𝑇 , ∇ · 𝑣 = 0, in Ω𝑇 , 𝑣 = 𝑞, in 𝜕𝐷, 𝑣 = 0, in 𝑆𝑇 , (8.3) where 𝑣 : Ω𝑇 → R𝑑 is the velocity, 𝑝 : Ω𝑇 → R is the pressure, 𝜈 is the viscosity, and 𝑓 , 𝑞 : R𝑑 → R𝑑 are the forcing and boundary functions. The condition that 𝑣 = 0 in 𝑆𝑇 is commonly known as a no slip boundary and is prevalent in many engineering applications. The function 𝑞, on the other hand, defines the inlet and outlet boundary conditions for the flow. We assume that 𝑓 ∈ 𝐻 −1 (R𝑑 ; R𝑑 ) and 𝑞 ∈ 𝐶 (R𝑑 ; R𝑑 ). We can then define our boundary function 𝑔 ∈ 𝐶 (R𝑑 ; R𝑑 ) such that 𝑔(𝑥) = 0 for any 𝑥 ∈ 𝐷 with dist(𝑥, 𝜕𝐷) ≥ 𝜖 and 𝑔(𝑥) = 𝑞(𝑥) for any 𝑥 ∈ 𝐷 with, dist(𝑥, 𝜕𝐷) > 𝜖/2 as well as any 𝑥 ∉ 𝐷. Continuity of 𝑔 can be ensured by an appropriate extension for any 𝑥 ∈ 𝐷 such that dist(𝑥, 𝜕𝐷) < 𝜖 and dist(𝑥, 𝜕𝐷) ≥ 𝜖/2 Whitney, 1934. We define 𝑢 : Ω𝑇 → R𝑑+1 by 𝑢 = (𝑣, 𝑝) as the unique weak solution of (8.3) with U𝑇 = 𝐻 1 (Ω𝑇 ; R𝑑 ) × 𝐿 2 (Ω𝑇 )/R Temam, 1984. We define U = 𝐻 1 (𝐷; R𝑑 ) × 𝐿 2 (𝐷)/R and the family of extension operators {𝐸𝑇 :
𝑝 U𝑇 → U} by 𝐸𝑇 (𝑢) = 𝐸𝑇𝑣 (𝑣), 𝐸𝑇 ( 𝑝) where 𝐸𝑇𝑣 : 𝐻 1 (Ω𝑇 ; R𝑑 ) → 𝐻 1 (𝐷; R𝑑 ) 𝑝 and 𝐸𝑇 : 𝐿 2 (Ω𝑇 )/R → 𝐿 2 (𝐷)/R are defined as the restriction onto 𝐷 of the extension operators defined in Stein, 1970, Chapter 6, Theorem 5. This establishes the existence of the operator Ψ : T × 𝐻 −1 (R𝑑 ; R𝑑 ) × 𝐶 (R𝑑 ; R𝑑 ) → 𝐻 1 (𝐷; R𝑑 ) × 𝐿 2 (𝐷)/R mapping the geometry, forcing, and boundary condition to the (extended) solution of the steady Navier-Stokes equation (8.3). Homomorphic extensions of deformation-based operators have been shown in Cohen, Schwab, and Zech, 2018. We leave for future work studying the regularity properties of the presently defined operator. 160 8.4 Geometric-Informed Neural Operator We propose a geometry-informed neural operator (GINO), a neural operator architecture for varying geometries and mesh regularities. GINO is a deep neural operator model consisting of three main components, (𝑖) multiple local kernel integration layers, (𝑖𝑖) a sequence of FNO layers for global kernel integration which precedes (𝑖𝑖𝑖) the final kernel integral layers. Each layer of GINO follows the form of generic kernel integral of the form (8.5). Local integration is computed using graphs, while global integration is done in Fourier space. Neural operator A neural operator Ψ Kovachki et al., 2021 maps the input functions 𝑎 = (𝑇, 𝑓 , 𝑔) to the solution function 𝑢. The neural operator Ψ is composed of multiple layers of point-wise and integral operators, Ψ = Q ◦ K𝐿 ◦ . . . ◦ K1 ◦ ¶. (8.4) The first layer ¶ is a pointwise operator parameterized by a neural network. It transforms the input function 𝑎 into a higher-dimensional latent space ¶ : 𝑎 ↦→ 𝑣 0 . Similarly, the last layer acts as a projection layer, which is a pointwise operator Q : 𝑣 𝑙 ↦→ 𝑢, parameterized by a neural network 𝑄. The model consists of 𝐿 layers of integral operators K𝑙 : 𝑣 𝑙−1 ↦→ 𝑣 𝑙 in between. ∫ 𝑣 𝑙 (𝑥) = 𝜅 𝑙 (𝑥, 𝑦)𝑣 𝑙−1 (𝑦)d𝑦 (8.5) 𝐷 where 𝜅 𝑙 is a learnable kernel function. Non-linear activation functions are incorporated between each layer. Graph operator block To efficiently compute the integral in equation (8.5), we truncate the integral to a local ball at 𝑥 with radius 𝑟 > 0, as done in (Zongyi Li, Kovachki, Azizzadenesheli, B. Liu, Bhattacharya, et al., 2020b), ∫ 𝑣 𝑙 (𝑥) = 𝜅(𝑥, 𝑦)𝑣 𝑙−1 (𝑦) d𝑦. (8.6) 𝐵𝑟 (𝑥) We discretize the space and use a Riemann sum to compute the integral. This process involves uniformly sampling the input mesh points and connecting them with a graph for efficient parallel computation. Specifically, for each point 𝑥 ∈ 𝐷, we randomly sample points {𝑦 1 , . . . , 𝑦 𝑀 } ⊂ 𝐵𝑟 (𝑥) and approximate equation (8.6) 161 as 𝑣 𝑙 (𝑥) ≈ 𝑀 ∑︁ 𝜅(𝑥, 𝑦𝑖 )𝑣 𝑙−1 (𝑦𝑖 )𝜇(𝑦𝑖 ), (8.7) 𝑖=1 where 𝜇 denotes the Riemannian sum weights corresponding to the ambient space of 𝐵𝑟 (𝑥). For a fixed input mesh of 𝑁 points, the computational cost of equation (8.7) scales with the number of edges, denoted as 𝑂 (𝐸) = 𝑂 (𝑀 𝑁). Here, the number of sampling points 𝑀 is the degree of the graph. It can be either fixed to a constant sampling size, or scale with the area of the ball. Encoder. Given an input point cloud {𝑥 1in , . . . , 𝑥 in 𝑁 } ⊂ 𝑆𝑇 , we employ a GNOgrid grid encoder to transform it to a function on a uniform latent grid {𝑥 1 , . . . , 𝑥 𝑆 } ⊂ 𝐷. The encoder is computed as discretization of an integral operator 𝑣 0 (𝑥 grid ) ≈ Í𝑀 grid , 𝑦 in )𝜇(𝑦 in ) over ball 𝐵 (𝑥 grid ). To inform the grid density, GINO 𝑟 in 𝑖=1 𝜅(𝑥 𝑖 𝑖 computes Riemannian sum weights 𝜇(𝑦𝑖in ). Further, we use Fourier features in the kernel Sitzmann et al., 2020. For simple geometries, this encoder can be omitted, see Section 8.5. grid grid Decoder. Similarly, given a function defined on the uniform latent grid {𝑥 1 , . . . , 𝑥 𝑆 } ⊂ 𝐷, we use a GNO-decoder to query arbitrary output points {𝑥 1out , . . . , 𝑥 out 𝑁 } ⊂ Ω𝑇 . Í𝑀 grid grid grid out out The output is evaluated as 𝑢(𝑥 ) ≈ 𝑖=1 𝜅(𝑥 , 𝑦𝑖 )𝑣 𝑙 (𝑦𝑖 )𝜇(𝑦𝑖 ) over ball grid 𝐵𝑟out (𝑥 out ). Here, the Riemannian weight, 𝜇(𝑦𝑖 ) = 1/𝑆 since we choose the latent space to be regular grid. Since the queries are independent, we divide the output points into small batches and run them in parallel, which enables us to use much larger models by saving memory. Efficient graph construction. The graph construction requires finding neighbors to each node that are within a certain radius. The simplest solution is to compute all possible distances between neighbors, which requires 𝑂 (𝑁 2 ) computation and memory. However, as the 𝑁 gets larger, e.g., 10 ∼ 100 million, computation and memory become prohibitive even on modern GPUs. Instead, we use a hash gridbased implementation to efficiently prune candidates that are outside of a ℓ ∞ -ball first and then compute the ℓ 2 distance between only the candidates that survive. This reduces the computational complexity to 𝑂 (𝑁 𝑑𝑟 3 ) where 𝑑 denotes unit density and 𝑟 is the radius. This can be efficiently done using first creating a hash table of voxels with size 𝑟. Then, for each node, we go over all immediate neighbors to the current voxel that the current node falls into and compute the distance between all points 162 in these neighboring voxels. Specifically, we use the CUDA implementation from Open3D Zhou, Park, and Koltun, 2018. Then, using the neighbors, we compute the kernel integration using gather-scatter operations from torch-scatter Fey et al., 2023. Further, if the degree of the graph gets larger, we can add Nyström approximation by sampling nodes (Zongyi Li, Kovachki, Azizzadenesheli, B. Liu, Bhattacharya, et al., 2020b). Fourier operator block The geometry encoding 𝑣 0 and the geometry specifying map 𝑇, both evaluated on a regular grid discretizing 𝐷 are passed to a FNO block. We describe the basic FNO block as first outlined in Zongyi Li, Kovachki, Azizzadenesheli, B. Liu, Bhattacharya, et al., 2020a. We will first define global convolution in the Fourier space and use it to build the full FNO operator block. To that end, we will work on the 𝑑-dimensional unit torus T𝑑 . We define an integral operator with kernel 𝜅 ∈ 𝐿 2 (T𝑑 ; R𝑛×𝑚 ) as the mapping C : 𝐿 2 (T𝑑 ; R𝑚 ) → 𝐿 2 (T𝑑 ; R𝑛 ) given by
C(𝑣) = F −1 F (𝜅) · F (𝑣) , ∀ 𝑣 ∈ 𝐿 2 (T𝑑 ; R𝑚 ) Here F , F −1 are the Fourier transform and its inverse respectively, defined for 𝐿 2 by the appropriate limiting procedure. The Fourier transform of the function 𝜅 will be parameterized directly by some fixed number of Fourier modes, denoted 𝛼 ∈ N. In particular, we assume ∑︁ 𝜅(𝑥) = 𝑐 𝛾 e𝑖⟨𝛾,𝑥⟩ , ∀ 𝑥 ∈ T𝑑 𝛾∈𝐼 for some index set 𝐼 ⊂ Z𝑑 with |𝐼 | = 𝛼 and coefficients 𝑐 𝛾 ∈ C𝑛×𝑚 . Then we may view F : 𝐿 2 (T𝑑 ; R𝑛×𝑚 ) → ℓ 2 (Z𝑑 ; C𝑛×𝑚 ) so that F (𝜅)(𝛾) = 𝑐 𝛾 if 𝛾 ∈ 𝐼 and F (𝜅)(𝛾) = 0 if 𝛾 ∉ 𝐼. We directly learn the coefficients 𝑐 𝛾 without ever having to evaluate 𝜅 in physical space. We then define the full operator block K : 𝐿 2 (T𝑑 ; R𝑚 ) → 𝐿 2 (T𝑑 ; R𝑛 ) by
K (𝑣)(𝑥) = 𝜎 𝑊𝑣(𝑥) + C(𝑣) , ∀ 𝑥 ∈ T𝑑 where 𝜎 is a pointwise non-linearity and 𝑊 ∈ R𝑛×𝑚 is a learnable matrix. We further modify the layer by learning the kernel coefficients in tensorized form, adding skip connections, normalization layers, and learnable activations as outlined in Kossaifi et al., 2023. We refer the reader to this work for further details. 163 Ground-truth pressure Predicted pressure -197. -96.7 3.54 104. 204. -2.82e+03 -1.90e+03 -991. -79.0 833. Relative error 0.00 0.200 0.400 0.600 0.800 0.00 0.370 0.739 1.11 1.48 Figure 8.3: Visualization of a ground-truth pressure and corresponding prediction by GINO from the Shape-Net Car (top) and Ahmed-body (bottom) datasets, as well as the absolute error. Adaptive instance normalization. For many engineering problems of interest, the boundary information is a fixed, scalar, inlet velocity specified on some portion of 𝜕𝐷. In order to efficiently incorporate this scalar information into our architecture, we use a learnable adaptive instance normalization X. Huang and Belongie, 2017 combined with a Fourier feature embedding Sitzmann et al., 2020. In particular, the scalar velocity is embedded into a vector with Fourier features. This vector then goes through a learnable MLP, which outputs the scale and shift parameters of an instance normalization layer Ulyanov, Vedaldi, and Lempitsky, 2016. In problems where the velocity information is not fixed, we replace the normalization layers of the FNO blocks with this adaptive normalization. We find this technique improves performance, since the magnitude of the output fields usually strongly depends on the magnitude of the inlet velocity. 8.5 Experiments We explore a range of models on two CFD datasets. The large-scale Ahmed-Body dataset, which we generated, and also the Shape-Net Car dataset from Umetani and Bickel, 2018. Both datasets contain simulations of the Reynold-Averaged NavierStokes (RANS) equations for a chosen turbulence model. The goal is to estimate the full pressure field given the shape of the vehicle as input. We consider GNO (Zongyi Li, Kovachki, Azizzadenesheli, B. Liu, Bhattacharya, et al., 2020b), MeshGraphNet (Pfaff et al., 2020), GeoFNO (Zongyi Li, D. Z. Huang, et al., 2022), 3D UNet (Wolny 164 Model latent resolution radius training error test error GINO 32 0.055 14.11% 13.59% GINO 48 0.055 8.99% 10.20% 64 0.055 6.00% 8.47% GINO GINO 80 0.055 5.77% 7.87% GINO 32 0.110 8.66% 10.10% GINO 48 0.073 7.25% 9.17% 64 0.055 6.00% 8.47% GINO GINO 80 0.044 6.22% 7.89% When fixing the radius, larger latent resolutions lead to better performance. The gaps become smaller when fixing the number of edges and scaling the radius correspondingly. Table 8.2: Ablation on the Ahmed-body with different sizes of the latent space et al., 2020) with linear interpolation, FNO (Zongyi Li, Kovachki, Azizzadenesheli, B. Liu, Bhattacharya, et al., 2020a), and GINO. We train each model for 100 epochs with Adam optimizer and step learning rate scheduler. The implementation details can be found in the Appendix. All models run on a single Nvidia V100 GPU. Ahmed-Body dataset We generate the industry-level vehicle aerodynamics simulation based on the Ahmedbody shapes (Ahmed, Ramm, and Faltin, 1984). The shapes are parameterized with six design parameters: length, width, height, ground clearance, slant angle, and fillet radius. We also vary the inlet velocity from 10m/s to 70m/s, leading to Reynolds numbers ranging from 4.35 × 105 to 6.82 × 106 . We use the GPU-accelerated OpenFOAM solver for steady state simulation using the SST 𝑘 − 𝜔 turbulence model (Menter, 1993) with 7.2 million mesh points in total with 100k mesh points on the surface. Each simulation takes 7-19 hours on 2 Nvidia v100 GPUs with 16 CPU cores. We generate 551 shapes in total and divide them into 500 for training and 51 for validation. Shape-Net Car dataset We also consider the Car dataset generated by Umetani and Bickel, 2018. The input shapes are from the ShapeNet Car category (Chang et al., 2015). In Umetani and Bickel, 2018, the shapes are manually modified to remove the side mirrors, spoilers, and tires. The RANS equations with the 𝑘 − 𝜖 turbulence model and SUPG stabilization are simulated to obtain the time-averaged velocity and pressure fields using a finite element solver (Zienkiewicz, Taylor, and Nithiarasu, 2013). The inlet 165 Model training error test error GNO 18.16% 18.77% Geo-FNO (sphere) 10.79% 15.85% UNet (interp) 12.48% 12.83% FNO (interp) 9.65% 9.42% GINO (encoder-decoder) 7.95% 9.47% GINO (decoder) 6.37% 7.12% We do a benchmark study with several standard machine-learning methods on the Shape-Net and Ahmed body datasets. The training error is normalized L2 error; the test error is de-normalized L2. Table 8.3: Shape-Net Car dataset (3.7k mesh points). velocity is fixed at 20m/s (72km/h) and the estimated Reynolds number is 5 × 106 . Each simulation takes approximately 50 minutes. The car surfaces are stored with 3.7k mesh points. We take the 611 water-tight shapes out of the 889 instances, and divide the 611 instances into 500 for training and 111 for validation. As shown in Table 8.3 8.4 and Figure 8.3, GINO achieves the best error rate with a large margin compared with previous methods. On the Ahmed-body dataset, GINO achieves 8.31% while the previous best method achieve 11.16%. On the Shape-Net Car, GINO achieves 7.12% error rate compared to 9.42% on FNO. It takes 0.1 seconds to evaluate, which is 100,000x faster than the GPU-parallel OpenFOAM solver that take 10 hours to generates the data. We further performance a full cost-accuracy analysis in the following section. For ablations, we consider channel dimensions [32, 48, 64, 80], latent space [32, 48, 64, 80], and radius from 0.025 to 0.055 (with the domain size normalized to [-1, 1]). As depicted in Figure 8.5a and Table 8.2, larger latent spaces and radii yield superior performance. Cost-accuracy analysis on drag coefficient To compare the performance of our model against the industry-standard OpenFOAM solver, we perform a full cost-accuracy trade-off analysis on computing the drag coefficient, which is the standard design objective of vehicles and airfoils. The result shows GINO is 26,000x faster at computing the drag coefficients. Figure 8.4b below shows the cost-accuracy curve, measured in terms of inference time needed for a relative error in the drag coefficient for GINO and OpenFOAM. The detailed setup is discussed in the appendix. 166 Model training error test error MeshGraphNet 9.08% 13.88% UNet (interp) 9.93% 11.16% FNO (interp) 12.97% 12.59% GINO (encoder-decoder) 9.36% 9.01%% GINO (decoder) 9.34% 8.31% Previous works such as GNO and Geo-FNO cannot scale to large meshes with 100k points. We instead add the MeshGraphNet for graph comparison. Again, for UNet, FNO, and GINO, we fix the latent grid to 64 × 64 × 64. The training error is normalized L2; the test error is de-normalized L2. Table 8.4: Ahmed-body dataset (100k mesh points). (b) Cost-accuracy trade-off curve for (a) Example drag coefficient computed OpenFOAM vs. GINO. Increasing cost after every iteration of the OpenFOAM of GINO is achieved by increasing the solver. The reference drag corresponds size of the latent space. Loss function for to the last point of the filtered signal. Tri- models on the red line includes only error angle indicates the time when the solver in pressure and wall shear stress, while, reaches a 3% relative error with respect for models on orange line, the loss also to the reference drag. includes drag. Figure 8.4: Cost-accuracy trade-off analysis for the drag coefficient. Discretization-convergence and ablation studies We investigate discretization-convergence by varying different parts of GINO. Specifically, we vary the latent grid resolution and the sampling rates for inputoutput meshes. In these experiments, we fixed the training and test samples to be the same, i.e., same latent grid resolution or sampling rate, but varied the shape and input conditions. Discretization-convergence wrt the latent grid. Here, each model is trained and tested on (the same) latent resolutions, specifically 32, 48, 64, 80, and 88, and the architecture is the same. As depicted in Figure 8.5a, GINO demonstrates a comparable error rate across all resolutions. A minor improvement in errors is 167 (b) Varying the sampling (a) Varying resolutions of the rates of the input-output mesh (c) Train with a low sampling latent grid (same resolution (same rate training and test- rate and test on full mesh for training and testing). ing). (zero-shot super-resolution). Figure 8.5: Discretization-convergence studies and zero-shot super-resolution. observed when employing a larger latent space. Conversely, the errors associated with the UNet model grow as the resolution is decreased due to the decreasing receptive field of its local convolution kernels. Discretization-convergence in the input-output mesh. Here, GINO is trained and tested with sub-sampled input-output meshes at various sampling rates (2x, 4x, 6x, 8x). As illustrated in Figure 8.5b, GINO exhibits a consistent error rate across all sampling rates. A slight increase in errors is observed on coarser meshes. Zero-shot super-resolution. GINO possesses the ability to perform zero-shot super-resolution. The model is trained on a coarse dataset, sub-sampled by 2x, 4x, 6x, and 8x, and subsequently tested on the full mesh, that is not seen during training. The error remains consistent across all sampling rates 8.5c. This characteristic enables the model to be trained at a coarse resolution when the mesh is dense, consequently reducing the computational requirements. Drag Coefficient Comparison For many engineering tasks, the goal is often to determine a single quantity of interest from a simulation which can then be used within an overall design process. In the design of automobiles, a sought after quantity is the drag coefficient of the vehicle. Intuitively, it is a number inversely proportional to the efficiency with which a vehicle passes through a fluid. Engineers are therefore often interested in designing geometries with minimal drag coefficients. For a fluid with unit density, the drag coefficient is defined as ∫  ∫
2 𝑝(𝑥) 𝑛(𝑥) ˆ · 𝑖ˆ(𝑥) 𝑑𝑥 + 𝑇𝑤 (𝑥) · 𝑖ˆ(𝑥) 𝑑𝑥 (8.8) 𝑐𝑑 = 2 𝑣 𝐴 𝜕Ω 𝜕Ω 168 where 𝜕Ω ⊂ R3 is the surface of the car, 𝑝 : R3 → R is the pressure, 𝑛ˆ : 𝜕Ω → S2 is the outward unit normal vector of the car surface, 𝑖ˆ : R3 → S2 is the unit direction of the inlet flow, 𝑇𝑤 : 𝜕Ω → R3 is the wall shear stress on the surface of the car, 𝑣 ∈ R is the speed of the inlet flow, and 𝐴 ∈ R is the area of the smallest rectangle enclosing the front of the car. For our Ahmed-body dataset, we train several GINO models (decoder) to predict the pressure on the surface of the car as well as the wall shear stress. Since the inlet flow is always parallel to the 𝑥-axis i.e. 𝑖ˆ(𝑥) = (−1, 0, 0), we predict only the first component of the wall shear stress. Since our results from Table 8.2 indicate that varying the size of the latent space has a significant effect on predictive performance, we make this the only hyper-parameter of the model and fix all others. We choose latent resolutions varying from 24 to 86. Furthermore, we consider two different loss functions. The first is simply the average of the relative 𝐿 2 errors for pressure and wall shear stress. The second includes in this average the relative error in the drag coefficient computed using equation (8.8). The drag is always computed from our full-field predictions and is never given as part of the output from the model. To compare the performance of our model against the industry-standard OpenFOAM solver, we perform a full cost-accuracy trade-off analysis. During data generation, we keep track of the drag coefficient predicted by OpenFOAM after every iteration. While the coefficient converges with more iterations, this convergence is not monotone and can often appear quite noisy. This makes computing the error from the raw data not possible. We therefore apply a box filter to the raw signal to compute a filtered version of the drag which acts as smoother. We take as the reference drag, the drag at the last iteration of the filtered signal. To compute the number of iterations it takes for the solver to predict a drag coefficient at a given relative error, we trace back the predictions from the filtered signal and return the first time at which this prediction incurs the given error with respect to the reference drag. An example of this methodology is shown in Figure 8.4a. The errors for our GINO model are computed with respect to the true drag coefficient from the last iteration of the solver. This is because we take as ground truth the pressure and wall shear stress from this last iteration and train our model to predict them. Figure 8.4b shows the cost-accuracy curve, measured in terms of inference time needed for a relative error in the drag coefficient for GINO and OpenFOAM. The cost of GINO is computed as the time, averaged over the test set, needed to predict the drag coefficient by running the model. This time includes both data pre-processing 169 (computing the SDF) as well as the model run-time and the drag calculation given the predicted fields. All models are ran on a single NVIDIA V100 GPU. The cost for OpenFOAM is computed as described in the previous paragraph and is averaged over the test set. The solver is ran on two NVIDIA V100 GPUs in parallel. We observe a four to five order of magnitude speed-up when using GINO. At a 3% relative error, we find the speed-up from our model which includes drag in the loss to be 26, 000×. As we increase the size of the latent space, the cost of GINO grows, however, we observe a plateau in the drag error. This is common in machine learning models as the error from using finite data starts to dominate the approximation error. Furthermore, we use only the size of the latent space as a hyper-parameter, keeping the number of learnable parameters fixed. It is interesting to explore further how parametrically scaling the model impacts predictive power. 8.6 Discussion and Conclusion In this work, we propose the GINO model for 3D PDEs with complex geometries. The GINO model consists of the graph-kernel blocks for the encoder and decoder that go to a latent uniform space, where the Fourier blocks run on the latent space to capture the global interaction. We experiment on two CFD datasets: Shape-Net car geometries and large-scale Ahmed’s body geometries, the latter encompassing over 600 car geometries featuring hundreds of thousands of mesh points. The evidence from these case studies illustrates that our method offers a substantial speed improvement, with a factor of 100,000 times acceleration in comparison to the GPU-based OpenFOAM solver. Concurrently, our approach has achieved onefourth to one-half the error rates compared to prevailing neural networks such as 3D U-Net. This underscores the potential of our method to significantly enhance computational efficiency while maintaining a competitive level of accuracy within the realm of CFD applications. Limitation: The trained surrogate model is limited to a specific category of shapes. The quality of the model depends on the quality of the training dataset. For CFD with more complex shapes, it is not easy to obtain a large training dataset. We will explore physics-informed approaches (Zongyi Li, Zheng, et al., 2021) and generate time-dependent high-fidelity simulations in the future. References Ahmed, Syed R, G Ramm, and Gunter Faltin (1984). “Some salient features of the time-averaged ground vehicle wake”. In: SAE transactions, pp. 473–503. 170 Figure 8.6: Illustrations of the Ahmed-body dataset Alet, Ferran et al. (2019). “Graph element networks: adaptive, structured computation and memory”. In: International Conference on Machine Learning. PMLR, pp. 212–222. Allen, Kelsey R et al. (2022). “Physical design using differentiable learned simulators”. In: arXiv preprint arXiv:2202.00728. Balu, Aditya et al. (2021). “Distributed multigrid neural solvers on megavoxel domains”. In: Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1–14. Battaglia, Peter et al. (2016). “Interaction networks for learning about objects, relations and physics”. In: Advances in neural information processing systems 29. 171 Battaglia, Peter W et al. (2018). “Relational inductive biases, deep learning, and graph networks”. In: arXiv preprint arXiv:1806.01261. Bhattacharya, Kaushik et al. (2021). “Model reduction and neural networks for parametric PDEs”. In: The SMAI journal of computational mathematics 7, pp. 121– 157. Brandstetter, Johannes et al. (2022). “Clifford neural layers for PDE modeling”. In: arXiv preprint arXiv:2209.04934. Chang, Angel X et al. (2015). “Shapenet: An information-rich 3d model repository”. In: arXiv preprint arXiv:1512.03012. Cohen, Albert, Christoph Schwab, and Jakob Zech (2018). “Shape Holomorphy of the Stationary Navier–Stokes Equations”. In: SIAM Journal on Mathematical Analysis 50.2, pp. 1720–1752. Cooley, James W and John W Tukey (1965). “An algorithm for the machine calculation of complex Fourier series”. In: Mathematics of computation 19.90, pp. 297– 301. Durasov, Nikita et al. (2021). “DEBOSH: Deep Bayesian Shape Optimization”. In: arXiv preprint arXiv:2109.13337. Fey, Matthias et al. (2023). PyTorch Extension Library of Optimized Scatter Operations. url: https://github.com/rusty1s/pytorch_scatter. Hamilton, Will, Zhitao Ying, and Jure Leskovec (2017). “Inductive representation learning on large graphs”. In: Advances in neural information processing systems 30. Hao, Zhongkai et al. (2023). “GNOT: A General Neural Operator Transformer for Operator Learning”. In: arXiv preprint arXiv:2302.14376. Hermosilla, Pedro et al. (2018). “Monte carlo convolution for learning on nonuniformly sampled point clouds”. In: ACM Transactions on Graphics (TOG) 37.6, pp. 1–12. Huang, Xun and Serge J. Belongie (2017). “Arbitrary Style Transfer in Real-time with Adaptive Instance Normalization”. In: CoRR abs/1703.06868. Jasak, Hrvoje, Aleksandar Jemcov, Zeljko Tukovic, et al. (2007). “OpenFOAM: A C++ library for complex physics simulations”. In: International workshop on coupled methods in numerical dynamics. Vol. 1000, pp. 1–20. Kipf, Thomas N and Max Welling (2016). “Semi-supervised classification with graph convolutional networks”. In: arXiv preprint arXiv:1609.02907. Kissas, Georgios et al. (2022). “Learning operators with coupled attention”. In: Journal of Machine Learning Research 23.215, pp. 1–63. Korzun, Ashley M et al. (2022). “Application of a Detached Eddy Simulation Approach with Finite-Rate Chemistry to Mars-Relevant Retropropulsion Operating Environments”. In: AIAA SciTech 2022 Forum, p. 2298. 172 Kossaifi, Jean et al. (2023). Multi-Grid Tensorized Fourier Neural Operator for High-Resolution PDEs. arXiv: 2310.00120. Kovachki, Nikola et al. (2021). “Neural operator: Learning maps between function spaces”. In: arXiv preprint arXiv:2108.08481. Lam, Remi et al. (2022). “GraphCast: Learning skillful medium-range global weather forecasting”. In: arXiv preprint arXiv:2212.12794. Lee, Seungjun (n.d.). “Mesh-Independent Operator Learning for Partial Differential Equations”. In: ICML 2022 2nd AI for Science Workshop. Li, Yangyan et al. (2018). “Pointcnn: Convolution on x-transformed points”. In: Advances in neural information processing systems 31. Li, Zijie, Kazem Meidani, and Amir Barati Farimani (2022). “Transformer for partial differential equations’ operator learning”. In: arXiv preprint arXiv:2205.13671. Li, Zongyi, Daniel Zhengyu Huang, et al. (2022). “Fourier neural operator with learned deformations for pdes on general geometries”. In: arXiv preprint arXiv:2207.05209. Li, Zongyi, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, et al. (2020a). “Fourier Neural Operator for Parametric Partial Differential Equations”. In: arXiv preprint arXiv:2010.08895. – (2020b). “Neural operator: Graph kernel network for partial differential equations”. In: arXiv preprint arXiv:2003.03485. Li, Zongyi, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Andrew Stuart, et al. (2020). “Multipole graph neural operator for parametric partial differential equations”. In: Advances in Neural Information Processing Systems 33. Li, Zongyi, Hongkai Zheng, et al. (2021). “Physics-informed neural operator for learning partial differential equations”. In: arXiv preprint arXiv:2111.03794. Lin, Haitao et al. (2022). “Non-equispaced Fourier Neural Solvers for PDEs”. In: arXiv preprint arXiv:2212.04689. Liu, Ning, Siavash Jafarzadeh, and Yue Yu (2023). “Domain Agnostic Fourier Neural Operators”. In: arXiv preprint arXiv:2305.00478. Menter, Florianr (1993). “Zonal two equation kw turbulence models for aerodynamic flows”. In: 23rd fluid dynamics, plasmadynamics, and lasers conference, p. 2906. Pathak, Jaideep et al. (2022). “Fourcastnet: A global data-driven high-resolution weather model using adaptive fourier neural operators”. In: arXiv preprint arXiv:2202.11214. Pfaff, Tobias et al. (2020). “Learning mesh-based simulation with graph networks”. In: arXiv preprint arXiv:2010.03409. Rahman, Md Ashiqur, Zachary E Ross, and Kamyar Azizzadenesheli (2022). “U-no: U-shaped neural operators”. In: arXiv preprint arXiv:2204.11127. 173 Remelli, Edoardo et al. (2020). “Meshsdf: Differentiable iso-surface extraction”. In: Advances in Neural Information Processing Systems 33, pp. 22468–22478. Sanchez-Gonzalez, Alvaro et al. (2020). “Learning to simulate complex physics with graph networks”. In: International conference on machine learning. PMLR, pp. 8459–8468. Sitzmann, Vincent et al. (2020). “Implicit neural representations with periodic activation functions”. In: Advances in Neural Information Processing Systems 33, pp. 7462–7473. Stein, Elias M. (1970). Singular Integrals and Differentiability Properties of Functions. Princeton University Press. Sun, Hongyu et al. (2022). “Accelerating Time-Reversal Imaging with Neural Operators for Real-time Earthquake Locations”. In: arXiv preprint arXiv:2210.06636. Temam, R. (1984). Navier-Stokes Equations: Theory and Numerical Analysis. Studies in mathematics and its applications. North-Holland. Ulyanov, Dmitry, Andrea Vedaldi, and Victor Lempitsky (2016). “Instance normalization: The missing ingredient for fast stylization”. In: arXiv preprint arXiv:1607.08022. Umetani, Nobuyuki and Bernd Bickel (2018). “Learning three-dimensional flow for interactive aerodynamic design”. In: ACM Transactions on Graphics (TOG) 37.4, pp. 1–10. Wandel, Nils, Michael Weinmann, and Reinhard Klein (2021). “Teaching the incompressible Navier–Stokes equations to fast neural surrogate models in three dimensions”. In: Physics of Fluids 33.4. Wen, Gege et al. (2023). “Real-time high-resolution CO 2 geological storage prediction using nested Fourier neural operators”. In: Energy & Environmental Science 16.4, pp. 1732–1741. Whitney, Hassler (1934). “Analytic Extensions of Differentiable Functions Defined in Closed Sets”. In: Transactions of the American Mathematical Society 36.1, pp. 63–89. Wolny, Adrian et al. (July 2020). “Accurate and versatile 3D segmentation of plant tissues at cellular resolution”. In: eLife 9. Ed. by Christian S Hardtke et al., e57613. issn: 2050-084X. doi: 10.7554/eLife.57613. url: https://doi.org/10. 7554/eLife.57613. Wu, Wenxuan, Zhongang Qi, and Li Fuxin (2019). “Pointconv: Deep convolutional networks on 3d point clouds”. In: Proceedings of the IEEE/CVF Conference on computer vision and pattern recognition, pp. 9621–9630. Yang, Yan et al. (2021). “Seismic wave propagation and inversion with neural operators”. In: The Seismic Record 1.3, pp. 126–134. Zhou, Qian-Yi, Jaesik Park, and Vladlen Koltun (2018). “Open3D: A Modern Library for 3D Data Processing”. In: arXiv:1801.09847. 174 Zienkiewicz, Olek C, Robert Leroy Taylor, and Perumal Nithiarasu (2013). The finite element method for fluid dynamics. Butterworth-Heinemann. 175 Chapter 9 GEOMETRY: NEURAL OPERATOR WITH LEARNED DEFORMATION Deep learning surrogate models have shown promise in solving partial differential equations (PDEs). Among them, the Fourier neural operator (FNO) achieves good accuracy, and is significantly faster compared to numerical solvers, on a variety of PDEs, such as fluid flows. However, the FNO uses the Fast Fourier transform (FFT), which is limited to rectangular domains with uniform grids. In this work, we propose a new framework, viz., Geo-FNO, to solve PDEs on arbitrary geometries. Geo-FNO learns to deform the input (physical) domain, which may be irregular, into a latent space with a uniform grid. The FNO model with the FFT is applied in the latent space. The resulting Geo-FNO model has both the computation efficiency of FFT and the flexibility of handling arbitrary geometries. Our Geo-FNO is also flexible in terms of its input formats, viz., point clouds, meshes, and design parameters are all valid inputs. We consider a variety of PDEs such as the Elasticity, Plasticity, Euler’s, and Navier-Stokes equations, and both forward modeling and inverse design problems. Comprehensive cost-accuracy experiments show that Geo-FNO is 105 times faster than the standard numerical solvers and twice more accurate compared to direct interpolation on existing ML-based PDE solvers such as the standard FNO. 9.1 Introduction Data-driven engineering design has the potential to accelerate the design process by orders of magnitude compared to conventional methods. It can enable extensive exploration of the design space and yield new designs with far greater efficiency. This is because the conventional design process requires repeated evaluations of partial differential equations (PDEs) for optimization, which can be time-consuming. Examples include computational fluid dynamics (CFD) based aerodynamic design and topology optimization for additive manufacturing. Deep learning approaches have shown promise in speeding up PDE evaluations and automatically computing derivatives, hence accelerating the overall design cycle. However, most deep learning approaches currently focus on predicting a few important quantities in the design process, e.g. lift and drag in aerodynamics, as opposed to completely emulating the entire simulation (i.e., pressure or Mach number fields in aerodynamics). This limits 176 the design space to be similar to the training data, and may not generalize to new geometries and scenarios. In contrast, we develop an efficient deep learning-based approach for design optimization, which emulates the entire PDE simulation on general geometries leading to a versatile tool for exploring the design space. Deformable meshes. Solutions of PDEs frequently have high variations across the problem domain, and hence some regions of the domain often require a finer mesh compared to other regions for obtaining accurate solutions. For example, in airfoil flows, the region near the airfoil requires a much higher resolution for accurately modeling the aerodynamics, as shown in Figure 9.8. To address such variations, deformable mesh methods such as adaptive finite element methods (FEM) have been developed. Two adaptive mesh methods are commonly used: the adaptive mesh refinement method that adds new nodes (and higher-order polynomials), and the adaptive moving mesh method that relocates the existing nodes Babuška and Suri, 1990; W. Huang and Russell, 2010. While mesh refinement is popular and easy to use with traditional numerical solvers, it changes the mesh topology and requires special dynamic data structures, which reduce speed and make it hard for parallel processing. On the other hand, the adaptive moving mesh methods retain the topology of the mesh, making it possible to integrate with spectral methods. Spectral methods solve the equation on the spectral space (i.e., Fourier, Chebyshev, and Bessel), which usually have exponential convergence guarantees when applicable Gottlieb and Orszag, 1977. However, these spectral methods are limited to simple computational domains with uniform grids. When the mesh becomes nonuniform, spectral basis functions are no longer orthogonal and the spectral transform is no longer invertible, so the system in the deformed Fourier space is not equivalent to the original one anymore. Thus, traditional numerical methods are slow on complex geometries due to computational constraints. Neural operators. Deep learning surrogate models have recently yielded promising results in solving PDEs. One class of such models is data-driven, and they directly approximate the input-output map through supervised learning. Such models have achieved significant speedups compared to traditional solvers in numerous applications (Zhu and Zabaras, 2018; Bhatnagar et al., 2019). Among them, the graph neural networks have not been studied for complex geometries (Allen et al., 2022; Sanchez-Gonzalez et al., 2020) Recently, a new class of data-driven models known as neural operators aim to directly learn the solution operator of the 177 PDE in a mesh-convergent manner. Unlike standard deep learning models, such as convolutional neural networks from computer vision, neural operators are consistent to discretization and hence, better suited for solving PDEs. They generalize the previous finite-dimensional neural networks to learn operator mapping between infinite-dimensional function spaces. Examples include Fourier neural operator (FNO) (Li, Kovachki, et al., 2020c), graph neural operator Li, Kovachki, et al., 2020b, and DeepONet (Lu, Jin, and Karniadakis, 2019). We consider FNO (Li, Kovachki, et al., 2020a) in this paper due to its superior cost-accuracy trade-off (De Hoop et al., 2022). FNO implements a series of layers computing global convolution operators with the fast Fourier transform (FFT) followed by mixing weights in the frequency domain and inverse Fourier transform. These global convolutional operators are interspersed with non-linear transformations such as GeLU. By composing global convolution operators and non-linear activations, FNO can approximate highly non-linear and non-local solution operators. FNO and its variants are able to simulate many PDEs such as the Navier-Stokes equation and seismic waves, do high-resolution weather forecasts, and predict CO2 migration with unprecedented cost-accuracy trade-off (Pathak et al., 2022; Yang et al., 2021; Wen et al., 2022). Limitations of FNO on irregular geometries. While FNO is fast and accurate, it has limitations on the input format and the problem domain. Since FNO is implemented with FFT, it can be only applied on rectangular domains with uniform meshes. When applying it to irregular domain shapes, previous works usually embed the domain into a larger rectangular domain (Lu, Meng, et al., 2022). However, such embeddings are less efficient and wasteful, especially for highly irregular geometries. Similarly, if the input data is in the form of non-uniform meshes such as triangular meshes, previous works use basic interpolation methods between the input non-uniform mesh and a uniform mesh on which FFT is computed Li, Kovachki, et al., 2020a. This can cause large interpolation errors, especially for non-linear PDEs. Our contributions. To address the above challenges, we develop a geometryaware Fourier neural operator (Geo-FNO) and we summarize our contributions below: 1. We propose a geometry-aware discretization-convergent FNO framework 178 Top: The Geo-FNO model deforms the input functions from irregular physical space to a uniform latent space. After the standard FNO Li, Kovachki, et al., 2020a is applied, the Geo-FNO will deform the latent solution function back to the physical space. 𝑃, 𝑄 are lifting and projection. F and F −1 are the Fourier transforms. 𝐾 is a linear transform (9.18). Bottom: The deformation, either given or learned, defines a correspondence between the physical space and computational space. It induces an adaptive mesh and a generalized Fourier basis on the physical space. Figure 9.1: Geometry-aware FNO (Geo-FNO) that works on arbitrary geometries and a variety of input formats such as point clouds, non-uniform meshes, and design parameters. 2. Geo-FNO deforms the irregular input domain into a uniform latent mesh on which the FFT can be applied. Such deformation can be fixed or learned with the FNO architecture in an end-to-end manner. For the latter case, we design a neural network to model the deformation and train it end-to-end with Geo-FNO. 3. We experiment on different geometries for the Elasticity, Plasticity, Advection, Euler’s, and Navier-Stokes equations on both forward modeling and inverse design tasks. The cost-accuracy experiments show that Geo-FNO has up to 105 acceleration compared to the numerical solver, as well as half the error rate compared to previous interpolation-based methods as shown in Figure 9.3. Further, Geo-FNO maintains discretization-convergence. The model trained on low resolution datasets can be directly evaluated at high resolution cases. Geo-FNO thus learns a deformable mesh along with the solution operator in an endto-end manner by applying a series of FFTs and inverse FFTs on a uniform latent mesh, as shown in Figure 9.1. Thus, Geo-FNO combines both the computational 179 efficiency of the FFT and the flexibility of learned deformations. It is as fast as the original FNO but can represent the variations in solutions and domain shapes more efficiently and accurately. The learned deformation from the physical irregular domain to a uniform computational domain is inspired by the traditional adaptive moving mesh method (W. Huang and Russell, 2010). However, the adaptive moving mesh method has not had much success with traditional numerical solvers due to the loss of orthogonality of the spectral transform. On a deformed mesh, the system in the Fourier space is no longer equivalent to the original system, so it does not make sense to solve the PDE in the deformed Fourier space. However, Geo-FNO does not have this limitation. It does not involve solving equations in the Fourier space. Instead, Geo-FNO approximates the solution operator via the computation on Fourier space in a data-driven manner. In principle, our Geo-FNO framework can be directly extended to general topologies. Given complex input topology we can apply the decomposition techniques such as triangle tessellation to divide the domain into sub-domains such as each sub-domain is regular. Similarly, for other input formats such as the signed distance functions (SDF), we can sample collocation points. Furthermore, it is possible to extend the Geo-FNO framework to the physics-informed neural operator setting which incorporates PDE constraints (Li, Zheng, et al., 2021). Since the deformation is parameterized by a neural network, we can obtain the exact derivatives and apply the chain rule to compute the residual error. 9.2 Problem Settings and Preliminaries Problem settings. In this work, we consider parametric partial differential equations defined on various domains. Assume the problem domain 𝐷 𝑎 is parameterized by design parameters 𝑎 ∈ A, which follows some distribution 𝑎 ∼ 𝜇. For simplicity, we assume all the domains are bounded, orientable manifolds embedded in some background Euclidean space Ω (e.g., R3 ). We consider both stationary problems and time-dependent problems of the following forms: 𝑑𝑢 = R (𝑢), 𝑑𝑡 𝑢 = 𝑢0, in 𝐷 𝑎 × 𝑇 𝑢 = 𝑢𝑏 , in 𝜕𝐷 𝑎 × 𝑇 in 𝐷 𝑎 × {0} (9.1) where 𝑢 0 ∈ U (0) is the initial condition; 𝑢 𝑏 is the boundary condition; and 𝑢(𝑡) ∈ U (𝑡) for 𝑡 > 0 is the target solution function. R is a possibly non-linear differential operator. We assume that 𝑢 exists and is bounded for all time and for every 𝑢 0 ∈ U (𝑡) 180 at every 𝑡 ∈ 𝑇. This formulation gives rise to the solution operator G † : (𝑎, 𝑢 0 , 𝑢 𝑏 ) ↦→ 𝑢. Prototypical examples include fluid mechanics problems such as the Burgers’ equation and the Navier-Stokes equation and solid mechanics problems such as elasticity and plasticity defined on various domain shapes such as airfoils and pipes. In this paper, we assume the initial condition 𝑢 0 and the boundary condition 𝑢 𝑏 are fixed. Specifically, we also consider stationary problem R (𝑢) = 0 where 𝑢 0 is not necessary. In this case, the solution operator simplifies to G † : 𝑎 ↦→ 𝑢. Input formats. In practice, the domain shape can be parameterized in various ways. It can be given as meshes, functions, and design parameters. In the most common cases, the domain space is given as meshes (point clouds) 𝑎 = T = {𝑥𝑖 } 𝑁 ⊂ Ω, such as triangular tessellation meshes used in many traditional numerical solvers. Alternatively, the shapes can be described by some functions 𝑎 = 𝑓 : Ω → R. For example, if the domain is a 2D surface, then it can be parameterized by its boundary function 𝑓 : [0, 1] → 𝜕𝐷 𝑎 . Similarly, the domain can be described as signed distance functions (SDF) or occupancy functions. These are common examples used in computer vision and graphics. At last, we also consider specific design parameters 𝑎 ∈ R𝑑 . For example, it can be the width, height, volume, angle, radius, and spline nodes used to model the boundary. The design parameters restrict the domain shape in a relatively smaller subspace of all possible choices. This format is most commonly used in a design problem, which is easy to optimize and manufacture. For the latter two cases, it is possible to sample a mesh T based on the shape. Learning the solution operator Given a PDE as defined in (9.1) and the corresponding solution operator G † , one can use a neural operator G𝜃 with parameters 𝜃 as a surrogate model to approximate G † . Usually we assume a dataset {𝑎 𝑗 , 𝑢 𝑗 } 𝑁𝑗=1 is available, where G † (𝑎 𝑗 ) = 𝑢 𝑗 and 𝑎 𝑗 ∼ 𝜇 are i.i.d. samples from some distribution 𝜇 supported on A. In this case, one can optimize the solution operator by minimizing the relative empirical data loss on a given data pair v u t∫ |𝑢(𝑥) − G𝜃 (𝑎)(𝑥)| 2 d𝑥 ∥𝑢 − G𝜃 (𝑎)∥ L 2 𝐷𝑎 ∫ = (9.2) Ldata (𝑢, G𝜃 (𝑎)) = ∥𝑢∥ L 2 𝑢(𝑥) 2 d𝑥 𝐷 𝑎 181 The operator data loss is defined as the average error across all possible inputs v u∫ 𝑁 t ∑︁ |𝑢 𝑗 (𝑥) − G𝜃 (𝑎 𝑗 )(𝑥)| 2 d𝑥 1 𝐷𝑎 † ∫ Jdata (G𝜃 ) = E𝑎∼𝜇 [Ldata (G (𝑎), G𝜃 (𝑎))] ≈ . 2 d𝑥 𝑁 𝑗=1 𝑢 (𝑥) 𝑗 𝐷 𝑎 (9.3) When the prediction 𝑢 is time-dependent, the L 2 integration should also account for the temporal dimension. Neural operator In this work, we will focus on the neural operator model designed for the operator learning problem. The neural operator, proposed in (Li, Kovachki, et al., 2020c), is formulated as a generalization of standard deep neural networks to operator setting. Neural operators compose linear integral operator K with pointwise non-linear activation function 𝜎 to approximate highly non-linear operators. Definition 9.2.1 (Neural operator G𝜃 ) Define the neural operator G𝜃 := Q ◦ (𝑊 𝐿 + K𝐿 + 𝑏 𝐿 ) ◦ · · · ◦ 𝜎(𝑊1 + K1 + 𝑏 1 ) ◦ P (9.4) where P : R𝑑 𝑎 → R𝑑1 , Q : R𝑑 𝐿 → R𝑑𝑢 are the pointwise neural networks that encode the lower dimension function into higher dimensional space and vice versa. The model stack 𝐿 layers of 𝜎(𝑊𝑙 + K𝑙 + 𝑏 𝑙 ) where 𝑊𝑙 ∈ R𝑑𝑙+1 ×𝑑𝑙 are pointwise linear operators (matrices), K𝑙 : {𝐷 → R𝑑𝑙 } → {𝐷 → R𝑑𝑙+1 } are integral kernel operators, 𝑏 𝑙 : 𝐷 → R𝑑𝑙+1 are the bias term made of a linear layer, and 𝜎 are fixed activation functions. The parameters 𝜃 consists of all the parameters in P, Q, 𝑊𝑙 , K𝑙 , 𝑏 𝑙 . Definition 9.2.2 (Fourier integral operator K) Define the Fourier integral operator  
−1 K (𝜙)𝑣 𝑡 (𝑥) = F 𝑅𝜙 · (F 𝑣 𝑡 ) (𝑥) ∀𝑥 ∈ 𝐷 (9.5) where 𝑅𝜙 is the Fourier transform of a periodic function 𝜅 : 𝐷¯ → R𝑑 𝑣 ×𝑑 𝑣 parameterized by 𝜙 ∈ ΘK . The Fourier transform F is defined as ∫ ∑︁ (F 𝑣)(𝑘) = ⟨𝑣, 𝜓(𝑘)⟩ 𝐿 (𝐷) = 𝑣(𝑥)𝜓(𝑥, 𝑘)𝜇(𝑥) ≈ 𝑣(𝑥)𝜓(𝑥, 𝑘) 𝑥 𝑥∈T (9.6) 182 where 𝜓(𝑥, 𝑘) = 𝑒 2𝑖𝜋⟨𝑥,𝑘⟩ ∈ 𝐿 (𝐷) is the Fourier basis and T is the mesh sampled from the distribution 𝜇. In Li, Kovachki, et al., 2020c, the domain 𝐷 is assumed to be a periodic, square torus and the mesh T is uniform, so the Fourier transform F can be implemented by the fast Fourier transform algorithm (FFT). In this work, we aim to extend the framework to non-uniform meshes and irregular domains. 9.3 Geometry-Aware Fourier Neural Operator The idea of the geometry-aware Fourier neural operator is to deform the physical space into a regular computational space so that the fast Fourier transform can be performed on the computational space. Formally, we want to find a diffeomorphic deformation 𝜙𝑎 between the input domain 𝐷 𝑎 and the unit torus 𝐷 𝑐 = [0, 1] 𝑑 . The computational mesh 𝐷 𝑐 is shared among all input space 𝐷 𝑎 . It is equipped with a uniform mesh and standard Fourier basis. Once the coordinate map 𝜙𝑎 is determined, the map induces an adaptive mesh and deformed Fourier basis on the physical space. In the community of numerical PDEs, the diffeomorphism corresponds to the adaptive moving mesh (W. Huang and Russell, 2010). Deformation from the physical space to the computational space. Let 𝐷 𝑎 be the physical domain and 𝐷 𝑐 be the computational domain. Let 𝑥 ∈ 𝐷 𝑎 and 𝜉 ∈ 𝐷 𝑐 be the corresponding mesh points. Denote the input meshes as T𝑎 = {𝑥 (𝑖) } ⊂ 𝐷 𝑎 and the computational meshes as T 𝑐 = {𝜉 (𝑖) } ⊂ 𝐷 𝑐 . For simplicity, we assume the output mesh is the same as the input mesh. The adaptive moving mesh is defined by a coordinate transformation 𝜙𝑎 that transforms the points from the computational mesh to the physical mesh, as shown in Figure 9.1 (b) 𝜙𝑎 : 𝐷 𝑐 → 𝐷 𝑎 𝜉 ↦→ 𝑥 (9.7) −1 Ideally, 𝜙𝑎 is a diffeomorphism, meaning it has an inverse 𝜙−1 𝑎 , and both 𝜙 𝑎 and 𝜙 𝑎 are smooth. When such a diffeomorphism exists, there is a correspondence between the physical space and the computational space. Let T 𝑐 ⊂ 𝐷 𝑐 be the uniform mesh and 𝜓 𝑐 ∈ 𝐿 (𝐷 𝑐 ) be the standard spectral basis defined on the computational space. The coordinate transformation 𝜙𝑎 ⊂ 𝐷 𝑎 induces an adaptive mesh T𝑎 and adaptive spectral basis 𝜓𝑎 ∈ 𝐿(𝐷 𝑎 ) (pushforward) T𝑎 := 𝜙𝑎 (T 𝑐 ) 𝜓𝑎 (𝑥) := 𝜓 𝑐 ◦ 𝜙−1 𝑎 (𝑥) (9.8) 183 It can be interpreted as solving the system R (𝑣) = 0 with adaptive mesh T𝑎 and generalized basis 𝜓𝑎 . Conversely, for any function defined on the physical domain 𝑣 ∈ 𝐿 (𝐷 𝑎 ), it can be transformed to the computational domain 𝑣 𝑐 ∈ 𝐿 (𝐷 𝑐 ) (pullback) 𝑣 𝑐 (𝜉) := 𝑣(𝜙𝑎 (𝜉)) (9.9) Similarly, for any system of equations R (𝑣) = 0 such as (9.1) defined on the physical domain 𝐷 𝑎 , the transformation 𝜙𝑎 induced a new deformed system of equations R 𝑐 (𝑣 𝑐 ) = 0. It can be interpreted as solving the deformed system R 𝑐 (𝑣 𝑐 ) = 0 with uniform mesh T 𝑐 and standard basis 𝜓 𝑐 . Examples: Chebyshev Basis. Chebyshev method is a standard spectral method for solving PDE with non-periodic boundary Driscoll, Hale, and Trefethen, 2014. It can be induced from the standard cosine basis (in discrete cosine transform) with a cosine grid. The Chebyshev polynomials on domain [−1, 1] has the form of 𝑘 𝜓cheb := cos(𝑘 · cos−1 (𝑥)), 𝑘 = 0, 1, 2 . . . It’s equivalent to apply the cosine deformation 𝜙𝑎 : [0, 𝜋] → [−1, 1] 𝜙𝑎 (𝑥) = cos(𝑥) to the basis of cosine series 𝑘 𝜓cos := cos(𝑘 · 𝑥), 𝑘 = 0, 1, 2 . . . in the same manner as a deformed basis (9.8) 𝑘 𝑘 𝜓cheb = 𝜓cos ◦ 𝜙−1 𝑎 (𝑥) Intuitively, the cosine deformation places more grid points around the boundary −1 and 1, which addresses the stability issue around the boundary. Notice, in practice the discrete cosine transform on [0, 𝜋] can be implemented by the faster Fourier transform using an extended domain [−𝜋, 𝜋] and restricts to real coefficients. We will discuss the use of extended domain in Section 9.3 Fourier continuation. Therefore, GeoFNO equipped with the cosine deformation can implement the Chebyshev method. However, it is hard to deform the Fourier basis into spherical basis, since the underlying domain is not homeomorphic. For these non-homeomorphic domain we will use domain decomposition as discussed in Section 9.3. 184 Geometric Fourier transform Based on the deformation, we can define the spectral transform in the computational space. Fourier transforms conducted in the computational domain are standard, since we have a uniform structured mesh in the computational domain. In this subsection, We will introduce a geometric spectral transform between the function 𝑣(𝑥) in the physical domain 𝐷 𝑎 and the function 𝑣ˆ (𝑘) in the spectral space of the computation domain 𝐷 𝑐 . To transform the function 𝑣(𝑥) from the physical domain to the spectral space of the computation domain, we define the forward transform: ∫ (F𝑎 𝑣)(𝑘) := 𝑣 𝑐 (𝜉)𝑒 −2𝑖𝜋⟨𝜉,𝑘⟩ 𝑑𝜉 (9.10) 𝑐 𝐷 ∫ −1 = 𝑣(𝑥)𝑒 −2𝑖𝜋⟨𝜙 𝑎 (𝑥),𝑘⟩ |det[∇𝑥 𝜙−1 (9.11) 𝑎 (𝑥)]|𝑑𝑥 𝐷 ≈ −1 1 ∑︁ 𝑚(𝑥)𝑣(𝑥)𝑒 −2𝑖𝜋⟨𝜙 𝑎 (𝑥),𝑘⟩ |T𝑎 | 𝑥∈T (9.12) 𝑎 where the weight 𝑚(𝑥) = |det[∇𝑥 𝜙−1 𝑎 (𝑥)]|/𝜌 𝑎 (𝑥) and 𝜌 𝑎 is a distribution from which the input mesh T𝑎 is sampled. Notice, in many cases, we have an input mesh T𝑎 so 𝑐 −1 that computational mesh 𝜙−1 𝑎 (T𝑎 ) = 𝑇 is uniform. In this case, |det[∇𝑥 𝜙 𝑎 (𝑥)]| = 𝜌𝑎 (𝑥), and 𝑚(𝑥) = 1 can be omitted. Other choices can be made, for example, we can define 𝑚(𝑥) using heuristics depending on the solution 𝑢(𝑥) or residual error. The weight 𝑚(𝑥) corresponds to the monitor functions in the literature of adaptive meshes, where the adaptive mesh is finer near 𝑥, when 𝑚(𝑥) is large. To transform the spectral function 𝑣ˆ (𝑘) from the spectral space of the computation domain to the physical domain, we define the inverse transform (F𝑎−1 𝑣ˆ )(𝑥) = (F −1 𝑣ˆ )(𝜙−1 𝑎 (𝑥)) = ∑︁ −1 𝑣ˆ (𝑘)𝑒 2𝑖𝜋⟨𝜙 𝑎 (𝑥),𝑘⟩ (9.13) 𝑘 It is worth mentioning that F𝑎 ◦ F𝑎−1 = 𝐼, since both are defined on the computational space. ∫ −1 (F𝑎 ◦ F𝑎 𝑣ˆ )(𝑘) = (F𝑎−1 𝑣ˆ ) 𝑐 (𝜉)𝑒 −2𝑖𝜋⟨𝜉,𝑘⟩ 𝑑𝜉 (9.14) 𝑐 𝐷 ∫ ∑︁ = 𝑣ˆ (𝑘)𝑒 2𝑖𝜋⟨𝜉,𝑘⟩ 𝑒 −2𝑖𝜋⟨𝜉,𝑘⟩ 𝑑𝜉 = 𝑣ˆ (𝑘) (9.15) 𝐷𝑐 𝑘 Notice both F𝑎 (9.12) and F𝑎−1 (9.13) only involve the inverse coordinate transform 𝑐 −1 𝜙−1 𝑎 : 𝐷 𝑎 → 𝐷 . Intuitively, in the forward Fourier transform F𝑎 , we use 𝜙 𝑎 to 185 transform the input function to the computational space, while in the inverse Fourier transform F𝑎−1 , we use 𝜙−1 𝑎 to transform the query point to the computational space to evaluate the Fourier basis. It makes the implementation easy in that we only need to define 𝜙−1 𝑎 . Architecture of Geo-FNO and the deformation neural network We consider two scenarios: (1) the coordinate map is given, and (2) learning a coordinate map. In many cases, the input mesh T𝑎 is non-uniform but structured. For example, the meshes of airfoils are usually generated in the cylindrical structure. If the input mesh is structured, meaning it can be indexed as a multi-dimensional array T𝑎 [𝑖1 , . . . , 𝑖 𝑑 ] with 0 ≤ 𝑖 𝑘 ≤ 𝑠 𝑘 ∀𝑘, then its indexing induces a canonical coordinate map 𝜙−1 𝑎 : T𝑎 [𝑖 1 , . . . , 𝑖 𝑑 ] ↦→ (𝑖 1 /𝑠1 , . . . , 𝑖 𝑑 /𝑠 𝑑 ) (9.16) In this case, (𝑖1 /𝑠1 , . . . , 𝑖 𝑑 /𝑠 𝑑 ) forms a uniform mesh within a unit cube, allowing for the direct application of the FFT. And Geo-FNO reduces to a standard FNO directly applied to the input meshes. When we need to learn a coordinate map, we parameterize the coordinate map 𝜙−1 𝑎 as a neural network and learn it along with the solution operator 𝐺 𝜃 in an end-to-end manner with loss (9.3). 𝜙−1 𝑎 : (𝑥 1 , 𝑥 2 , · · · , 𝑥 𝑑 , 𝑎) ↦→ (𝜉1 , 𝜉2 , · · · , 𝜉 𝑑 ) (9.17) We parameterize the deformation as 𝜙−1 𝑎 (𝑥, 𝑎) = 𝑥 + 𝑓 (𝑥, 𝑎), where 𝑓 is a standard feed-forward neural network. Specially, 𝑓 takes input (𝑥1 , 𝑥2 , · · · , 𝑥 𝑑 , 𝑎), where 𝑎 is the geometry parameters. This formulation initializes 𝜙−1 𝑎 around an identity map, which is a good initial choice. Following the works of implicit representation (Mildenhall et al., 2020), (Sitzmann et al., 2020), we use sinusoidal features in the first layers sin(𝐵𝑥), 𝐵 = 2𝑖 to improve the expressiveness of the network. Concretely, we define 𝑓 as a three-layer feed-forward neural network with a width of 32. The input to this network comprises a combination of components, namely 𝑥 and 𝑎, alongside a series of trigonometric terms such as sin(21 𝑥), cos(21 𝑥), sin(22 𝑥), cos(22 𝑥), · · · , sin(2𝐾 𝑥) and cos(2𝐾 𝑥), which are applied coordinatewisely. And 𝑎 can represent the coordinates of the point clouds that encode the design geometry. −1 Finally, 𝜙−1 𝑎 is used to compute F𝑎 and F𝑎 (See (9.12) and (9.13)). As shown in Figure 9.1, Geo-FNO has a similar architecture as the standard FNO, 186 but with a deformation in the input and output G𝜃 := Q ◦ (𝑊 𝐿 + K𝐿 (𝜙𝑎 ) + 𝑏 𝐿 ) ◦ · · · ◦ 𝜎(𝑊1 + K1 (𝜙𝑎 ) + 𝑏 1 ) ◦ P (9.18) Here, we begin by employing P to lift (encode) the input to a higher channel dimension. Subsequently, the first Fourier convolution operator K1 employs the geometric Fourier transform F𝑎 (9.12), while the last layer K𝐿 employs the inverse geometric Fourier transform F𝑎−1 (9.13). To finalize the process, we utilize Q to project (decode) the output back to the desired dimension. In the case where the input is presented as a structured mesh, we define the coordinates as (9.16). Under this circumstance, both F𝑎 and F𝑎−1 reduce to the standard FFT, leading to Geo-FNO being reduced to the standard FNO. Fourier continuation For some PDE problems, the input domain has an irregular topology that is not homeomorphic to a disk or a torus. Consequentially, there does not exist a diffeomorphism between 𝐷 𝑎 and 𝐷 𝑐 . In this case, we will first embed the input domain into a larger regular domain 𝑖 𝐷 𝑎 ↩− → 𝐷¯ 𝑎 so that 𝐷¯ 𝑎 is diffeomorphic to 𝐷 𝑐 . For example, the elasticity problem section 9.4 has a hole in its square domain. We can first embed the domain into a full square by completing the hole. Such embedding corresponds to the Fourier continuation (Bruno, Han, and Pohlman, 2007) techniques used in conventional spectral solvers. Usually it requires to extend the function 𝑢 ∈ U (𝐷 𝑎 ) to 𝑢¯ ∈ U ( 𝐷¯ 𝑎 ) by fitting polynomials. However, in the data-driven setting, a such extension can be done implicitly during the training. We only need to compute the loss on the original space 𝐷 𝑎 and discard the extra output from the underlying space 𝐷¯ 𝑎 . This continuation technique is universal. According to the Whitney embedding theorem, any mdimensional manifold can be smoothly embedded into a Euclidean space R2𝑚 . Domain decomposition For complex topologies, another technique is to decompose the domain into multiple sub-domains such that each of the domains is equipped with a standard topology. We mainly consider 2-dimensional manifolds in our experiments, although the technique is applicable for any dimension. Given any 2-manifold M, we apply the decomposition M = 𝐷 1 # · · · #𝐷 𝑛 187 Figure 9.2: Elasticity (a) and plasticity problems (b) introduced in section 9.4. The comparison is shown between the reference obtained using a traditional solver (left) and the Geo-FNO result (right). where # is the connect-sum, and each 𝐷 𝑖 is homeomorphic to either a sphere 𝑆 2 or torus 𝑇 2 . We then train 𝑛 sub-models G1 · · · G𝑛 , each equipped with a deformation map 𝜙𝑎1 · · · 𝜙𝑎 𝑛 . This domain decomposition is also universal since any orientable compact 2-dimensional manifold is homeomorphic to a sphere or a 𝑛-genus torus, where 𝑛-genus torus can be naturally decomposed into 𝑛 tori. Compared to continuation, the decomposition method does not require raising the dimensionality, but it needs multiple sub-models. A numerical example is included in Section 9.4. 9.4 Numerical Examples In this section, we study Geo-FNO numerically, and compare the Geo-FNO against other machine-learning models on PDEs with various geometries, including hyperelastic problem in section 9.4, plastic problem in section 9.4, advection equation on a sphere problem in section 9.4, airfoil problem in section 9.4, and pipe problem in section 9.4. We demonstrate that Geo-FNO can be applied on irregular domains and non-uniform meshes (See figs. 9.2 and 9.8) and show it is more accurate than combining direct interpolation with existing ML-based PDE solvers such as the standard FNO (Li, Kovachki, et al., 2020a) and UNet(Ronneberger, Fischer, and Brox, 2015), as well as mesh-free methods such as Graph neural operators (GNO) (Li, Kovachki, et al., 2020c) and DeepONet (Lu, Jin, and Karniadakis, 2019) (See table 9.1 and table 9.2). Meanwhile, Geo-FNO perseveres the speed of standard FNO. It can accelerate the numerical solvers up to 105 times on the airfoil problem (See section 9.4). All experiments are performed on a single Nvidia 3090 GPU. With no special mention, we train all the models with 500 epochs with an initial learning rate of 0.001 which decays by half every 100 epochs. We use the relative 𝐿 2 error for both the training and testing. And the standard Adam optimizer is used. The code is available at https://github.com/neuraloperator/Geo-FNO. 188 Implicit scheme (CPU) Geo-FNO (GPU) Interp-UNet (GPU) Explicit scheme (CPU) Geo-FNO (CPU) Interp-UNet (CPU) Rel. L error 10 1 10 2 10 2 10 1 100 101 102 103 Time (s) Figure 9.3: The cost-accuracy trade-off for Geo-FNO/UNet and traditional numerical solver with implicit/explicit temporal integrators on the airfoil problem. Hyper-elastic problem The governing equation of a solid body can be written as 2 𝑠𝜕 𝒖 +∇·𝝈 =0 𝜌 𝜕𝑡 2 where 𝜌 𝑠 is the mass density, 𝒖 is the displacement vector, 𝝈 is the stress tensor. Constitutive models, which relate the strain tensor 𝝐 to the stress tensor, are required to close the system. We consider the unit cell problem Ω = [0, 1] × [0, 1] with an arbitrary shape void at the center, which is depicted in Figure 9.2 (a). The prior of 0.2 2 2 −1 the void radius is 𝑟 = 0.2 + 1+exp( 𝑟) ˜ with 𝑟˜ ∼ N (0, 4 (−∇ + 3 ) ), which embeds the constraint 0.2 ≤ 𝑟 ≤ 0.4. The unit cell is clamped on the bottom edge and tension traction 𝒕 = [0, 100] is applied on the top edge. The material is the incompressible Rivlin-Saunders material Pascon, 2019 and the constitutive model of the material is given by 𝜕𝑤(𝝐) 𝝈= 𝜕𝝐 𝑤(𝝐) = 𝐶1 (𝐼1 − 3) + 𝐶2 (𝐼2 − 3) where 𝐼1 = tr(𝐶) and 𝐼2 = 21 [(tr(𝐶) 2 − tr(𝐶 2 )] are scalar invariants of the right Cauchy Green stretch tensor 𝐶 = 2𝝐 + I. And energy density function parameters are 𝐶1 = 1.863 × 105 and 𝐶1 = 9.79 × 103 . We generate 1000 training data and 200 test data with a finite element solver D. Z. Huang, Xu, et al., 2020 with about 100 quadratic quadrilateral elements. It takes about 5 CPU seconds for each simulation. 189 The inputs 𝑎 are given as point clouds with a size of around 1000. The target output is the stress field. The elasticity problem has an unstructured input format. We compare the meshconvergent methods such as Geo-FNO, Graph neural operator (GNO), and Deep operator network (DeepONet), as well as interpolation-based ML methods including uniform meshes (FNO, UNet) and heuristic adaptive meshes (R-mesh, O-mesh). The details of these methods are listed below. • GNO: Since the graph structure is flexible on the problem geometry, graph neural networks have been widely used for complex geometries (SanchezGonzalez et al., 2020; Pfaff et al., 2020). In this work, we compare with the graph kernel operator (Li, Kovachki, et al., 2020c; Li, Kovachki, et al., 2020b), which is a graph neural network-based neural operator model. It implements the linear operator K by message passing on graphs. We build a full graph with edge connection radius 𝑟 = 0.2, width 64, and kernel with 128. • DeepONet: the deep operator network (Lu, Jin, and Karniadakis, 2019) is a line of works designed with respect to the linear approximation of operator as shown in (T. Chen and H. Chen, 1995). It has two neural networks a trunk net and a branch net to represent the basis and coefficients of the operator. Both DeepONets and neural operators aim to parameterize infinitely dimensional solution operators, but neural operators directly output the full field solution functions, while DeepONets output the solution at specific query points. This design difference gives neural operators, especially FNO, an advantage when the data are fully observed while DeepONet has an advantage when the data are partially observed. In this work, we use five layers for both the trunk net and branch net, each with a width of 256. • FNO interpolation: as a baseline, we simply interpolate the input point cloud into a 41-by-41 uniform grid and train a plain FNO model (Li, Kovachki, et al., 2020a). As shown in figure 9.4 (c), the interpolation causes an interpolation error, which loses information on the singularities (the red regions). As a result, the testing error is constantly higher than 5%. • UNet interpolation: Similarly, we train a UNet model (Ronneberger, Fischer, and Brox, 2015) on the interpolated uniform grid. As FNO-interpolation, the error is constantly higher than 5% as shown in the figure 9.4 (d). 190 • Geo-FNO (R mesh): We consider a heuristic adaptive mesh method for each input geometry by deforming a uniform 41 by 41 grid, as shown in the figure 9.4 (a). The stretching is applied in the radial direction originated at the void center (0.5, 0.5) to attain a finer mesh near the interface of the void. Let 𝑟 denote the radial distance of mesh points, the deformation map is 2  𝑠)  𝑟 𝑠 + 𝛼(𝑟 − 𝑟 𝑠 ) + (1 − 𝛼) (𝑟−𝑟  𝑟 𝑒 −𝑟 𝑠 𝑟→ 2  𝑟 𝑠 − 𝛼(𝑟 𝑠 − 𝑟) − (1 − 𝛼) (𝑟 𝑠𝑟−𝑟) 𝑠  𝑟 ≥ 𝑟𝑠 𝑟 ≤ 𝑟𝑠 where 𝑟 𝑠 and 𝑟 𝑒 denote the void radius and the unit cell radius in this radial direction, 𝛼 = 0.2 denotes the stretching factor, where the ratio of mesh sizes between these near the void interface and these near the unit cell boundary 𝛼 is about 2−𝛼 . It is worth mentioning that the deformation map remains void interface, the origin, and the unit cell boundary. Once the adaptive meshes are generated, we interpolate the input data to the adaptive meshes. The surrogate model is learned on the adaptive meshes. We used four Fourier layers with mode 12 and width 32. • Geo-FNO (O mesh): Similarly, we design another heuristic adaptive mesh method with cylindrical meshing, as shown in the figure 9.4 (b). A body-fitted O-type mesh is generated for each geometry with 64 points in the azimuth direction and 41 points in the radial direction. The points in the azimuth direction are uniformly located between 0 and 2𝜋, and the points in the radial direction are uniformly located between 𝑟 𝑠 and 𝑟 𝑒 , where 𝑟 𝑠 and 𝑟 𝑒 denote the void radius and the unit cell radius. Again, we interpolate the input data to the adaptive meshes. The surrogate model is learned on the adaptive meshes. We used four Fourier layers with mode 12 and width 32. As shown in Table 9.1, for the elasticity problem, the Geo-FNO model has a significantly lower error compared to combining direct interpolation with existing ML-based PDE solvers such as the standard FNO(Li, Kovachki, et al., 2020a) and UNet(Ronneberger, Fischer, and Brox, 2015). Both FNO+interpolation and UNet+interpolation methods have a test error larger that 5%, which is likely caused by the interpolation error. Geo-FNO also has a lower error compared to mesh-free methods such as Graph neural operators (GNO) (Li, Kovachki, et al., 2020c) and DeepONet (Lu, Jin, and Karniadakis, 2019) due to the advantages of the spectral transform, similar to standard FNO (De Hoop et al., 2022). Notice that the Geo-FNO 191 Model Geo-FNO (learned) GraphNO DeepONet Geo-FNO (R mesh) Geo-FNO (O mesh) FNO interpolation UNet interpolation mesh size 972 972 972 1681 1353 1681 1681 model size training time 1546404 1s 571617 32s 1025025 45s 1188417 0.6s 1188385 0.5s 1188353 0.5s 7752961 0.9s training error testing error 0.0125 0.0229 0.1271 0.1260 0.0528 0.0965 0.0308 0.0536 0.0344 0.0363 0.0314 0.0508 0.0089 0.0531 Table 9.1: Benchmark on the elasticity problem. Inputs are point clouds. The table presents the mesh size as the number of input mesh points, the model size as the quantity of training parameters, the training time for each epoch, and the relative training and test errors. with a learned deformation has better accuracy compared to Geo-FNO with fixed heuristic deformations (R-mesh) and (O-mesh). Remark 9.4.1 (Visualization of 𝜙−1 𝑎 ) Figure 9.5 shows the effects of the coordinate transform, where (a) is the prediction G𝜃 (𝑎)(𝑥) on the original physical mesh and (b) is the prediction on the computational mesh (without transforming back to the physical space) G𝜃 (𝑎)(𝜙−1 𝑎 (𝑥)). The top row visualizes the solution on the original 𝑖 −1 𝑖 mesh 𝑇 and 𝜙𝑎 (𝑇 ), while the bottom row is the full field solution on 𝑇 𝑐 directly evaluated on the Fourier coefficients. The solution on the latent space has more "white region", which is latent encoding that does not show up in the final solution. This latent encoding is similar to the Fourier continuation, making the solution easier to be represented with the Fourier basis. As shown in the figure, the solution function on the computational mesh has a cleaner wave structure and is more periodic compared to the physical space, allowing it to be better represented by the Fourier basis. Interestingly, there exists a vertical discontinuity on the latent space around 𝑥 = 0.5 in the latent encoding. Since most features are horizontal, the vertical discontinuity does not affect the output. The gap can be seen as a result of encoding in the high-dimensional channel space. Plastic problem We consider the plastic forging problem where a block of material Ω = [0, 𝐿]×[0, 𝐻] is impacted by a frictionless, rigid die at time 𝑡 = 0. The die is parameterized by an arbitrary function 𝑆 𝑑 ∈ 𝐻 1 ([0, 𝐿]; R) and traveled at a constant speed 𝑣. The block is clamped on the bottom edge and the displacement boundary condition is imposed 192 The first column is the original unstructured input; the second column is the interpolated input; the third column is the prediction; the last column is the error. As shown in the figures, interpolation causes an error, which is less accurate than geometry-aware methods. Figure 9.4: Interpolation into different meshes for the elasticity problem 193 The left column is the prediction of Geo-FNO on the physical space G 𝜃 (𝑎) (𝑥); the right column is the prediction of Geo-FNO on the computational space G 𝜃 (𝑎) (𝜉) = G 𝜃 (𝑎) (𝜙 𝑎−1 (𝑥)). The top row is the solution on the input mesh. The bottom row is the full-field solution. The latent space has a cleaner wave structure. Figure 9.5: Visualization of the deformation map on the top edge. The governing equation is the same as the previous example but with an elasto-plastic constitutive model given by 𝝈 = C : (𝝐 − 𝝐 𝑝 ) 𝝐¤𝑝 = 𝜆∇𝝈 𝑓 (𝝈) √︂ 3 1 |𝝈 − tr(𝝈) · 𝐼 | 𝐹 − 𝝈𝑌 𝑓 (𝝈) = 2 3 where 𝜆 is the plastic multiplier constrained by 𝜆 ≥ 0, 𝑓 (𝝈) ≤ 0, and 𝜆 · 𝑓 (𝝈) = 0. The isotropic stiffness tensor C is with Young’s modulus 𝐸 = 200 GPa and Poisson’s ratio 0.3. The yield strength 𝝈𝑌 is set to 70 MPa with the mass density 𝜌 𝑠 = 7850kg·m−3 . We generate 900 training data and 80 test data by using the commercial Finite Element software ABAQUS Smith, 2009, using 3000 4-node quadrilateral bi-linear elements (CPS4R in ABAQUS’s notation). Without lose of generality, we fix 𝐿 = 50mm, 𝐻 = 15mm, and 𝑣 = 3 ms−1 . For each sample, we generate a random function 𝑆 𝑑 by randomly sampling 𝑆 𝑑 (𝑥 𝑘 ) on {𝑥 𝑘 = 𝑘 𝐿/7; 𝑘 = 0, .., 7} with 194 The top row is the truth and the bottom row is the prediction. The five columns represent the changes in time. The color represents the norm of displacement. Figure 9.6: Geo-FNO on Plasticity a uniform probability measure. The die geometry is then generated by interpolating 𝑆 𝑑 (𝑥 𝑘 ) using piecewise Cubic Hermit Splines. It takes about 600 CPU seconds for each simulation. The target solution operator maps from the shape of the die to the time-dependent mesh grid and deformation. The data is given on the 101 × 31 structured mesh with 20 time steps. The plastic forging problem is a time-dependent problem, so we used the FNO3d model as the backbone to do convolution in both spatial dimension and temporal dimensions. Since the data is given on a structured mesh, the deformation (9.16) has an analytical form, so there is no need to learn a deformation. Geo-FNO is equivalent to directly applying the standard FNO on the structured mesh, and hence it perseveres the speed of standard FNO with an inference time of around 0.01 second per. In this experiment, Geo-FNO outputs both the position of the grid as well as the displacement at each grid point. As shown in fig. 9.6, Geo-FNO correctly predicts the evolution of the shape and the displacement. It has a moderate error on the top of the material, which is flat in the ground truth but curved in the prediction. Overall, Geo-FNO serves as an efficient surrogate model with test error 0.0074 (See table 9.2). Advection equation on unit sphere We consider the cos bell test introduced in Rasch, 1994, where a cos bell 𝑐(0, 𝑥) = 𝑐) )) is randomly generated and centered at 𝑥 𝑐 = (𝜆 𝑐 , 𝜃 𝑐 ) with ℎ𝑐 (1.0 + cos(𝜋 dist(𝑥,𝑥 𝑟 20𝜋 radius 𝑟 ∼ U[ 10𝜋 128 , 128 ] and height ℎ 𝑐 ∼ U[0.5, 1.5] on the unit sphere. Here dist 𝜋 denotes great-circle distance, 𝜆 𝑐 , 𝜃 𝑐 ∼ U[ −𝜋 3 , 3 ] denote longitude and latitude. The cos bell is advected following 𝜕𝑐 + ∇(𝒗𝑐) = 0, 𝜕𝑡 195 Model Geo-FNO FNO interpolation UNet interpolation Airfoil training testing 0.0134 0.0138 0.0287 0.0421 0.0039 0.0519 Pipe training testing 0.0047 0.0067 0.0083 0.0151 0.0109 0.0182 Plasticity training testing 0.0071 0.0074 − − − − The Plasticity requires the mesh as a target output, so interpolation methods do not apply. Table 9.2: Benchmark on airfoils, pipe flows, and plasticity. Inputs are structured meshes. Model topology mesh size FNO2D 𝑇2 128x64 FNO3D 𝐷3 64x64x64 Geo-FNO2D 𝐷 2 #𝐷 2 128x64 2 128x64 UNet2D 𝐷 128x64 DeepONet - model size 2368001 35397665 4736002 7752961 2624769 training error testing error 0.0119 0.0381 > 100% > 100% 0.0119 0.0332 0.1964 0.3132 0.1113 0.1599 Table 9.3: Advection equation on sphere
here the advective speed 𝒗(𝜆, 𝜃) = cos 𝛽 + sin 𝛽 tan 𝜃 cos 𝜆, −𝑈 sin 𝛽 sin 𝜆 with 𝛽 = 𝜋/2 and 𝑈 = 2𝜋/256. We generate 1000 training data and 200 test data with a finite volume solver. The inputs 𝑐(0, 𝑥) are given as point clouds with a size of around 8000 on the unit sphere. The target output is the solution 𝑐(𝑡, 𝑥) at 𝑡 = 256/3. As shown in Figure 9.7 and table 9.3, Geo-FNO can work on topology differently than torus or disk. Through domain decomposition, particularly by splitting the sphere into the northern and southern hemispheres, we can apply FFT2D for the spherical problem without raising dimensionality. It outperforms baseline models such as the original FNO and UNet models applied on the 2D spherical map projection, whose topology is not a sphere. We also compare against the FNO3D model that embeds the sphere into 𝑅 3 . It exhibits a remarkably high level of test error, which we attribute to the considerably larger training data and resolution demands associated with 3D learning. This underscores the significance of selecting an appropriate latent computation domain. Airfoil problem with Euler’s equation We consider the transonic flow over an airfoil, where the governing equation is the Euler equation, as follows, 𝜕𝜌 𝑓 + ∇ · (𝜌 𝑓 𝒗) = 0, 𝜕𝑡 𝜕𝜌 𝑓 𝒗 + ∇ · (𝜌 𝑓 𝒗 ⊗ 𝒗 + 𝑝I) = 0, 𝜕𝑡   𝜕𝐸 + ∇ · (𝐸 + 𝑝)𝒗 = 0 𝜕𝑡 196 Figure 9.7: Advection equation on the unit sphere. It shows that 2D Geo-FNO with domain decomposition can simulate PDE on general topologies. Figure 9.8: The airfoil flows (a) and pipe flows (b) introduced in section 9.4; The comparison is shown between the reference obtained using a traditional solver (top) and the Geo-FNO result (bottom). where 𝜌 𝑓 is the fluid density, 𝒗 is the velocity vector, 𝑝 is the pressure, and 𝐸 is the total energy. The viscous effect is ignored. The far-field boundary condition is 𝜌∞ = 1, 𝑝 ∞ = 1.0, 𝑀∞ = 0.8, 𝐴𝑜 𝐴 = 0 where 𝑀∞ is the Mach number and 𝐴𝑜 𝐴 is the angle of attack, and at the airfoil, no-penetration condition is imposed. The shape parameterization of the airfoil follows the design element approach Farin, 2014. The initial NACA-0012 shape is mapped onto a ‘cubic’ design element with 8 control nodes, and the initial shape is morphed to a different one following the displacement field of the control nodes of the design element. The displacements of control nodes are restricted to vertical direction only with prior 𝑑 ∼ U[−0.05, 0.05]. We generate 1000 training data and 200 test data with a second-order implicit finite volume solver. The C-grid mesh with about (200 × 50) quadrilateral elements is used, and the mesh is adapted near the airfoil but not around the shock. It takes about 1 CPU-hour for each simulation. The mesh point locations and Mach number on these mesh points are used as input and output data. 197 Similar with the plastic problem, the data for the airfoil problem is given on a structured mesh (Omesh), the deformation (9.16) is prescribed, so there is no need to learn a deformation. Geo-FNO is equivalent to directly applying the standard FNO on the structured mesh. The prediction of Geo-FNO on a sample test is presented in fig. 9.8-a, although there are two shocks induced by the airfoil, the Geo-FNO prediction matches well with the reference. We also compare Geo-FNO with the interpolation-based ML methods, which directly interpolates the target on a larger rectangular space. As shown in table 9.2, Geo-FNO outperforms both FNO and UNet with interpolation. In the following discusses, we explore discretization convergence of Geo-FNO model, throughly compare the cost-accuracy trade-off between Geo-FNO and traditional Euler solvers. And finally we demonstrate that we can conduct real-time design optimization with the trained Geo-FNO model. Discretization Convergence Discretization-convergence is an important property for surrogate models (Kovachki, Li, et al., 2021). It means the model can be applied to any discretizations and resolutions, and further, the same set of parameters can be transferred to different discretizations and resolutions. Such a design philosophy guides the research on neural operators to model the target solution operator as a mapping between function spaces, not just a specific model at one testing resolution. The graph neural operator and Fourier neural operator (with discrete Fourier transform) are both discretization-convergent. However, when implemented with the Fast Fourier transform, Fourier neural operator is restricted to the uniform grids and therefore it loses the discretization-convergent property. The proposed Geo-FNO model, on the other hand, extends FNO (with FFT) to non-uniform meshes, and retains discretization-convergence. Geo-FNO can be trained on a low-resolution dataset and evaluated at a higher resolution. On the Airfoil problem, we train GeoFNO on a 56 × 51 mesh. It achieves test errors of 0.0147, 0.0329, and 0.0428 on 56 × 51, 111 × 51, and 221 × 51 meshes, respectively. Conventional deep learning methods such as U-Nets are not capable of transferring among different resolutions. Cost-Accuracy Study Geo-FNO is used as the surrogate model to efficiently approximate expensive PDE simulations. Such surrogate modeling is an enabling methodology for many-query computations in science and engineering, e.g., design optimization. In principle, the relative merits of different surrogate models can be evaluated by understanding, for each one, the cost required to achieve a given level of accuracy. De Hoop et. al. demonstrates in (De Hoop et al., 2022), that FNO 198 has a superior cost-accuracy trade-off among different neural operator approaches. In this section, we study the cost-accuracy trade-off in comparison with traditional numerical solvers with different resolutions. Specifically, we consider the Euler equation airfoil test 9.4, where the reference data are generated by a 220×50 grid with 2000 implicit backward-Euler pseudo-time-stepping iterations. For comparison, we use the same solver but with different spatial resolutions: 220 × 50, 110 × 20, and 44 × 5 grids, and different time integrators, including implicit backward-Euler and explicit Runge-Kutta-2 schemes, with different pseudo time-stepping iterations. For each setup, we estimate the error and CPU time by an average of 10 runs sampled from the data set. The cost-accuracy trade-off is depicted in Figure 9.3. Geo-FNO is at least 104 x faster while having the same accuracy, for the presented test. We should also mention, the speed-up of Geo-FNO is due to the avoidance of time-stepping iterations, Euler flux computation, and GPU acceleration. Inverse design Once the Geo-FNO model is trained, it can be used to do the inverse design. We can directly optimize the design parameters to achieve the design goal. For example, as shown in fig. 9.9, the shape of the airfoil is parameterized by the vertical displacements of seven spline nodes. We set the design goal to minimize the drag lift ratio. We first train the model mapping from the input mesh to the output pressure field, then program the maps from the vertical displacement of spline nodes to input mesh and from the output pressure field to the drag lift ratio and finally optimize the vertical displacement of spline nodes in an end-to-end manner. As shown in the figure, the resulting airfoil becomes asymmetry with larger upper camber over the optimization iteration, which increases the lift coefficient and matches the physical intuition. Finally, we use the traditional numerical solver to verify optimized design shape. For the optimized design shape, both Geo-FNO and traditional numerical solver lead to a drag coefficient of 0.04 and a lift coefficient of 0.29. Pipe problem with Navier-Stokes equation Finally, we consider the incompressible flow in a pipe, where the governing equation is the incompressible Navier-Stokes equation, as follows, 𝜕𝒗 + (𝒗∇)𝒗 = −∇𝑝 + 𝜈∇2 𝒗, 𝜕𝑡 ∇·𝒗 =0 where 𝒗 is the velocity vector, 𝑝 is the pressure, and 𝜈 = 0.005 is the viscosity. The parabolic velocity profile with maximum velocity 𝒗 = [1, 0] is imposed at the 199 Figure 9.9: The inverse design for the airfoil flow problem with end-to-end optimization. The optimal design using the simulation from the numerical solver matches the prediction from Geo-FNO. inlet, the free boundary condition is imposed at the outlet, and no-slip boundary condition is imposed at the pipe surface. The pipe is of length 10 and width 1, and the centerline of the pipe is parameterized by 4 piecewise cubic polynomials, which are determined by the vertical positions and slopes on 5 spatially uniform control nodes. The vertical position at these control nodes obeys 𝑑 ∼ U[−2, 2] and the slop at these control nodes obeys 𝑑 ∼ U[−1, 1]. We generate 1000 training data and 200 test data with an implicit finite element solver with about 4000 Taylor-Hood Q2-Q1 mixed elements D. Z. Huang, Pazner, et al., 2020. It takes about 70 CPU seconds for each simulation. The mesh point locations (129 × 129) and horizontal velocity on these mesh points are used as input and output data. Similar with the plastic problem, the data for the pipe problem is given on a structured mesh, the deformation eq. (9.16) has an analytical form, so there is no need to learn a deformation. Geo-FNO is equivalent to directly applying the standard FNO on the structured mesh. The prediction of Geo-FNO on a sample test is presented in fig. 9.8-b, the Geo-FNO is able to capture the boundary layer, and the prediction matches well with the reference. Geo-FNO also outperforms these interpolationbased ML ethods, which directly interpolates the target on a larger rectangular space, as shown in table 9.2. 9.5 Discussion and Conclusion In this work, we propose a geometry-aware FNO framework (Geo-FNO) that applies to arbitrary geometries and a variety of input formats. The Geo-FNO deforms the irregular input domain into a uniform latent mesh on which the FFT can be applied. Such deformation can be fixed or learned with the FNO architecture in an end-toend manner. The Geo-FNO combines both the computational efficiency of the FFT 200 and the flexibility of learned deformations. It is as fast as the original FNO but can represent the variations in solutions and domain shapes more efficiently and accurately. In the end, we also discussed several potential extensions of Geo-FNO. Physics-informed settings. In this work, we mainly consider learning surrogate models in the data-driven setting. However, the dataset may not always be available. In this case, we can use the physics-informed setting by optimizing the physicsinformed equation loss. Given a fixed input 𝑎, the output function 𝑢 = G𝜃 (𝑎) can be explicitly written out using the formula (9.18) and (9.13). The derivative of the −1 deformed basis 𝜓𝑎 = 𝑒 2𝑖𝜋⟨𝜙 𝑎 (𝑥),𝑘⟩ can be exactly computed using chain rule with the auto-differentiation of the neural network 𝜙−1 𝑎 . Using the exact gradient, one can minimize the residual error R(G𝜃 (𝑎)) to find out the parameter G𝜃 (𝑎) that represents the solution function. The Geo-PINO method will be an optimization-based spectral method for general geometry. General topologies. In this work, we mainly studied simple topologies of 2D disks or 2D disks with holes. If the problem topology is more challenging, there does not exist a diffeomorphism from the physical space to the uniform computational space. Thankfully, we can use the Fourier continuation and decomposition to convert the problem domain into simpler ones. It is known that 2D connected orientable surfaces can be classified as either a sphere or an n-genus torus. For spheres, it is natural to use the unit sphere as the computational space and the spherical harmonics as the computational basis. For n-genus torus (𝑛 ≥ 2), usually, there do not exist useful harmonics series, but we can decompose the domain, which requires training multiple FNO models on each of the sub-domain in a coupled manner. We leave the domain decomposition as exciting future work. Theoretical guarantees. In the end, it will be interesting to extend the universal approximation bound of Fourier neural operator (Kovachki, Lanthaler, and Mishra, 2021) to the solution operator of PDEs defined on general geometries. In (Kovachki, Lanthaler, and Mishra, 2021), the approximation is achieved by using existing pseudo-spectral solvers. Since FNO’s architecture can approximate the operation in the pseudo-spectral solvers, FNO can approximate the target solution operators. For general domains, usually, there does not exist a pseudo-spectral solver. However, we can transform the problem into a computational space. By applying the universal approximation bound on the deformed equation R𝑐 = 0, as well as the approximation 201 bound for the neural network 𝜙−1 𝑎 , it is possible to obtain a bound for Geo-FNO. We also leave this direction as future work. References Allen, Kelsey R et al. (2022). “Physical Design using Differentiable Learned Simulators”. In: arXiv preprint arXiv:2202.00728. Babuška, Ivo and Manil Suri (1990). “The p-and hp versions of the finite element method, an overview”. In: Computer methods in applied mechanics and engineering 80.1-3, pp. 5–26. Bhatnagar, Saakaar et al. (2019). “Prediction of aerodynamic flow fields using convolutional neural networks”. In: Computational Mechanics, pp. 1–21. Bruno, Oscar P, Youngae Han, and Matthew M Pohlman (2007). “Accurate, highorder representation of complex three-dimensional surfaces via Fourier continuation analysis”. In: Journal of computational Physics 227.2, pp. 1094–1125. Chen, Tianping and Hong Chen (1995). “Universal approximation to nonlinear operators by neural networks with arbitrary activation functions and its application to dynamical systems”. In: IEEE Transactions on Neural Networks 6.4, pp. 911– 917. De Hoop, Maarten et al. (2022). “The Cost-Accuracy Trade-Off In Operator Learning With Neural Networks”. In: arXiv preprint arXiv:2203.13181. Driscoll, Tobin A, Nicholas Hale, and Lloyd N Trefethen (2014). Chebfun guide. Farin, Gerald (2014). Curves and surfaces for computer-aided geometric design: a practical guide. Elsevier. Gottlieb, David and Steven A Orszag (1977). Numerical analysis of spectral methods: theory and applications. SIAM. Huang, Daniel Z, Will Pazner, et al. (2020). “High-order partitioned spectral deferred correction solvers for multiphysics problems”. In: Journal of Computational Physics 412, p. 109441. Huang, Daniel Z, Kailai Xu, et al. (2020). “Learning constitutive relations from indirect observations using deep neural networks”. In: Journal of Computational Physics 416, p. 109491. Huang, Weizhang and Robert D Russell (2010). Adaptive moving mesh methods. Vol. 174. Springer Science & Business Media. Kovachki, Nikola, Samuel Lanthaler, and Siddhartha Mishra (2021). “On universal approximation and error bounds for Fourier Neural Operators”. In: Journal of Machine Learning Research 22, Art–No. Kovachki, Nikola, Zongyi Li, et al. (2021). “Neural operator: Learning maps between function spaces”. In: arXiv preprint arXiv:2108.08481. 202 Li, Zongyi, Nikola Kovachki, et al. (2020a). Fourier Neural Operator for Parametric Partial Differential Equations. arXiv: 2010.08895 [cs.LG]. – (2020b). Multipole Graph Neural Operator for Parametric Partial Differential Equations. arXiv: 2006.09535 [cs.LG]. – (2020c). “Neural operator: Graph kernel network for partial differential equations”. In: arXiv preprint arXiv:2003.03485. Li, Zongyi, Hongkai Zheng, et al. (2021). “Physics-informed neural operator for learning partial differential equations”. In: arXiv preprint arXiv:2111.03794. Lu, Lu, Pengzhan Jin, and George Em Karniadakis (2019). “DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators”. In: arXiv preprint arXiv:1910.03193. Lu, Lu, Xuhui Meng, et al. (2022). “A comprehensive and fair comparison of two neural operators (with practical extensions) based on fair data”. In: Computer Methods in Applied Mechanics and Engineering 393, p. 114778. Mildenhall, Ben et al. (2020). “Nerf: Representing scenes as neural radiance fields for view synthesis”. In: European conference on computer vision. Springer, pp. 405– 421. Pascon, João Paulo (2019). “Large deformation analysis of plane-stress hyperelastic problems via triangular membrane finite elements”. In: International Journal of Advanced Structural Engineering 11.3, pp. 331–350. Pathak, Jaideep et al. (2022). “FourCastNet: A Global Data-driven High-resolution Weather Model using Adaptive Fourier Neural Operators”. In: arXiv preprint arXiv:2202.11214. Pfaff, Tobias et al. (2020). Learning Mesh-Based Simulation with Graph Networks. arXiv: 2010.03409 [cs.LG]. Rasch, Philip J (1994). “Conservative shape-preserving two-dimensional transport on a spherical reduced grid”. In: Monthly weather review 122.6, pp. 1337–1350. Ronneberger, Olaf, Philipp Fischer, and Thomas Brox (2015). “U-net: Convolutional networks for biomedical image segmentation”. In: International Conference on Medical image computing and computer-assisted intervention. Springer, pp. 234– 241. Sanchez-Gonzalez, Alvaro et al. (2020). “Learning to simulate complex physics with graph networks”. In: arXiv preprint arXiv:2002.09405. Sitzmann, Vincent et al. (2020). “Implicit neural representations with periodic activation functions”. In: arXiv preprint arXiv:2006.09661. Smith, Michael (2009). ABAQUS/Standard User’s Manual, Version 6.9. English. United States: Dassault Systèmes Simulia Corp. 203 Wen, Gege et al. (2022). “U-FNO–An enhanced Fourier neural operator-based deeplearning model for multiphase flow”. In: Advances in Water Resources, p. 104180. Yang, Yan et al. (2021). “Seismic wave propagation and inversion with Neural Operators”. In: The Seismic Record 1.3, pp. 126–134. Zhu, Yinhao and Nicholas Zabaras (2018). “Bayesian deep convolutional encoder–decoder networks for surrogate modeling and uncertainty quantification”. In: Journal of Computational Physics. issn: 0021-9991. doi: https://doi.org/10.1016/ j.jcp.2018.04.018. url: http://www.sciencedirect.com/science/ article/pii/S0021999118302341. 204 C h a p t e r 10 GEOMETRY: NEURAL OPERATOR WITH OPTIMAL TRANSPORT Classical geometry learning problems typically focus on discretized meshes or point clouds. In this work, we extend geometry learning to an operator learning problem by generalizing discretized meshes to mesh density functions. Our approach formulates geometry embedding as an optimal transport (OT) problem, mapping the mesh density function to the uniform density function in reference space. Compared to previous methods with interpolation or shared deformation, our OT-based method has instance-dependent deformation, which is flexible and efficient. Further, for 3D simulations focused on surfaces, our OT-based neural operator embeds the surface geometry into 2D latent space. By constraining computations to the 2D surface manifold, we achieve significant computational efficiency gains compared to volumetric simulation. Experiments utilizing Reynolds-averaged Navier–Stokes equations (RANS) on the ShapeNet-Car and DrivAerNet datasets reveal that our method not only achieves superior accuracy but also significantly reduces computational expenses in terms of both time and memory usage compared to earlier machine learning models. Additionally, our model demonstrates improved accuracy on the FlowBench dataset, underscoring the benefits of utilizing instance-dependent deformation for datasets with highly variable geometries. 10.1 Introduction Handling complex 3D geometries remains one of the fundamental challenges in scientific computing. While standard numerical solvers based on finite element or spectral methods have been successful on simple regular domains, they struggle with complex geometries due to computationally expensive meshing processes that often require iterative refinement. The challenge of geometric modeling poses a significant obstacle across multiple domains, including fluid dynamics, solid mechanics, and earth science applications. This challenge is particularly evident in 3D aerodynamic simulations of automobiles, where a single shape takes over three hundred hours on CPU (Elrefaie, Dai, and Ahmed, 2024) or ten hours on GPU (Zongyi Li, N. B. Kovachki, et al., 2023). Machine learning methods have emerged as promising alternatives for solving partial 205 differential equations (PDEs) on complex geometries, offering dramatic improvements in computational efficiency (Bhatnagar et al., 2019; Pfaff et al., 2021; Thuerey et al., 2020; Hennigh et al., 2021). These approaches can operate effectively at lower resolutions compared to traditional numerical solvers, significantly reducing computational overhead. However, most existing ML-based methods are constrained to specific resolutions, limiting their flexibility and broader applicability. To address this limitation, we focus on neural operators, a recent breakthrough in scientific computing that offers a resolution-independent approach to solving PDEs. Neural operators for complex geometries. Neural operators represent an innovative class of data-driven models designed to directly learn the mapping of solution operators for PDEs in a mesh-free manner (Zongyi Li, N. Kovachki, et al., 2020a; N. Kovachki et al., 2023; Lu et al., 2021). Unlike conventional deep learning models, neural operators are designed to be invariant to discretization, making them particularly effective for solving PDEs. Significant advances in neural operator research have focused on addressing PDEs with complex geometries (Zongyi Li, D. Z. Huang, et al., 2022; M. Yin et al., 2024; Ahmad et al., 2024), primarily through embedding the geometries into uniform latent spaces where efficient spectral methods such as the Fast Fourier Transform (FFT) (Cooley and Tukey, 1965) can be applied. The Geometry-Aware Fourier Neural Operator (Geo-FNO) (Zongyi Li, D. Z. Huang, et al., 2022) marked a significant advancement in this direction by introducing a diffeomorphic mapping from physical to computational domains structured as regular grids. This innovation enabled the application of the FFT in latent computational spaces, dramatically improving computational efficiency. However, Geo-FNO faces two major limitations. It learns a shared deformation map for a class of shapes, which cannot address case-dependent geometry features. In addition, GeoFNO encodes and decodes the geometry using a Fourier transform based on computationally expensive matrix-vector multiplication, which restricts scaling to large 3D simulations. Building on these ideas, the Geometry-Informed Neural Operator (GINO) (Zongyi Li, N. B. Kovachki, et al., 2023) combined Graph Neural Operators (GNO) (Zongyi Li, N. Kovachki, et al., 2020b) with Fourier Neural Operators (FNO) (Zongyi Li, N. Kovachki, et al., 2020a). By leveraging the adaptability of graphs for local interactions and the computational efficiency of FFT for global physics, it became the first neural operator capable of tackling large-scale 3D aerodynamics challenges. 206 Despite its potential, the method struggles with the inherent limitations of graph embeddings’ locality and the high computational cost of 3D latent spaces. These challenges are especially evident in large-scale scenarios. Another line of works encode the geometry into structure-free tokens based on transformer (Hao et al., 2023; Wu et al., 2024; Alkin et al., 2024) or implicit neural representation (Y. Yin et al., 2022; Serrano et al., 2023; H. Chen et al., 2023; P. Y. Chen et al., 2022). These methods are flexible and generic, but they generally do not preserve geometric properties in their encoding. Meanwhile, their encoders are generally not invertible, which limits their applications in inverse meshing optimization and shape design. Despite these advances, current neural operator approaches continue to grapple with the computational burden of 3D PDEs and the challenge of learning operators across diverse geometries. To address these challenges, we propose reformulating geometry embedding as an optimal transport problem for each instance. This allows solution operators to be learned directly on the surface manifold, with deformation tailored to each instance. This novel approach fundamentally transforms the handling of complex geometries in neural operators. Geometry encoding with optimal transport Optimal transport provides a rigorous mathematical framework for determining the most efficient transformation between densitys. We leverage this framework by interpreting surface meshes as continuous density functions, where mesh density reflects the underlying geometric complexity and surface curvature. Our key insight is to formulate the geometry embedding problem as an optimal transport problem that maps these mesh density functions to uniform density functions in a latent space. Unlike traditional approaches that rely on direct projection and interpolation—which often result in problematic point clustering and density distortions—optimal transport inherently preserves the structural properties of the mesh while ensuring a smooth, physically meaningful transformation, as illustrated in Figure 10.1. This preservation property shares conceptual similarities with adaptive moving mesh methods (Budd, Russell, and Walsh, 2015), but offers flexibility for different topologies. Recent computational advances, particularly the Sinkhorn algorithm (Cuturi, 2013), have made optimal transport practically feasible by providing efficient approximations to the transport problem. Building on these developments, we demonstrate how optimal transport can effectively embed surface mesh sub-manifolds into latent 207 Figure 10.1: OT plans as bi-partite graphs. In the figure, the green nodes represent the input shape and the red nodes represent the latent grid. Compared to other graph mapping strategies such as ball connection and nearest neighbor connection, OT preserves the global measure, which is essential for computing integral operators. space while maintaining their essential geometric properties, as illustrated in Figure 10.2. Conformal mapping (Gu et al., 2004; Choi, Lam, and Lui, 2015) represents a common approach for mapping irregular meshes to canonical manifolds. These mappings preserve local angles, making them particularly suitable for Laplaciantype equations. The smoothness and harmonic properties of conformal mappings work exceptionally well with numerical solvers such as finite element methods, as they have less local distortion on each cell within a mesh. However, despite these advantages, conformal mapping is suboptimal for evaluating the integral operators widely employed in neural operators. For this reason, our work adopts optimal transport as the preferred geometric transformation. While optimal transport may not be as smooth locally, they preserve global measure, which is essential for integral operators. Our Contribution: In this work, we generalize geometry learning from discretized mesh points to mesh density function. Our key innovation lies in formulating geometry embedding as a pre-determined transform that maps the input mesh density function to the uniform density function in the canonical reference space. Such geometric transform is computed as optimal transport. We explore both the Monge formulation (transport maps) and Kantorovich formulation (transport plans), showing how this unified framework naturally encompasses previous approaches: transport maps generalize the deformation maps in Geo-FNO, while transport plans extend the graph representations in GINO. Building on this theoretical foundation, we introduce the optimal transport neural 208 Figure 10.2: Illustration of the optimal transport neural operator (OTNO). (a) OT Encoder: The surface mesh is encoded onto a latent computational mesh, and the OT coupling is visualized by representing the coordinates of points as RGB colors. (b) Fourier Neural Operator (FNO): Within the latent space, we apply S/FNO to calculate solutions on the latent mesh. (c) OT Decoder: We decode the solutions from the latent space back to the original surface mesh; here, the colors indicate solution values. operator (OTNO), which combines OT-based geometry encoding/decoding with (Spherical) Fourier Neural Operators (Zongyi Li, N. Kovachki, et al., 2020a; Bonev et al., 2023) (Figure 10.2). We use Kantorovich formulation to achieve sparse transport plan and implement it with Sinkhorn algorithm(Cuturi, 2013), and use Monge formulation to obtain bijective map and implement it with Projection pursuit Monge map (PPMM)(Meng et al., 2019). Our optimal transport technique enables crucial dimension reduction by embedding the surface manifolds from the 𝑑-dimensional ambient space into the (𝑑 − 1)dimensional latent spaces. This capability is particularly valuable for automotive and aerospace applications, where many critical simulations, including Reynoldsaveraged Navier-Stokes (RANS) and Large-Eddy Simulation (LES), fundamentally operate on boundary value problems. The input is a 2D surface design, and the desired outputs are surface quantities like pressure and shear velocity that determine total drag. We validate our method through simulations based on RANS equations on the ShapeNet-Car (Umetani and Bickel, 2018) and DrivAerNet-Car (Elrefaie, Morar, et al., 2024) datasets. Results under different sampling rates (Figure 10.3) show that our approach achieves the fastest convergence rate compared to baseline methods and achieves both the smallest error and shortest time cost when using full dataset. Given different sampling size, our model maintains robust performance in geometry representation and embedding. Moreover, our OTNO conducts geometry embedding for each shape individually, 209 Figure 10.3: Convergence Plot: This figure presents a comparison of convergence rates among different models. We vary the physical mesh size and scale the latent mesh size along the physical size. The data indicates that our model, OTNO, exhibits the fastest exponential convergence rate of 1.85, surpassing both GINO at 1.37 and GeoFNO at 1.32. distinguishing it from previous methods such as GeoFNO and GINO, which learned a shared deformation network across all geometries. As confirmed by our experiments on the FlowBench dataset, this feature enables our model to better handle datasets composed of a wider variety of shapes. Our main contributions are summarized as follows: 1. A novel optimal transport framework for mesh embedding that unifies and generalizes previous approaches by mapping mesh density functions to latent uniform density functions, bridging the gap between Geo-FNO and GINO methodologies. 2. The Sub-Manifold Method, a dimension-reduction technique for high-dimensional PDEs that restricts solution operators to 2D surface manifolds, implemented with optimal transport and coupled with latent spectral neural operators for efficient PDE resolution in reduced dimensional space, which achieve significant reduction in computational expense. 3. Comprehensive validation on industry-standard datasets—ShapeNet (3.7k points) (Chang et al., 2015) and DrivAerNet (400k points) (Elrefaie, Dai, and Ahmed, 2024)—demonstrates unprecedented efficiency in RANS pressure field prediction. Our method achieves a performance improvement of 2x-8x faster processing and 2x-8x smaller memory usage compared to current machine learning methods, while also slightly enhancing accuracy. Moreover, it is approximately 7,000 times faster than traditional approaches. 210 4. Optimal transport provides instance-dependent geometry embeddings. Our model notably excels across diverse geometries, as evidenced on the FlowBench dataset, which contains shapes with greater variability. 10.2 Problem Setting: Geometric Operator Learning In this work, we consider boundary value problems arising from partial differential equations (PDEs). Our aim is to learn the operator from the boundary geometries of PDEs to their boundary solutions. In previous works, the boundary shapes have been parameterized as discrete design parameters (Timmer, 2009), occupancy functions or signed distance functions Zongyi Li, N. B. Kovachki, et al., 2023. In this work, we model the geometries as supports of mesh density functions defined on an ambient Euclidean space. Suppose we have density functions 𝑓 ∈ F defined in domain R𝑑 . Let Ω 𝑓 = supp( 𝑓 ) for 𝑓 ∈ F be the manifolds in which the PDEs are defined. We assume that 𝜕Ω 𝑓 ∈ 𝑉 for any 𝑓 ∈ F , where 𝑉 is a compact subset of R𝑑 . Denote L and B as a partial differential operator and a boundary operator. The primary high-dimensional PDEs are as follows: L (𝑢) = ℎ, in Ω 𝑓 , B (𝑢) = 𝑏, in 𝜕Ω 𝑓 , where 𝜕Ω 𝑓 denote the boundary manifold of Ω 𝑓 , which is a sub-manifold of Ω 𝑓 . ℎ and 𝑏 are some fixed functions on R𝑑 . We assume (L, B) satisfies that for any (ℎ, 𝑏), the PDE (10.1) has a unique solution 𝑢 on Ω 𝑓 . Denote U as the set of solution functions. Then the solution operator of these high-dimensional PDEs is given by: G:F →U (10.1) 𝑓 ↦→ 𝑢 in Ω 𝑓 , Sub-Manifold Operator For the associated boundary value problems, the solution operator on the sub-manifold 𝜕Ω 𝑓 is given by: G ∗ : 𝑓 sub ↦→ 𝑢 sub where 𝑓 sub = ∫ 𝑓 | 𝜕Ω 𝑓 𝜕Ω 𝑓 𝑓 (𝑥)𝑑𝑆 in 𝜕Ω 𝑓 , (10.2) (𝑑𝑆 is a surface measure) denotes the normalized function of 𝑓 constrained to 𝜕Ω 𝑓 which is a density function on 𝜕Ω 𝑓 , and 𝑢 sub denotes our target solution function on 𝜕Ω 𝑓 . Our aim is to solve the operator G ∗ in Eq (10.2) on the sub-manifolds {𝜕Ω 𝑓 : 𝑓 ∈ F }. 211 For linear elliptic PDEs such as the Helmholtz equations on regular domains (Ihlenburg and Babuška, 1995; Hubbert, 1956), solution operators can be explicitly constructed on boundary sub-manifolds through a well-established process: defining basis functions for boundary inputs and solutions via singular value decomposition (SVD), then establishing a linear mapping between these bases. However, this approach breaks down for nonlinear problems like RANS and LES with complex geometries, where explicit linear mappings become mathematically impossible. To overcome this limitation, we introduce a novel approach that combines optimal transport for geometry embedding with neural operators learned directly on the surface sub-manifold, enabling efficient handling of nonlinear boundary value problems. Reynolds-Averaged Navier-Stokes Equations One common example in computational fluid dynamics is the shape design problem. Given a mesh with mesh density function 𝑓 ∈ 𝑉 ⊂ R3 , we aim to solve the Reynolds-averaged Navier–Stokes Equations in Ω 𝑓 = supp( 𝑓 ), which is the open volume outside of automotive or airfoil. 1 − Δv + (v · ∇)v + ∇p = h in Ω 𝑓 , 𝑅𝑒 (10.3) ∇v = 0 in Ω 𝑓 , v = b in 𝜕Ω 𝑓 , where v represents the time-averaged velocity field, b is boundary condition of velocity, and p denotes the time-averaged pressure field. The 𝑅𝑒 represents the Reynolds number, which accounts for turbulence effects in the averaged flow field. The vector h includes any external forces acting on the fluid. Given that the boundary condition b and vector h are fixed, these equations on various manifolds give rise to the solution operator G : 𝑓 ↦→ 𝑢 = (v, p). Our objective is to solve the pressure field p on the boundary manifold 𝜕Ω 𝑓 which is the surface of automotive or airfoil in practice. Correspondingly, the solution operator we target is a sub-manifold operator of the original 3D PDE solution operator G: G ∗ : 𝑓 sub ↦→ 𝑢 sub = p in Ω. (10.4) where 𝑓 sub is the density function of the surface mesh of the automotive or airfoil and a mass function on it in practice. Although this setting presents certain limitations, it remains a general framework applicable to many realistic 3D geometry design problems, where the desired solutions are typically concentrated on the surfaces of objects such as cars and airfoils. 212 10.3 Geometry Embedding as Optimal Transport In this section, our goal is to embed a complex geometric domain into a simpler latent geometric domain while simultaneously embedding the associated density function. To address this challenge, we construct a computational framework based on optimal transport, as described below. For convenience, we use 𝑓 , 𝑢 and Ω to denote the 𝑓 sub , 𝑢 sub and 𝜕Ω 𝑓 in Eq (10.2). According to the settings from Sec 10.2, 𝑓 is a density function defined on the complex geometric domain Ω and Ω = supp( 𝑓 ). We define a measure 𝜇 built from the density 𝑓 as follows: 𝑑𝜇 = 𝑓 (𝑥) 𝑑𝑥 on Ω. (10.5) The task is then to embed the physical density function 𝑓 on Ω into a latent density function 𝑔 on a simple geometric domain Ω∗ within the same metric space R𝑑 . This is equivalent to finding a transformation between measures 𝜇 and 𝜆, where 𝑑𝜆 = 𝑔(𝜉)𝑑𝜉 represents a uniform measure on a canonical geometric domain Ω∗ , such as a unit sphere or torus. And we use 𝑥 and 𝜉 to represent positions in Ω and Ω∗ . Thereby, we can model the geometries as the density functions (probability measures) and then encode these density functions using transport maps/plans, finally cooperate with the latent operators such as FNO to solve the PDEs, detailed as follows. Transport Map Given a transport map 𝑇 from latent measure 𝜆 to physical measure 𝜇, which is a function defined on latent domain Ω∗ : 𝑇 : (Ω∗ , 𝜆) → (Ω, 𝜇), 𝜉 ↦→ 𝑥, (10.6) we define the latent neural operator such that maps the transport map 𝑇 to the latent solution function 𝑣 on the latent domain Ω∗ : G : 𝑇 ↦→ 𝑣 on Ω∗ . (10.7) The final solution is 𝑢(𝑥) = 𝑣 ◦ 𝑇 −1 (𝑥) ∀𝑥 ∈ Ω. (10.8) Moreover, for any function 𝑛(𝑥) on the physical domain Ω, the transport map 𝑇 enables encoding it onto the latent domain Ω∗ through: 𝑛∗ (𝜉) = 𝑛(𝑇 (𝜉)), ∀𝜉 ∈ Ω∗ . (10.9) 213 Transport Plan Given a transport Plan 𝑃 from latent measure 𝜆 to physical measure 𝜇, which is a probability measure on Ω∗ × Ω with marginals 𝜇 on Ω and 𝜆 on Ω∗ : 𝑃 : (Ω∗ , 𝜆) × (Ω, 𝜇) → [0, 1], (10.10) (𝜉, 𝑥) ↦→ 𝑃(𝜉, 𝑥), we define the latent neural operator such that maps the marginal map of 𝑃 to the latent solution function 𝑣 on the latent domain Ω∗ ): ∫ 𝑃(·, 𝑥)𝑥𝑑𝜇(𝑥) ↦→ 𝑣 on Ω∗ . (10.11) G : ∫Ω 𝑃(·, 𝑥)𝑑𝜇(𝑥) Ω The final solution is ∫ 𝑢(𝑥) = 𝑃(𝜉, 𝑥)𝑣(𝜉)𝑑𝜆(𝜉) ∀𝑥 ∈ Ω. (10.12) Ω∗ Moreover, for any function 𝑛(𝑥) on the physical domain Ω, the transport plan 𝑃 enables encoding it onto the latent domain Ω∗ through: ∫ ∗ 𝑛 (𝜉) = 𝑃(𝜉, 𝑥)𝑛(𝑥)𝑑𝜇(𝑥). (10.13) Ω Monge Problem: OT Map Monge originally formulated the OT problem as finding the most economical map to transfer one measure to another (Monge, 1781). Later, Kantorovich introduced a relaxation of these strict transportation maps to more flexible transportation plans, solved using linear programming techniques (Kantorovich, 2006). In the following sections, we explore both types of OT methods—map and plan. Additionally, we analyze state-of-the-art methods for the deformation layer in operator learning, categorizing them into these two types. The transportation map 𝑇 : Ω∗ → Ω is measure preserving, if for any Borel set 𝐵 ⊆ Ω, ∫ ∫ 𝑑𝜆(𝜉) = 𝑇 −1 (𝐵) 𝑑𝜇(𝑥), (10.14) 𝐵 denoted as 𝑇#𝜆 = 𝜇. Monge raised the OT problem: given the transportation cost function 𝑐 : Ω∗ × Ω → R+ , we aim to find a transportation map 𝑇 : Ω∗ → Ω that minimizes the total transportation cost, ∫ (MP) inf{ 𝑐(𝜉, 𝑇 (𝜉))𝑑𝜆(𝜉) | 𝑇#𝜆 = 𝜇}. (10.15) Ω∗ 214 Moreover, the following theorem and continuity property hold for the Monge formulation. Theorem 10.3.1 (Existence and uniqueness of transport map (Brenier, 1991)) Suppose the measures 𝜇 and 𝜆 are with compact supports Ω, Ω∗ ⊆ R𝑑 respectively, and they have equal total mass 𝜇(Ω) = 𝜆(Ω∗ ). Assume the corresponding density functions satisfy 𝑓 , 𝑔 ∈ 𝐿 1 (R𝑑 ), and the cost function is 𝑐(𝜉, 𝑥) = 12 |𝜉 − 𝑥| 2 , then the OT map from 𝜆 to 𝜇 exists and is unique. It can be expressed as 𝑇 (𝜉) = 𝜉 + ∇𝜙(𝜉), where 𝜙 : Ω∗ → R is a convex function, and 𝜙 is unique up to adding a constant. Lemma 10.3.1 (Continuity of transport map) Given that cost function is the squared Euclidean distance and 𝜆 is a measure with uniform density function with a compact support, if 𝜇 is absolutely continuous and strictly positive also with a compact support, then the OT map 𝑇 is continuous almost everywhere. (This lemma can be easily derived from Theorem 10.3.1.) In practice, the measures are discretized on the grids and represented by the discrete measure: 𝑛 𝑛 ∑︁ ∑︁ 𝜇(𝑥) = 𝜇𝑖 𝛿(𝑥 − 𝑥𝑖 ), 𝜆(𝜉) = 𝜆𝑖 𝛿(𝜉 − 𝜉𝑖 ), (10.16) 𝑖=1 𝑖=1 ) , 𝜆𝑖 = 𝑛1 , X = {𝑥𝑖 } are the locations of the vertex in complex where 𝜇𝑖 = Í𝑛𝑓 (𝑥𝑓𝑖(𝑥 𝑖) 𝑖=1 geometric domain Ω, 𝜉 = {𝜉𝑖 } are the positions of the computational mesh in Ω∗ and Í𝑛 Í𝑛 𝜆𝑖 satisfy measure perserving. 𝜇𝑖 = 𝑖=1 𝛿 is the Dirac delta function. Thus 𝑖=1 The corresponding discrete Monge problem is as follows: (MP) min 𝑛 ∑︁ 𝑐(𝜉𝑖 , 𝑇 (𝜉𝑖 )) · 𝜆𝑖 . (10.17) 𝑇:𝜉→X 𝑖=1 As for Geo-FNO(Zongyi Li, D. Z. Huang, et al., 2022), the deformation map can be viewed as generalized OT map with a special choice of cost function: ∫ ∫ −1 𝑐(𝑥, 𝑇 (𝑥))𝑑𝜇(𝑥) := 𝐺 (𝑇 −1 (𝑥)) − 𝑢 sub (𝑥)𝑑𝜇(𝑥) (10.18) Ω Ω where 𝑢 sub is our target solution function as presented in Eq (10.2). However, this is a non-standard, generalized cost function, since it depends not on (𝑥, 𝑇 −1 (𝑥)), but the global solution operator 𝐺 and domain Ω. 215 Similar to learning the function 𝜙 (defined in Theorem 10.3.1) in Monge’s formulation, Geo-FNO (Zongyi Li, D. Z. Huang, et al., 2022) implements a skip connection 𝜉 = 𝜓(𝑥) + 𝑥 within the deformation network to learn only 𝜓. The difference is that Geo-FNO adopts an end-to-end approach, optimizing the deformation map based on the final solution error and learning a shared network to implement geometry deformation. In contrast, the Monge problem learns the deformation map separately from the operator learning and optimizing based on the instance-specific transportation cost for each geometry respectively. Despite these differences in optimization strategies, both approaches can be categorized as employing a map-type for the deformation layer. Kantorovich Problem: OT Plan Relaxing the map constraint, the Kantorovich formulation seeks probability measures 𝑃 on Ω∗ × Ω that attains the infimum, ∫ (KP) inf{ 𝑐(𝜉, 𝑥)𝑑𝑃(𝜉, 𝑥) | 𝑃 ∈ Γ(𝜆, 𝜇)}. (10.19) Ω∗ ×Ω where Γ(𝜆, 𝜇) denotes the collection of all probability measures on Ω∗ × Ω with marginals 𝜆 on Ω∗ and 𝜇 on Ω, and 𝑐 : Ω∗ × Ω → R+ is transportation cost function. For its discrete implementation, we define 𝜉 = {𝜉1 , . . . , 𝜉𝑛1 } and X = {𝑥 1 , . . . , 𝑥 𝑛2 } as vertex locations in the domains Ω∗ and Ω, where 𝜉 is a uniform mesh in computational space. Denote 𝑎 = (𝜆(𝜉1 ), . . . , 𝜆(𝜉𝑛1 )) = 𝑛11 1𝑛1 and 𝑏 = (𝜇(𝑥 1 ), . . . , 𝜇(𝑥 𝑛2 )) as the density mass vectors of the discrete probability measure 𝜆 and uniform measure 𝜇, respectively. Define 𝑀 ∈ R𝑛1 ×𝑛2 as the cost matrix, where each element represents the cost function value 𝑐(𝜉𝑖 , 𝑥 𝑗 ). The discrete OT problem then seeks to minimize the total transport cost, known as the Wasserstein distance: (KP) 𝑊 𝑀 (𝑎, 𝑏) = min ⟨𝑃, 𝑀⟩ 𝑃∈Γ(𝑎,𝑏) = min 𝑛1 ∑︁ 𝑛2 ∑︁ 𝑀𝑖 𝑗 · 𝑃𝑖 𝑗 𝑃∈Γ(𝑎,𝑏) = min 𝑖=1 𝑗=1 𝑛1 ∑︁ 𝑛2 ∑︁ 𝑐(𝜉𝑖 , 𝑥 𝑗 ) · 𝑃𝑖 𝑗 , (10.20) 𝑃∈Γ(𝑎,𝑏) 𝑖=1 𝑗=1 where Γ(𝑎, 𝑏) represents the set of all feasible coupling matrices, defined as the discrete probability measures Γ(𝑎, 𝑏) = {𝑃 ≥ 0 | 𝑃1𝑛2 = 𝑎, 𝑃𝑇 1𝑛1 = 𝑏}. Moreover, the following sparity property holds for the discrete implementation of the Kantorovich formulation (Peyré, Cuturi, et al., 2019): 216 Lemma 10.3.2 (Sparsity of Transport Plan) The solution to the linear programming is sparse; the number of non-zero entries in the transport plan 𝑃 is at most 𝑛1 + 𝑛2 − 1. Similar to learning the optimal coupling matrix in the Kantorovich formulation, GINO (Zongyi Li, N. B. Kovachki, et al., 2023) employs a Graph Neural Operator (GNO), which can be interpreted as a learnable Graph Laplacian. GNO constructs the adjacency matrix by performing neighborhood searches within a fixed radius and learns the edge weights using kernel functions. In contrast, the Kantorovich formulation solves a global transport plan by optimizing the Wasserstein distance, which accounts for the pairwise distances between all points across the two domains. Despite these differences in optimization strategies, both approaches can be categorized as employing a plan-type transformation layer. Cost function We choose the squared Euclidean distance as the cost function in optimal transport due to both its mathematical convenience and its relevance in geometric applications. In the context of transporting probability measures between two geometric domains embedded in 3D space, the squared Euclidean distance 𝑐(𝑥, 𝑦) = ∥𝑥 − 𝑦∥ 2 naturally captures the geometric cost of moving mass from one location to another. It emphasizes longer displacements more heavily, encouraging smoother and more localized transport plans. Moreover, it is widely used in the literature and aligns with the assumptions of classical results like Brenier’s theorem, which guarantees the existence of a unique optimal map under this cost. Thus, it serves as a principled and standard choice for geometric transformation tasks in 3D Euclid space. 10.4 Optimal Transport Neural Operator (OTNO) In this paper, we introduce a novel model, the optimal transport neural operator (OTNO), which efficiently integrates optimal transport with neural operators. Our model employs the Projection pursuit Monge map to obtain an approximate solution of the Monge OT map 𝑇, and employs Sinkhorn method (Cuturi, 2013) to obtain an approximate solution of the Kantorovich OT plan 𝑇. The resulting Map/plan is then utilized to construct the OT encoder/decoder for neural operator. The methodology and implementation are detailed below. 217 OTNO Algorithm Given a dataset {(X𝑗 , 𝑢 𝑗 )} 𝑁𝑗=1 of surface sampling meshes and corresponding PDEs’   𝑛1𝑗   𝑛1𝑗 1 ∈ ∈ R𝑛 𝑗 ×3 and 𝑢 𝑗 = 𝑢 𝑗,𝑘 𝑘=1 solution values on surface, where X𝑗 = 𝑥 𝑗,𝑘 𝑘=1 2   2 1 𝑛 𝑗 ∈ R𝑛 𝑗 ×3 using a R𝑛 𝑗 ×𝑠 . For each X𝑗 , generate a latent surface mesh Ξ 𝑗 = 𝜉 𝑗,𝑙 𝑙=1 2D parametric mapping from a square grid, where 𝑛2𝑗 is a perfect square. Compute 1 2 the normals N 𝑗 ∈ R𝑛 𝑗 ×3 on X𝑗 and H 𝑗 ∈ R𝑛 𝑗 ×3 on Ξ 𝑗 . By solving OT problem from each latent mesh Ξ 𝑗 to each physical sampling mesh h i 𝑛2𝑗 2 𝑁 ′ ′ ∈ R𝑛 𝑗 ×3 . X𝑗 , we obtain a set of transported mesh {X𝑗 } 𝑗=1 , where X𝑗 = 𝑥 ′𝑗,𝑙 𝑙=1 Note that X𝑗′ serves as a representation of the physical surface from which 𝑋 𝑗 is sampled. Algorithm 3 Optimal Transport Neural Operator (OTNO) 𝑁 ′ 𝑁 1: Given physical mesh {X 𝑗 } 𝑁 𝑗=1 , latent mesh {Ξ 𝑗 } 𝑗=1 , transported mesh {X 𝑗 } 𝑗=1 and solution values {𝑢 𝑗 } 𝑁𝑗=1 . 2: Initialize a FNO G𝜃 . 3: for 𝑗 = 1 to 𝑁 do 4: 1. Build Index Mapping: ! 5: Encoder indices E = arg min ∥𝑥 ′𝑗,𝑙 − 𝑥 𝑗,𝑘 ∥ 2 : 𝑙 = 1, . . . , 𝑛2𝑗 𝑘=1,...,𝑛1𝑗 ! 6: Decoder indices D = arg min ∥𝑥 ′𝑗,𝑙 − 𝑥 𝑗,𝑘 ∥ 2 : 𝑘 = 1, . . . , 𝑛1𝑗 𝑙=1,...,𝑛2𝑗 2 2. OT encoder: M 𝑗 = X𝑗 (E) ∈ R𝑛 𝑗 ×3 , where M 𝑗 selects rows from X𝑗 according to the indices specified in E. 2 8: 3. Latent FNO: 𝑣 𝑗 = G𝜃 (T𝑗 ), where T𝑗 = (Ξ 𝑗 , M 𝑗 , H 𝑗 × N 𝑗 (E)) ∈ R𝑛 𝑗 ×9 9: (H 𝑗 × N 𝑗 (E) computes the cross product between rows). 1 10: 4. OT decoder: 𝑢′𝑗 = 𝑣 𝑗 (D) ∈ R𝑛 𝑗 ×𝑠 , where 𝑢 𝑗 selects rows from 𝑣 𝑗 according to the indices specified in D. 11: end for 𝑁 ∑︁ ∥𝑢′𝑗 − 𝑢 𝑗 ∥ U . 12: Compute the empirical loss over all dataset instances: 7: 𝑗=1 Encoder For each point 𝑥 ′𝑗,𝑙 (𝑙 = 1, . . . , 𝑛2𝑗 ) in the transported mesh X𝑗′ , we find the closet point in X𝑗 and denote 𝑒 𝑘 as the index of this closest point in !X𝑗 . Thus, we obtain an index sequence E = arg min ∥𝑥 ′𝑗,𝑙 − 𝑥 𝑗,𝑘 ∥ : 𝑙 = 1, . . . , 𝑛2𝑗 = 𝑘=1,...,𝑛1𝑗 218 2 (𝑒 1 , . . . , 𝑒 𝑛2 ). Using these indices, we encode the mesh X𝑗 to M 𝑗 = X𝑗 (E) ∈ R𝑛 𝑗 ×3 , 𝑗 where M 𝑗 selects rows from X𝑗 according to the indices specified in E. Latent Operator In the latent space, we deploy the FNO G𝜃 to execute the latent neural operator G as introduced in Eq (10.7)(10.11). Denote 𝑇 as the transport map in Eq (10.7) or marginal map of transport plan in Eq (10.11). Then it can be fundamentally represented by pairs of latent mesh and transported mesh (Ξ 𝑗 , 𝑇 (Ξ 𝑗 )) = (Ξ 𝑗 , X𝑗′ ). Considering 𝑇 as a deformation map, we include deformation-related features by using the cross-product over normals H 𝑗 × 𝑇 (H 𝑗 ) (further discussion on normal feature is in Sec 10.6). Thus, we configure T𝑗 as (Ξ 𝑗 , X𝑗′ , H 𝑗 × 𝑇 (H 𝑗 )) to fully encapsulate the map’s properties. To enhance the quantity of surface representation, we substitute points in X𝑗′ with their closest counterparts in X𝑗 , i.e. use M 𝑗 = X𝑗 (E) to replace X𝑗′ . Similarly, we use N 𝑗 (E) to replace the 𝑇 (H 𝑗 ). Consequently, T𝑗 = (Ξ 𝑗 , M 𝑗 , H 𝑗 × N 𝑗 (E)) serves as a comprehensive representation of the deformation map 𝑇, as well as the input for the latent FNO. Decoder Corresponding to the encoder, we compute the index 𝑑 𝑘 of the closest point in the transported mesh X𝑗′ for each point 𝑥 𝑗,𝑘 (𝑘 = 1, . . . , 𝑛1𝑗 ) in X𝑗 and ! build a decoder index sequence D = arg min ∥𝑥 ′𝑗,𝑙 − 𝑥 𝑗,𝑘 ∥ 2 : 𝑘 = 1, . . . , 𝑛1𝑗 = 𝑙=1,...,𝑛2𝑗 (𝑑1 , . . . , 𝑑𝑛1 ). Using these indices, we decode the solutions 𝑣 𝑗 back to physical 𝑗 surface. 2 Note that the latent surface mesh Ξ 𝑗 ∈ R𝑛 𝑗 ×3 is constructed from a 2D parametric 2 grid. This ensures that the input features 𝑇 𝑗 ∈ R𝑛 𝑗 ×9 for the FNO are organized on a 2D grid. Consequently, the FNO computations occur in 2D space, leveraging the structured layout of Ξ 𝑗 , rather than directly handling the unstructured 3D point cloud. OTNO with OT Plan Transported Mesh Using the Sinkhorn method (Cuturi, 2013) detailed in the following Sec 10.4, we obtain dense coupling matrices 𝑃 𝑗 that represent the OT plans from the latent computational mesh Ξ 𝑗 to the boundary sampling mesh X𝑗 . How to effectively utilize these dense matrices within a neural operator framework become a problem worth discussing. 219 Directly applying the dense matrix by multiplying it with the mesh will lead to a lower accuracy of the final predictions because it only approximates the solution plan for the Kantorovich problem (10.20). And saving large-scale dense matrices is costly. Therefore we discuss how to use these optimal coupling matrices more effectively. As the matrix 𝑃 𝑗 represents a discrete probability measure on X𝑗 × Ξ 𝑗 , a practical approach is to focus on the maximum probability elements, referred to as the "Max" strategy. Additionally, a more effective "Mean" strategy involves replacing each point 𝑥 ′ ∈ X𝑗′ = 𝑃 𝑗 X𝑗 with the nearest point in X𝑗 , significantly reducing approximation ambiguity caused by Sinkhorn method. Further, we can encode/decode by the top-k nearest neighbors in X𝑗 /X𝑗′ instead of only find the single nearest one. Details of these different strategies are further discussed in Section 10.6. Latent Mesh In practice, we design the computational grid to have more vertices than the boundary sampling mesh to ensure maximal information retention during the encoding and decoding processes. Let 𝛼 be √︃ an expansion factor, and the number of points in 2 latent space is then set to 𝑛 𝑗 = ⌈ 𝛼 × 𝑛1𝑗 ⌉ 2 . This approach implies that the encoder functions as an interpolator, while the decoder acts as a selective querier. A simple illustration of the encoder and decoder processes is shown in Fig.10.4. And detailed ablation studies for latent mesh shape and latent mesh size are in Sec 10.6 and Sec 10.6. Figure 10.4: Illustration of OT encoder and OT decoder. The curve represents the boundary sampling points and the line denotes the latent computational grid Voxel Downsampling To achieve a well-balanced input density function, we employ voxel downsampling as a normalization process that constrains the density function within a predefined range [0, 𝑎], where 𝑎 is determined by the voxel size. This approach mitigates excessive clustering in regions with high point density, 220 such as around the wheels in some some car data, therefore ensuring a more uniform spatial distribution of points. Sinkhorn Algorithm Sinkhorn (Cuturi, 2013) method added an entropy regularizer to the Kantorovich potential and greatly improved efficiency. They first propose the Sinkhorn distance that added entropy constraint: 𝑑 𝑀,𝛼 (𝑎, 𝑏) = min ⟨𝑃, 𝑀⟩, 𝑃∈Γ 𝛼 (𝑎,𝑏) (10.21) where Γ𝛼 (𝑎, 𝑏) = {𝑃 ∈ Γ(𝑎, 𝑏) | KL(𝑃, 𝑎𝑏𝑇 ) ≤ 𝛼} = {𝑃 ∈ Γ(𝑎, 𝑏) | ℎ(𝑃) ≥ ℎ(𝑎) + ℎ(𝑏) − 𝛼}. (10.22) Then they consider the dual problem that arises from Lagrange multiplier: 𝑃 𝛽 = argmin 𝑃∈Γ(𝑎,𝑏) ⟨𝑃, 𝑀⟩ − 1 ℎ(𝑃), 𝛽 (10.23) This formulation leads to a problem where 𝛽 adjusts the trade-off between the transport cost and the entropy of the transport plan 𝑃. When 𝛽 increases, the influence of the entropy regularization decreases, making 𝑃 𝛽 converge closer to the solution of the original Kantorovich problem (10.20). This implies that a larger 𝛽 leads to a solution that is more accurate and economically efficient. By introducing the entropy constraint, the Sinkhorn distance not only regularizes the OT problem but also ensures that the solution is computationally feasible even for large-scale problems. This regularization dramatically improves the numerical stability and convergence speed of the algorithm. Implementation In the implementation, we set 𝑎 = 𝑛11 1𝑛1 given that the sampling mesh X = {𝑥1 , . . . , 𝑥 𝑛1 } from the computational fluid dynamics (CFD) simulations typically demonstrates advantages such as increased point density in regions with sharp changes. This uniform mass vector on the mesh X provides a denser distribution in these critical regions, enhancing aerodynamics prediction. For cases with excessively dense regions, we employ voxel downsampling to prevent excessive density variations, thereby maintaining the feasibility of our uniform mass vector. We set 𝛽 = 1 × 106 to ensure that the solution is economical without incurring 221 excessive computational costs. For computational support, we utilize the geomloss implementation from the Python Optimal Transport (POT) library (Flamary et al., 2021), which supports GPU acceleration and use lazy tensors to store dense coupling matrix obtained from Sinkhorn method. OTNO with OT Map Transported Mesh Using the PPMM, detailed in Section 10.4, we obtain the transported mesh 𝑋 ′𝑗 , which is mapped from the latent mesh Ξ 𝑗 to the physical mesh 𝑋 𝑗 . This results in the same number of points as the latent mesh while representing the distribution of the physical mesh. Since PPMM operates directly on the latent mesh, projecting it towards the physical mesh (as shown in Algorithm 4), the output is the transported mesh itself rather than a mapping that can operate on any function on the latent mesh, and it cannot offer an inverse direction itself. (An OT plan, discretized as a coupling matrix, can handle both.) Therefore, we also compute the indices for encoding and decoding, as outlined in Algorithm 3. Latent Mesh Since the latent mesh has the same number of points as the physical mesh, and the 2-dimensional FNO on the latent space requires a square mesh, the number of points in both the latent mesh and the physical mesh should be a perfect square. Therefore, we first apply voxel downsampling to normalize the mesh density, followed by random sampling to further downsample the mesh so that ensure the number of points is a perfect square. Finally, we generate the latent mesh to match the square number of points in the downsampled physical mesh. We choose a spherical mesh as the latent mesh, given its suitability for the Projection Pursuit Monge Map (PPMM). Projection pursuit Monge map (PPMM) The PPMM proposes an estimation method for large-scale OT maps by combining the concepts of projection pursuit regression and sufficient dimension reduction. As summarized in Algorithm 4, in each iteration, the PPMM applies a one-dimensional OT map along the most "informative" projection direction. The direction 𝑒 𝑘 is considered the most "informative" in the sense that the projected samples 𝑋 𝑘 𝑒 𝑘 222 Algorithm 4 Projection pursuit Monge map Input: two matrix 𝑋 ∈ R𝑛×𝑑 and 𝑌 ∈ R𝑛×𝑑 𝑘 ← 0, 𝑋0 ← 𝑋 repeat (a) calculate the most ’informative’ projection direction 𝑒 𝑘 ∈ R𝑑 between 𝑋 𝑘 and 𝑌 (b) find the one-dimensional OT Map 𝜙 (𝑘) that matches 𝑋 𝑘 𝑒 𝑘 to 𝑌 𝑒 𝑘 (using look-up table) (c) 𝑋 𝑘+1 ← 𝑋 𝑘 + (𝜙 (𝑘) (𝑋 𝑘 𝑒 𝑘 ) − 𝑋 𝑘 𝑒𝑇𝑘 )𝑒𝑇𝑘 and 𝑘 ← 𝑘 + 1 until converge The final mapping is given by 𝜙ˆ : 𝑋 → 𝑋 𝑘 and 𝑌 𝑒 𝑘 exhibit the greatest "discrepancy." The specific method for calculating this direction is detailed in Algorithm 1 in paper Meng et al., 2019. 10.5 Experiments We conducted experiments on three CFD datasets. Two of them are 3D car datasets, where the target prediction is the pressure field or drag coefficient, which depends solely on the car surface—a 2D manifold. The ShapeNet dataset includes 611 car designs, each with 3.7k vertices and corresponding average pressure values, following the setup from (Zongyi Li, N. B. Kovachki, et al., 2023). The DrivAerNet dataset, sourced from (Elrefaie, Dai, and Ahmed, 2024), contains 4k meshes, each with 200k vertices, along with results from CFD simulations that measure the drag coefficient. Additionally, we further evaluate our model on the 2D FlowBench dataset (Tali et al., 2024), which features a wider variety of shapes, including three groups, each containing 1k shapes with a resolution of 512 × 512. Although the PDE solutions are not on the boundary sub-manifold, preventing our model from reducing dimensions by embedding boundary geometries, this dataset provides an excellent opportunity to explore our model’s capabilities across diverse shapes. Experiments: 3D Car Datasets ShapeNet Car Dataset To ensure a fair comparison, we maintained identical experimental settings to those used in GINO paper (Zongyi Li, N. B. Kovachki, et al., 2023). We compared the relative error, total time (including data processing and training), and GPU memory usage against key baselines from (Zongyi Li, N. B. Kovachki, et al., 2023). As detailed in Table.10.1, our method, OTNO, achieved a relative error of 6.70%, 223 Model Geo-FNO UNet GINO OTNO(Plan) Relative L2 Total Time (hr) GPU Memory (MB) 13.50% 13.14% 7.21% 6.70% 1.96 6.86 10.45 1.52 12668 7402 12734 1890 Table 10.1: Predict pressure field on ShapeNet Car Dataset which is a slight improvement over the lowest error reported in the baseline—7.21% by GINO. Across all baselines, our method demonstrated reduced total time and lower GPU memory usage. Compared to GINO, our method significantly reduced both time costs and memory expenses by factors of eight and seven, respectively. The visual results for pressure prediction are presented in Fig.10.5a. DrivAerNet Car Dataset Herein, we follow all the experimental settings from DrivAerNet paper (Elrefaie, Dai, and Ahmed, 2024) and compare our model to RegDGCNN, as proposed in (Elrefaie, Dai, and Ahmed, 2024), and the state-of-the-art neural operator model, GINO (Zongyi Li, N. B. Kovachki, et al., 2023). Note that we do not use pressure data for training; instead, we only use the drag coefficient (Cd). The visual result for Cd prediction is presented in Fig.10.5b. For OTNO, we employ voxel downsampling with a size of 0.05, corresponding to approximately 60k samples (see ablation in 10.6). As detailed in Table.10.2, OTNO(Plan) demonstrates significantly superior accuracy and reduced computational costs compared to RegDGCNN, halving the MSE, speeding up computations by a factor of 5, and reducing memory usage by a factor of 24. Compared to GINO, OTNO(Plan) achieves a slightly lower MSE, a marginally higher R2 score, and notably reduces total computation time and memory usage by factors of 4 and 5, respectively. Given that PPMM is designed for large-scale OT maps, we evaluate OTNO(Map) on this large-scale dataset to assess its performance. Unfortunately, both the error and time cost are worse than OTNO(Plan). However, the memory cost is significantly lower. This is primarily because OTNO(Plan) expands the latent space, while OTNO(Map) does not. 224 Model RegDGCNN GINO OTNO(Map) OTNO(Plan) MSE (e-05) R2 Score Total Time (hr) GPU Memory (MB) 6.63 3.33 3.93 3.28 0.887 0.955 0.947 0.956 10.78 7.73 10.63 5.26 72392 14696 2896 9702 Table 10.2: Predict drag coefficient (Cd) on DrivAerNet Car Dataset (single A100) (a) Results of pressure on ShapeNet dataset. (b) Results of drag coefficient on DrivAerNet dataset Figure 10.5: Results visualization for OTNO on Car Datasets Showcase of Dataset with Diverse Geometries - 2D Flow Datasets To assess the performance of instance-dependent deformation in our model across diverse geometries, we conducted further evaluations using the FlowBench dataset (Tali et al., 2024). We use two metrics for training: 1. M1: Global metrics: The errors of in velocity and pressure fields over the 225 entire domain. 2. M2: Boundary layer metrics: The errors in velocity and pressure fields over a narrow region around the object. We define the boundary layer by considering the solution conditioned on the Signed Distance Field (0 ≤ 𝑆𝐷𝐹 ≤ 0.2). The results of training using the M1 metric (global) and M2 metric (boundary) are presented in Table 10.3, and Table 10.4, respectively. We present a subset of the visualization results in Fig. 10.6. The results demonstrate that OTNO significantly outperforms in accuracy under both the M1 (global) and M2 (boundary) metrics, particularly for M1. Furthermore. And OTNO achieves notably better accuracy across all three groups, especially for G2 (harmonics). However, cost reduction is not observed in the FlowBench dataset. It is important to note that the cost reductions seen in car datasets stem from the sub-manifold method, which employs optimal transport to generate 2D representations and perform computations in 2D latent space instead of 3D. In the FlowBench dataset, however, the solutions are not restricted to the boundary submanifold. Even for the M2 metric, although the relevant data is close to the boundary with a width of 0.2, it does not reduce to a 1D line. As a result, computations cannot be confined to the sub-manifold, and the associated cost reduction benefits are consequently absent. We do not present Geo-FNO results on this dataset as the relative L2 errors consistently exceed 60%. Our analysis suggests that Geo-FNO, which uses an end-to-end approach to learn a shared deformation map and the latent operator, is less suited for the diverse shapes present in the FlowBench dataset. In contrast, our OTNO model, which solves the OT plan/map for each shape separately, exhibits superior performance on diverse geometries. 10.6 Ablation Studies In this section, we discuss ablation experiments for the optimal transport neural operators. We will first cover transport plan formulation, and then transport map based on PPMM. 226 Group Model Relative L2 Time per Epoch (sec) GPU Memory (MB) G1 FNO GINO OTNO 16.27% 8.62% 3.06% 43 371 578 4548 57870 26324 G2 FNO GINO OTNO 56.67% 43.16% 7.16% 43 390 603 4548 58970 22868 G3 FNO GINO OTNO 23.20% 13.27% 4.02% 44 383 606 4548 73140 26008 Table 10.3: Prediction under M1 Metric (global) on FlowBench Dataset (single A100) Group Model Relative L2 Time per Epoch (sec) GPU Memory (MB) G1 FNO GINO OTNO 5.65% 5.83% 3.91% 43 190 135 4536 67632 7300 G2 FNO GINO OTNO 29.37% 19.74% 14.36% 43 177 124 4534 73792 7012 G3 FNO GINO OTNO 10.47% 10.69% 7.18% 43 219 153 4538 67838 11406 Table 10.4: Prediction under M2 Metric (boundary) on FlowBench Dataset OTNO with OT Plan (Sinkhorn) Encoder and Decoder Strategy Using the Sinkhorn algorithm, we solve the Kantorovich optimal transport problem between the latent mesh Ξ ∈ R𝑛2 and the physical mesh X ∈ R𝑛1 , resulting in a large, dense coupling matrix 𝑃 ∈ R𝑛2 ×𝑛1 that approximates the OT plan. The direct way to get transported mesh is as X ′ = 𝑃 · X, we refer to as the Matrix Strategy. However, storing these large, dense matrices for each physical mesh is too costly, and the approximate matrices obtained from the Sinkhorn Method introduce a degree of imprecision in the results. Therefore, this section discusses strategies to implement these approximate matrices more efficiently, aiming to conserve memory and reduce imprecision. 227 (a) Results of velocity in y direction under M1 metric (global). (b) Results of velocity in y direction under M2 metric (boundary). Figure 10.6: Results Visualization for OTNO on FlowBench Dataset 1. Max vs Mean As the matrix 𝑃 ∈ R𝑛2 ×𝑛1 represents a discrete probability measure on Ξ × X, a natural approach to take use of the plan is to transport by the maximum probability, termed as "Max" strategy. Let X = (𝑥 1 , . . . , 𝑥 𝑛1 ) ∈ R𝑛1 ×𝑑 . Denote 𝑃 = ( 𝑝 𝑘,𝑙 )𝑛2 ×𝑛1 . We calculate the indices of the max element in each row 𝑖 𝑘 = arg min{𝑝 𝑘, 𝑗 : 𝑗 = 1, . . . , 𝑛1 }, for 𝑘 = 1, . . . , 𝑛2 , (10.24) 𝑗 and then use these indices to get a refined transported mesh: M = (𝑥𝑖1 , . . . , 𝑥𝑖 𝑛2 ) ∈ R𝑛2 ×𝑑 . Another approach, perhaps more intuitive, involves finding the element in mesh X that is closest to the directly transported mesh X ′ = (𝑥1′ , . . . , 𝑥 𝑛′ 2 ) ∈ R𝑛2 ×𝑑 . We name this approach as "Mean" strategy due to the weight product across rows in X ′ = 𝑃 · X. The specific process is described as follows: first compute the indices 𝑖 𝑘 = arg min{|𝑥 ′𝑘 − 𝑥 𝑗 | : 𝑗 = 1, . . . , 𝑛1 }, for 𝑘 = 1, . . . , 𝑛2 . 𝑗 The refined transported mesh is M = (𝑥𝑖1 , . . . , 𝑥𝑖 𝑛2 ) ∈ R𝑛2 ×𝑑 . As shown in Table.10.5, "Mean" strategy has a better performance. (10.25) 228 Dataset Matrix Max Mean ShapeNet (Relative L2) DrivAerNet (MSE e-5) 7.21% 5.6 7.01% 6.70% 4.8 3.4 Table 10.5: Comparison of Max vs. Mean Strategy Dataset Multi Enc Multi Dec Single ShapeNet (Relative L2) DrivAerNet (MSE e-5) 20.55% 4.2 72.30% 4.5 6.70% 3.4 Table 10.6: Comparison of Multi vs. Single Encoder/Decoder Strategy 2. Single vs Multiple Besides the implementation of OT, integrating OT with FNO poses another significant question. For FNO, the requirement of inputs is just need to be features on a latent grid, suggesting that the encoder’s output can combine multiple transported distributions, same for the decoder. Thus, under the "Mean" strategy, we further investigate utilizing indices of the top 𝑘 closest elements, termed as "Multi-Enc" and "Multi-Dec". We choose 𝑘 = 8 for "Multi-End" and 𝑘 = 3 for "Multi-Dec" setups. The comparative results are presented in Table 10.6, indicating that "Multi-Enc" and "Multi-Dec" configurations underperform relative to "Single" strategy which utilize the index of the closest point instead of top 𝑘 closest points. Normal Features Our model, which incorporates latent FNO, learns the operator mapping from 𝑇 to the latent solution function 𝑣, where 𝑇 represents the transport map or the marginal map of the transport plan, as described in equations (10.7) and (10.11). The basic representation of the map 𝑇 in our experiments can be defined as T = (Ξ, 𝑇 (Ξ)). The previous section discussed how to refine the transported mesh 𝑇 (Ξ), and in this section, we further explore the addition of normal features to enhance the representation T , leveraging the normal’s ability to describe a surface. We propose three different approaches to integrate normals and compare them with the case where no normal features are added. The three methods are as follows: "Car": Only add car surface normals. "Concatenate": Add the concatenation of latent surface normals and car surface normals as normal features. "Cross Product": 229 Add the cross product of latent surface normals and car surface normals as normal features to capture the deformation information of the transport. As shown in Table 10.7, the "Cross Product" method performs the best. Dataset Non Car Concatenate Cross Product ShapeNet (Rel L2) 7.19% DrivAerNet (MSE e-5) 3.4 6.82% 3.9 6.83% 3.8 6.70% 3.4 Table 10.7: Comparison of Different Normal Features Shape of mesh in latent space We investigated the effects of different shapes for the latent mesh, i.e., the target shapes for optimal transport. The results presented in Table 10.8 indicate that the torus provides the best performance due to its alignment with the periodic Fourier function. Although the capped hemisphere and the spherical surface share the same topological structure as the car surface, their performance is suboptimal. Additionally, we experimented with the double sphere method (Mildenberger and Quellmalz, 2023), which unfortunately yielded worse accuracy compared to the sphere and doubled the time costs. Dataset ShapeNet (Rel L2) DrivAerNet (MSE e-5) Hemisphere Plane Double Sphere Sphere Torus 8.9% 4.7 8.67% 4.4 7.41% 4.6 7.09% 4.1 6.70% 3.4 Table 10.8: Comparison of Different Mesh Shapes in Latent Space Expand Factor For OTNO(Plan), we found that expanding the size of the latent mesh within a certain range leads to better results. The ablation experiments for different expansion factors 𝛼 = 1, 2, 3, 4are presented. As shown in Fig. 10.7, on both the DrivAerNet and ShapeNet datasets, 𝛼 = 3 achieves the best accuracy. OTNO with OT Map (PPMM) Number of Iterations The theoretical time complexity of the PPMM is 𝑂 (𝐾𝑛 log(𝑛)), where 𝐾 is the number of iterations and 𝑛 is the number of points. While the original study 230 (a) Tests of different expand factor for OTNO(Plan) on ShapeNet dataset (b) Tests of different expand factor for OTNO(Plan) on DrivAerNet dataset Figure 10.7: Ablation Studies of expand factor for OTNO(Plan) claims that empirically 𝐾 = 𝑂 (𝑑) works reasonably well, our experience with 3D car datasets tells a different story. In this section, we conduct ablation studies to explore the relationship between the optimal number of iterations and the number of points. We employ voxel downsampling to generate subsets of the car surface mesh, varying in the number of points. The results for different voxel sizes (𝑟) and iteration numbers (𝐾) are presented in Table 10.9. From these results, we observe that 𝐾 ∝ 𝑟 −1 . Given that the car surface mesh is a 2D manifold embedded in 3D space, the number of points 𝑛 should be inversely proportional to the square of the voxel size 𝑟, i.e., 𝑛 ∝ 𝑟 −2 . Consequently, 𝐾 ∝ 𝑛1/2 , and the experimental time complexity of PPMM in our model on the car dataset is accordingly 𝑂 (𝑛3/2 log(𝑛)). Voxel Size / Itr 500 1000 2000 4000 0.2 0.1 0.05 6.2 5.3 4.8 6.6 4.6 4.2 6.9 5.5 3.9 8.0 5.4 4.5 Table 10.9: Results on the DrivAerNet Dataset for Different Voxel Sizes and Different Iteration Numbers of PPMM (MSE (e-5)) Sampling Mesh Size It is a challenging problem to efficiently and effectively solve large-scale OT problems. We investigate the Sinkhorn method for large-scale OT plans and the PPMM algorithm for large-scale OT maps, and accordingly build two models, OTNO(Plan) and OTNO(Map). In this section, we conduct experiments on the large-scale DrivAerNet dataset to explore our models’ ability to handle large sampling meshes. We 231 (a) Test errors associate with the number of sampling points. We set the voxel size 𝑟 = 0.2, 0.1, 0.05, 0.025, corresponding to about 1k, 4.5k, 18k, 60k points. (b) Time complexity associate with the number of sampling points. We set the voxel size 𝑟 = 0.1, 0.05, 0.025, corresponding to about 4.5k, 18k, 60k points. Figure 10.8: Ablation Studies of Sampling Mesh Size for OTNO(Map) and OTNO(Plan) on DrivAerNet Dataset use voxel downsampling to reduce the mesh with 200k points from the DrivAerNet dataset into a normalized sampling mesh of varying sizes. As shown in Fig. 10.8a, both OTNO(Plan) and OTNO(Map) achieve their lowest accuracy at a voxel size of 0.05, corresponding to approximately 18k points. But when the sampling mesh size increases more, the accuracy does not improves. This suggests that our model has a range of limitations when dealing with large-scale problems, primarily due to the difficulty of solving large-scale OT problems with high precision. Regarding the comparison between OTNO(Plan) and OTNO(Map), we find that OTNO(Plan) consistently outperforms OTNO(Map) in terms of accuracy, regardless of the sampling mesh size. From the experimental results, we observe that the training time of OTNO(Plan) is much larger than OTNO(Map) as shown in Fig. 10.8b. This is because we expand the latent space by a factor of 3 for OTNO(Plan), and the FNO on the latent space has linear complexity, attributed to the linear FFT. Note that the time complexity of OTNO(Plan) and OTNO(Map) are 𝑂 (𝑛2 ) and 𝑂 (𝑛3/2 log(𝑛)), respectively, both of which are larger than the training complexity of 𝑂 (𝑛). Therefore, the overall time complexity of OTNO(Plan) and OTNO(Map) relative to the number of points 𝑛 should be 𝑂 (𝑛2 ) and 𝑂 (𝑛3/2 log(𝑛)), respectively. These results align with the plots shown in Fig. 10.8b that the time cost of OTNO(Plan) increases faster than that of OTNO(Map) as the sampling mesh size grows. 232 10.7 Discussion and Conclusion Recently, several studies have focused on integrating diffeomorphic transformations within PDE solution operators to extract geometric information and ultimately solve PDEs. Examples include DIMON (M. Yin et al., 2024), which employs LDDMM to learn diffeomorphic deformations, and DNO (Zhao et al., 2025), which utilizes harmonic mapping for the same purpose. A recent work, Diffeomorphic Latent Neural Operators (Ahmad et al., 2024), compared conformal mapping and LDDMM, which are both diffeomorphic mappings, with discrete optimal transport via the Hungarian algorithm, which is a non-diffeomorphic mapping. Their findings suggest that, for invertible spatial transformations, diffeomorphic approaches generally offer better performance. However, the Hungarian algorithm is designed primarily for discrete assignment problems, which is not suited for spatial transformations between meshes. In contrast, our method employs continuous OT algorithms, which are better equipped for distribution transformations that encode/decode geometric information. We observe several advantages of our method employing continuous OT algorithms compared to diffeomorphic methods. First, we found that smoothness is an unnecessarily strict requirement—piecewise continuity is often sufficient for effective geometric embeddings. For example, using a sphere surface as a latent shape results in a globally continuous transported mesh, whereas for a torus, the transported mesh is only piecewise continuous. Yet, as shown in Table 10.8, the torus-based representation outperforms the sphere-based one. Secondly invertibility is not always beneficial. Invertibility requires the size of latent mesh should be the same as the size of phyical mesh. However, As illustrated in Fig. 10.7a, expanding the size of latent mesh within a certain range leads to improved performance. Moreover, the topology of the mapping warrants attention. While diffeomorphisms preserve topology, our OT framework does not require topological consistency. As demonstrated in Table 10.8, the torus representation achieves the best performance despite a change in topology. We also tested this on a 2D elasticity example to assess its generalization capability. The results, shown in Fig. 10.10, confirm that enforcing topological consistency does not yield improved outcomes. In this work, we propose an OT framework for geometry embedding, which maps density functions from physical geometric domains to latent uniform density functions on regular latent geometric domains. We developed the OTNO model by 233 Latent grid Ring Square Topology Relative Error 3.8% 2.7% Table 10.10: Demonstrate topology independence of OTNO on elasticity integrating neural operator with optimal transport to solve PDEs on complex geometries. Specifically, there are two implementations: OTNO(Plan) and OTNO(Map), using Kantorovich and Monge formulations, respectively. Our model achieves particularly good performance in automotive and aerospace applications, where the inputs are 2D surface designs, and the outputs are surface pressure and velocity. These applications provide the opportunity to use optimal transport to embed the physical surface mesh in 3D space into a 2D parameterized latent mesh, allowing computations in a lower-dimensional space. The effectiveness and efficiency of our model in these scenarios are confirmed by our experiments on the ShapeNet-Car and DrivAerNet-Car datasets. Moreover, by leveraging the advantage of instance-dependent OT Map/Plan, our model handles diverse geometries effectively. As demonstrated on the FlowBench dataset, which includes a wider variety of shapes, OTNO significantly outperforms other models in terms of accuracy. While our proposed methods demonstrate promise, they present certain limitations that pave the way for future research. A primary challenge lies in the trade-off between computational complexity and accuracy. The Sinkhorn algorithm, underpinning our OTNO(Plan) model, exhibits a complexity of 𝑂 (𝑛2 ), posing scalability challenges for large-scale point clouds. Our alternative, the OTNO(Map) model leveraging the PPMM method, achieves a reduced experimental complexity of 𝑂 (𝑛3/2 log(𝑛)) in a lower-dimensional Euclidean space, yet at the cost of diminished accuracy compared to OTNO(Plan). Consequently, a key direction for future work is the development of novel, sub-quadratic algorithms for optimal transport that can achieve higher accuracy, thereby bridging this performance gap. A second promising avenue involves a deeper investigation into the optimal selection of the latent sub-manifold. Our experimental results indicate that the optimal latent manifold is non-canonical, diverging from the intuitively favored spherical topology typically chosen for its topological preservation properties. Future efforts will focus on systematically exploring criteria and methods for determining the most suitable 234 latent manifold structure to enhance model performance. References Ahmad, Zan et al. (2024). Diffeomorphic Latent Neural Operators for Data-Efficient Learning of Solutions to Partial Differential Equations. arXiv: 2411 . 18014 [cs.LG]. url: https://arxiv.org/abs/2411.18014. Alkin, Benedikt et al. (2024). “Universal physics transformers”. In: arXiv preprint arXiv:2402.12365. Bhatnagar, Saakaar et al. (2019). “Prediction of aerodynamic flow fields using convolutional neural networks”. In: Computational Mechanics 64, pp. 525–545. Bonev, Boris et al. (2023). “Spherical fourier neural operators: Learning stable dynamics on the sphere”. In: International conference on machine learning. PMLR, pp. 2806–2823. Brenier, Yann (1991). “Polar factorization and monotone rearrangement of vectorvalued functions”. In: Communications on pure and applied mathematics 44.4, pp. 375–417. Budd, CJ, Robert D Russell, and E Walsh (2015). “The geometry of r-adaptive meshes generated using optimal transport methods”. In: Journal of Computational Physics 282, pp. 113–137. Chang, Angel X et al. (2015). “ShapeNet: An Information-Rich 3D Model Repository”. In: arXiv:1512.03012 [cs.GR]. Chen, Honglin et al. (2023). “Implicit neural spatial representations for timedependent pdes”. In: International Conference on Machine Learning. PMLR, pp. 5162–5177. Chen, Peter Yichen et al. (2022). “Crom: Continuous reduced-order modeling of pdes using implicit neural representations”. In: arXiv preprint arXiv:2206.02607. Choi, Pui Tung, Ka Chun Lam, and Lok Ming Lui (2015). “FLASH: Fast landmark aligned spherical harmonic parameterization for genus-0 closed brain surfaces”. In: SIAM Journal on Imaging Sciences 8.1, pp. 67–94. Cooley, James W and John W Tukey (1965). “An algorithm for the machine calculation of complex Fourier series”. In: Mathematics of computation 19.90, pp. 297– 301. Cuturi, Marco (2013). “Sinkhorn distances: Lightspeed computation of optimal transport”. In: Advances in neural information processing systems 26. Elrefaie, Mohamed, Angela Dai, and Faez Ahmed (2024). DrivAerNet: A Parametric Car Dataset for Data-Driven Aerodynamic Design and Graph-Based Drag Prediction. arXiv: 2403.08055 [cs.LG]. url: https://arxiv.org/abs/ 2403.08055. 235 Elrefaie, Mohamed, Florin Morar, et al. (2024). “DrivAerNet++: A Large-Scale Multimodal Car Dataset with Computational Fluid Dynamics Simulations and Deep Learning Benchmarks”. In: arXiv preprint arXiv:2406.09624. Flamary, Rémi et al. (2021). “POT: Python Optimal Transport”. In: Journal of Machine Learning Research 22.78, pp. 1–8. url: http://jmlr.org/papers/ v22/20-451.html. Gu, Xianfeng et al. (2004). “Genus zero surface conformal mapping and its application to brain surface mapping”. In: IEEE transactions on medical imaging 23.8, pp. 949–958. Hao, Zhongkai et al. (2023). “Gnot: A general neural operator transformer for operator learning”. In: International Conference on Machine Learning. PMLR, pp. 12556–12569. Hennigh, Oliver et al. (2021). “NVIDIA SimNet™: An AI-accelerated multi-physics simulation framework”. In: International conference on computational science. Springer, pp. 447–461. Hubbert, M King (1956). “Darcy’s law and the field equations of the flow of underground fluids”. In: Transactions of the AIME 207.01, pp. 222–239. Ihlenburg, Frank and Ivo Babuška (1995). “Finite element solution of the Helmholtz equation with high wave number Part I: The h-version of the FEM”. In: Computers & Mathematics with Applications 30.9, pp. 9–37. Kantorovich, LK (2006). “On a Problem of Monge.” In: Journal of Mathematical Sciences 133.4. Kovachki, Nikola et al. (2023). “Neural operator: Learning maps between function spaces with applications to pdes”. In: Journal of Machine Learning Research 24.89, pp. 1–97. Li, Zongyi, Daniel Zhengyu Huang, et al. (2022). “Fourier Neural Operator with Learned Deformations for PDEs on General Geometries”. In: ArXiv abs/2207.05209. url: http://arxiv.org/abs/2207.05209. Li, Zongyi, Nikola Kovachki, et al. (2020a). “Fourier neural operator for parametric partial differential equations”. In: arXiv preprint arXiv:2010.08895. – (2020b). “Neural operator: Graph kernel network for partial differential equations”. In: arXiv preprint arXiv:2003.03485. Li, Zongyi, Nikola Borislavov Kovachki, et al. (2023). “Geometry-Informed Neural Operator for Large-Scale 3D PDEs”. In: NIPS. Lu, Lu et al. (2021). “Learning nonlinear operators via DeepONet based on the universal approximation theorem of operators”. In: Nature machine intelligence 3.3, pp. 218–229. Meng, Cheng et al. (2019). “Large-scale optimal transport map estimation using projection pursuit”. In: Advances in Neural Information Processing Systems 32. 236 Mildenberger, Sophie and Michael Quellmalz (2023). “A double Fourier sphere method for d-dimensional manifolds”. In: Sampling Theory, Signal Processing, and Data Analysis 21.2, p. 23. Monge, Gaspard (1781). “Mémoire sur la théorie des déblais et des remblais”. In: Mem. Math. Phys. Acad. Royale Sci., pp. 666–704. Peyré, Gabriel, Marco Cuturi, et al. (2019). “Computational optimal transport: With applications to data science”. In: Foundations and Trends® in Machine Learning 11.5-6, pp. 355–607. Pfaff, Tobias et al. (2021). “Learning mesh-based simulation with graph networks”. In: ICLR. Serrano, Louis et al. (2023). “Operator learning with neural fields: Tackling pdes on general geometries”. In: Advances in Neural Information Processing Systems 36, pp. 70581–70611. Tali, Ronak et al. (2024). “FlowBench: A Large Scale Benchmark for Flow Simulation over Complex Geometries”. In. Thuerey, Nils et al. (2020). “Deep learning methods for Reynolds-averaged Navier– Stokes simulations of airfoil flows”. In: AIAA Journal 58.1, pp. 25–36. Timmer, W (2009). “An overview of NACA 6-digit airfoil series characteristics with reference to airfoils for large wind turbine blades”. In: 47th AIAA aerospace sciences meeting including the new horizons forum and aerospace exposition, p. 268. Umetani, Nobuyuki and Bernd Bickel (2018). “Learning three-dimensional flow for interactive aerodynamic design”. In: ACM Transactions on Graphics (TOG) 37.4, pp. 1–10. Wu, Haixu et al. (2024). “Transolver: A fast transformer solver for pdes on general geometries”. In: arXiv preprint arXiv:2402.02366. Yin, Minglang et al. (2024). “Dimon: Learning solution operators of partial differential equations on a diffeomorphic family of domains”. In: arXiv preprint arXiv:2402.07250. Yin, Yuan et al. (2022). “Continuous pde dynamics forecasting with implicit neural representations”. In: arXiv preprint arXiv:2209.14855. Zhao, Zhiwei et al. (2025). “Diffeomorphism neural operator for various domains and parameters of partial differential equations”. In: Communications Physics 8.1, p. 15. 237 C h a p t e r 11 APPLICATION: WEATHER FORECAST FourCastNet, short for Fourier ForeCasting Neural Network, is a global data-driven weather forecasting model that provides accurate short to medium-range global predictions at 0.25◦ resolution. FourCastNet accurately forecasts high-resolution, fast-timescale variables such as the surface wind speed, precipitation, and atmospheric water vapor. It has important implications for planning wind energy resources, predicting extreme weather events such as tropical cyclones, extra-tropical cyclones, and atmospheric rivers. FourCastNet matches the forecasting accuracy of the ECMWF Integrated Forecasting System (IFS), a state-of-the-art Numerical Weather Prediction (NWP) model, at short lead times for large-scale variables, while outperforming IFS for variables with complex fine-scale structure, including precipitation. FourCastNet generates a week-long forecast in less than 2 seconds, orders of magnitude faster than IFS. The speed of FourCastNet enables the creation of rapid and inexpensive large-ensemble forecasts with thousands of ensemble-members for improving probabilistic forecasting. We discuss how data-driven deep learning models such as FourCastNet are a valuable addition to the meteorology toolkit to aid and augment NWP models. 11.1 Introduction The beginnings of modern numerical weather prediction (NWP) can be traced to the 1920s. Now ubiquitous, they contribute to economic planning in key sectors such as transport, logistics, agriculture, and energy production. Accurate weather forecasts have saved countless human lives by providing advance notice of extreme events. The quality of weather forecasts has been steadily improving over the past decades (c.f. Bauer, Thorpe, and Brunet (2015) and Alley, Emanuel, and Zhang (2019)). The earliest dynamically-modeled numerical weather forecast for a single point was computed using a slide rule and table of logarithms by Lewis Fry Richardson in 1922 (Richardson, 2007) and took six weeks to compute a 6-hour forecast of the atmosphere. By the 1950s, early electronic computers greatly improved the speed of forecasting, allowing operational forecasts to be calculated fast enough to be useful for future prediction. In addition to better computing capabilities, improvements in weather forecasting have been achieved through better parameterization of fine- 238 scale processes through deeper understanding of their physics and higher-quality atmospheric observations. The latter has resulted in improved model initializations via data assimilation. There is now increasing interest around developing data-driven DL-based models for weather forecasting owing to their orders of magnitude lower computational cost as compared to state-of-the-art NWP models (Schultz et al., 2021; Balaji, 2021; Irrgang et al., 2021; Reichstein et al., 2019). Many studies have attempted to build data-driven models for forecasting the large-scale circulation of the atmosphere, either trained on climate model outputs, general circulation models (GCM) (Scher and Messori, 2018; Scher and Messori, 2019; Ashesh Chattopadhyay, Nabizadeh, and Hassanzadeh, 2020), reanalysis products (Weyn, Durran, and Caruana, 2019; Weyn, Durran, and Caruana, 2020; Weyn, Durran, Caruana, and Cresswell-Clay, 2021; Rasp, Dueben, et al., 2020; Rasp and Thuerey, 2021a; Rasp and Thuerey, 2020; Ashesh Chattopadhyay, Mustafa Mustafa, Hassanzadeh, Bach, et al., 2021; Arcomano, Szunyogh, Pathak, et al., 2020; Chantry et al., 2021; Grönquist et al., 2021), or a blend of climate model outputs and reanalysis products (Rasp and Thuerey, 2021a). Data-driven models have great potential to improve weather predictions by overcoming model biases present in NWP models and by enabling the generation of large ensembles at low computational cost for probabilistic forecasting and data assimilation. By training on reanalysis data or observations, data-driven models can avoid limitations that exist in NWP models (Schultz et al., 2021; Balaji, 2021), such as biases in convection parameterization schemes that strongly affect precipitation forecasts. Once trained, data-driven models are orders of magnitude faster than traditional NWP models in generating forecasts via inference, thus enabling the generation of very large ensembles (Ashesh Chattopadhyay, Mustafa Mustafa, Hassanzadeh, Bach, et al., 2021; Weyn, Durran, Caruana, and Cresswell-Clay, 2021). In this regard, Weyn, Durran, Caruana, and Cresswell-Clay (2021) have shown that large data-driven ensembles improve subseasonal-to-seasonal (S2S) forecasts over operational NWP models that can only incorporate a small number of ensemble members. Furthermore, a large ensemble helps improve data-driven predictions of extreme weather events in short- and long-term forecasts (Ashesh Chattopadhyay, Nabizadeh, and Hassanzadeh, 2020). Most data-driven weather models, however, use low-resolution data for training, 239 usually at the 5.625◦ resolution as in Rasp and Thuerey (2021b) or 2◦ as in Weyn, Durran, and Caruana (2020). These prior attempts have achieved good results on forecasting some of the coarse, low-resolution atmospheric variables. However, the coarsening procedure leads to the loss of crucial, fine-scale physical information. For data-driven models to be truly impactful, it is essential that they generate forecasts at the same or greater resolution than current state-of-the-art numerical weather models, which are run at ≈ 0.1◦ resolution. Forecasts at 5.625◦ spatial resolution, for instance, result in a mere 32 × 64 pixels grid representing the entire globe. Such a forecast is not able to resolve features smaller than ≈ 500 km. Such coarse forecasts fail to account for the important effects of small-scale dynamics on the large scales and the impact of topographic features such as mountain ranges and lakes on small-scale dynamics. This limits the practical utility of low-resolution forecasts. While low-resolution forecasts may be justified for variables that do not possess complex fine-scale structure, such as the geopotential height at 500 hPa (𝑍500 ), higher-resolution data (e.g., at 0.25◦ resolution) can substantially improve the predictions of data-driven models for variables like low-level winds (𝑈10 and 𝑉10 ) that have complex fine-scale structures. Moreover, high-resolution models can resolve the formation and dynamics of high-impact extreme events such as tropical cyclones, which would otherwise be inadequately represented on a coarser grid. Our approach. We develop FourCastNet, a Fourier-based neural network forecasting model, to generate global data-driven forecasts of key atmospheric variables at a resolution of 0.25°, which corresponds to a spatial resolution of roughly 30 km × 30 km near the equator and a global grid size of 720 × 1440 pixels. This allows us, for the first time, to make a direct comparison with the high-resolution Integrated Forecasting System (IFS) model of the European Center for Medium-Range Weather Forecasting (ECMWF). Figure 11.1 shows an illustrative global near-surface wind speed forecast at a 96-hour lead time generated using FourCastNet. We highlight key high-resolution details that are resolved and accurately tracked by our forecast, including Super Typhoon Mangkhut and three named cyclones heading towards the eastern coast of the United States (Florence, Issac, and Helene). FourCastNet is about 45,000 times faster than traditional NWP models on a nodehour basis. This orders of magnitude speedup, along with the unprecedented accuracy of FourCastNet at high resolution, enables inexpensive generation of extremely 240 To prepare this figure, we initialized FourCastNet with an initial condition from the out-ofsample test dataset with the calendar timestamp September 8, 2018 at 00:00 UTC. Starting from this initial condition, the model was allowed to run freely for 16 time-steps of six hours each in inference mode (Figure 11.2(c)) corresponding to a 96-hour forecast. Panel (a) shows the wind speed at model initialization. Panel (b) shows the model forecasts at forecast lead time of 96 hours (upper panel) and the corresponding true wind speeds at that time (lower panel). FourCastNet is able to forecast the wind speeds 96 hours in advance with remarkable fidelity and correct fine-scale features. The forecast accurately captures the formation and track of Super Typhoon Mangkhut that begins to form at roughly 10° 𝑁, 210°𝑊 (see Inset 1). Further, the model captures the intensification and track of the typhoon over a period of four days. During the period of this forecast, the model reveals three named hurricanes (Florence, Issac, and Helene) forming in the Atlantic Ocean and approaching the eastern coast of North America. Figure 11.1: Illustrative example of a global near-surface wind forecast generated by FourCastNet over the entire globe at a resolution of 0.25° . 241 large ensemble forecasts. This dramatically improves probabilistic weather forecasting. Massive large-ensemble forecasts of events such as hurricanes, atmospheric rivers, and extreme precipitation can be generated in seconds using FourCastNet. This could lead to better-informed disaster response. Furthermore, FourCastNet’s reliable, rapid, and cheap forecasts of near-surface wind speeds can improve wind energy resource planning at onshore and offshore wind farms. The energy required to train FourCastNet is approximately equal to the energy required to generate a 10-day forecast with 50 ensemble members using the IFS model. Once trained, however, FourCastNet uses about 12,000 times less energy to generate a forecast than the IFS model. We expect FourCastNet to be only trained once; the energy consumption of subsequent fine tuning is negligible. FourCastNet uses a Fourier transform-based token-mixing scheme (Guibas et al., 2022) with a vision transformer (ViT) backbone (Dosovitskiy et al., 2021). This approach is based on the recent Fourier neural operator that learns in a resolutioninvariant manner and has shown success in modeling challenging partial differential equations (PDE) such as fluid dynamics (Li, Kovachki, et al., 2021). We chose a ViT backbone since it is capable of modeling long-range dependencies well. Combining ViT with Fourier-based token mixing yields a state-of-the-art highresolution model that resolves fine-grained features and scales well with resolution and size of dataset. This approach enables training high-fidelity data-driven models at truly unprecedented resolution.1 In summary, FourCastNet makes four significant contributions to data-driven weather forecasting: 1. FourCastNet predicts, with unparalleled accuracy at forecast lead times of up to one week, challenging variables with complex fine-scale structure, such as surface winds and precipitation. Until now, no deep learning model was able to forecast surface winds on global scales. Additionally, DL models for precipitation on global scales have been inadequate for resolving fine-scale structures. This has important implications for disaster mitigation and wind energy resource planning. 2. FourCastNet, at 0.25◦ resolution, has eight times greater resolution than stateof-the-art DL-based global weather models. Due to its high resolution and 1 We estimate that FourCastNet could be trained on currently available GPU hardware in about two months with 40 years of global 5-km data, if such data were available. 242 accuracy, FourCastNet accurately resolves extreme weather patterns such as tropical cyclones and atmospheric rivers that have been inadequately represented by prior DL models owing to their coarser grids. 3. FourCastNet’s predictions are comparable to the IFS model on metrics of Root Mean Squared Error (RMSE) and Anomaly Correlation Coefficient (ACC) at lead times of up to three days. After, predictions of all modeled variables lag close behind IFS at lead times of up to a week. Whereas the IFS model has been developed over decades, contains greater than 150 variables at more than 50 vertical levels in the atmosphere, and is guided by physics, FourCastNet models 20 variables at five vertical levels, and is purely data driven. This comparison points to the enormous potential of data-driven modeling in complementing and eventually replacing NWP. 4. FourCastNet’s reliable, rapid, and computationally inexpensive forecasts facilitate the generation of very large ensembles, thus enabling estimation of well-calibrated and constrained uncertainties in extremes with higher confidence than current NWP ensembles that have at most 50 members owing to their high computational cost. Fast generation of 1,000-member ensembles dramatically changes what is possible in probabilistic weather forecasting, including improving reliability of early warnings of extreme weather events and enabling rapid assessment of their impacts. 11.2 Training Methods The ECMWF provides a publicly available, comprehensive dataset called ERA5 (Hersbach et al., 2020) which consists of hourly estimates of several atmospheric variables at a latitude and longitude resolution of 0.25° from the surface of the earth to roughly 100 km altitude from 1979 to the present day. ERA5 is an atmospheric reanalysis (Kalnay et al., 1996) dataset and is the result of an optimal combination of observations from various measurement sources and the output of a numerical model using a Bayesian estimation process called data-assimilation (Kalnay, 2003). The dataset is essentially a reconstruction of the optimal estimate of the observed history of the Earth’s atmosphere. We use the ERA5 dataset to train FourCastNet. While the ERA5 dataset has several prognostic variables available at 37 vertical levels with an hourly resolution, computational and data limitations along with other operational considerations for DL models restricts our choice, based on physical reasoning, to a subset of these available variables to train our model on. 243 In this work, we focus on forecasting two important and challenging atmospheric variables, namely, (1) the wind velocities at a distance of 10m from the surface of the earth and (2) the 6-hourly total precipitation. There are a few reasons for our focus on these variables. First, surface wind velocities and precipitation require highresolution models to resolve and forecast accurately because they contain and are influenced by many fine-scale features. Due to computational and model architectural limitations, previous efforts in DL-based weather prediction have not been able to produce global forecasts for these variables at full ERA5 resolution. Near-surface wind velocity forecasts have a tremendous amount of utility due to their key role in planning energy storage, grid transmission, and other operational considerations at on-shore and off-shore wind farms. As we show in Section 11.3, near-surface wind forecasts (along with wind forecasts above the atmospheric boundary layer) can help track extreme wind events such as hurricanes and can be used for disaster preparedness. Our second focus is on forecasting total precipitation where DL models can potentially show great promise. NWP models, such as the operational IFS, have several parameterization schemes to tractably forecast precipitation and since neural networks are known to have impressive capabilities at deducing parameterizations from high-resolution observational data, they are well-suited for this task. Although we focus on forecasting near-surface wind-speed and precipitation, our model also forecasts with remarkable accuracy several other variables. In our forecast, we include the geopotential height, temperature, wind velocity, and relative humidity at a few different vertical levels, a few near-surface variables such as surface pressure and mean sea-level pressure as well as the integrated total column of water vapor. FourCastNet: Model Description To produce our high-resolution forecasts, we choose the Adaptive Fourier Neural Operator (AFNO) model (Guibas et al., 2022). This particular neural network architecture is appealing as it is specifically designed for high-resolution inputs and synthesizes several key recent advances in DL into one model. Namely, it combines the Fourier Neural Operator (FNO) learning approach of Li, Kovachki, et al., 2021, which has been shown to perform well in modeling challenging PDE systems, with a powerful ViT backbone. The vision transformer (ViT) architecture and its variants have emerged as the state-of-the-art in computer vision over the previous years, showing remarkable 244 performance on a number of tasks and scaling well with increased model and dataset sizes. Such performance is attributed mainly to the multi-head self-attention mechanism in these networks, which allows the network to model interactions between features (called tokens in ViT representation terms) globally at each layer in the network. However, spatial mixing via self-attention is quadratic in the number of tokens, and thus quickly becomes infeasible for high-resolution inputs. Several ViT variants with reduced computational complexity have been proposed, with various alternate mechanisms for spatial token mixing employed in each. However, the AFNO model is unique in that it frames the mixing operation as continuous global convolution, implemented efficiently in the Fourier domain with FFTs, which allows modeling dependencies across spatial and channel dimensions flexibly and scalably. With such a design, the spatial mixing complexity is reduced to O (𝑁 log 𝑁), where 𝑁 is the number of image patches or tokens. This scaling allows the AFNO model to be well-suited to high-resolution data at the current 0.25° resolution considered in this paper as well as potential future work at an even higher resolution. In the original FNO formulation, the operator learning approach showed impressive results solving turbulent Navier-Stokes systems, so incorporating this into a data-driven atmospheric model is a natural choice. Given the general popularity of convolutional network architectures, and particularly their usage in previous works forecasting ERA5 variables (Rasp and Thuerey, 2021b; Weyn, Durran, and Caruana, 2020), it is worth contrasting our AFNO model with these more conventional architectures. For one, the ability of AFNO to scale well with resolution yields immediate practical benefits – at our 720x1440 resolution, the FourCastNet model memory footprint is about 10GB with a batch size of 1. To contrast this, we can look at the 19-layer ResNet architecture from a prior result on WeatherBench (Rasp and Thuerey, 2021b), which was trained at a very coarse resolution (32×64 pixels). Naively transferring this architecture to our dataset and training at 720×1440 resolution would require 83GB for a batch size of 1. This is prohibitive, and is compounded by the fact that it is somewhat of a lower bound – with order-of-magnitude increases in resolution, a convolution-based network’s receptive field would similarly need to grow via the addition of even more layers. Beyond practical considerations, our preliminary non-exhaustive experiments suggested that convolutional architectures showed poor performance on capturing small scales over many time steps in auto-regressive inference. These observations along with our knowledge of the current state of the art for high-resolution image process- 245 ing in image de-noising, super-resolution and de-blurring are a strong motivation for our choice of a ViT architecture over a convolutional architecture. While we refer the reader to the original AFNO paper (Guibas et al., 2022) for more details, we briefly describe the flow of computation in our model here. First, the input variables on the 720 × 1440 lat-lon grid are projected to a 2D grid (ℎ × 𝑤) of patches (with a small patch size 𝑝 × 𝑝, where e.g., 𝑝 = 8), with each patch represented as a 𝑑-dimensional token. Then, the sequence of patches are fed, along with a positional encoding, to a series of AFNO layers. Each layer, given an input tensor of patches 𝑋 ∈ Rℎ×𝑤×𝑑 , performs spatial mixing followed by channel mixing. Spatial mixing happens in the Fourier domain as follows: Step 1. Transform tokens to the Fourier domain with 𝑧 𝑚,𝑛 = [DFT(𝑋)] 𝑚,𝑛 , (11.1) where 𝑚, 𝑛 index the patch location and DFT denotes a 2D discrete Fourier transform. Step 2. Apply token weighting in the Fourier domain, and promote sparsity with a Soft-Thresholding and Shrinkage operation as 𝑧˜𝑚,𝑛 = 𝑆𝜆 (MLP(𝑧 𝑚,𝑛 )), (11.2) where 𝑆𝜆 (𝑥) = sign(𝑥) max(|𝑥| − 𝜆, 0) with the sparsity controlling parameter 𝜆, and MLP() is a 2-layer multi-layer perceptron with block-diagonal weight matrices which are shared across all patches. Step 3. Inverse Fourier to transform back to the patch domain and add a residual connection as ˜ 𝑚,𝑛 + 𝑋𝑚,𝑛 . 𝑦 𝑚,𝑛 = [IDFT( 𝑍)] (11.3) Training While our primary interest lies in forecasting the surface wind velocities and precipitation, the complex atmospheric system contains strong nonlinear interactions across several variables such as temperatures, surface pressures, humidity, moisture content from the surface of the earth to the stratosphere, etc. In order to model these interactions, we choose a few variables (Table 11.1) to represent the instantaneous 246 (a) The multi-layer transformer architecture that utilizes the Adaptive Fourier Neural Operator with shared MLP and frequency soft-thresholding for spatial token mixing. The input frame is first divided into a ℎ × 𝑤 grid of patches, where each patch has a small size 𝑝 × 𝑝 × 𝑐. Each patch is then embedded in a higher dimensional space with high number of latent channels and position embedding is added to form a sequence of tokens. Tokens are then mixed spatially using AFNO, and subsequently for each token the latent channels are mixed. This process is repeated for 𝐿 layers, and finally a linear decoder reconstructs the patches for the next frame from the final embedding. The right-hand panels describe the FourCastNet model’s additional training and inference modes: (b) two-step fine-tuning, (c) backbone model that forecasts the 20 variables in Table 11.1 with secondary precipitation diagnostic model (note that p(𝑘 + 1) denotes the 6 hour accumulated total precipitation that falls between 𝑘 + 1 and 𝑘 + 2 time steps) (d) forecast model in free-running autoregressive inference mode. Figure 11.2: Architecture of FourCastNet 247 state of the atmosphere. These variables are specifically chosen to model important processes that influence low-level winds and precipitation. As such, we treat all the prognostic variables equally and the model architecture or optimization scheme does not afford special treatment to any of the prognostic variables. Each of the variables in Table 11.1 is re-gridded from a Gaussian grid to a regular Euclidean grid using the standard interpolation scheme provided by the Copernicus Climate Data Store (CDS) Application Programming Interface (API). Following the re-gridding process, each of the 20 variables is represented as a 2D field of shape (721 × 1440) pixels. Thus, a single training data point at an instant in time containing all 20 variables is represented by a tensor of shape (721 × 1440 × 20). While the ERA5 dataset is available at a temporal resolution of 1 hour, we choose to sub-sample the dataset and use snapshots spaced 6 hours apart to train our model. Within each 24 hour day, we choose to sample the 20 variable subset of the ERA5 dataset at 0000 hrs, 0600 hrs, 1200 hrs and 1800 hrs. We divide the dataset into three sets, namely training, validation and out-of-sample testing datasets. The training dataset consists of data from the year 1979 to 2015 (both included). The validation dataset contains data from the years 2016 and 2017. The out-of-sample testing dataset consists of the years 2018 and beyond. ⃗⃗⃗⃗ We collectively denote the modeled variables by the tensor 𝑋 (𝑘Δ𝑡), where 𝑘 denotes the time index and Δ𝑡 is the temporal spacing between consecutive snapshots in the training dataset. We will consider the ERA5 dataset as the truth and denote the ⃗⃗⃗⃗ true variables by 𝑋 true (𝑘Δ𝑡). With the understanding that Δ𝑡 is fixed at 6 hours throughout this work, we omit Δ𝑡 in our notation for convenience where appropriate. The training procedure consists of two steps, pre-training and fine-tuning. In the pre-training step, we train the AFNO model using the training dataset in a supervised ⃗⃗⃗⃗ ⃗⃗⃗⃗ fashion to learn the mapping from 𝑋 (𝑘) to 𝑋 (𝑘 + 1). In the fine-tuning step, we start from the previously pre-trained model and optimize the model to predict two time ⃗⃗⃗⃗ ⃗⃗⃗⃗ steps, i.e., The model first generates the output 𝑋 (𝑘 + 1) from the input 𝑋 (𝑘). The ⃗⃗⃗⃗ ⃗⃗⃗⃗ model then uses its own output 𝑋 (𝑘 +1) as an input and generates the output 𝑋 (𝑘 +2). ⃗⃗⃗⃗ ⃗⃗⃗⃗ We then compute a training loss by comparing each of 𝑋 (𝑘 + 1) and 𝑋 (𝑘 + 2) to the respective ground truth from the training data and use the sum of the two training losses for optimizing the model. In both, the pre-training and fine-tuning steps, the training dataset is used to optimize the model and the validation dataset is used to estimate the model skill during hyper-parameter optimization. The out-of-sample testing dataset is untouched. The training dataset consists of 54020 samples while 248 Vertical Level Surface 1000ℎ𝑃𝑎 850ℎ𝑃𝑎 500ℎ𝑃𝑎 50ℎ𝑃𝑎 Integrated Variables 𝑈10 , 𝑉10 , 𝑇2𝑚 , 𝑠𝑝, 𝑚𝑠𝑙 𝑝 𝑈, 𝑉, 𝑍 𝑇, 𝑈, 𝑉, 𝑍, 𝑅𝐻 𝑇, 𝑈, 𝑉, 𝑍, 𝑅𝐻 𝑍 𝑇𝐶𝑊𝑉 Table 11.1: Prognostic variables modeled by the DL model. Abbreviations are as follows. 𝑈10 (𝑉10 ): zonal (meridional) wind at 10 m; 𝑇2𝑚 : temperature at 2 m above ground; 𝑇, 𝑈, 𝑉, 𝑍, 𝑅𝐻: temperature, zonal velocity, meridional velocity, geopotential, and relative humidity (at the specified pressure level); 𝑇𝐶𝑊𝑉: total column water vapor. the validation dataset contains 2920 samples. We refer to the trained and fine-tuned model as the ‘backbone’. The model is pre-trained using a cosine learning-rate schedule with a starting learning rate ℓ1 for 80 epochs. Following the pre-training, the model is fine-tuned for a further 50 epochs using a cosine learning-rate schedule and a lower learning rate ℓ2 . The precipitation model (described in Section 11.2) is then added to the trained backbone and trained for 25 epochs using a cosine learning rate schedule with an initial learning rate ℓ3 . The end to end training takes about 16 hours wall-clock time on a cluster of 64 Nvidia A100 GPUs. Precipitation Model The total precipitation (𝑇 𝑃) in the ERA5 re-analysis dataset is a variable that represents the the accumulated liquid and frozen water that falls to the Earth’s surface through rainfall and snow. It is defined in units of length as the depth of water that would accumulate if spread evenly over a unit grid box of the model. Compared to the variables handled by our backbone model, 𝑇 𝑃 exhibits certain features that complicate the task of forecasting it—the probability distribution of 𝑇 𝑃 is strongly peaked at zero with a long tail towards positive values. Hence, 𝑇 𝑃 exhibits more sparse spatial features than the other prognostic variables. In addition, 𝑇 𝑃 does not have significant impact on the variables that guide the dynamical evolution of the atmosphere (e.g. winds, pressures, and temperatures), and capturing it accurately in NWP involves complex parameterizations for processes like phase changes. For these reasons, we treat the total precipitation (𝑇 𝑃) as a diagnostic variable ⃗⃗⃗ and denote it by 𝑝(𝑘Δ𝑡). Total precipitation is not included in the 20 variable 249 dataset used to train the backbone model2 . Rather, we train a separate AFNO model to diagnose 𝑇 𝑃 using the outputs of the backbone model, as indicated in Figure 11.2(c). This approach decouples the difficulties of modeling precipitation (which typically deteriorates in accuracy fairly quickly) from the general task of forecasting the atmospheric state. In addition, once trained, our diagnostic 𝑇 𝑃 model could potentially be used in conjunction with other forecast models (either traditional NWP or data-driven forecasts). The model used to diagnose precipitation from the output of the backbone has the same base AFNO architecture, with an additional 2D convolutional layer (with periodic padding) and a ReLU activation as the last layer, used to enforce nonnegative precipitation outputs. Since the backbone model makes predictions in 6-hour increments, we train our diagnostic precipitation model to predict the 6hourly accumulated total precipitation (rather than the 1 hour precipitation in the raw ERA5 data). This also enables easy comparison with the IFS model, which is archived in 6-hour increments and thus also predicts 6-hourly accumulated precipitation. Following (Rasp, Dueben, et al., 2020), we additionally log-transform the precipitation field: 𝑇˜𝑃 = log (1 + 𝑇 𝑃/𝜖), with 𝜖 = 1𝐸 − 5. Since total precipitation values are highly sparse, this transformation discourages the network from predicting zeros and ensures a less skewed distribution of values. For any comparisons with the IFS model or ERA5 ground truth, we transform 𝑇 𝑃 back to units of length. Inference We generate forecasts of the core atmospheric variables in Table 11.1 and the total precipitation by using our trained models in autoregressive inference mode as shown in Figure 11.2(d). The model is initialized with an initial condition ⃗⃗⃗⃗ ( 𝑋 true ( 𝑗)) from the year 2018 3 out-of-sample held out dataset for 𝑁 𝑓 different initial conditions and allowed to freely run iteratively for 𝜏 time-steps to generate forecasts ⃗⃗⃗⃗ ⃗⃗⃗⃗ 𝜏 . The initial conditions 𝑋 { 𝑋 pred ( 𝑗 +𝑖Δ𝑡)}𝑖=1 true ( 𝑗) are spaced apart by 𝐷 days based on a rough estimate of the temporal de-correlation time for each of the variables being forecast. The value of 𝐷 and 𝑁 𝑓 is thus different for each of the forecast variables. We also use the IFS forecasts for the year 2018 from The International Grand Global Ensemble (TIGGE) archive for comparative analysis. The archived IFS forecasts, with initial conditions matching the times of corresponding initial 2 This approach is similar to previous work (Rasp and Thuerey, 2021b), which trained a separate model for precipitation than for the other atmospheric variables. 3 The year 2018 was chosen from the out-of-sample dataset due to ready availability of IFS forecasts for that year from the TIGGE archive. 250 conditions for the FourCastNet model forecast, are used for comparing our model’s accuracy to that of the IFS model. 11.3 Results Figure 11.1 qualitatively shows the forecast skill of our FourCastNet model on forecasting the surface wind speeds over the entire globe at a resolution of 0.25° -latlong. The wind speeds are computed as the magnitude of the surface √︂ wind velocity   2 + 𝑉2 using the zonal and meridonal components of the wind velocity i.e., 𝑈10 10 To prepare this figure, we initialized the FourCastNet model with an initial condition from the out-of-sample test dataset. Starting from this initial condition, the model was allowed to run freely for 16 time-steps in inference mode (Figure 11.2(d)). The calendar time-stamp of the initial condition used to generate this forecast was September 8, 2018 at 00:00 UTC. Figure 11.1(a) shows the wind speed at model initialization. Figure 11.1(b) shows the model forecasts at a lead time of 96 hours (upper-panel) and the corresponding true wind speeds at that time (lower-panel). We note that the FourCastNet model is able to forecast the wind speeds upto 96 hours in advance with remarkable fidelity with correct fine-scale features. Notably, this figure illustrates the forecast of the formation and track of a super-typhoon named Mangkhut that is beginning to form in the initialization frame at roughly 10° 𝑁 latitude, 210°𝑊 longitude. The model qualitatively tracks with remarkable fidelity the intensification of the typhoon and its track over a period of 4 days. Also of note are three simultaneous named hurricanes (Florence, Issac and Helene) forming in the Atlantic ocean and approaching the eastern coast of North America during the period of this forecast. The FourCastNet model appears to be able to forecast the formation and track of these phenomena remarkably well. We provide a further discussion of hurricane forecasts with a few quantitative results and case studies in Section 11.3. In Fig 11.3, we show the forecast skill of our model in diagnosing total precipitation over the entire globe. Using the free running FourCastNet model predictions (from above) for the 20 prognostic variables as input to the precipitation model, we diagnose total precipitation at the same time steps. Fig 11.3(a) shows the precipitation at the initial time, Fig 11.3(b) shows the model predictions at lead time 36 hours along with the corresponding ground truth. The inset panels show the precipitation fields over a local region along the western coast of the United States, highlighting the ability of the FourCastNet model to resolve and forecast localized areas of high 251 precipitation with remarkable accuracy. Forecasting precipitation is known to be an extremely difficult task due to its intermittent and stochastic nature. Despite these challenges, we observe that the FourCastNet diagnosis shows excellent skill in capturing short-term high-resolution precipitation features, which can have significant impact in predicting extreme events. We also note that this is the first time a DL model has been successfully utilized to provide competitive precipitation diagnosis at this scale. Hurricanes In this section, we explore the potential utility of developing DL models for forecasting hurricanes, a category of extreme events with tremendous destructive potential. A rapidly available, computationally inexpensive atmospheric model that could could forewarn the possibility of hurricane formation and track the path of the hurricane would be of great utility for mitigating loss of life and property damage. As the stakes for mis-forecasting such extreme weather phenomena are very high, more rigorous studies need to be undertaken before DL can be considered a mature technology to forecast hurricanes. The results herein should be considered a preliminary and exploratory dive for inspiring future research into the potential of DL models to provide valuable models of this phenomenon. Prior to this work, DL models were trained on data that was too coarse and thus incapable of resolving atmospheric variables finely enough. Prior models could not generate accurate predictions of wind speed and other important prognostic variables with long enough forecast lead times to consider hurricane forecasts. Our model has reasonably good resolution and generates accurate medium-range forecasts of variables that allow us to track the generation and path of hurricanes. For a case-study we consider a hurricane that occurred in 2018 (a year that is part of our out-of-sample dataset), namely hurricane Michael. Michael was a category 5 hurricane on the Saffir -Simpson Hurricane Wind Scale that made landfall in Florida causing catastrophic damage (Beven II, Berg, and Hagen, 2019). Michael started as a tropical depression around October 7, 2018. Within a day, the depression intensified into a hurricane. After undergoing rapid intensification in the gulf of Mexico, Michael reached category 5 status. Soon after, Michael made landfall in Florida on October 10, 2018. Thus within a short period of roughly 72 hours, Michael went from a tropical depression to a category 5 hurricane to landfall. 252 Land-sea borders are shown using a thin white trace. For ease of visualization, the precipitation field is plotted as a log-transformed field in all panels. Panel (a) shows the TP fields at the time of forecast initialization. Panel (b) shows the TP forecast generated by the FourCastNet model (upper panel) over the entire globe at 0.25° -lat-long resolution with the corresponding truth (lower panel). Inset 1 shows the I.C., forecast and true precipitation fields at a lead time of 36 hours over a local region along the western coast of the United States. This highlights the ability of the FourCastNet model to resolve and predict localized regions of high precipitation, in this case due to an atmospheric river. Inset 2 shows the I.C., forecast, and true precipitation fields near the coast of the U.K. and highlights an extreme precipitation event due to an extra-tropical cyclone that is predicted very well by the FourCastNet model. The precipitation is diagnosed from the FourCastNet predicted prognostic variables as described in Figure 11.2(d). The calendar time-stamp of the initial condition used to generate this forecast was 00:00 UTC on April 4, 2018. The high-resolution FourCastNet model demonstrates excellent skill in capturing small scale features that are key to precipitation forecasting. Figure 11.3: Illustration of a global Total Precipitation (TP) forecast using the FourCastNet model. 253 We use our trained model as described in Section 11.2 (with no further changes) to study the potential of our model for forecasting the formation, rapid intensification and tracking of hurricane Michael. The FourCastNet model is capable of rapidly generating large ensemble forecasts. We start from the initial condition at the calendar time 00:00 hours on October 7, 2018 UTC. The initial condition was perturbed with Gaussian noise to generate an ensemble of 𝐸 = 100 perturbed initial conditions. We provide further discussion of ensemble forecasting using FourCastNet in Section 11.3. Figure 11.4 shows the track of the hurricane and the intensification as forecast by the 100-member FourCastNet ensemble using the Mean Sea Level Pressure to estimate the eye of the hurricane and the minimum pressure at the eye. Figure 11.4(a) shows the mean position of the minima of Mean Sea Level Pressure using a 100 member ensemble forecast generated by FourCastNet (red circles). The corresponding ground truth according to ERA5 reanalysis is indicated on the same plot (blue squares) over a trajectory spanning 108 hours. The shaded ellipses in the figure have a width and height equal to the 90th percentile spread in the longitudinal and latitudinal positions respectively of the hurricane eye as indicated by the MSLP minima in the 100-member FourCastNet ensemble. Figure 11.4(b) quantitatively demonstrates that the FourCastNet model is able to predict the intensification of the hurricane as the hurricane eye pressure drops rapidly in the first 72 hours. The minimum MSLP at the eye of hurricane Michael as forecast by FourCastNet is indicated by red circles and the corresponding true minimum from the ERA5 reanalysis is shown by blue circles. The red shaded region shows the region between the first and third quartiles of minimum MSLP in the 100-member ensemble. While this is an impressive result for a model trained on 0.25° resolution data, the model fails to fully forecast the extent of the sharp drop in pressure between 36 and 48 hours. We hypothesize that this is likely due to the fact that the current version of the FourCastNet model does not account for a number of convective and radiative processes that would be crucial to such a forecast. Additionally we expect an AFNO model trained on even higher resolution data to improve such a forecast. Figures 11.4(c),(d) respectively provide a qualitative visualization of three prognostic variables that are useful for tracking the formation, intensification and path of a hurricane, namely the wind speed at the surface and at 850hPa level (calculated as the magnitude of the velocity from the meridional and zonal components of the respective velocity – 𝑈10 , 𝑉10 , 𝑈850 , 𝑉850 ), and the Mean Sea Level Pressure. We believe there is tremendous potential to improve these forecasts by training even 254 higher resolution DL weather models using the AFNO architecture. Our forecasts of the wind speeds and the mean sea level pressure qualitatively match the ground truth remarkably well over a period of 72 hours. Figures 11.4(a),(b), (c), (d) clearly show that the DL model is able to forecast the formation, intensification and track of the hurricane from a tropical depression to landfall on the coast of Florida. Further research is warranted to quantitatively study the potential of our DL model to accurately forecast hurricanes and similar extreme phenomena but these results show great promise in the ability of DL models to aid in the forecasting of one of the most destructive phenomena affecting human life. Atmospheric Rivers Atmospheric rivers are columns of moisture that are transported by atmospheric circulation currents and carry large amounts of water vapor from the tropics to the extra-tropical regions. They are called ‘rivers’ as they often carry an amount of water equivalent to that of the flow rate of major rivers. Large atmospheric rivers can cause extreme precipitation upon landfall, with the potential to cause flooding and extensive damage. More moderately-sized atmospheric rivers are crucial to the water supply of the western United States. Thus, forecasting atmospheric rivers and their landfall locations is crucial for early warning of flooding in low-lying coastal areas as well as for water resource planning. Figure 11.5 shows the use of our FourCastNet model for predicting the formation and evolution of an atmospheric river (using the Total Column of Water Vapor variable) in April 2018 as it made eventual landfall in Northern California. This type of river which passes through Hawaii is often called the Pineapple Express. Atmospheric rivers show up very clearly in the ‘Total Column Water Vapor’ field that is forecast by the FourCastNet backbone model. The FourCastNet model has very good prediction accuracy for 𝑇𝐶𝑊𝑉, with ACC> 0.6 out beyond 8 days. For this atmospheric river, the FourCastNet model was initialized using an intial condition on April 4, 2018 at 00:00 hours UTC, which we display in Figure 11.5(a). Figures 11.5(b) and 11.5(c) show the forecast of the 𝑇𝐶𝑊𝑉 fields generated by the FourCastNet model (top panels) at a lead time of 36 hours and 72 hours respectively and the corresponding ground truth (bottom panels). Whie TCWV is a reasonable proxy for atmospheric rivers, we expect future iterations of our model to include Integrated Vapor Transport and Total Column of Liquid 255 As an illustrative example, we have chosen Hurricane Michael which underwent rapid intensification during the course of its four day trajectory. Panel (a) shows the mean position of the minima of Mean Sea Level Pressure (indicating the eye of hurricane Michael) as forecast by a 100 member ensemble forecast using FourCastNet (red circles) and the corresponding ground truth according to ERA5 reanalysis (blue squares) for 108 hours starting from the initial condition at 00:00 hours on October 7, 2018 UTC. To generate an ensemble forecast, the initial condition was perturbed with Gaussian noise as described in Section 11.3 and 100 forecast trajectories were computed. The shaded ellipses have a width and height equal to the 90th percentile spread of the longitudinal and latitudinal positions respectively of the hurricane eye as indicated by the MSLP minima in the 100member FourCastNet ensemble. Panel (b) shows the minimum MSLP at the eye of hurricane Michael as forecast by FourCastNet (red filled circles) along with the corresponding true minimum from the ERA5 reanalysis (blue filled circles). The red shaded region shows the 90 percent confidence region in the 100-member ensemble forecast. Panels (c) and (d) respectively show the surface wind speed and 850hPa wind speed predictions at lead times of 18 hours, 36 hours, 54 hours and 72 hours generated by FourCastNet along with the corresponding true wind speeds at those times. The surface wind speed and the 850hPa speed in the initial condition (Oct. 7, 2018 00:00 UTC) that was used to initialize this forecast is shown in the leftmost column. Collectively, the minimum MSLP tracks, surface wind speed and the 850hPa wind speed forecasts show the formation, intensification and path of Hurricane Michael as it goes from a tropical depression to a category 5 hurricane with landfall on the west coast of Florida. Figure 11.4: The FourCastNet model has excellent skill on forecasting fine-scale, rapidly changing variables relevant to a hurricane forecast. 256 Water as additional variables to aid in the forecast of atmospheric rivers. Atmospheric rivers are important phenomena that can cause extreme precipitation and contribute significantly to the supply of precipitable water in several parts of the world. Panels (a)-(c) visualize the Total Column Water Vapor (𝑇𝐶𝑊𝑉) in a FourCastNet model forecast initialized at 00:00 UTC on April 4, 2018. Panel (a) Shows the 𝑇𝐶𝑊𝑉 field in the initial condition that was used to initialize the FourCastNet model. Panels (b) and (c) show the forecasts of the 𝑇𝐶𝑊𝑉 field produced by the FourCastNet model (top panels) at lead times of 36 and 72 hours respectively along with the corresponding true 𝑇𝐶𝑊𝑉 fields at those instants of time. The forecast shows an atmospheric river building up and making landfall on the northern California coastline. Figure 11.5: Illustrative example of the utility of the FourCastNet model for forecasting atmospheric rivers. Quantitative Skill of FourCastNet We illustrate the forecast skill of our model for 𝑁 𝑓 initial conditions from the outof-sample dataset (consisting of the year 2018) and generate a forecast for each initial condition. For each forecast, we evaluate the latitude-weighted Anomaly Correlation Coefficient (ACC) and Root Mean Squared Error (RMSE) for all of the variables included in the forecast. We report the mean ACC and RMSE for each of the variables along with the first and third quartile values of the ACC and RMSE at each forecast time step, to show the dispersion of these metrics over different initial conditions. As a comparison, we also compute the same ACC and RMSE metrics for the corresponding IFS forecast with time-matched initial conditions. Figure 11.6(a-f) shows the latitude weighted ACC for the FourCastNet model fore- 257 averaged over several forecasts initialized using initial conditions in the out-of-sample testing dataset corresponding to the calendar year 2018 for the variables (a) 𝑈10 , (b) 𝑇 𝑃, (c) 𝑇2𝑚 , (d) 𝑍500 , (e) 𝑇850 , and (f) 𝑉10 . The ACC values are averaged over 𝑁 𝑓 initial conditions over a full year with an interval of 𝐷 days between consecutive initial conditions to account for seasonal variability in forecast skill. The appropriately colored shaded regions around the ACC curves indicate the region between the first and third quartile values of the corresponding quantity at each time step. Figure 11.6: Latitude weighted ACC for the FourCastNet model forecasts (red line with markers) and the corresponding matched IFS forecasts (blue line with markers) casts (Red line with markers) and the corresponding matched IFS forecasts (Blue line with markers) for the variables (a) 𝑈10 , (b) 𝑇 𝑃, (c) 𝑇2𝑚 , (d) 𝑍500 , (e) 𝑇850 , (f) 𝑉10 . The ACC and RMSE values are averaged over 𝑁 𝑓 initial conditions with an interval of 𝐷 days between consecutive initial conditions. The shaded regions around the ACC curves indicate the region between the first and third quartile values of the corresponding quantity at each time step. In general, the FourCastNet predictions are very competitive with IFS, with our model achieving similar ACC and RMSE over a horizon of several days. At shorter lead times (∼ 48hrs or less), we actually outperform the IFS model in ACC and/or RMSE for key variables like precipitation, winds, and temperature. Remarkably, we achieve this accuracy using only part of the full variable set available to the IFS model, and we do so at a fraction of the compute cost (see section 11.4 for a detailed speed comparison between models). 258 Ensemble Forecasts Using FourCastNet Ensemble forecasts have become a crucial component of numerical weather prediction (Palmer, 2019), and consume the largest share of compute costs at operational weather forecasting centers (Bauer, Quintino, et al., 2020). An ensemble forecast improves upon a single deterministic forecast by modeling multiple possible trajectories of a system. For a chaotic atmosphere with uncertain initial conditions, ensemble forecasting helps quantify the likelihood of extreme events and improves the accuracy of long-term predictions. Thus, rapidly generating large ensemble forecasts is an extremely promising direction for DL-based weather models (Weyn, Durran, Caruana, and Cresswell-Clay, 2021), which can provide immense speedups over traditional NWP models. NWP models such as the IFS perform ensemble forecasting with up to 51 ensemble members. The initial conditions for the ensemble forecasts are obtained by perturbing the analysis state obtained from data assimilation. As seen in Section 11.3, ensemble forecasting is useful for generating probabilistic forecasts of extreme events such as hurricanes. While the individual perturbed ensemble members typically show lower forecast skill than the unperturbed ‘control’ forecast, the mean of a large number of such perturbed ensemble members has better forecast skill than the control. In Section 11.4, we estimate that FourCastNet is roughly 45,000 times faster than a traditional NWP model. This speed allows us to consider probabilistic ensemble forecasting with massive ensemble sizes. Ensemble weather forecasts using FourCastNet are highly computationally efficient because (1.) Inference time for a single forecast on a GPU is very fast and (2.) An ensemble of initial conditions can be folded into the the ‘batch’ dimension in a tensor and as such, inference on a large batch (𝑂 (100) or more) of initial conditions using a few GPUs is straightforward. As a simple test of ensemble forecasting, we generate an ensemble forecast using FourCastNet from a given ERA5 initial condition by perturbing the initial condition using Gaussian random noise. This allows us to simulate initial condition uncertainty due to errors in the estimate of the starting state of the forecast. This method of ensemble generation is the same as methods used in Ensemble Kalman Filtering (EnKF) (Evensen, 2003) for background forecast covariance estimation and not too dissimilar from the way operational NWP models generate perturbed initial ⃗⃗⃗⃗ conditions. Thus, given an initial condition 𝑋 true (𝑘) from our out-of-sample testing ⃗⃗⃗⃗ dataset, we generate an ensemble of 𝐸 perturbed initial conditions { 𝑋 (𝑒) (𝑘) = 259 ⃗⃗⃗⃗ ⃗⃗⃗⃗ ˆ (𝑘) is the standardized initial condition with zero 𝐸 , where 𝑋 𝑋ˆ true (𝑘) + 𝜎𝜉} 𝑒=1 true ⃗⃗ ⃗⃗ mean and unit variance and 𝜉 ∼ N ( 0, 1) is a normally distributed random variable ⃗⃗⃗⃗ of the same shape as 𝑋 true and with unit mean and variance. The perturbations are scaled by a factor 𝜎 = 0.3. We refer to the forecast starting from the unperturbed initial condition as the control forecast. We generate an ensemble of perturbed forecasts each starting from a perturbed initial conditions and compute the ensemble mean of the perturbed forecasts at every forecast time step. We compute a control forecast and an ensemble mean forecast for 𝑁 𝑓 initial conditions separated by 𝐷 days. We report the mean ACC and RMSE over all 𝑁 𝑓 initial conditions for both the control and the mean forecast in Figure 11.7. Figure 11.7 shows the ACC and RMSE of the FourCastNet ensemble mean (magenta line with markers) and FourCastNet unperturbed control (red line with markers) forecasts along with the unperturbed control IFS model (blue line with markers) forecasts for reference. It is challenging to unambiguously visualize in a single plot, both the spread due to simulated initial condition uncertainty in an ensemble forecast and the spread due to seasonal and day-to-day variability. As such, we do not visualize the spread in ACC and RMSE over the 𝑁 𝑓 forecasts and simply report the mean. Indeed, in Figure 11.7, we see that the ensemble mean from our 100-member FourCastNet ensemble results in a net improvement in ACC and RMSE at longer timescales over the unperturbed control. We do observe a marginal degradation in skill for the ensemble mean at short (< 48hr) lead times, as averaging over the individual ensemble members likely averages over relevant fine-scale features. Nevertheless, these ensemble forecasts are impressive, and warrant further work in how to optimally choose ensemble members. In addition to perturbing initial conditions with Gaussian noise, as we do here, it is possible and likely worthwhile to introduce more nuanced perturbations to both the initial conditions as well as the model itself. This is a promising direction of research for future work. Forecast Skill Over Land For Near-surface Wind Speed Most wind farms are located on land or just off of coastlines, so accurately modeling near-surface wind speed over these regions is of critical importance to wind energy resource planning. To demonstrate the accuracy of FourCastNet predictions over √︃ 2 + 𝑉 2 ) forecast and ground truth over landmasses, we plot the 10m wind speed ( 𝑈10 10 North America in Figure 11.8. We find that FourCastNet can qualitatively capture 260 We compare the forecast skill of the unperturbed ‘control’ forecasts using FourCastNet (red) with the mean of a 100-member ensemble forecast using FourCastNet (magenta) for 𝑍500 (panel a) and 𝑈10 (panel b). The IFS unperturbed control forecast is included for reference (blue). We find that the 100-member FourCastNet ensemble mean is more skillful than the FourCastNet control at longer forecast lead times. The 100-member FourCastNet ensemble mean shows significant improvement over the unperturbed FourCastNet control forecast beyond 70 hours for 𝑈10 and 100 hours for 𝑍500 . Due to the challenge of clearly disambiguating in a single plot the forecast spread arising from simulated initial condition uncertainty and the forecast spread due to seasonal and day-to-day variability, we choose not to visualize the spread in ACC and RMSE over the 𝑁 𝑓 forecasts and simply report the mean. Figure 11.7: Illustration of the improvement in forecast skill of FourCastNet by utilizing large ensembles. the spatial patterns and intensities of surface winds with impressive accuracy up to several days in advance. Moreover, the visualizations emphasize the importance of running forecasts at high resolution, as the surface wind speed exhibits significant fine-scale spatial variations which would be lost with a coarser grid. We evaluate the forecast skill of our model over land versus over oceans quantitatively. By computing a separate land-masked ACC and a sea-masked ACC for the surface wind velocity components, we find that the forecast quality of our model for surface wind speed over landmass is almost as good as it is over the ocean. This is significant, as surface wind speed over land is strongly affected by orographic features such as mountains, making it in general harder to forecast surface winds 261 This is a significant result for wind energy resource planning,√︃as windfarms are located on 2 + 𝑉 2 ) forecast (top four land or just offshore. The figure shows The 10m wind speed ( 𝑈10 10 panels) generated by FourCastNet and corresponding ground truth (bottom four panels) for forecast lead times of 18 hours, 36 hours, 54 hours and 72 hours. The forecast was initialized with an initial condition at calendar time 06:00:00 on July 4 2018 UTC. Figure 11.8: The FourCastNet model shows excellent skill on forecasting overland wind speed, a challenging problem due to topographic features such as mountains and lakes. over land than over the oceans. Extremes We assess the ability of the FourCastNet model to capture instantaneous extremes by looking at the top quantiles of each field at a given time step. Similar to the approach in Fildier, W. D. Collins, and Muller, 2021, we use 50 logarithmically-spaced quantile bins 𝑄 = 1 − {10−1 , ..., 10−4 } (corresponding to percentiles {90%, ..., 99.99%}) to emphasize the most extreme values (generally, the FourCastNet predictions and ERA5 targets match closely up to around the 98th percentile). We choose the 99.99th as the top percentile bin because percentiles beyond there sample less than 1000 pixels in each image and are subject to more variability. We show example plots of the top quantiles for 𝑈10 and 𝑇 𝑃 at 24-hour forecast times in the left panel of Figure 11.9 (these particular forecasts were initialized at 00:00 UTC Jan 1 2018). At this particular time, both the FourCastNet and IFS models under-predict extreme precipitation, while for extreme winds in 𝑈10 the IFS model over-predicts and FourCastNet under-predicts. To get a more comprehensive picture, we need to evaluate the model performance at multiple forecast times over multiple initial conditions in order to ascertain if there is a systematic bias in the model’s predictions for extreme values. 262 Figure 11.9: Comparison of extreme percentiles between ERA5, FourCastNet, and IFS. The left panel shows the top percentiles of the 𝑇 𝑃 and 𝑈10 distribution at a forecast time of 24 hours, for a randomly sampled initial condition. The right panel shows the 𝑇 𝑃 and 𝑈10 relative quantile error (RQE, defined in the text) as a function of forecast time, averaged over 𝑁 𝑓 initial conditions in the calendar year 2018 (filled region spans the 1st and 3rd quartiles). On average, RQE trends slightly negative for both models as they under-predict the most extreme values for these variables, especially for 𝑇 𝑃. To this end, we define the relative quantile error (RQE) at each time step 𝑙 as RQE(𝑙) = ∑︁ ⃗⃗⃗⃗ ⃗⃗⃗⃗ 𝑞 ⃗⃗⃗⃗ 𝑞 𝑞 ( 𝑋 pred (𝑙) − 𝑋 true (𝑙))/ 𝑋 true (𝑙), (11.4) 𝑞∈𝑄 ⃗⃗⃗⃗ ⃗⃗⃗⃗ where 𝑋 𝑞 (𝑙) is the 𝑞 th -quantile of 𝑋 (𝑙). RQE trends negative for a given variable if a model systematically under-predicts that variable’s extremes, and we indeed find that both the FourCastNet and IFS models show a slight negative RQE over different forecast times and initial conditions for both 𝑇 𝑃 and 𝑈10 . This can be seen in the right-hand panel of Figure 11.9. For 𝑈10 , the difference between FourCastNet and IFS is negligible and, on average, both models underestimate the extreme percentiles by just a few percentage points in RQE. For 𝑇 𝑃, the difference with respect to IFS is more pronounced, and FourCastNet underestimates the extreme percentiles by ∼ 35% in RQE, compared to ∼ 15% for IFS. This is not surprising given the forecasts visualized in Figure 11.3, which show 263 the FourCastNet predictions being generally smoother than the ERA5 targets. As the extreme values tend to be concentrated in extremely small regions (sometimes down to the gridbox/pixel scale), a model that fails to fully resolve these scales will have a harder time capturing 𝑇 𝑃 extremes. Given the noise and uncertainties, predicting precipitation extremes is well-known to be a challenging problem, but we believe our model could be improved further by focusing more on such fine-scale features. We leave this for future work. 11.4 Computational Cost of FourCastNet In comparing the speed of forecast generation between FourCastNet and IFS, we have to deal with the rather difficult problem of comparing a forecast computed using a CPU cluster (in the case of the IFS model) and a forecast that is computed on a single (or perhaps a few) GPU(s) (FourCastNet). We take a nuanced approach to reporting this comparison. Our motivation is not to create a definitive apples to apples comparison and tout a single numerical factor advantage for our model, but merely to illustrate the order-of-magnitude differences in forecast generation time and also highlight the radically different perspectives of computation when comparing traditional NWP models with DL models. Through this comparison, we also wish to highlight the significant potential of FourCastNet and future DL models to offer an important addition to the toolkit of a meteorologist. To estimate the forecast speed of the IFS model, we use figures provided in Bauer, Quintino, et al. (2020) as a baseline. In Ref. (Bauer, Quintino, et al., 2020), we see that the IFS model computes a 15-day, 51-member ensemble forecast using the “L91” 18km resolution grid on 1530 Cray XC40 nodes with dual socket Intel Haswell processors in 82 minutes. The IFS model archived in TIGGE, which we compare the FourCastNet predictions with in Section 11.3, also uses the L91 18km grid for computation (but is archived at the ERA5 resolution of 30km). Based on this information, we estimate that to compute a 24-hour 100-member ensemble forecast, the reference IFS model would require 984,000 node-seconds. We estimate the energy consumption for computing such a 100-member forecast to be 271MJ4 . We now estimate the latency and energy consumption of the FourCastNet model. The FourCastNet model can compute a 100-member 24-hour forecast in 7 seconds by using a single node on the Perlmutter HPC cluster which contains 4 A100 GPUs per node. This is achieved by performing batched inference on the 4 A100 GPUs 4 A dual-socket Intel Haswell node draws a Thermal Design Power (TDP) of 270 Watts 264 using a batch size of 25. Thus, the FourCastNet model takes 7 node-seconds per forecast day for a 100-member ensemble. With a peak power consumption of 1kW per node, we estimate this 100-member 24-hour forecast to use 8kJ of energy. We attempt to account for the resolution difference between the 18km L91 model and the 30km FourCastNet model by additionally reporting the inference time for an 18km FourCastNet model. Since we did not train the FourCastNet with 18km resolution data (due to the lack of such a publicly available dataset), the reported numbers simply estimate the computational costs for such a hypothetical model by performing inference on data and model parameters interpolated to 18km resolution from the original 30km resolution. Table 11.2 provides a comparison of computational speed and energy consumption of the IFS L91 18km model and the FourCastNet model at 30km resolution, as well as the extrapolated 18km resolution. These results suggest that the FourCastNet model can compute a 100-member ensemble forecast using vastly fewer nodes, at a speed that is between 45,000 times faster (at the 18 km resolution) and 145,000 times faster (at the 30 km resolution) on a node-to-node comparison. By the same estimates, FourCastNet has an energy consumption that is between 12,000 (18 km) and 24,000 (30 km) times lower than that of the IFS model. Latency and Energy consumption for a 24-hour 100-member ensemble forecast FCN - 30km FCN - 18km IFS IFS / FCN(18km) Ratio (actual) (extrapolated) Nodes required 3060 1 2 1530 Latency 984000 7 22 44727 (Node-seconds) Energy Consumed 271000 7 22 12318 (kJ) In comparison, the IFS model needs 3060 nodes for such a forecast. In this table, we provide information about latency and energy consumption for the FourCastNet model in comparison with the IFS model. The FourCastNet model at a 30km resolution is about 145,000 times faster on a single-node basis than the IFS model. We can also estimate the cost of generating an 18km resolution forecast using FourCastNet. Such a hypothetical 18km model would be about 45,000 times faster than the IFS on a single-node basis. The FourCastNet model at 30km resolution uses 24,000 times less energy to compute the ensemble forecast than the IFS model, while a hypothetical FourCastNet model at 18km resolution would use 12000 times less energy. Table 11.2: The FourCastNet model can compute a 100-member ensemble forecast on a single 4GPU A100 node. This comparison comes with several caveats. The IFS model generates forecasts that 265 are provably physically consistent, while FourCastNet in its current iteration does not impose physics constraints. The IFS model also outputs an order of magnitude more variables at as many as 100 vertical levels. Notably, the IFS model is generally more accurate than FourCastNet (although in several variables, the DL model approaches the accuracy of the IFS model and exceeds it for precipitation in certain cases). On the other hand, it is worth noting our rudimentary speed assessments of FourCastNet do not employ any of the common optimizations used for inference of DL models (e.g., model distillation, pruning, quantization, or reduced precision). We expect implementing these would greatly accelerate our speed of inference, and lead to further gains in computational efficiency over IFS. While the above caveats are important, it is fair to say that if one were only interested in limited-purpose forecasting (e.g., a wind farm operator interested in short-term surface wind speed forecasts), FourCastNet would be a very attractive option as the infrastructure requirements are minimal. FourCastNet can generate a 10-day, global forecast at full ERA5 resolution using a single device, which is simply not possible with IFS, and such a forecast completes in seconds. This means one could generate reasonably accurate forecasts using a tabletop computer with a single GPU, rather than needing a substantial portion of a compute cluster. Similarly, only a handful of GPUs are needed for generating ensemble forecasts with 100s of ensemble members, and such ensembles run quickly and efficiently using batched inference. This greatly lowers the barrier to entry for doing data-assimilation and uncertainty quantification, and future work in this direction is warranted to explore these possibilities. 11.5 Comparison Against State-of-the-art DL Weather Prediction To the best of our knowledge, the current state-of-the-art DL weather prediction model is the DLWP model of Weyn, Durran, and Caruana (2020)—they employ a deep convolutional network with a cubed-sphere remapped coordinate system to predict important weather forecast variables. The authors work with a coarser resolution of 2◦ and forecast variables relating to geopotential heights, geopotential thickness, and 2-m temperature (see (Weyn, Durran, and Caruana, 2020) for further details). The FourCastNet model predicts more variables than the DLWP model at a resolution that is higher than the DLWP model by a factor of 8. The significantly higher resolution of the FourCastNet model resolves fine-scale features present in variables such as wind velocities and precipitation allowing us to resolve important phenomena such as hurricanes, extreme precipitation and atmospheric rivers. This would not be possible at a lower resolution such as 2° (and almost entirely a futile 266 exercise at a 5° resolution.) For reference, we have visualized the MSLP over the trajectory of hurricane Michael at a resolution of 2° . Thus, the FourCastNet model has many characteristics that make it superior to the prior SOTA DLWP model. Nonetheless, we undertake a comparison of forecasts generated by the FourCastNet model with those of the DLWP model by coarsening the FourCastNet outputs to bring them to a resolution comparable to that of the DLWP model. We emphasize that this comparison has been provided only for the sake of completeness. Coarsening our forecasts and making them less effective in order to accommodate a prior benchmark at a lower resolution is not fair to our model. We downsample our predictions eight times (in each direction, using bilinear interpolation) to coarsen them to a resolution that is comparable to that of the DLWP model. Since the two variables reported in the DLWP results are 𝑍500 and 𝑇2𝑚 , we re-compute our ACC and RMSE metrics for those two variables. We also note that the ACC metric in the DLWP baseline was computed using daily climatology (we use a time-averaged climatology in this work, motivated by (Rasp, Dueben, et al., 2020)) and, hence, we modify our ACC computation using the same definition for a fair comparison. We show our comparisons for ACC and RMSE in Figure 11.10. We observe that even at the lower resolution of the DLWP work, the FourCastNet model predictions show significant improvement over the current state-of-the-art DLWP model in both variables. Additionally, the FourCastNet model operates at a resolution that is 8 times higher than the DLWP model allowing it to resolve many important fine-scale phenomena. 11.6 Implications, Discussion, and Future Work FourCastNet is a novel global data-driven DL-based weather forecasting model based on the FNO and AFNO (Li, Kovachki, et al., 2021; Guibas et al., 2022). FourCastNet’s speed, computational cost, energy footprint, and capacity for generating large ensembles has several important implications for science and society. In particular, FourCastNet’s high-resolution, high-fidelity wind and precipitation forecasts are of tremendous value. Even though FourCastNet was developed in less than a year and has only a fraction of the number of variables and vertical levels compared to NWP, its accuracy is comparable to the IFS model and better than state-of-the-art DL weather prediction models (Weyn, Durran, Caruana, and Cresswell-Clay, 2021; Rasp, Dueben, et al., 2020) on short timescales. We anticipate that with additional resources and further development, FourCastNet could match the capabilities of current NWP models on all timescales and at all vertical levels of the atmosphere. 267 (a) 𝑍500 and (b) 𝑇2𝑚 . We observe that the FourCastNet predictions show significant improvement over the baseline model. We also note that the FourCastNet generates predictions that have a higher resolution by a factor of 8, and is thus able to resolve many more important fine-scale features than the DLWP model. Figure 11.10: Comparison of ACC and RMSE metrics between the (downsampled) FourCastNet predictions, (downsampled) IFS, and baseline state-of-the-art DLWP model Implications FourCastNet’s predictions are four to five orders of magnitude faster than traditional NWP models. This has two important implications. First, large ensembles of thousands of members can be generated in seconds, thus enabling estimation of well-calibrated and constrained uncertainties in extremes with higher confidence than current NWP ensembles that have at most approximately 50 members owing to their high computational cost. Fast generation of 1,000-member ensembles dramatically changes what is possible in probabilistic weather forecasting, including improving reliability of early warnings of extreme weather events and enabling rapid assessment of their impacts. Second, FourCastNet is suitable for rapidly testing hypotheses about mechanisms of weather variability and their predictability. The unprecedented accuracy in short-range forecasts of precipitation and its extremes has potentially massive benefits for society such as enabling rapid responses for disaster mitigation. Furthermore, a highly accurate DL-based diagnostic precipitation model provides the flexibility to input prognostic variables from different 268 models or observational sources. For the wind energy industry, FourCastNet’s rapid and reliable high-resolution wind forecasts can help mitigate disasters from extreme wind events and enables planning for fluctuations in wind power output. Wind farm designers can benefit from fast and reliable high-resolution wind forecasts to optimize wind farm layouts that account for a wide variety of wind and weather conditions. Discussion FourCastNet’s skill improves with increasing number of modeled variables. A larger model trained on more variables, perhaps even on entire 3D atmospheric fields, may extend prediction horizons even further and with better uncertainty estimates. Not far in the future, FourCastNet could be trained on all nine petabytes of the ERA5 dataset to predict all currently predicted variables in NWP at all atmospheric levels. Although the cost of training such a model will be huge, fast inference will enable rapid predictions of entire 3D fields in a few seconds. Such an advancement will likely revolutionize weather prediction. Due to the current absence of a data-assimilation component, FourCastNet cannot yet generate up-to-the-minute weather forecasts. If observations are available, however, such a component could be readily added given the ease of generating large ensembles for methods such as Ensemble Kalman Filtering with data-driven background covariance estimation (Ashesh Chattopadhyay, Mustafa Mustafa, Hassanzadeh, and Karthik Kashinath, 2020). Therefore, in principle, future iterations of FourCastNet could be trained on observational data. This will enable real-time weather prediction by initializing the model with real-time observations. With ever-increasing demands for very high-resolution forecasts, NWP has seen a steady growth in resolution. The increase in computational cost of NWP for a doubling of resolution is nearly 12-fold (23.5 ). Current IFS forecasts are at 9-km resolution but we require forecasts at sub-km resolution for improvements in a wide variety of applications, such as energy and agricultural planning, transportation, and disaster mitigation. Simultaneously, DL continues to have ever-increasing accuracy and predictive power with larger models that have hundreds of billions of parameters (Rajbhandari et al., 2020). With advances in large-scale DL we expect that FourCastNet can be trained to predict weather on sub-km scales. Even though training such a large DL model will be computationally expensive, since inference of large DL models can still be done rapidly (DeepSpeed: Accelerating large-scale 269 model inference and training via system optimizations and compression 2021), a sub-km resolution version of FourCastNet will have even more dramatic speedup over sub-km resolution NWP, likely more than six orders of magnitude. DLWP has shown good skill on S2S timescales (Weyn, Durran, Caruana, and Cresswell-Clay, 2021). FourCastNet has better skill at short timescales (up to two weeks). We envision a coupled model using a two-timescale approach that combines DLWP and FourCastNet with two-way interactions to achieve unprecedented accuracies on short-, medium-, and long-range weather forecasts. FourCastNet is a purely data-driven DL weather model. The physical systems of weather and climate are governed by the laws of nature, some of which are wellunderstood, such as Navier-Stokes equations for the fluid dynamics of atmosphere and oceans. Weather forecasts and climate predictions that obey known physical laws are more trustworthy than those that do not. Furthermore, models that obey the laws of physics are more likely to be robust under climate change. An emerging field in AI applications in the sciences is Physics-informed Machine Learning (Kashinath et al., 2021). The Fourier Neural Operator has been extended to be physics-informed (Li, Zheng, et al., 2021). Future versions of FourCastNet will incorporate physical laws. A physics-informed version of FourCastNet could be trained with fewer datapoints. This benefit is particularly valuable at higher resolutions in order to reduce the data volume requirements for training. FourCastNet could also combine with a physics-based NWP model Arcomano, Szunyogh, Wikner, et al. (2021), to generate long-term stable forecasts over S2S timescales. An important question that remains unanswered is whether FourCastNet generalizes under climate change. FourCastNet was trained on data from 1979 to 2015 and tested on data from 2016 to 2020. We know that Earth’s climate has changed over this period of time. Therefore, FourCastNet has been trained on data from a changing climate. However, FourCastNet may not predict weather reliably under extreme climate change expected in the decades to come. A future version will initialize with climate model output to evaluate FourCastNet’s performance under different warming scenarios. A grand challenge for the climate community is to predict the changing behavior of extreme weather events under climate change, such as their frequency, intensity, and spatio-temporal nature. Once FourCastNet achieves high fidelity under extreme climate change, it can address this grand challenge. 270 References Alley, Richard B, Kerry A Emanuel, and Fuqing Zhang (2019). “Advances in weather prediction”. In: Science 363.6425, pp. 342–344. Arcomano, Troy, Istvan Szunyogh, Jaideep Pathak, et al. (2020). “A machine learning-based global atmospheric forecast model”. In: Geophysical Research Letters 47.9, e2020GL087776. Arcomano, Troy, Istvan Szunyogh, Alexander Wikner, et al. (2021). “A Hybrid Approach to Atmospheric Modeling that Combines Machine Learning with a Physics-Based Numerical Model”. In: Journal of Advances in Modeling Earth Systems. Balaji, V (2021). “Climbing down Charney’s ladder: machine learning and the postDennard era of computational climate science”. In: Philosophical Transactions of the Royal Society A 379.2194, p. 20200085. Bauer, Peter, Tiago Quintino, et al. (2020). The ecmwf scalability programme: Progress and plans. European Centre for Medium Range Weather Forecasts. doi: 10.21957/gdit22ulm. url: https://www.ecmwf.int/node/19380. Bauer, Peter, Alan Thorpe, and Gilbert Brunet (2015). “The quiet revolution of numerical weather prediction”. In: Nature 525.7567, pp. 47–55. Beven II, J.L., R. Berg, and A. Hagen (Apr. 2019). Tropical Cyclone Report Hurricane Michael. Chantry, Matthew et al. (2021). “Opportunities and challenges for machine learning in weather and climate modelling: hard, medium and soft AI”. In: Philosophical Transactions of the Royal Society A 379.2194, p. 20200083. Chattopadhyay, Ashesh, Mustafa Mustafa, Pedram Hassanzadeh, Eviatar Bach, et al. (2021). “Towards physically consistent data-driven weather forecasting: Integrating data assimilation with equivariance-preserving spatial transformers in a case study with ERA5”. In: Geoscientific Model Development Discussions, pp. 1–23. Chattopadhyay, Ashesh, Mustafa Mustafa, Pedram Hassanzadeh, and Karthik Kashinath (2020). “Deep spatial transformers for autoregressive data-driven forecasting of geophysical turbulence”. In: Proceedings of the 10th International Conference on Climate Informatics, pp. 106–112. Chattopadhyay, Ashesh, Ebrahim Nabizadeh, and Pedram Hassanzadeh (2020). “Analog forecasting of extreme-causing weather patterns using deep learning”. In: Journal of Advances in Modeling Earth Systems 12.2, e2019MS001958. DeepSpeed: Accelerating large-scale model inference and training via system optimizations and compression (2021). url: https : / / www . microsoft . com / en - us / research / blog / deepspeed - accelerating - large - scale model-inference-and-training-via-system-optimizatio/ns-andcompression/ (visited on 05/24/2021). 271 Dosovitskiy, Alexey et al. (2021). An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. arXiv: 2010.11929 [cs.CV]. Evensen, Geir (2003). “The ensemble Kalman filter: Theoretical formulation and practical implementation”. In: Ocean dynamics 53.4, pp. 343–367. Fildier, Benjamin, William D. Collins, and Caroline Muller (2021). “Distortions of the Rain Distribution With Warming, With and Without Self-Aggregation”. In: Journal of Advances in Modeling Earth Systems 13.2. e2020MS002256 2020MS002256, e2020MS002256. doi: https://doi.org/10.1029/2020MS002256. eprint: https : / / agupubs . onlinelibrary . wiley . com / doi / pdf / 10 . 1029 / 2020MS002256. url: https : / / agupubs . onlinelibrary . wiley . com/doi/abs/10.1029/2020MS002256. Grönquist, Peter et al. (2021). “Deep learning for post-processing ensemble weather forecasts”. In: Philosophical Transactions of the Royal Society A 379.2194, p. 20200092. Guibas, John et al. (Apr. 2022). “Adaptive Fourier Neural Operators: Efficient Token Mixers for Transformers”. In: International Conference on Representation Learning (to appear). Hersbach, Hans et al. (2020). “The ERA5 global reanalysis”. en. In: Quarterly Journal of the Royal Meteorological Society 146.730, pp. 1999–2049. issn: 1477870X. Irrgang, Christopher et al. (2021). “Towards neural Earth system modelling by integrating artificial intelligence in Earth system science”. In: Nature Machine Intelligence 3.8, pp. 667–674. Kalnay, Eugenia (2003). Atmospheric modeling, data assimilation and predictability. Cambridge University Press. Kalnay, Eugenia et al. (1996). “The NCEP/NCAR 40-year reanalysis project”. In: Bulletin of the American meteorological Society 77.3, pp. 437–472. Kashinath, K et al. (2021). “Physics-informed machine learning: case studies for weather and climate modelling”. In: Philosophical Transactions of the Royal Society A 379.2194, p. 20200093. Li, Zongyi, Nikola Kovachki, et al. (2021). “Fourier Neural Operator for Parametric Partial Differential Equations”. In: International Conference on Learning Representations (ICLR). Li, Zongyi, Hongkai Zheng, et al. (2021). Physics-Informed Neural Operator for Learning Partial Differential Equations. arXiv: 2111.03794 [cs.LG]. Palmer, Tim (2019). “The ECMWF ensemble prediction system: Looking back (more than) 25 years and projecting forward 25 years”. In: Quarterly Journal of the Royal Meteorological Society 145.S1, pp. 12–24. doi: https : / / doi . org/10.1002/qj.3383. eprint: https://rmets.onlinelibrary.wiley. 272 com/doi/pdf/10.1002/qj.3383. url: https://rmets.onlinelibrary. wiley.com/doi/abs/10.1002/qj.3383. Rajbhandari, Samyam et al. (2020). “Zero: Memory optimizations toward training trillion parameter models”. In: SC20: International Conference for High Performance Computing, Networking, Storage and Analysis. IEEE, pp. 1–16. Rasp, Stephan, Peter D Dueben, et al. (2020). “WeatherBench: a benchmark data set for data-driven weather forecasting”. In: Journal of Advances in Modeling Earth Systems 12.11, e2020MS002203. Rasp, Stephan and Nils Thuerey (2020). “Purely data-driven medium-range weather forecasting achieves comparable skill to physical models at similar resolution”. In: arXiv preprint arXiv:2008.08626. – (2021a). “Data-Driven Medium-Range Weather Prediction With a Resnet Pretrained on Climate Simulations: A New Model for WeatherBench”. In: Journal of Advances in Modeling Earth Systems 13.2, e2020MS002405. – (2021b). “Data-driven medium-range weather prediction with a Resnet pretrained on climate simulations: A new model for WeatherBench”. In: Journal of Advances in Modeling Earth Systems, e2020MS002405. Reichstein, Markus et al. (2019). “Deep learning and process understanding for data-driven Earth system science”. In: Nature 566.7743, pp. 195–204. Richardson, Lewis Fry (2007). Weather prediction by numerical process. Cambridge university press. Scher, Sebastian and Gabriele Messori (2018). “Predicting weather forecast uncertainty with machine learning”. In: Quarterly Journal of the Royal Meteorological Society 144.717, pp. 2830–2841. – (2019). “Weather and climate forecasting with neural networks: using general circulation models (GCMs) with different complexity as a study ground”. In: Geoscientific Model Development 12.7, pp. 2797–2809. Schultz, MG et al. (2021). “Can deep learning beat numerical weather prediction?” In: Philosophical Transactions of the Royal Society A 379.2194, p. 20200097. Weyn, Jonathan A, Dale R Durran, and Rich Caruana (2019). “Can machines learn to predict weather? Using deep learning to predict gridded 500-hPa geopotential height from historical weather data”. In: Journal of Advances in Modeling Earth Systems 11.8, pp. 2680–2693. – (2020). “Improving data-driven global weather prediction using deep convolutional neural networks on a cubed sphere”. In: Journal of Advances in Modeling Earth Systems 12.9, e2020MS002109. Weyn, Jonathan A, Dale R Durran, Rich Caruana, and Nathaniel Cresswell-Clay (2021). “Sub-seasonal forecasting with a large ensemble of deep-learning weather prediction models”. In: arXiv preprint arXiv:2102.05107. 273 C h a p t e r 12 APPLICATION: CARBON CAPTURE AND STORAGE Carbon capture and storage (CCS) is an essential climate change mitigation strategy for reducing carbon dioxide emissions. We consider the storage aspect of CCS, which involves injecting carbon dioxide into underground reservoirs. This requires accurate and high-resolution predictions of carbon dioxide plume migration and reservoir pressure buildup. However, such modeling is challenging at scale due to the high computational costs of existing numerical methods. We introduce a novel machine learning approach for four-dimensional spatial-temporal modeling, which speeds up predictions nearly 700,000 times compared to existing methods. It provides a general-purpose numerical simulator alternative for highly accurate flow and pressure predictions under diverse reservoir conditions, geological heterogeneity, and injection schemes. Our framework, Nested Fourier Neural Operator (FNO), uses a hierarchy of FNO models to produce outputs at different refinement levels. It enables unprecedented real-time high-resolution modeling for basin-scale carbon dioxide storage. 12.1 Introduction Carbon capture and storage (CCS) is the process of capturing carbon dioxide and permanently storing it in subsurface geological formations to mitigate climate change (NAS, 2018). We consider the storage modeling of CCS, which involves the multiphase flow of carbon dioxide and water through porous media. Numerical simulations are used to guide critical engineering decisions in CCS projects by forecasting carbon dioxide gaseous plume migration in the formation and pressure buildup caused by the injection. Current numerical simulations are very expensive due to the multi-scale and multiphysics modeling involved, the high-resolution grids required, and the large spatiotemporal ranges needed for real-world scenarios. The governing partial differential equations (PDEs) involving the multiphase variation of Darcy’s law are expensive to solve (Pruess, Oldenburg, and Moridis, 1999; Blunt, 2017). Carbon dioxide and water are immiscible and mutually soluble, requiring multi-physics simulation coupled with thermodynamics (Pruess and Garcia, 2002). High-resolution spatial grids (1-2 m resolution) are necessary to provide accurate estimates of the carbon 274 dioxide plume migration, which is controlled by the complex interplay of buoyancy, viscous, and capillary forces (Doughty, 2010; Wen and Benson, 2019). Pressure buildup, as well as the dry-out effect, i.e., evaporation of formation fluid into the gas phase (Pruess and Müller, 2009; André, Peysson, and Azaroual, 2014), also demand particularly high resolutions around the injection well. In real-world scenarios, CCS projects often involve large reservoir spatial domains of hundreds of kilometers (Chadwick et al., 2004) and long time frames that span from decades to hundreds of years (NETL, 2017). One approach for reducing the computational costs of numerical simulations is to use non-uniform grids to capture different responses at different grid resolutions. A popular method, known as local grid refinement (Bramble et al., 1988) (LGR), has enabled scaling simulations to real-world three-dimensional (3D) carbon dioxide storage projects, where the fine-grid responses capture the plume migration while the coarser grid responses capture the far-field pressure buildup (Eigestad et al., 2009; Faigle et al., 2014; Kamashev and Amanbek, 2021). However, even with nonuniform grid approaches, these numerical models are still too expensive for tasks that require repetitive forward simulations, e.g., site selection (Callas et al., 2022), inversion (Strandli, Mehnert, and Benson, 2014), or optimization (Nghiem et al., 2010; Zhang and Agarwal, 2012). In practice, these numerical simulation methods are forced to reduce computational costs by coarsening the grid resolution (Kou et al., 2022) and/or simplifying the physics (Cavanagh and Ringrose, 2011), which reduces the accuracy of modeling. In recent years, machine learning (ML) approaches are emerging as a promising alternative to numerical simulation for subsurface flow problems (Zhu and Zabaras, 2018; Mo et al., 2019; M. Tang, Liu, and Durlofsky, 2020; Wen, M. Tang, and Benson, 2021; Wen, Hay, and Benson, 2021). ML methods are usually much faster than numerical simulators because the inference is very cheap once the ML models are trained. However, standard ML methods suffer from the lack of generalization and fail to provide accurate estimates away from the domain of training data. This limits the usage of ML in challenging applications such as carbon dioxide storage modeling because it requires generalization under diverse geology, reservoir conditions, and injection schemes. A recent machine-learning framework, termed neural operators (Kovachki et al., 2021; Li et al., 2020b; Li et al., 2020c), overcomes these generalization challenges by directly learning the solution operator for the PDE system family instead of just learning a single instance. By learning the solution 275 (a-b) Permeability for a dipped 3D reservoir with four injection wells; white and black lines indicate level 0 to 4’s boundary; the black dotted lines in the zoomed-in circles show the locations of injection perforation intervals. (c) Each grey block represents an FNO model; light grey arrows point to the input and output’s level; dark grey arrows show when one model’s output is used as another model’s input. Notice we feed level 0 pressure buildup output to level 1 gas saturation model because carbon dioxide plumes never migrate to level 0. (d) Pressure buildup and gas saturation at 30 years. Figure 12.1: Introduction to Nested-FNO. operator, neural operators generalize well to different conditions in the PDE system as well as discretizations without the need for re-training the ML model. Fourier neural operator (FNO) is a class of neural operators that uses Fourier transform to learn the solution operator (Li et al., 2020a) efficiently. A variant of FNO was previously used in carbon dioxide storage modeling (Wen, Li, et al., 2022) that showed 60,000 times speedup and excellent generalization. This previous model (Wen, Li, et al., 2022) is limited to a 2D spatial domain that can only represent flat reservoirs with a single injection well. However, In real-world scenarios, CCS projects often involve multiple injection wells and dipped reservoirs. 276 Due to the buoyancy effect, reservoir dipping angles can significantly influence the gaseous plume migration. These processes can only be accurately captured by high-resolution 3D spatial domains, where the costs of collecting training data from numerical solvers become prohibitive. We incorporate these practical modeling features in the current work. Here, we present an ML approach with unprecedented capabilities of high-resolution 4D spatial-temporal modeling of pressure buildup and gas saturation for realistic carbon dioxide storage projects. We integrate the FNO architecture with the LGR modeling approach to obtain a Nested Fourier Neural Operator (Nested FNO) architecture. As shown in Figure 12.1, five levels of FNOs are used to predict outputs in five levels of grid refinements; each coarser-level model’s output is used as the finer-level model’s input. This nested approach vastly reduces the computational cost needed during data collection and overcomes the memory constraints in model training. For instance, Nested FNO only needs less than 2,500 training data at the coarsest resolution (level 0) and about 6,000 samples for the finer levels (1-4). Despite the small training size, Nested FNO generalizes well to large problem dimensions with millions of cells. The spatial resolutions provided by Nested FNO are finer than most current simulations run with existing numerical models (e.g., Sleipner benchmark model (Data: Sleipner CO2 reference dataset, published via the CO2 DataShare online portal administrated by SINTEF AS n.d.) and Decatur model (Data: llinois State Geological Survey (ISGS), Illinois Basin Decatur Project (IBDP) CO2 Injection Monitoring Data, April 30, 2021. Midwest Geological Sequestration Consortium (MGSC) Phase III Data Sets. DOE Cooperative Agreement No. DE-FC26-05NT42588 n.d.)). Nested FNO provides 700,000 times speedup. The fast inference enables rigorous probabilistic assessments for maximum pressure buildup and carbon dioxide plume footprint. Such assessments are important for decisions on injection design and land lease acquisition. Running such assessments take nearly two years with traditional simulators, and it took only 2.8 seconds with the Nested FNO model. 12.2 Methods Data overview The data set is generated using a semi-adaptive LGR approach to ensure both high fidelity and computational tractability. We use the global (level 0) resolution grid in the large spatial domain to mimic typical saline storage formations with infinite boundary conditions. Next, we apply four levels of local refinements (levels 1 to 277 (a) Visualizations of gas saturation predictions at 30 years for a 3-well case. Each row shows permeability, gas saturation ground truth, prediction, and error. The white lines indicate the boundary between each level. (b) Reservoir permeability and the location of each well. (c) Testing set plume saturation error versus time for 250 random cases. The red dotted line shows the 95% prediction bands of the error. (d) Error histograms for 250 cases in the training and test set. The solid red column indicates the error for the shown example. Figure 12.2: Gas saturation prediction. 4) around each well to gradually increase the grid resolutions. Going from levels 0 to 4, we reduce the cell size by 80x on the 𝑥, 𝑦 dimensions and 10x on the 𝑧 dimension to resolve near-well plume migration, dry-out, and pressure buildup. See Supplementary, Local Grid Refinement for details on the LGR design. Nested FNO architecture Nested FNO uses a sequence of FNO models to predict the data set with multiple refinement levels. The computational domain of the Nested FNO is a 3D space with time, 𝐷 = Ω × 𝑇, where 𝑇 is the time interval of 30 years and Ω is the reservoir domain. As shown in Figure 12.1, the 3D reservoir domain consists of subdomains 𝑃 Ω𝑖 at levels 0 to 4 for each grid refinement. We use nine FNO models: G0...4 for 𝑆 pressure buildup (𝑃), and G1...4 for gas saturation (𝑆), to predict outputs at each level. We extend the original FNO (Li et al., 2020a) architecture into 4D to produce outputs in the 3D space-time domain. See Supplementary, Fourier Neural Operator for detailed architecture and parameters. The input for each model includes the permeability field, initial hydro-static pressure, 278 reservoir temperature, injection scheme, as well as spatial and temporal encoding. In carbon dioxide-water multiphase flow, pressure buildup travels significantly faster than gas saturation. Therefore, we use G0𝑃 to predict level 0 (global) pressure buildup as well as the pressure interaction between wells. We then feed level 0 pressure buildup prediction around each injection well (𝑃ˆ0 | 𝑤𝑒𝑙𝑙 𝑗 ) to models on level 1. Each subsequent model takes the input on domain Ω𝑖 together with the coarser-level prediction of 𝑃ˆ or 𝑆ˆ on Ω𝑖−1 , and outputs the prediction of 𝑃ˆ or 𝑆ˆ on Ω𝑖 . By giving the coarser-level prediction to the finer-level model as an input, we also provide the boundary conditions of the finer-level subdomain, which significantly improves the finer-level predictions. Training procedure To train the Nested FNO, we first prepare the input-output pairs for each subdomain and train each of the nine models independently. For each model, we use the ground truth, i.e., numerical simulation, pressure buildup and gas saturation on the coarselevel training domain to construct the input. This approach is time efficient because it allows us to train all models concurrently instead of sequentially going from coarser-level to finer-level models. Refer to Supplementary, Training procedure for more details. Inference procedure Once we train the nine models in the Nested FNO, we can predict the gas saturation and pressure buildup according to Algorithm 5. Notice that the number of subdomains in Ω depends on the number of injection wells. For example, a reservoir with three injection wells has 13 subdomains Ω = {Ω0 , Ω𝑙𝑒𝑣𝑒𝑙1...4,𝑤𝑒𝑙𝑙1 , Ω𝑙𝑒𝑣𝑒𝑙1...4,𝑤𝑒𝑙𝑙2 , Ω𝑙𝑒𝑣𝑒𝑙1...4,𝑤𝑒𝑙𝑙3 }. Therefore, we repeat the inference for each injection well. The input can be constructed given any random combination of reservoir condition (depth, temperature, and dip angle), injection scheme (number of wells, rate, location, perforation interval), and permeability field, as long as the variables are within the training data sampling ranges. Evaluation metrics To evaluate the gas saturation prediction accuracy in reservoirs with multiple levels of refinements, we introduce the plume saturation error 𝛿 𝑆 , defined as: 279 Algorithm 5 Predict gas saturation and pressure buildup in a reservoir with 𝑛 injection wells. G denotes the a model, 𝑃 denotes pressure buildup, 𝑆 denotes gas situation, and 𝑎 denotes input. Use G0𝑃 to predict 𝑃ˆ0 given 𝑎 0 for each well 𝑗 = 1, . . . , 𝑛 do Construct input (𝑎 1, 𝑗 , 𝑃ˆ0 | 𝑗 ) Use G1𝑃 and above input to predict 𝑃ˆ1, 𝑗 Use G1𝑆 and above input to predict 𝑆ˆ1, 𝑗 for each level 𝑖 = 2, . . . , 4 do Construct input (𝑎𝑖, 𝑗 , 𝑆ˆ𝑖−1, 𝑗 ) Use G𝑖𝑆 and above input to predict 𝑆ˆ𝑖, 𝑗 Construct input (𝑎𝑖, 𝑗 , 𝑃ˆ𝑖−1, 𝑗 ) Use G𝑖𝑃 and above input to predict 𝑃ˆ𝑖, 𝑗 end for end for 1 ∑︁ ∑︁ 𝐼𝑡,𝑖 |𝑆𝑡,𝑖 − 𝑆ˆ𝑡,𝑖 |, 𝛿𝑆 = Í 𝐼𝑡,𝑖 𝑡∈𝑇 𝑖∈Ω 𝐼𝑡,𝑖 = 1 𝑖 𝑓 (𝑆𝑡,𝑖 > 0.01) ∪ (| 𝑆ˆ𝑡,𝑖 | > 0.01). (12.1) 𝑆 is the ground truth gas saturation, 𝑆ˆ is the predicted gas saturation, 𝑇 includes all times snapshots over the 30 years, and Ω includes all the cells as in the original domain of the numerical simulator; refer Supplementary, Training procedure for more details. We use this metric because the reservoir domain includes many cells with zero gas saturation; taking an average with these zero predictions leads to an overestimation of the gas saturation accuracy. 𝛿 𝑆 is a more strict metric focusing on the error within the plume. For pressure buildup, we introduce relative error 𝛿 𝑃 : 𝛿𝑃 = 1 ∑︁ ∑︁ |𝑃𝑡,𝑖 − 𝑃ˆ𝑡,𝑖 | 𝑛Ω 𝑛𝑇 𝑡∈𝑇 𝑖∈Ω 𝑃𝑡,𝑚𝑎𝑥 (12.2) Here 𝑃 is the ground truth pressure buildup given by numerical simulation, 𝑃ˆ is the predicted pressure buildup, 𝑃𝑡,𝑚𝑎𝑥 is the maximum reservoir pressure buildup at time 𝑡, 𝑛Ω is the number of cells in Ω, and 𝑛𝑇 is the number time steps. This metric is commonly used for evaluating reservoir pressure buildup (H. Tang et al., 2021; Wen, Li, et al., 2022). 280 Each model’s separate and sequential error for (a) pressure buildup before fine-tuning, (b) pressure buildup after fine-tuning, (d) gas saturation before fine-tuning, and (e) gas saturation after fine-tuning. On the legend, ‘seq’ denotes sequential prediction, ‘sep’ denotes separate prediction. The transparent lines indicate the before fine-tune error. (e) Training and validation set 𝛿4𝑃 of fine-tuning using Option 1 to 3. (f) Principle component number and cumulative percentage of the 40 strongest rank for G3𝑃 ’s error. Figure 12.3: Fine-tuning. Fine-tuning procedure Separate vs. sequential prediction. As described in Algorithm 5, during inference, the input for each model in levels 1 to 4 includes the 𝑆ˆ or 𝑃ˆ predicted by their corresponding coarser-level model. However, during training, the inputs are constructed by ground truth numerical simulation data. The discrepancy in training and inference leads to error accumulation, especially for the models that appear later in the prediction sequence. To investigate this effect, we introduce two ways to evaluate each model: (1) separate prediction using the ground truth input taken from the numerical simulation (as in training), and (2) sequential prediction using predicted values from the coarser level as input (as in inference). Figure 12.3 a compares the average relative pressure 𝑃 for each model using both separate and sequential prediction methods. buildup 𝛿Ω 𝑗 𝑃 focuses on the ability of each model to produce outputs similar to Unlike 𝛿 𝑃 , 𝛿Ω 𝑗 the training data, defined as: 281 𝑃 𝛿Ω = 𝑗 1 ∑︁ ∑︁ |𝑃𝑡,𝑖 − 𝑃ˆ𝑡,𝑖 | . 𝑛𝑇 𝑛Ω 𝑗 𝑡∈𝑇 𝑖∈Ω 𝑃𝑡,𝑚𝑎𝑥 (12.3) 𝑗 Figure 12.3 a shows that all models have low errors and negligible overfitting 𝑃 quickly when using separate predictions. However, with sequential prediction, 𝛿Ω 𝑗 accumulates, going from coarser to finer-level models. The validation error of level 4 using sequential prediction increased by 13 times compared to separate predictions. 𝑆 for each model Similarly, for the gas saturation, the plume gas saturation error 𝛿Ω 𝑗 is defined as: 1 ∑︁ ∑︁ 𝑆 Í |𝑆𝑡,𝑖 − 𝑆ˆ𝑡,𝑖 | 𝛿Ω = 𝑗 𝐼𝑡,𝑖 𝑡∈𝑇 𝑖∈Ω 𝑗 𝐼𝑡,𝑖 = 1 𝑖 𝑓 (𝑆𝑡,𝑖 > 0.01) ∪ ( |𝑆ˆ𝑡,𝑖 | > 0.01) (12.4) 𝑆 using separate verses sequential prediction. We obFigure 12.3 d compares 𝛿Ω 𝑗 served less error accumulation for gas saturation than pressure buildup, which indicates that the prediction of gas saturation does not rely as heavily on coarser-level models. Random perturbation. To reduce the error accumulation, we explored several finetuning techniques to improve generalizability using the level 4 pressure prediction as an example. To fine-tune G4𝑃 , we add a perturbation to the ground truth input, ′ = 𝑃 + 𝜁 where 𝑖 represents a sample taken from the training set. We defined 𝑃3,𝑖 3,𝑖 𝑖 the coarser-level model’s error in the training set as 𝜖 3 = 𝑃ˆ3 − 𝑃3 , and explore three configurations of perturbation 𝜁𝑖 . • Option 1: 𝜁𝑖 = 𝜖3, 𝑗 - randomly sample an instance from 𝜖3 . • Option 2: 𝜁𝑖 = 𝜖3,𝑖 - choose the error corresponding to the specific training sample (i.e., fine-tune with the predicted label 𝑃ˆ3,𝑖 ). • Option 3: 𝜁𝑖 ∼ N (𝜇𝜖3 , 𝜎𝜖3 ) - generate a random Gaussian error using the mean and standard deviation of 𝜖3 . As shown in Figure 12.3 c, Option 1 provides the best validation set performance with the smallest overfitting. By providing a randomly sampled noise instance from 𝜖3 with each training data, we let the finer-level models become aware of the 282 presence of a structured error and learn to filter it out. Option 2 gives the best training set error but is significantly overfitted. Interestingly, Option 3 leads to the largest errors in both the training and validation set despite being a commonly used machine learning technique for introducing randomness. We applied Option 1 to G1𝑃 , G4𝑃 , G1𝑆 , and G2𝑆 . After fine-tuning, sequential prediction errors are reduced for both pressure buildup and gas saturation (Figure 12.3 b and e). For pressure buildup, we observe a dramatic improvement on level 4, where the validation error decreased by more than 50%. Structure of prediction error. We hypothesize that the FNO model’s predictions and errors lie in a perturbed low-dimensional manifold in the output function spaces due to its structure. To verify our hypothesis, we analyzed the functional principle components on error 𝜖3 (Blanchard, Bousquet, and Zwald, 2007). As shown in Figure 12.3 f, only a few principal components are needed to describe nearly a third of the error. Gaussian noised functions did not improve prediction because the Gaussian noise resides in an infinite-dimensional space, whereas the actual error only lives in a small linear sub-space. 12.3 Results We consider carbon dioxide injection into dipped 3D reservoirs through multiple wells over 30 years, as shown in Figure 12.1 a. Our data set considers a comprehensive collection of input variables for practical carbon dioxide storage projects, covering most realistic scenarios of potential sites. Each input variable is sampled according to expert knowledge, including reservoir conditions (depth, temperature, dip angle), injection schemes (number of injection wells, rates, perforation intervals), and permeability heterogeneity (mean, standard deviation, correlation lengths). See Supplementary, Data set generation for details for the input variable sampling. Prediction accuracy Gas saturation The migration of carbon dioxide plumes is governed by the complex interplay of viscous, capillary, and gravity forces. Due to the presence of the dip angle, the reservoir condition around each well can lead to various fluid properties in the same reservoir. In addition, carbon dioxide plumes can form distinctively different shapes and sizes according to different injection schemes and permeability heterogeneity. As shown in Figure 12.2 a, Nested FNO captures these complex processes with 283 (a) Global and (b) well pressure buildup predictions at 30 years. Each row shows pressure buildup ground truth, prediction, and relative error. The white lines indicate the boundary between each level. (c) Testing set pressure relative error versus time for 250 random cases. The red dotted line shows the 95% prediction bands of the error. (d) Error histograms for 250 cases in the training and test set. The solid red column indicates the error for the visualized example. Figure 12.4: Pressure buildup prediction. excellent accuracy. Near the injection perforation interval, we observe dry-out zones where the gas saturation is near one, e.g., at the bottom half of well 1. This is a highly nonlinear response caused by the vaporization of the formation fluids into the gas phase (Pruess and Müller, 2009). This important physical process is neglected by many numerical models because it can only be captured with high grid resolutions and high-fidelity simulations. Nevertheless, Nested FNO predicts the dry-out with high accuracy. We observe slightly more errors at the plumes and dry-out zones’ edges. This is because the simulation data around discontinuous saturation transitions consist of inherent numerical artifacts that are less systematic. Overall, Nested FNO displays small overfitting considering the problem’s high dimensionality (Figure 12.2 d). The average saturation error (𝛿 𝑆 ) for the gaseous carbon dioxide plume is 1.2% for the training set and 1.8% for the testing set. See Methods for 𝛿 𝑆 ’s definition. This accuracy is sufficient for practical applications such as forecasting plume footprints and estimating sweep efficiencies. 284 Pressure buildup As demonstrated in Figure 12.4 a, Nested FNO precisely captures the local pressure buildup responses around each well, as well as the global interaction among them. The high resolution provides accurate estimates of the maximum pressure buildup, which is an essential indicator of reservoir integrity. The global level prediction provides the spatial extent of the region of pressure buildup influence, another important parameter required for regulatory purposes (EPA, 2013). Additionally, Figure 12.4 a-c shows that accuracy is consistent across different prediction domains and throughout the injection period. These predictions are sufficient to guide important engineering decisions such as injection schemes and monitoring program designs. The relative pressure buildup error 𝛿 𝑃 (as defined in Methods) for the training and the testing set are 0.3% and 0.5%, respectively. Similar to the gas saturation, we observe small overfitting from the error histogram (Figure 12.4 d) for the training and testing set. This generalization is remarkable, considering the training data size for this high-dimensional problem. The generalizability is achieved through a fine-tuning technique which we will introduce in Methods. (a) Histogram of CO2 plume footprint predictions given 1,000 permeability realizations from the same geological parameters. The result satisfies a log-normal distribution; P5, P50, and P95 are marked on the distribution. (b) Ten realizations of CO2 plume at 30 years. (c) Histogram of CO2 pressure buildup predictions given the same 1,000 permeability realizations. The result satisfies a log-normal distribution; P5, P50, and P95 are marked on the distribution. (d) Ten realizations of pressure buildup at 30 years. Figure 12.5: Probabilistic assessment. 285 Computational Speed-up Once the Nested FNO model is trained, we can use it as a general-purpose numerical simulator alternative in realistic 3D reservoirs. This approach allows traditional users to skip numerical simulations for an entire class of problems and directly obtain efficient, high-fidelity, and high-resolution predictions (Wen, Hay, and Benson, 2021). Our approach differs from the task-specific “surrogate” modeling approach, which often involves a specific set of reservoir conditions. As a result, we calculate the computational speedup by comparing the Nested FNO’s prediction time to the numerical simulation run time of a state-of-the-art full-physics simulator ECLIPSE (Schlumberger, 2014). Each model in Nested FNO takes about 0.005s to infer, while the total prediction time depends on the number of injection wells. On average, the Nested FNO provides 400,000 (1-well case) to 700,000 (4-well case) times speedup compared to traditional simulation. Refer to Supplementary, Speedup analysis for detailed specifications for the simulator and ML models. Probabilistic assessment Nested FNO’s fast prediction speed enables rigorous probabilistic assessments that were previously unattainable. As an example, we investigate the maximum pressure buildup and carbon dioxide plume footprint for a four-well reservoir where each well injects at a 1MT/year rate. Refer to Supplementary, Probabilistic assessment for detailed setups. We generate 1,000 permeability realizations using the same distribution and spatial correlations, then use Nested FNO to predict gas saturation plumes and pressure buildup for each realization. As shown in Figure 12.5, the probabilistic estimates of the carbon dioxide plume footprint and maximum pressure buildup can support project developers and regulators in making important engineering decisions (Pawar et al., 2016). For example, the plume footprint helps determine the area of the land lease acquisition required and the monitoring program’s design (NETL, 2017); the maximum reservoir pressure buildup helps evaluate the safety of an injection scheme and ensure reservoir integrity. Running such assessments takes nearly two years with traditional numerical simulators, and it takes only 2.8 seconds with Nested FNO. 12.4 Discussion We present the Nested FNO model for high-resolution 4D gas saturation and pressure buildup in CO2 storage problems. The trained model provides exceptionally 286 fast predictions and can support many engineering tasks that require repetitive forward simulations, including but not limited to (1) probabilistic assessment - as demonstrated above, (2) site selection (Callas et al., 2022) - quick screening for a large number of potential reservoirs, (3) storage optimizations (Kumar and Bryant, 2008; Zhang and Agarwal, 2012) - exhaustive search in the parameter space, and (4) seismic inversion (Yin et al., 2022) - provide forward simulation outputs and gradients. Nested FNO can readily serve as a high-fidelity and high-resolution numerical simulator alternative and facilitate rigorous analyses for these tasks. A highlight of the Nested FNO is its excellent generalizability. The training sizes for Nested FNO are small (2,408 for the level 0 model and 5,916 for level 1-4 models), considering the large problem dimension with millions of cells. We achieve this generalizability through (1) a novel fine-tuning technique for nested architecture as introduced in Methods and (2) the utilization of the FNO architecture. Most existing machine learning approaches for subsurface flow use convolutional neural network (CNN)-based models with local kernels and deep architectures that are prone to overfitting (Jiang, Tahmasebi, and Mao, 2021; Wu and Qiao, 2020; Kadeethum et al., 2021; M. Tang, Liu, and Durlofsky, 2021; H. Tang et al., 2021). Unlike CNN, FNO uses global kernels to learn an infinite-dimensional input and output mapping in the function space (Li et al., 2020a). Our previous study shows that FNO-based models require only one-third of the training data compared to the CNN-based model in a 2D spatial-temporal setting (Wen, Li, et al., 2022). As a result, using FNO greatly reduces the demand for training data; combining FNO with the LGR approach makes this high-resolution 3D problem tractable. Besides data-driven approaches, another line of work, often referred to as a physicsinformed neural network, attempts to solve the governing PDE by parameterizing governing relations and initial/boundary conditions using neural networks (Raissi, Perdikaris, and Karniadakis, 2019). However, these approaches have not yet shown significant advantages in computational efficiency for subsurface multiphase flow through heterogeneous media (Fuks and Tchelepi, 2020; Almajid and Abu-Alsaud, 2021; Fraces and Hamdi Tchelepi, 2021; Haghighat, Amini, and Juanes, 2022). On the contrary, as a data-driven approach, Nested FNO provides a framework with great potential not only for CO2 storage but also for other scientific problems that involve multi-level modeling. For instance, in weather forecast modeling, different cyclones can develop locally while interacting with each other on a global level (Dong and Neumann, 1983); in nuclear fusion, the collision of multiple nuclei 287 in a particle involves long-distance interactions as well as inner-nucleus manyparticle physics (Jin et al., 2021). References Almajid, Muhammad M and Moataz O Abu-Alsaud (2021). “Prediction of porous media fluid flow using physics informed neural networks”. In: Journal of Petroleum Science and Engineering, p. 109205. André, Laurent, Yannick Peysson, and Mohamed Azaroual (2014). “Well injectivity during CO2 storage operations in deep saline aquifers–Part 2: Numerical simulations of drying, salt deposit mechanisms and role of capillary forces”. In: International Journal of Greenhouse Gas Control 22, pp. 301–312. Blanchard, Gilles, Olivier Bousquet, and Laurent Zwald (2007). “Statistical properties of kernel principal component analysis”. In: Machine Learning 66.2, pp. 259– 294. Blunt, Martin J (2017). Multiphase flow in permeable media: A pore-scale perspective. Cambridge university press. Bramble, James H et al. (1988). “A preconditioning technique for the efficient solution of problems with local grid refinement”. In: Computer Methods in Applied Mechanics and Engineering 67.2, pp. 149–159. Callas, Catherine et al. (2022). “Criteria and workflow for selecting depleted hydrocarbon reservoirs for carbon storage”. In: Applied Energy 324, p. 119668. Cavanagh, Andrew and Philip Ringrose (2011). “Simulation of CO2 distribution at the In Salah storage site using high-resolution field-scale models”. In: Energy Procedia 4, pp. 3730–3737. Chadwick, RA et al. (2004). “Geological reservoir characterization of a CO2 storage site: The Utsira Sand, Sleipner, northern North Sea”. In: Energy 29.9-10, pp. 1371–1381. Data: llinois State Geological Survey (ISGS), Illinois Basin - Decatur Project (IBDP) CO2 Injection Monitoring Data, April 30, 2021. Midwest Geological Sequestration Consortium (MGSC) Phase III Data Sets. DOE Cooperative Agreement No. DE-FC26-05NT42588 (n.d.). doi: 10 . 11582 / 2022 . 00017. url: https : / / co2datashare . org / illinois - basin - decatur - project dataset / static / IBDP % 5C % 20 - %5C % 20CO2 % 5C % 20DataShare % 5C % 20Terms%5C%20Final.pdf. Data: Sleipner CO2 reference dataset, published via the CO2 DataShare online portal administrated by SINTEF AS (n.d.). doi: 10.11582/2020.00004. url: https://co2datashare.org/. Dong, Keqin and Charles J Neumann (1983). “On the relative motion of binary tropical cyclones”. In: Monthly weather review 111.5, pp. 945–953. 288 Doughty, Christine (2010). “Investigation of CO2 plume behavior for a large-scale pilot test of geologic carbon storage in a saline formation”. In: Transport in porous media 82.1, pp. 49–76. Eigestad, Geir Terje et al. (2009). “Geological modeling and simulation of CO2 injection in the Johansen formation”. In: Computational Geosciences 13.4, pp. 435– 450. EPA (2013). Geologic Sequestration of Carbon Dioxide - Underground Injection Control (UIC) Program Class VI Well Area of Review Evaluation and Corrective Action Guidance. 816-R-13-005. Faigle, Benjamin et al. (2014). “Efficient multiphysics modelling with adaptive grid refinement using a MPFA method”. In: Computational Geosciences 18.5, pp. 625–636. Fraces, Cedric G and Hamdi Tchelepi (2021). “Physics Informed Deep Learning for Flow and Transport in Porous Media”. In: SPE Reservoir Simulation Conference. OnePetro. Fuks, O and H Tchelepi (2020). “Physics Based Deep Learning for Nonlinear TwoPhase Flow in Porous Media”. In: ECMOR XVII. Vol. 2020. European Association of Geoscientists & Engineers, pp. 1–10. Haghighat, Ehsan, Danial Amini, and Ruben Juanes (2022). “Physics-informed neural network simulation of multiphase poroelasticity using stress-split sequential training”. In: Computer Methods in Applied Mechanics and Engineering 397, p. 115141. Jiang, Zhihao, Pejman Tahmasebi, and Zhiqiang Mao (2021). “Deep Residual U-net Convolution Neural Networks with Autoregressive Strategy for Fluid Flow Predictions in Large-scale Geosystems”. In: Advances in Water Resources, p. 103878. Jin, Shi et al. (2021). “The LISE package: Solvers for static and time-dependent superfluid local density approximation equations in three dimensions”. In: Computer Physics Communications 269, p. 108130. Kadeethum, Teeratorn et al. (2021). “A framework for data-driven solution and parameter estimation of PDEs using conditional generative adversarial networks”. In: Nature Computational Science 1.12, pp. 819–829. Kamashev, Aibar and Yerlan Amanbek (2021). “Reservoir Simulation of CO2 Storage Using Compositional Flow Model for Geological Formations in Frio Field and Precaspian Basin”. In: Energies 14.23, p. 8023. Kou, Zuhao et al. (2022). “Method for upscaling of CO2 migration in 3D heterogeneous geological models”. In: Journal of Hydrology, p. 128361. Kovachki, Nikola et al. (2021). “Neural operator: Learning maps between function spaces”. In: arXiv preprint arXiv:2108.08481. 289 Kumar, Navanit and Steven Bryant (2008). “Optimizing injection intervals in vertical and horizontal wells for CO2 sequestration”. In: SPE Annual Technical Conference and Exhibition. OnePetro. Li, Zongyi et al. (2020a). “Fourier neural operator for parametric partial differential equations”. In: arXiv preprint arXiv:2010.08895. – (2020b). “Multipole graph neural operator for parametric partial differential equations”. In: arXiv preprint arXiv:2006.09535. – (2020c). “Neural operator: Graph kernel network for partial differential equations”. In: arXiv preprint arXiv:2003.03485. Mo, Shaoxing et al. (2019). “Deep convolutional encoder-decoder networks for uncertainty quantification of dynamic multiphase flow in heterogeneous media”. In: Water Resources Research 55.1, pp. 703–728. NAS (2018). Negative Emissions Technologies and Reliable Sequestration. National Academies Press. isbn: 9780309484527. doi: 10.17226/25259. NETL (2017). Best Practices: Risk Management and Simulation for Geologic Storage Projects. url: https://www.netl.doe.gov/sites/default/files/ 2018-10/BPM%7B%5C_%7DRiskAnalysisSimulation.pdf. Nghiem, Long et al. (2010). “Simulation and optimization of trapping processes for CO2 storage in saline aquifers”. In: Journal of Canadian Petroleum Technology 49.08, pp. 15–22. Pawar, Rajesh J et al. (2016). “The National Risk Assessment Partnership’s integrated assessment model for carbon storage: A tool to support decision making amidst uncertainty”. In: International Journal of Greenhouse Gas Control 52, pp. 175–189. Pruess, Karsten and Julio Garcia (2002). “Multiphase flow dynamics during CO2 disposal into saline aquifers”. In: Environmental Geology 42.2, pp. 282–295. Pruess, Karsten and Nadja Müller (2009). “Formation dry-out from CO2 injection into saline aquifers: 1. Effects of solids precipitation and their mitigation”. In: Water Resources Research 45.3. Pruess, Karsten, Curtis M Oldenburg, and GJ Moridis (1999). TOUGH2 user’s guide version 2. Tech. rep. Lawrence Berkeley National Lab.(LBNL), Berkeley, CA (United States). Raissi, Maziar, Paris Perdikaris, and George E Karniadakis (2019). “Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations”. In: Journal of Computational Physics 378, pp. 686–707. Schlumberger (2014). ECLIPSE reservoir simulation software Reference Manual. 290 Strandli, Christin W., Edward Mehnert, and Sally M. Benson (2014). “CO2 plume tracking and history matching using multilevel pressure monitoring at the Illinois basin - Decatur project”. In: Energy Procedia 63, pp. 4473–4484. issn: 18766102. doi: 10.1016/j.egypro.2014.11.483. Tang, Hewei et al. (2021). “A deep learning-accelerated data assimilation and forecasting workflow for commercial-scale geologic carbon storage”. In: International Journal of Greenhouse Gas Control 112, p. 103488. Tang, Meng, Yimin Liu, and Louis J Durlofsky (2020). “A deep-learning-based surrogate model for data assimilation in dynamic subsurface flow problems”. In: Journal of Computational Physics 413, p. 109456. – (2021). “Deep-learning-based surrogate flow modeling and geological parameterization for data assimilation in 3D subsurface flow”. In: Computer Methods in Applied Mechanics and Engineering 376, p. 113636. Wen, Gege and Sally M. Benson (2019). “CO2 plume migration and dissolution in layered reservoirs”. In: International Journal of Greenhouse Gas Control 87.May, pp. 66–79. issn: 17505836. doi: 10.1016/j.ijggc.2019.05.012. Wen, Gege, Catherine Hay, and Sally M. Benson (2021). “CCSNet: a deep learning modeling suite for CO2 storage”. In: Advances in Water Resources, p. 104009. issn: 0309-1708. doi: https://doi.org/10.1016/j.advwatres.2021. 104009. Wen, Gege, Zongyi Li, et al. (2022). “U-FNO—An enhanced Fourier neural operatorbased deep-learning model for multiphase flow”. In: Advances in Water Resources 163, p. 104180. Wen, Gege, Meng Tang, and Sally M. Benson (2021). “Towards a predictor for CO2 plume migration using deep neural networks”. In: International Journal of Greenhouse Gas Control 105, p. 103223. issn: 1750-5836. doi: 10.1016/j. ijggc.2020.103223. Wu, Haiyi and Rui Qiao (2020). “Physics-Constrained Deep Learning for Data Assimilation of Subsurface Transport”. In: Energy and AI 3, p. 100044. Yin, Ziyi et al. (2022). “Learned coupled inversion for carbon sequestration monitoring and forecasting with Fourier neural operators”. In: arXiv preprint arXiv:2203.14396. Zhang, Zheming and Ramesh K Agarwal (2012). “Numerical simulation and optimization of CO2 sequestration in saline aquifers for vertical and horizontal well injection”. In: Computational Geosciences 16.4, pp. 891–899. Zhu, Yinhao and Nicholas Zabaras (2018). “Bayesian deep convolutional encoder– decoder networks for surrogate modeling and uncertainty quantification”. In: Journal of Computational Physics 366, pp. 415–447. 291