50.3 Neural Operators for PDEs

1. From Functions to Operators

What if you could hand a neural network an entire pressure field, temperature map, or velocity profile and get back the solution to a completely new configuration in milliseconds, without ever re-solving the underlying equations? That is exactly what a neural operator does: instead of learning a mapping between finite-dimensional vectors (\(f: \mathbb{R}^n \to \mathbb{R}^m\)), it learns a mapping between function spaces, \(\mathcal{G}: \mathcal{A} \to \mathcal{U}\), where \(\mathcal{A}\) and \(\mathcal{U}\) are spaces of functions defined on some domain \(D \subset \mathbb{R}^d\).

A neural operator takes an entire function as input and produces an entire function as output. Unlike conventional networks that map fixed-size vectors, it operates on function spaces. Most physics and engineering problems require solving the same equation repeatedly for different configurations. A neural operator amortizes this cost by learning the solution map itself: each new configuration requires only a single forward pass (milliseconds) instead of a full numerical solve (seconds to hours). The core mechanism parameterizes the integral kernel of the solution operator with a neural network. Input-output function pairs in the training data teach the network to reproduce the operator that the PDE defines. Use a neural operator instead of a classical solver when you need thousands of evaluations across varying inputs (optimization loops, uncertainty quantification, real-time control); use a classical solver when you need a single high-fidelity answer or when training data is unavailable.

Consider the steady-state Darcy flow equation, which models groundwater flow through porous media:

Here \(a(x)\) is the permeability field (how easily fluid flows through the material at each point), \(u(x)\) is the pressure field (the solution), and \(f(x)\) is the source term. The operator we want to learn maps the permeability field to the pressure field: \(\mathcal{G}: a(\cdot) \mapsto u(\cdot)\). This is an operator because both the input \(a\) and the output \(u\) are functions of spatial position, not finite-dimensional vectors.

The key advantage is zero-shot generalization: once \(\mathcal{G}\) is learned, it maps any permeability field to its corresponding pressure solution, including fields never seen during training. A classical solver must re-solve the PDE from scratch for each new \(a(x)\). A neural operator evaluation takes milliseconds. In short: a neural operator learns the solution map once so that every future query costs a forward pass, not a fresh solve.

2. The Universal Approximation Theorem for Operators

Chen and Chen (1995) proved the universal approximation theorem for operators, the theoretical backbone of neural operator methods. Lu et al. (2021) later extended the result. The theorem guarantees that for any continuous nonlinear operator \(\mathcal{G}: \mathcal{A} \to \mathcal{U}\) and any \(\epsilon > 0\), a neural network architecture exists that approximates \(\mathcal{G}\) to within \(\epsilon\) in the appropriate norm.

More precisely, let \(\mathcal{A}\) be a compact set in a Banach space (a complete normed vector space, generalizing Euclidean space to infinite dimensions) and \(\mathcal{G}: \mathcal{A} \to C(K)\) be a continuous operator mapping into continuous functions on a compact set \(K\). Then \(\mathcal{G}\) can be approximated by a network of the form:

This is the DeepONet architecture: a branch network that processes the input function \(a\) (evaluated at sensor locations \(x_1, \dots, x_m\)) and a trunk network that processes the evaluation point \(y\). The output is the inner product of the branch and trunk outputs.

DeepONet is theoretically elegant, but for PDE problems on regular grids, the FNO achieves better accuracy with simpler training. The FNO is therefore the primary architecture below; DeepONet returns in Section 6.

3. Fourier Neural Operator Architecture

The key observation behind the FNO is that many PDE operators are naturally expressed in Fourier space. Convolution in physical space becomes pointwise multiplication in frequency space. The FNO exploits this by learning the operator kernel directly in the Fourier domain.

An FNO layer performs four operations, illustrated in Figure 50.3.1 below:

1. Fourier transform: apply the Fast Fourier Transform (FFT) to the input function, mapping it from physical space to frequency space.
2. Spectral filtering: multiply the Fourier coefficients by a learnable weight tensor \(R_\phi\), keeping only the lowest \(k_{\max}\) modes. This is the learned kernel in Fourier space.
3. Inverse Fourier transform: apply the inverse FFT to return to physical space.
4. Local bypass: add a pointwise linear transformation \(Wv(x)\) (a 1x1 convolution) that captures local, high-frequency interactions missed by the truncated Fourier representation.

Mathematically, one FNO layer maps input \(v_t\) to output \(v_{t+1}\):

where \(\mathcal{F}\) denotes the FFT, \(R_\phi\) is the learnable spectral weight tensor, \(W\) is a pointwise weight matrix, and \(\sigma\) is a nonlinear activation (typically Gaussian Error Linear Unit, or GELU). Stacking \(L\) such layers (typically \(L = 4\)) with a lifting layer (a pointwise linear map that projects the low-dimensional input, such as coordinates plus the coefficient field, into a higher-dimensional channel space) at the input and a projection layer (a pointwise linear map that collapses the high-dimensional channel representation back to the scalar solution field) at the output gives the complete FNO architecture.

import torch
import torch.nn as nn
import torch.fft

class SpectralConv2d(nn.Module):
    """Spectral convolution layer for 2D FNO."""

    def __init__(self, in_channels, out_channels, modes1, modes2):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.modes1 = modes1  # number of Fourier modes (x)
        self.modes2 = modes2  # number of Fourier modes (y)

        # Learnable weights in Fourier space (complex-valued)
        scale = 1.0 / (in_channels * out_channels)
        self.weights1 = nn.Parameter(
            scale * torch.rand(in_channels, out_channels,
                               modes1, modes2, dtype=torch.cfloat)
        )
        self.weights2 = nn.Parameter(
            scale * torch.rand(in_channels, out_channels,
                               modes1, modes2, dtype=torch.cfloat)
        )

    def compl_mul2d(self, input, weights):
        """Complex multiplication in Fourier space."""
        # input: (batch, in_ch, x, y)
        # weights: (in_ch, out_ch, x, y)
        return torch.einsum("bixy,ioxy->boxy", input, weights)

    def forward(self, x):
        batch_size = x.shape[0]

        # FFT along spatial dimensions
        x_ft = torch.fft.rfft2(x)

        # Multiply relevant Fourier modes with learnable weights
        out_ft = torch.zeros(
            batch_size, self.out_channels,
            x.size(-2), x.size(-1) // 2 + 1,
            dtype=torch.cfloat, device=x.device,
        )
        # Low-frequency modes (positive frequencies)
        out_ft[:, :, :self.modes1, :self.modes2] = \
            self.compl_mul2d(
                x_ft[:, :, :self.modes1, :self.modes2],
                self.weights1,
            )
        # Low-frequency modes (negative frequencies)
        out_ft[:, :, -self.modes1:, :self.modes2] = \
            self.compl_mul2d(
                x_ft[:, :, -self.modes1:, :self.modes2],
                self.weights2,
            )

        # Inverse FFT
        x = torch.fft.irfft2(out_ft, s=(x.size(-2), x.size(-1)))
        return x


class FNO2d(nn.Module):
    """
    Fourier Neural Operator for 2D PDEs.

    Architecture:
        Lift -> [SpectralConv + Linear + Activation] x L -> Project
    """

    def __init__(self, modes1, modes2, width, n_layers=4):
        super().__init__()
        self.modes1 = modes1
        self.modes2 = modes2
        self.width = width
        self.n_layers = n_layers

        # Lifting layer: input channels -> width
        self.lift = nn.Linear(3, width)  # (x, y, a(x,y)) -> width

        # Spectral convolution layers
        self.spectral_convs = nn.ModuleList([
            SpectralConv2d(width, width, modes1, modes2)
            for _ in range(n_layers)
        ])

        # Local linear bypass (1x1 convolution)
        self.linear_layers = nn.ModuleList([
            nn.Conv2d(width, width, 1)
            for _ in range(n_layers)
        ])

        # Projection layer: width -> 1 (solution field)
        self.project = nn.Sequential(
            nn.Linear(width, 128),
            nn.GELU(),
            nn.Linear(128, 1),
        )

    def forward(self, x):
        # x shape: (batch, nx, ny, 3) = (batch, nx, ny, [x, y, a])
        batch, nx, ny, _ = x.shape

        # Lift to higher dimension
        h = self.lift(x)                    # (batch, nx, ny, width)
        h = h.permute(0, 3, 1, 2)          # (batch, width, nx, ny)

        # FNO layers
        for spectral, linear in zip(
            self.spectral_convs, self.linear_layers
        ):
            h1 = spectral(h)                # global Fourier path
            h2 = linear(h)                  # local bypass path
            h = torch.nn.functional.gelu(h1 + h2)

        # Project back to solution space
        h = h.permute(0, 2, 3, 1)          # (batch, nx, ny, width)
        h = self.project(h)                 # (batch, nx, ny, 1)

        return h.squeeze(-1)               # (batch, nx, ny)

4. Training on Darcy Flow

The Darcy flow benchmark is the standard test problem for neural operators. The input is a permeability field \(a(x, y)\) (typically a piecewise constant "checkerboard" sampled from a Gaussian random field, where the value at each spatial point is drawn from a multivariate Gaussian distribution with a prescribed spatial correlation structure), and the target is the pressure solution \(u(x, y)\) computed by a high-fidelity finite element solver. The training set contains 1,000 input-output pairs on a 64x64 grid; the test set contains 200 unseen pairs.

import torch
from torch.utils.data import DataLoader, TensorDataset

def load_darcy_flow(data_path, n_train=1000, n_test=200, resolution=64):
    """
    Load the Darcy flow benchmark dataset.

    Each sample contains:
      - a(x,y): permeability coefficient field (64x64)
      - u(x,y): pressure solution field (64x64)
    """
    data = torch.load(data_path)
    a_train = data['a'][:n_train, :resolution, :resolution]
    u_train = data['u'][:n_train, :resolution, :resolution]
    a_test = data['a'][n_train:n_train+n_test, :resolution, :resolution]
    u_test = data['u'][n_train:n_train+n_test, :resolution, :resolution]

    # Create coordinate grid
    grid_x = torch.linspace(0, 1, resolution)
    grid_y = torch.linspace(0, 1, resolution)
    grid_x, grid_y = torch.meshgrid(grid_x, grid_y, indexing='ij')

    # Stack inputs: (x, y, a(x,y))
    def make_input(a_batch):
        batch_size = a_batch.shape[0]
        gx = grid_x.unsqueeze(0).expand(batch_size, -1, -1)
        gy = grid_y.unsqueeze(0).expand(batch_size, -1, -1)
        return torch.stack([gx, gy, a_batch], dim=-1)

    x_train = make_input(a_train)
    x_test = make_input(a_test)

    return (
        DataLoader(TensorDataset(x_train, u_train),
                   batch_size=32, shuffle=True),
        DataLoader(TensorDataset(x_test, u_test),
                   batch_size=32, shuffle=False),
    )


def train_fno(model, train_loader, test_loader, n_epochs=100, lr=1e-3):
    """Train FNO with relative L2 loss (the ratio of the prediction
    error norm to the true solution norm, so 0.01 means 1% error
    regardless of solution scale)."""
    optimizer = torch.optim.Adam(model.parameters(), lr=lr,
                                  weight_decay=1e-4)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer, T_max=n_epochs
    )

    for epoch in range(n_epochs):
        model.train()
        train_loss = 0.0

        for x_batch, u_batch in train_loader:
            optimizer.zero_grad()
            u_pred = model(x_batch)

            # Relative L2 loss: ||u_pred - u_true|| / ||u_true||
            loss = torch.mean(
                torch.norm(u_pred - u_batch, dim=(-2, -1))
                / torch.norm(u_batch, dim=(-2, -1))
            )
            loss.backward()
            optimizer.step()
            train_loss += loss.item()

        scheduler.step()

        # Evaluate
        if (epoch + 1) % 10 == 0:
            model.eval()
            test_loss = 0.0
            with torch.no_grad():
                for x_batch, u_batch in test_loader:
                    u_pred = model(x_batch)
                    test_loss += torch.mean(
                        torch.norm(u_pred - u_batch, dim=(-2, -1))
                        / torch.norm(u_batch, dim=(-2, -1))
                    ).item()

            n_train_batches = len(train_loader)
            n_test_batches = len(test_loader)
            print(
                f"Epoch {epoch+1:3d} | "
                f"Train: {train_loss/n_train_batches:.4f} | "
                f"Test: {test_loss/n_test_batches:.4f}"
            )

    return model

# Instantiate and train
model = FNO2d(modes1=12, modes2=12, width=32, n_layers=4)
print(f"FNO parameters: {sum(p.numel() for p in model.parameters()):,}")
# Typically ~2.4M parameters for this configuration

On the Darcy flow benchmark, an FNO with 12 Fourier modes and 4 layers achieves a relative L2 test error (the prediction error norm divided by the true solution norm, so that 0.01 means 1% error regardless of solution scale) typically below 1.5% after 100 epochs of training. Each forward pass takes approximately 3 milliseconds on a graphics processing unit (GPU), compared to roughly 10 seconds for the finite element solver, approximately a 3000x speedup. That speed gap is the value proposition of neural operators for engineering applications: once trained, they enable real-time evaluation of PDE solutions across the design space.

5. Navier-Stokes Surrogate

Darcy flow is a steady-state, linear PDE, so the 3000x speedup above is impressive but not yet a stress test; real engineering problems demand that neural operators handle time-dependent, nonlinear dynamics as well.

The more challenging and scientifically significant benchmark is the 2D Navier-Stokes equation for incompressible viscous flow:

where \(w = \nabla \times u\) is the vorticity (the curl of the velocity field, measuring local rotation in the fluid), \(\nu\) is the viscosity, and \(f\) is an external forcing. The operator learning task maps the vorticity at time steps \([0, T_{\text{in}}]\) to future vorticity at times \([T_{\text{in}}+1, T_{\text{out}}]\).

For this time-dependent problem, the FNO processes a stack of input time steps as channels and predicts a stack of output time steps. At viscosity \(\nu = 10^{-3}\), the flow is mildly turbulent, and the FNO typically achieves relative errors of approximately 5% for 10-step-ahead prediction. At \(\nu = 10^{-4}\) (more turbulent), errors typically rise to approximately 15%, reflecting the increasing complexity of the solution manifold.

6. DeepONet: Branch-Trunk Architecture

While FNO excels on regular grids, DeepONet handles irregular sensor locations and heterogeneous input representations. The architecture consists of two sub-networks:

Branch and Trunk Sub-Networks

The branch network processes the input function \(a\), evaluated at fixed sensor locations \((x_1, \dots, x_m)\). It takes the vector \((a(x_1), \dots, a(x_m))\) and produces a feature vector \(b \in \mathbb{R}^p\).

The trunk network processes the query location \(y\) where we want to evaluate the output function. It takes \(y\) and produces a feature vector \(\tau \in \mathbb{R}^p\).

The operator output is the inner product: \(\mathcal{G}(a)(y) \approx b \cdot \tau\).

import torch
import torch.nn as nn

class DeepONet(nn.Module):
    """
    DeepONet: branch-trunk architecture for operator learning.

    Branch: processes input function at sensor locations.
    Trunk: processes query evaluation point.
    Output: dot product of branch and trunk features.
    """

    def __init__(self, n_sensors, branch_width, trunk_width, p):
        """
        Args:
            n_sensors: number of sensor locations for input function
            branch_width: hidden layer width for branch network
            trunk_width: hidden layer width for trunk network
            p: dimension of the feature vectors (branch/trunk output)
        """
        super().__init__()

        # Branch network: sensor values -> feature vector
        self.branch = nn.Sequential(
            nn.Linear(n_sensors, branch_width),
            nn.Tanh(),
            nn.Linear(branch_width, branch_width),
            nn.Tanh(),
            nn.Linear(branch_width, p),
        )

        # Trunk network: query point -> feature vector
        self.trunk = nn.Sequential(
            nn.Linear(2, trunk_width),    # 2D spatial input
            nn.Tanh(),
            nn.Linear(trunk_width, trunk_width),
            nn.Tanh(),
            nn.Linear(trunk_width, p),
        )

        # Learnable bias
        self.bias = nn.Parameter(torch.zeros(1))

    def forward(self, a_sensors, y_query):
        """
        Args:
            a_sensors: (batch, n_sensors) input function values
            y_query: (batch, n_query, 2) query locations

        Returns:
            u_pred: (batch, n_query) predicted output function values
        """
        # Branch: (batch, p)
        b = self.branch(a_sensors)

        # Trunk: (batch, n_query, p)
        tau = self.trunk(y_query)

        # Dot product: (batch, n_query)
        u_pred = torch.einsum("bp,bqp->bq", b, tau) + self.bias

        return u_pred

7. Comparing Neural Operators, PINNs, and Classical Solvers

With both FNO (optimized for regular grids) and DeepONet (flexible on irregular geometries) now in hand, the natural question is when to reach for each, and when a classical solver or a PINN remains the better tool.

The table below compares classical solvers, PINNs, FNO, and DeepONet on the Darcy flow and Navier-Stokes benchmarks.

8. Physics-Informed Neural Operators

Physics-Informed Neural Operators (PINO, Li et al., 2021) combine the physics-informed losses of PINNs with the operator learning of FNO. PINO adds PDE residual losses to the FNO training objective, enabling training with fewer data pairs and better generalization to out-of-distribution inputs.

def pino_loss(model, x_batch, u_batch, x_colloc, pde_weight=1.0):
    """
    Combined data + physics loss for PINO.

    The data loss fits available solved PDE instances.
    The physics loss enforces the PDE at collocation points
    without requiring solved reference data.
    """
    # Data loss: fit known solutions
    u_pred = model(x_batch)
    loss_data = relative_l2_loss(u_pred, u_batch)

    # Physics loss: PDE residual at unsupervised collocation points,
    # where collocation points are spatial locations sampled across
    # the domain at which the PDE residual is evaluated without
    # requiring reference solution data
    x_colloc.requires_grad_(True)
    u_colloc = model(x_colloc)

    # Compute PDE residual (Darcy flow: -div(a * grad(u)) = f)
    # a is embedded in x_colloc[:, :, :, 2]
    a_field = x_colloc[:, :, :, 2]

    # Spatial gradients via autograd
    du_dx = torch.autograd.grad(
        u_colloc.sum(), x_colloc, create_graph=True
    )[0]
    du_dx_spatial = du_dx[:, :, :, :2]  # gradients w.r.t. (x, y)

    # Divergence of (a * grad u)
    flux = a_field.unsqueeze(-1) * du_dx_spatial
    div_flux = torch.autograd.grad(
        flux.sum(), x_colloc, create_graph=True
    )[0][:, :, :, :2].sum(dim=-1)

    # PDE residual: should equal source term f
    residual = -div_flux  # simplified; f handled separately
    loss_physics = torch.mean(residual**2)

    return loss_data + pde_weight * loss_physics

PINO is particularly valuable when generating training data is expensive. Instead of solving 1,000 PDE instances, you can solve 100 and supplement with physics-informed collocation at 10,000 unsupervised points. The physics loss acts as a regularizer that guides the operator toward physically consistent solutions in data-sparse regions of the input space.

9. Connecting to Engineering Design

These architectural advances are compelling on their own, but the real payoff comes when neural operators slot into a larger engineering workflow where speed per evaluation directly determines how many design candidates can be explored.

Neural operators transform engineering design by enabling real-time PDE evaluation within optimization loops. The workflow is:

1. Generate training data: run the classical solver for a representative sample of design configurations.
2. Train the operator: fit an FNO or DeepONet to map design parameters to performance metrics.
3. Optimize: run gradient-based or Bayesian optimization over the trained surrogate, evaluating thousands of candidates in seconds.
4. Validate: run the classical solver on the top candidates to verify the surrogate's predictions.

This four-step loop is the foundation of the physics discovery pipeline in Section 50.4, where we combine PySR for law discovery, PINNs for parameter estimation, FNO for surrogate modeling, and BoTorch for optimization into a single integrated pipeline.