Neural Operators and Function Space Learning

Prerequisites

• Module 3.1: Neural Differential Equations (MLP basics, gradient descent)
• Module 2.2: Numerical PDE Methods (finite differences, Poisson solvers)
• Fourier analysis basics (DFT, frequency domain)

Why This Matters

A classical PDE solver computes one solution at a time: change the forcing function, initial condition, or geometry and you must re-solve from scratch. Neural operators learn the entire solution operator—a mapping from input functions to output functions. Once trained, they predict solutions for new inputs in milliseconds, enabling real-time simulation, rapid design optimization, and uncertainty quantification that would be infeasible with traditional solvers.

Two architectures dominate the field. DeepONet (Lu et al. 2021) decomposes the operator into a branch network (encodes the input function) and a trunk network (encodes the query location), combining them via a learned basis expansion. The Fourier Neural Operator (FNO, Li et al. 2021) performs convolution in the spectral domain, achieving resolution invariance and computational efficiency. Both architectures are backed by universal approximation theorems for nonlinear operators.

Learning Objectives

1. Distinguish between a function $u(x)$ and an operator $\mathcal{G}[f](x) = u(x)$ that maps functions to functions.
2. Describe the DeepONet architecture: branch network, trunk network, and their combination via inner product.
3. Describe the FNO architecture: lifting, spectral convolution layers, and projection.
4. State the universal approximation theorem for operators (Chen & Chen 1995).
5. Generate training data for operator learning by solving PDEs with multiple input functions.
6. Train a toy DeepONet in NumPy on the Poisson equation operator.
7. Explain resolution invariance and why FNOs can generalize across discretizations.
8. Compare operator learning with pointwise prediction (e.g., PINNs) in terms of amortized cost.
9. Frame architecture and hyperparameter selection as an agent design problem.

Core Concepts

Common Misconceptions

Worked Examples

Interactive Code Lab

import numpy as np
import matplotlib.pyplot as plt

# ── 1. Data generation: solve -u'' = f with u(0)=u(1)=0 ─────
def solve_poisson(f_vals, x_grid):
    """Solve -u''=f via tridiagonal finite differences."""
    n = len(x_grid)
    dx = x_grid[1] - x_grid[0]
    # Interior points only (exclude boundaries)
    A_diag = 2.0 * np.ones(n - 2)
    A_off = -1.0 * np.ones(n - 3)
    # Solve tridiagonal system: (1/dx^2) * A * u_int = f_int
    rhs = dx**2 * f_vals[1:-1]
    # Thomas algorithm
    n_int = n - 2
    c = np.zeros(n_int); d = np.zeros(n_int)
    c[0] = A_off[0] / A_diag[0] if n_int > 1 else 0
    d[0] = rhs[0] / A_diag[0]
    for i in range(1, n_int):
        m = A_diag[i] - A_off[min(i-1, len(A_off)-1)] * c[i-1]
        c[i] = A_off[min(i, len(A_off)-1)] / m if i < n_int - 1 else 0
        d[i] = (rhs[i] - A_off[min(i-1, len(A_off)-1)] * d[i-1]) / m
    u_int = np.zeros(n_int)
    u_int[-1] = d[-1]
    for i in range(n_int - 2, -1, -1):
        u_int[i] = d[i] - c[i] * u_int[i+1]
    u = np.zeros(n)
    u[1:-1] = u_int
    return u

n_points = 32
x_grid = np.linspace(0, 1, n_points)

def random_forcing(n_samples):
    """Random Fourier forcing functions."""
    F = np.zeros((n_samples, n_points))
    for i in range(n_samples):
        n_modes = np.random.randint(1, 5)
        for k in range(1, n_modes + 1):
            a = np.random.randn() * 2.0
            F[i] += a * np.sin(k * np.pi * x_grid)
    return F

N_train, N_test = 500, 100
f_train = random_forcing(N_train)
u_train = np.array([solve_poisson(f, x_grid) for f in f_train])
f_test = random_forcing(N_test)
u_test = np.array([solve_poisson(f, x_grid) for f in f_test])

# ── 2. Tiny DeepONet in NumPy ────────────────────────────────
class TinyDeepONet:
    def __init__(self, n_sensors, trunk_in, hidden, p, seed=42):
        rng = np.random.RandomState(seed)
        sc = lambda m, n: rng.randn(m, n) * np.sqrt(2.0/(m+n))
        # Branch: n_sensors -> hidden -> p
        self.Wb1 = sc(n_sensors, hidden)
        self.bb1 = np.zeros(hidden)
        self.Wb2 = sc(hidden, p)
        self.bb2 = np.zeros(p)
        # Trunk: trunk_in -> hidden -> p
        self.Wt1 = sc(trunk_in, hidden)
        self.bt1 = np.zeros(hidden)
        self.Wt2 = sc(hidden, p)
        self.bt2 = np.zeros(p)
        # Bias
        self.b0 = np.zeros(1)
        self.params = [self.Wb1, self.bb1, self.Wb2, self.bb2,
                       self.Wt1, self.bt1, self.Wt2, self.bt2, self.b0]

    def branch(self, f_sensors):
        """f_sensors: (batch, n_sensors) -> (batch, p)"""
        h = np.tanh(f_sensors @ self.Wb1 + self.bb1)
        return h @ self.Wb2 + self.bb2

    def trunk(self, x_query):
        """x_query: (n_query, 1) -> (n_query, p)"""
        h = np.tanh(x_query @ self.Wt1 + self.bt1)
        return h @ self.Wt2 + self.bt2

    def forward(self, f_sensors, x_query):
        """Output: (batch, n_query)"""
        b = self.branch(f_sensors)   # (batch, p)
        t = self.trunk(x_query)      # (n_query, p)
        return b @ t.T + self.b0     # (batch, n_query)

# ── 3. Training ──────────────────────────────────────────────
p = 20  # basis dimension
net = TinyDeepONet(n_points, 1, 32, p)
x_query = x_grid.reshape(-1, 1)  # (n_points, 1)

lr = 0.001
batch_size = 32
n_epochs = 2000
losses = []

for epoch in range(n_epochs):
    idx = np.random.choice(N_train, batch_size, replace=False)
    f_batch = f_train[idx]    # (batch, n_points)
    u_batch = u_train[idx]    # (batch, n_points)

    # Forward
    u_pred = net.forward(f_batch, x_query)  # (batch, n_points)
    loss = np.mean((u_pred - u_batch)**2)
    losses.append(loss)

    # Numerical gradient update (parameter perturbation)
    eps = 1e-5
    for pi in range(len(net.params)):
        grad = np.zeros_like(net.params[pi])
        flat = net.params[pi].ravel()
        grad_flat = grad.ravel()
        # Sample subset of parameters to update (for speed)
        n_param = len(flat)
        sample_idx = np.random.choice(n_param, min(50, n_param), replace=False)
        for j in sample_idx:
            old_val = flat[j]
            flat[j] = old_val + eps
            net.params[pi] = flat.reshape(net.params[pi].shape)
            u_p = net.forward(f_batch, x_query)
            loss_p = np.mean((u_p - u_batch)**2)
            flat[j] = old_val - eps
            net.params[pi] = flat.reshape(net.params[pi].shape)
            u_m = net.forward(f_batch, x_query)
            loss_m = np.mean((u_m - u_batch)**2)
            grad_flat[j] = (loss_p - loss_m) / (2 * eps)
            flat[j] = old_val
        net.params[pi] = flat.reshape(net.params[pi].shape)
        net.params[pi] -= lr * grad.reshape(net.params[pi].shape)

    # Keep references in sync
    (net.Wb1, net.bb1, net.Wb2, net.bb2,
     net.Wt1, net.bt1, net.Wt2, net.bt2, net.b0) = net.params

    if (epoch + 1) % 500 == 0:
        print(f"Epoch {epoch+1}: MSE = {loss:.6f}")

# ── 4. Evaluation ────────────────────────────────────────────
u_pred_test = net.forward(f_test, x_query)
errors = np.linalg.norm(u_pred_test - u_test, axis=1) / (np.linalg.norm(u_test, axis=1) + 1e-10)
print(f"\nMean relative L2 error: {errors.mean():.4f}")
print(f"Max relative L2 error:  {errors.max():.4f}")

fig, axes = plt.subplots(2, 3, figsize=(14, 8))
for col in range(3):
    i = np.random.randint(N_test)
    axes[0, col].plot(x_grid, f_test[i], 'k-', lw=1.5)
    axes[0, col].set_title(f'Forcing f(x) #{i}')
    axes[0, col].set_xlabel('x'); axes[0, col].grid(True, alpha=0.3)
    axes[1, col].plot(x_grid, u_test[i], 'b-', lw=2, label='True')
    axes[1, col].plot(x_grid, u_pred_test[i], 'r--', lw=2, label='DeepONet')
    axes[1, col].set_title(f'Solution, err={errors[i]:.3e}')
    axes[1, col].set_xlabel('x'); axes[1, col].legend()
    axes[1, col].grid(True, alpha=0.3)
plt.suptitle('DeepONet: Poisson Operator Learning', fontsize=14)
plt.tight_layout()
plt.show()

Agent Lens: The Operator-Learning Agent

Active data acquisition. Generating training data requires running a classical PDE solver, which is expensive. A smart agent chooses which input functions to solve next based on where the current model is most uncertain. This reduces the total number of expensive simulations needed. Techniques from Bayesian optimization and active learning apply directly.

Architecture search as agent design. Choosing between DeepONet and FNO, selecting $k_{\max}$, choosing network width/depth, and setting the basis dimension $p$ are all design decisions. A meta-learning agent could try different configurations on small-scale problems and transfer the best architecture to the full-scale problem.

Multi-fidelity learning. An agent can combine cheap low-resolution data with expensive high-resolution data, learning to correct the low-fidelity predictions rather than learning from scratch. This is analogous to curriculum learning: start with easy (low-res) tasks, then refine with hard (high-res) tasks.

Exercises

Analytical

Exercise 1. Write out the DeepONet prediction formula for the antiderivative operator $\mathcal{G}[f](x) = \int_0^x f(s)\,ds$ with $p = 3$ basis functions, 4 sensor points, and query point $x = 0.7$.

Exercise 2. Explain why the FNO spectral convolution is more efficient than a pointwise convolution in physical space for long-range interactions.

Exercise 3. State the conditions under which Chen & Chen’s universal approximation theorem applies. Does it require the operator to be linear?

Exercise 4. For a 1D FNO with $k_{\max} = 5$ Fourier modes on a grid of $n = 64$ points, how many learnable parameters does one spectral convolution layer have (for single input/output channel)?

Computational

Exercise 5. Modify the code lab to learn the antiderivative operator $\mathcal{G}[f](x) = \int_0^x f(s)\,ds$ instead of the Poisson operator. Generate data analytically and compare training convergence.

Exercise 6. Implement a 1D spectral convolution layer in NumPy: take a vector of length 32, apply FFT, multiply the first 8 modes by random weights, apply IFFT, and plot the result.

Exercise 7. Train the toy DeepONet on forcing functions with amplitudes scaled to $[-1, 1]$, then test on forcing functions with amplitudes in $[-3, 3]$. How well does the operator extrapolate?

Exercise 8. Implement a simple resolution-transfer experiment: train the DeepONet on $n = 32$ point grids, then evaluate at $n = 64$ by querying the trunk network at 64 locations. Report the error compared to the 64-point FD solution.

Agentic

Exercise 9. Design an active data acquisition agent: maintain a pool of candidate forcing functions, train the DeepONet on a small initial set, then iteratively select the candidate with highest predicted uncertainty (approximated as ensemble disagreement) and add its solution to the training set. Compare with random acquisition.

Exercise 10. Implement a model selection agent that trains both a DeepONet and a simple fully-connected network on the same data and automatically selects the architecture with lower validation error. Run on 3 different PDE families and report which architecture wins each time.

Assessment

Quiz

Mini-Project: Build an Antiderivative Operator

Train a DeepONet in NumPy to learn $\mathcal{G}[f](x) = \int_0^x f(s)\,ds$.

1. Generate training data: 500 random polynomial inputs $f(s) = \sum_{k=0}^{4} a_k s^k$ with $a_k \sim \mathcal{N}(0,1)$
2. Compute antiderivatives analytically
3. Train a DeepONet with $p = 30$ basis functions
4. Test on 100 unseen polynomial inputs and 50 unseen sinusoidal inputs (out-of-distribution)
5. Report in-distribution and out-of-distribution errors

References & Next Steps

References

1. Lu, L., Jin, P., Pang, G., Zhang, Z., & Karniadakis, G. E. (2021). “Learning nonlinear operators via DeepONet based on the universal approximation theorem of operators.” Nature Machine Intelligence, 3, 218–229.
2. Li, Z., Kovachki, N., Azizzadenesheli, K., Liu, B., Bhatt, K., Stuart, A., & Anandkumar, A. (2021). “Fourier neural operator for parametric partial differential equations.” ICLR 2021.
3. Chen, T. & Chen, H. (1995). “Universal approximation to nonlinear operators by neural networks with arbitrary activation functions and its application to dynamical systems.” IEEE Transactions on Neural Networks, 6(4), 911–917.
4. Kovachki, N., Li, Z., Liu, B., Azizzadenesheli, K., Bhatt, K., Stuart, A., & Anandkumar, A. (2023). “Neural operator: Learning maps between function spaces with applications to PDEs.” JMLR, 24(89), 1–97.
5. Lanthaler, S., Mishra, S., & Karniadakis, G. E. (2022). “Error estimates for DeepONets: A deep learning framework in infinite dimensions.” Transactions of Mathematics and Its Applications, 6(1).

Full PyTorch Notebook

For complete DeepONet and FNO implementations: notebooks/neural_operator_full.ipynb

Next Steps

• Module 3.4: Agent-Based Control and RL — apply RL to control ODE systems directly.
• Module 3.2 ← revisit PINNs and compare the pointwise vs operator-learning paradigms.