Section 33.2: Neural Operators

1. From Functions to Operators

Every time a climate scientist changes an initial condition or an engineer tweaks a design parameter, a traditional PDE solver must restart from zero, burning hours of compute for each variation. Neural operators eliminate that bottleneck by learning the solver itself.

A neural operator accepts an entire function as input and returns another function as output. Instead of the familiar mapping between finite-dimensional vectors (\(f: \mathbb{R}^n \to \mathbb{R}^m\)), it learns a mapping between function spaces (collections of functions that share a common domain and structure, such as all smooth functions on an interval), \(\mathcal{G}: \mathcal{U} \to \mathcal{V}\), where \(\mathcal{U}\) and \(\mathcal{V}\) are (typically infinite-dimensional) spaces of functions. For a PDE, the operator might map an initial condition \(u_0(x) \in \mathcal{U}\) to the solution at time \(T\), \(u(x, T) \in \mathcal{V}\).

Because these function spaces can be infinite-dimensional, learning such a mapping requires architectures specifically designed for operator approximation, not just larger versions of standard neural networks. The theoretical foundation is the universal approximation theorem for operators (Chen & Chen, 1995): a neural network with a specific architecture can approximate any continuous nonlinear operator to arbitrary accuracy. DeepONet and the FNO are two different architectures that realize this theorem, each with distinct strengths.

A neural operator is a parametric map between infinite-dimensional function spaces. It trains on paired examples of input and output functions, then predicts the output for any new input in a single forward pass. Traditional PDE solvers re-solve from scratch whenever the input (initial condition, boundary, or coefficient field) changes. A neural operator amortizes that cost across all future queries. The network encodes the input function into a finite-dimensional latent representation, transforms it through learnable layers that respect the function-space structure, and decodes it back to a function. Choose a neural operator over a PINN when you need to evaluate the same PDE family for hundreds or thousands of different inputs; choose a traditional solver when you need guaranteed error bounds or have only a single instance to solve. In short: A neural operator trades one expensive training run for unlimited cheap predictions across an entire family of equations.

2. DeepONet: Branch and Trunk

With the universal approximation guarantee established, the next question is architectural. DeepONet provides the first major answer.

DeepONet (Deep Operator Network) factorizes the operator into two sub-networks:

• The branch network encodes the input function \(u(x)\), evaluated at a fixed set of sensor locations \(\{x_1, \ldots, x_m\}\), into a latent representation \(\mathbf{b} = [b_1, \ldots, b_p] \in \mathbb{R}^p\).
• The trunk network takes the query location \(y\) (where we want to evaluate the output) and produces a basis representation \(\boldsymbol{\tau} = [\tau_1(y), \ldots, \tau_p(y)] \in \mathbb{R}^p\).

The output is their inner product: Figure 33.2.1 illustrates DeepONet branch-trunk architecture.

This factorization has elegant properties: the branch network learns a data-dependent representation of the input function, while the trunk network learns a set of basis functions for the output space. The inner product combines them, analogous to how a traditional Galerkin method (a technique that approximates a continuous solution by projecting onto a finite set of basis functions) projects onto a finite basis. Figure 33.2 illustrates this two-network architecture and the dot-product combination step.

import torch
import torch.nn as nn

class DeepONet(nn.Module):
    """DeepONet: branch encodes input function, trunk encodes query location."""

    def __init__(self, branch_input_dim, trunk_input_dim, hidden_dim=128,
                 num_basis=64):
        super().__init__()
        # Branch network: input function values at sensor locations -> latent
        self.branch = nn.Sequential(
            nn.Linear(branch_input_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, num_basis),
        )
        # Trunk network: query location -> basis functions
        self.trunk = nn.Sequential(
            nn.Linear(trunk_input_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, num_basis),
        )
        self.bias = nn.Parameter(torch.zeros(1))

    def forward(self, u_sensors, y_query):
        """
        u_sensors: (batch, m) - input function at m sensor locations
        y_query: (batch, d) - query locations in output domain
        Returns: (batch,) - operator output at query locations
        """
        b = self.branch(u_sensors)          # (batch, num_basis)
        tau = self.trunk(y_query)            # (batch, num_basis)
        return torch.sum(b * tau, dim=-1) + self.bias

# Example: map initial condition u0(x) to solution u(x, T=0.5)
# m=100 sensor locations, 1D query points
model = DeepONet(branch_input_dim=100, trunk_input_dim=1)

3. Fourier Neural Operator: Spectral Convolutions

DeepONet's branch-trunk factorization works well for point queries and irregular sensor data, but it treats the spatial domain as an unstructured collection of locations. When the data lives on a regular grid, a second architecture exploits that structure directly by working in the frequency domain.

The Fourier Neural Operator (FNO) takes a fundamentally different approach. Instead of factorizing the operator, it processes the input function through a sequence of layers, each of which applies a linear transformation in Fourier space (the frequency-domain representation obtained by decomposing a signal into its constituent sinusoidal components) followed by a nonlinear activation in physical space:

where \(\mathcal{F}\) denotes the Fourier transform, \(R^{(l)}\) is a learnable weight matrix in Fourier space (applied to the first \(k_{\max}\) modes), \(W^{(l)}\) is a local linear transformation, and \(\sigma\) is a pointwise nonlinearity (typically Gaussian Error Linear Unit, or GELU).

The key properties of the spectral convolution are:

• Resolution invariance (the ability of a trained model to accept inputs discretized at any spatial resolution, not just the resolution seen during training): since convolutions in physical space become pointwise multiplications in Fourier space, and the Fourier transform is defined for any discretization, an FNO trained at one resolution can be evaluated at a different resolution without retraining. Standard convolutional networks cannot do this because their learned filters have a fixed pixel width tied to the training grid spacing.
• Global receptive field: each Fourier layer captures long-range dependencies across the entire domain, unlike local convolutional filters that require stacking many layers for global context.
• Efficient parameterization: by truncating to \(k_{\max}\) Fourier modes, the model captures the dominant spectral components of the solution while keeping the parameter count manageable.

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

class SpectralConv1d(nn.Module):
    """1D Fourier layer: learnable filter in spectral domain."""

    def __init__(self, in_channels, out_channels, modes):
        super().__init__()
        self.modes = modes  # number of Fourier modes to keep
        scale = 1.0 / (in_channels * out_channels)
        # Complex-valued weights for Fourier modes
        self.weights = nn.Parameter(
            scale * torch.randn(in_channels, out_channels, modes,
                                dtype=torch.cfloat)
        )

    def forward(self, x):
        """x: (batch, channels, spatial_points)."""
        # Transform to Fourier space
        x_ft = torch.fft.rfft(x)

        # Multiply relevant modes by learnable weights
        out_ft = torch.zeros_like(x_ft[:, :1, :].repeat(1, self.weights.shape[1], 1))
        out_ft[:, :, :self.modes] = torch.einsum(
            "bim,iom->bom", x_ft[:, :, :self.modes], self.weights
        )

        # Transform back to physical space
        return torch.fft.irfft(out_ft, n=x.size(-1))


class FNO1d(nn.Module):
    """1D Fourier Neural Operator with 4 Fourier layers."""

    def __init__(self, modes=16, width=64):
        super().__init__()
        self.fc0 = nn.Linear(2, width)  # lift: (x, u) -> high-dim

        # Four Fourier layers
        self.convs = nn.ModuleList([
            SpectralConv1d(width, width, modes) for _ in range(4)
        ])
        self.ws = nn.ModuleList([
            nn.Conv1d(width, width, 1) for _ in range(4)
        ])

        # Project back to output
        self.fc1 = nn.Linear(width, 128)
        self.fc2 = nn.Linear(128, 1)

    def forward(self, x):
        """x: (batch, spatial_points, 2) where channels are [coord, u0]."""
        x = self.fc0(x)                     # (batch, S, width)
        x = x.permute(0, 2, 1)              # (batch, width, S)

        for conv, w in zip(self.convs, self.ws):
            x1 = conv(x)                     # spectral pathway
            x2 = w(x)                        # local pathway
            x = torch.nn.functional.gelu(x1 + x2)

        x = x.permute(0, 2, 1)              # (batch, S, width)
        x = torch.nn.functional.gelu(self.fc1(x))
        return self.fc2(x)                   # (batch, S, 1)

4. DeepONet vs. FNO: When to Use Which

The two architectures have complementary strengths:

As a rule of thumb: use DeepONet when your data comes from irregular sensor networks or you need continuous output evaluation, and use FNO when your data lives on a regular grid and the problem has periodic or smooth boundary conditions. For problems that need both irregular inputs and resolution invariance, hybrid architectures (e.g., geo-FNO) bridge the gap.

5. The NeuralOperator Library

from neuralop.models import FNO
from neuralop.datasets import load_darcy_flow_small
from neuralop.training import Trainer
from neuralop.losses import LpLoss

# Load the Darcy flow benchmark dataset
train_loader, test_loaders, output_encoder = load_darcy_flow_small(
    n_train=1000, batch_size=32,
    test_resolutions=[16, 32],  # test at multiple resolutions
    test_batch_sizes=[32, 32],
)

# Create FNO model
model = FNO(
    n_modes=(16, 16),           # Fourier modes in each dimension
    in_channels=1,              # input: permeability field
    out_channels=1,             # output: pressure field
    hidden_channels=32,
    n_layers=4,
)

# Train with relative L2 loss
trainer = Trainer(model=model, n_epochs=50)
trainer.train(
    train_loader=train_loader,
    test_loaders=test_loaders,
    output_encoder=output_encoder,
)

6. Training Neural Operators: Data Generation and Best Practices

Choosing an architecture and configuring a library are only half the challenge; the other half is producing the training data that neural operators require.

Neural operators are supervised: they need pairs of (input function, output function) for training. PINNs self-supervise through PDE residuals, but neural operators require pre-computed solutions from traditional solvers. This creates a chicken-and-egg situation: you need a solver to generate training data for the model that will replace it.

The resolution is that training data generation is a one-time cost. Best practices include:

• Diverse training distribution: sample input functions from a rich distribution (e.g., Gaussian random fields, where each sample is a smooth function whose values at any finite set of points follow a multivariate Gaussian distribution, with varying correlation lengths) to ensure the operator generalizes across the expected range of inputs.
• Multi-resolution training: generate training data at coarse resolution (cheap) and validate at fine resolution, leveraging the FNO's resolution invariance.
• Physics-informed neural operators (PINO): combine the data-driven loss with a PDE residual loss, reducing the amount of training data needed. This hybrid marries the strengths of PINNs (no labeled data needed) and neural operators (fast inference for new inputs).

Exercises

1. Conceptual: Explain why a standard CNN cannot generalize across resolutions while an FNO can. What property of the Fourier transform makes this possible? How does the choice of \(k_{\max}\) (number of retained modes) affect accuracy and generalization?
2. Coding: Using the NeuralOperator library, train an FNO on the Darcy flow dataset at resolution 16 x 16 and evaluate it at 32 x 32. Report the relative \(L^2\) error at both resolutions. Then retrain at 32 x 32 and compare. How much does training resolution matter?
3. Analysis: Implement a simple DeepONet for the 1D heat equation (input: initial condition sampled from \(\sin(k\pi x)\) for \(k = 1, \ldots, 5\); output: solution at \(t = 0.5\)). Compare its accuracy and inference speed against: (a) the PINN from Section 33.1 (retrained for each initial condition), and (b) a traditional finite-difference solver. At what number of evaluations does the DeepONet break even in total wall-clock time?