SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zongyi Li Nikola Borislavov Kovachki Chris Choy Boyi Li Jean Kossaifi
Shourya Prakash Otta Mohammad Amin Nabian Maximilian Stadler
Christian Hundt Kamyar Azizzadenesheli Anima Anandkumar
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 point-cloud 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 discretization-convergent, 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 \times$  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.
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 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 [1].
To overcome these computational challenges, recent works propose deep learning-based methods, particularly neural operators [2], to speed up the simulation and inverse design. Neural operators
Figure 1: The architecture of GINO. 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.
generalize neural networks and learn operators, which are mappings between infinite-dimensional function spaces [2]. Neural operators are discretization convergent and can approximate general operators [3]. 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) [3] with numerous applications in scientific computing, including weather forecasting [4], carbon dioxide storage and reservoir engineering [5], 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) [2], and Fourier neural operators (FNO) [6] have been popular in various applications.
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) [7], which is also used in scientific computing [8-10]. 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 5. 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 [11, 12] 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) [13], 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 [14] 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 shown in Table 1, which we discussed in Section 5.
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
Table 1: Computational complexity of standard deep learning models.  $N$  is the number of mesh points;  $d$  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.
Model	Range	Complexity	Irregular grid	Discretization convergent
GNN	local	O(Ndegree)	✓	✗
CNN	local	O(N)	✗	✗
UNet	global	O(N)	✗	✗
Transformer	global	O(N2)	✓	✓
GNO (kernel)	radius r	O(Ndegree)	✓	✓
FNO (FFT)	global	O(N log N)	✗	✓
GINO [Ours]	global	O(N log N + Ndegree)	✓	✓
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, (i) Geometry encoder: multiple local kernel integration layers through GNO with graph operations, (ii) Global model: a sequence of FNO layers for global kernel integration, and (iii) Geometry decoder: the final kernel integral layers, as shown in Figure 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 GNO-based 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 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 [15], composed of  $500+$  car geometries with  $O(10^{5})$  mesh points on the surface and  $O(10^{7})$  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 [16]. 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 \times$  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.
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  $D \subset \mathbb{R}^d$  be a Lipschitz domain and  $\mathcal{A}$  a Banach space of real-valued functions on  $D$ . We consider the set of distance functions  $\mathcal{T} \subset \mathcal{A}$  so that, for each function  $T \in \mathcal{T}$ , its zero set  $S_T = \{x \in D : T(x) = 0\}$  defines a  $(d - 1)$ -dimensional sub-manifold. We assume  $S_T$  is simply connected, closed, smooth, and that there exists  $\epsilon > 0$  such that  $B_{\epsilon}(x) \cap \partial D = \emptyset$  for every  $x \in S_T$  and  $T \in \mathcal{T}$ . We denote by  $Q_T \subset D$ , the open volume enclosed by the sub-manifold  $S_T$  and assume that  $Q_T$  is a Lipschitz domain with  $\partial Q_T = S_T$ . We define the Lipschitz domain  $\Omega_T := D \setminus \bar{Q}_T$  so that,  $\partial \Omega_T = \partial D \cup S_T$ . Let  $\mathcal{L}$  denote a partial differential operator and consider the problem
$$
\mathcal {L} (u) = f, \quad \text {i n} \Omega_ {T}, \tag {1}
$$
$$
u = g, \qquad \text {i n} \partial \Omega_ {T},
$$
for some  $f \in \mathcal{F}, g \in \mathcal{B}$  where  $\mathcal{B}, \mathcal{F}$  denote Banach spaces of functions on  $\mathbb{R}^d$  with the assumption that the evaluation functional is continuous in  $\mathcal{B}$ . We assume that  $\mathcal{L}$  is such that, for any triplet  $(T, f, g)$ , the PDE (1) has a unique solution  $u \in \mathcal{U}_T$  where  $\mathcal{U}_T$  denotes a Banach space of functions on  $\Omega_T$ . Let  $\mathcal{U}$  denote a Banach space of functions on  $D$  and let  $\{E_T : \mathcal{U}_T \to \mathcal{U} : T \in \mathcal{T}\}$  be a family of extension operators which are linear and bounded. We define the mapping from the distance function to the solution function
$$
\Psi : \mathcal {T} \times \mathcal {F} \times \mathcal {B} \rightarrow \mathcal {U} \tag {2}
$$
by  $(T,f,g)\mapsto E_T(u)$  which is our operator of interest.
Navier-Stokes Equation. We illustrate the above abstract formulation with the following example. Let  $D = (0,1)^d$  be the unit cube and let  $\mathcal{A} = C(\bar{D})$ . We take  $\mathcal{T} \subset \mathcal{A}$  to be some subset such that the zero level set of every element defines a  $(d - 1)$ -dimensional closed surface which can be realized as the graph of a Lipschitz function and that there exists  $\epsilon > 0$  such that each surface is at least distance  $\epsilon$  away from the boundary of  $D$ . We now consider the steady Naiver-Stokes equations,
$$
- \nu \Delta v + (v \cdot \nabla) v + \nabla p = f, \quad \text {i n} \Omega_ {T},
$$
$$
\nabla \cdot v = 0, \quad \text {i n} \Omega_ {T}, \tag {3}
$$
$$
v = q, \qquad \text {i n} \partial D,
$$
$$
v = 0, \qquad \text {i n} S _ {T},
$$
where  $v: \Omega_T \to \mathbb{R}^d$  is the velocity,  $p: \Omega_T \to \mathbb{R}$  is the pressure,  $\nu$  is the viscosity, and  $f, q: \mathbb{R}^d \to \mathbb{R}^d$  are the forcing and boundary functions. The condition that  $v = 0$  in  $S_T$  is commonly known as a "no slip" boundary and is prevalent in many engineering applications. The function  $q$ , on the other hand, defines the inlet and outlet boundary conditions for the flow. We assume that  $f \in H^{-1}(\mathbb{R}^d; \mathbb{R}^d)$  and  $q \in C(\mathbb{R}^d; \mathbb{R}^d)$ . We can then define our boundary function  $g \in C(\mathbb{R}^d; \mathbb{R}^d)$  such that  $g(x) = 0$  for any  $x \in D$  with  $\mathrm{dist}(x, \partial D) \geq \epsilon$  and  $g(x) = q(x)$  for any  $x \in D$  with  $\mathrm{dist}(x, \partial D) > \epsilon / 2$  as well as any  $x \notin D$ . Continuity of  $g$  can be ensured by an appropriate extension for any  $x \in D$  such that  $\mathrm{dist}(x, \partial D) < \epsilon$  and  $\mathrm{dist}(x, \partial D) \geq \epsilon / 2$  [17]. We define  $u: \Omega_T \to \mathbb{R}^{d+1}$  by  $u = (v, p)$  as the unique weak solution of (3) with  $\mathcal{U}_T = H^1(\Omega_T; \mathbb{R}^d) \times L^2(\Omega_T)/\mathbb{R}$  [18]. We define  $\mathcal{U} = H^1(D; \mathbb{R}^d) \times L^2(D)/\mathbb{R}$  and the family of extension operators  $\{E_T: \mathcal{U}_T \to \mathcal{U}\}$  by  $E_T(u) = (E_T^v(v), E_T^p(p))$  where  $E_T^v: H^1(\Omega_T; \mathbb{R}^d) \to H^1(D; \mathbb{R}^d)$  and  $E_T^p: L^2(\Omega_T)/\mathbb{R} \to L^2(D)/\mathbb{R}$  are defined as the restriction onto  $D$  of the extension operators defined in [19, Chapter VI, Theorem 5]. This establishes the existence of the operator  $\Psi: T \times H^{-1}(\mathbb{R}^d; \mathbb{R}^d) \times C(\mathbb{R}^d; \mathbb{R}^d) \to H^1(D; \mathbb{R}^d) \times L^2(D)/\mathbb{R}$  mapping the geometry, forcing, and boundary condition to the (extended) solution of the steady Navier-Stokes equation (3). Homomorphic extensions of deformation-based operators have been shown in [20]. We leave for future work studying the regularity properties of the presently defined operator.
3 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,  $(i)$  multiple local kernel integration layers,  $(ii)$  a sequence of FNO layers for global kernel integration which precedes  $(iii)$  the final kernel integral layers. Each layer of GINO follows the form of generic kernel integral of the form (5). Local integration is computed using graphs, while global integration is done in Fourier space.
3.1 Neural operator
A neural operator  $\Psi$  [3] maps the input functions  $a = (T, f, g)$  to the solution function  $u$ . The neural operator  $\Psi$  is composed of multiple layers of point-wise and integral operators,
$$
\Psi = \mathcal {Q} \circ \mathcal {K} _ {L} \circ \dots \circ \mathcal {K} _ {1} \circ \mathcal {P}. \tag {4}
$$
The first layer  $\mathcal{P}$  is a pointwise operator parameterized by a neural network. It transforms the input function  $a$  into a higher-dimensional latent space  $\mathcal{P}:a\mapsto v_0$ . Similarly, the last layer acts as a projection layer, which is a pointwise operator  $\mathcal{Q}:v_l\mapsto u$ , parameterized by a neural network  $Q$ . The model consists of  $L$  layers of integral operators  $\mathcal{K}_l:v_{l - 1}\mapsto v_l$  in between.
$$
v _ {l} (x) = \int_ {D} \kappa_ {l} (x, y) v _ {l - 1} (y) d y \tag {5}
$$
where  $\kappa_{l}$  is a learnable kernel function. Non-linear activation functions are incorporated between each layer.
3.2 Graph operator block
To efficiently compute the integral in equation (5), we truncate the integral to a local ball at  $x$  with radius  $r > 0$ , as done in [2],
$$
v _ {l} (x) = \int_ {B _ {r} (x)} \kappa (x, y) v _ {l - 1} (y) d y. \tag {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  $x \in D$ , we randomly sample points  $\{y_1, \ldots, y_M\} \subset B_r(x)$  and approximate equation (6) as
$$
v _ {l} (x) \approx \sum_ {i = 1} ^ {M} \kappa \left(x, y _ {i}\right) v _ {l - 1} \left(y _ {i}\right) \mu \left(y _ {i}\right), \tag {7}
$$
where  $\mu$  denotes the Riemannian sum weights corresponding to the ambient space of  $B_r(x)$ . For a fixed input mesh of  $N$  points, the computational cost of equation (7) scales with the number of edges, denoted as  $O(E) = O(MN)$ . Here, the number of sampling points  $M$  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  $\{x_1^{\mathrm{in}},\ldots ,x_N^{\mathrm{in}}\} \subset S_T$ , we employ a GNO-encoder to transform it to a function on a uniform latent grid  $\{x_{1}^{\mathrm{grid}},\dots,x_{S}^{\mathrm{grid}}\} \subset D$ . The encoder is computed as discretization of an integral operator  $v_{0}(x^{\mathrm{grid}})\approx \sum_{i = 1}^{M}\kappa (x^{\mathrm{grid}},y_{i}^{\mathrm{in}})\mu (y_{i}^{\mathrm{in}})$  over ball  $B_{r_{\mathrm{in}}}(\boldsymbol{x}^{\mathrm{grid}})$ . To inform the grid density, GINO computes Riemannian sum weights  $\mu (y_i^{\mathrm{in}})$ . Further, we use Fourier features in the kernel [21]. For simple geometries, this encoder can be omitted, see Section 4.
Decoder. Similarly, given a function defined on the uniform latent grid  $\{x_1^{\mathrm{grid}},\ldots ,x_S^{\mathrm{grid}}\} \subset D$  we use a GNO-decoder to query arbitrary output points  $\{x_{1}^{\mathrm{out}},\dots,x_{N}^{\mathrm{out}}\} \subset \Omega_{T}$ . The output is evaluated as  $u(x^{\mathrm{out}})\approx \sum_{i = 1}^{M}\kappa (x^{\mathrm{out}},y_i^{\mathrm{grid}})v_l(y_i^{\mathrm{grid}})\mu (y_i^{\mathrm{grid}})$  over ball  $B_{r_{\mathrm{out}}}(\boldsymbol{x}^{\mathrm{out}})$ . Here, the Riemannian weight,  $\mu (y_i^{\mathrm{grid}}) = 1 / S$  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  $O(N^2)$  computation and memory. However, as the  $N$  gets larger, e.g., 10 ~ 100 million, computation and memory become prohibitive even on modern GPUs. Instead, we use a hash grid-based implementation to efficiently prune candidates that are outside of a  $\ell^\infty$ -ball first and then compute the  $\ell^2$  distance between only the candidates that survive. This reduces the computational complexity to  $O(Ndr^3)$  where  $d$  denotes unit density and  $r$  is the radius. This can be efficiently done using first creating a hash table of voxels with size  $r$ . 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 in these neighboring voxels. Specifically, we use the CUDA implementation from Open3D [22]. Then, using the neighbors, we compute the kernel integration using gather-scatter operations from torch-scatter [23]. Further, if the degree of the graph gets larger, we can add Nyström approximation by sampling nodes [2].
Figure 2: 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.
3.3 Fourier operator block
The geometry encoding  $v_{0}$  and the geometry specifying map  $T$ , both evaluated on a regular grid discretizing  $D$  are passed to a FNO block. We describe the basic FNO block as first outlined in [6]. 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  $d$ -dimensional unit torus  $\mathbb{T}^d$ . We define an integral operator with kernel  $\kappa \in L^{2}(\mathbb{T}^{d};\mathbb{R}^{n\times m})$  as the mapping  $\mathcal{C}:L^{2}(\mathbb{T}^{d};\mathbb{R}^{m})\to L^{2}(\mathbb{T}^{d};\mathbb{R}^{n})$  given by
$$
\mathcal {C} (v) = \mathcal {F} ^ {- 1} \big (\mathcal {F} (\kappa) \cdot \mathcal {F} (v) \big), \quad \forall v \in L ^ {2} \left(\mathbb {T} ^ {d}; \mathbb {R} ^ {m}\right)
$$
Here  $\mathcal{F},\mathcal{F}^{-1}$  are the Fourier transform and its inverse respectively, defined for  $L^2$  by the appropriate limiting procedure. The Fourier transform of the function  $\kappa$  will be parameterized directly by some fixed number of Fourier modes, denoted  $\alpha \in \mathbb{N}$ . In particular, we assume
$$
\kappa (x) = \sum_ {\gamma \in I} c _ {\gamma} \mathrm {e} ^ {i \langle \gamma , x \rangle}, \quad \forall x \in \mathbb {T} ^ {d}
$$
for some index set  $I \subset \mathbb{Z}^d$  with  $|I| = \alpha$  and coefficients  $c_{\gamma} \in \mathbb{C}^{n \times m}$ . Then we may view  $\mathcal{F}: L^2(\mathbb{T}^d; \mathbb{R}^{n \times m}) \to \ell^2(\mathbb{Z}^d; \mathbb{C}^{n \times m})$  so that  $\mathcal{F}(\kappa)(\gamma) = c_{\gamma}$  if  $\gamma \in I$  and  $\mathcal{F}(\kappa)(\gamma) = 0$  if  $\gamma \notin I$ . We directly learn the coefficients  $c_{\gamma}$  without ever having to evaluate  $\kappa$  in physical space. We then define the full operator block  $\mathcal{K}: L^2(\mathbb{T}^d; \mathbb{R}^m) \to L^2(\mathbb{T}^d; \mathbb{R}^n)$  by