DeepONet — parametric Poisson 2D

Download .py Back to chapter

Train a DeepONet to solve a 1-parameter family of Poisson problems via PDE-residual learning. The network sees no ground-truth solutions — only the physics — and learns the operator k → u(·) for the entire range k ∈ [0.5, 1.5].

Problem Setup

k Δu + 1 = 0   on [0, 2] × [0, 1],   u = 0 on ∂Ω,   k ~ U(0.5, 1.5)

50 random k values are sampled at the start of training; the solver replicates the spatial mesh across all 50 samples and computes the residual for every (k, x, y) triple in one forward pass.

Step 1: Parametric Domain

Multiplying a domain by an integer B replicates it across B independent samples. This is the operator-learning pattern:

N_SAMPLES = 50
dom = N_SAMPLES * jno.Shape.rect(0, 0, 2, 1, size=0.05).domain()
x, y, _ = dom.variable("interior")

k_values = jax.random.uniform(jax.random.PRNGKey(0), shape=(N_SAMPLES, 1, 1), minval=0.5, maxval=1.5)
k = dom.variable("k", k_values)

k_values has one scalar k per sample; attaching it as a tensor variable on the domain makes it accessible inside the symbolic expression.

Step 2: DeepONet Network

The branch input is the scalar k (a "function evaluated at one sensor"); the trunk input is the query coordinate (x, y). The output is the dot product of the two encoded vectors:

net = jno.nn(
    foundax.deeponet(
        n_sensors=1,         # branch input dimensionality
        coord_dim=2,         # trunk input dimensionality
        basis_functions=32,
        hidden_dim=128,
        activation=jax.numpy.tanh,
        key=jax.random.PRNGKey(0),
    )
)
net.optimizer(optax.adam(optax.cosine_decay_schedule(1e-3, 2_000, alpha=1e-5 / 1e-3)))

Step 3: Hard BCs + PDE Residual

u = net(k, jno.np.concat([x, y], axis=-1)) * x * (2 - x) * y * (1 - y)
pde = k * (u.d2(x) + u.d2(y)) + 1.0

The multiplicative ansatz x(2-x)y(1-y) vanishes on all four edges and enforces the homogeneous Dirichlet BC for every sample, so the boundary doesn't need a loss term.

Step 4: Solve

crux = jno.core(constraints=[pde.mse])
crux.solve(epochs=2_000, batchsize=32)

batchsize=32 means each gradient step uses 32 of the 50 parametric samples — a stochastic minibatch in k-space.

Result

Queried at a held-out coefficient k = 1.234 (never drawn during training), the trained operator's own output reproduces a finite-difference reference solution of k Δu + 1 = 0 to rel-\(L^2 \approx 0.007\). The DeepONet has learned the entire k → u(·) map from the PDE residual alone — no ground-truth solution was ever supplied.

What To Notice

• One network, one training run, 50 PDE solutions. After convergence, crux.eval(u) returns the solution field for every sampled k without any retraining.
• Branch/trunk factorisation is the operator-learning interpretation of "separation of variables in parameter space". It's cheap to scale (the trunk is the same for all samples), which makes DeepONet much faster than training one PINN per k.
• Pure PDE-residual training. No solution data is supplied — the network learns from physics alone. Compare with the FNO tutorials, which use a precomputed (f, u) dataset.

Download full script Back to 11 Operator Learning

Script Snippet

"""11 — DeepONet 2D for parametric Poisson"""

import foundax
import jax
import numpy as np
import optax
import scipy.sparse as sp
import scipy.sparse.linalg as spla

import jno

KEY = jax.random.PRNGKey(0)
N_SAMPLES = 50
EPOCHS = 2_000

# ── Parametric domain — replicate one mesh across N_SAMPLES random k values ──
dom = N_SAMPLES * jno.Shape.rect(0, 0, 2, 1, size=0.05).domain()
x, y, _ = dom.variable("interior")

k_values = jax.random.uniform(KEY, shape=(N_SAMPLES, 1, 1), minval=0.5, maxval=1.5)
k = dom.variable("k", k_values)

# ── Network ──────────────────────────────────────────────────────────────────
net = jno.nn(
    foundax.deeponet(
        n_sensors=1,  # branch input is the scalar k
        coord_dim=2,  # trunk input is (x, y)
        basis_functions=32,
        hidden_dim=128,
        activation=jax.numpy.tanh,
        key=KEY,
    )
)
net.optimizer(optax.adam(optax.cosine_decay_schedule(1e-3, EPOCHS, alpha=1e-5 / 1e-3)))

# ── Hard BC ansatz + PDE residual ────────────────────────────────────────────
u = net(k, jno.np.concat([x, y], axis=-1)) * x * (2 - x) * y * (1 - y)
pde = k * (u.d2(x) + u.d2(y)) + 1.0

# ── Solve ────────────────────────────────────────────────────────────────────
crux = jno.core(constraints=[pde.mse])
crux.solve(epochs=EPOCHS, batchsize=32)

# ── Held-out evaluation: query the trained operator at an unseen k, compare to
#    a finite-difference reference solution of  k Δu + 1 = 0,  u = 0 on ∂Ω. ────
K_TEST = 1.234  # a coefficient NOT in the training draw
GX, GY = 80, 40  # interior grid on [0, 2] x [0, 1]
xs = np.linspace(0.0, 2.0, GX + 1)[1:-1]
ys = np.linspace(0.0, 1.0, GY + 1)[1:-1]
XX, YY = np.meshgrid(xs, ys, indexing="xy")
coords = np.stack([XX.ravel(), YY.ravel()], axis=-1).astype(np.float32)

# the trained operator's OWN output at the held-out k (hard-BC ansatz applied)
raw = net.module(jax.numpy.array([K_TEST], dtype=jax.numpy.float32), jax.numpy.asarray(coords))
ansatz = coords[:, 0] * (2 - coords[:, 0]) * coords[:, 1] * (1 - coords[:, 1])
u_pred = (np.asarray(raw) * ansatz).reshape(XX.shape)

# finite-difference reference (5-point Laplacian, homogeneous Dirichlet)
nx, ny = len(xs), len(ys)
hx, hy = xs[1] - xs[0], ys[1] - ys[0]
Lx = sp.diags([1.0, -2.0, 1.0], [-1, 0, 1], shape=(nx, nx)) / hx**2
Ly = sp.diags([1.0, -2.0, 1.0], [-1, 0, 1], shape=(ny, ny)) / hy**2
A = sp.kron(Ly, sp.identity(nx)) + sp.kron(sp.identity(ny), Lx)
u_ref = spla.spsolve(A.tocsr(), np.full(nx * ny, -1.0 / K_TEST)).reshape(XX.shape)

rel_l2 = float(np.linalg.norm(u_pred - u_ref) / np.linalg.norm(u_ref))
print(f"held-out k={K_TEST}: rel-L2 vs finite-difference reference = {rel_l2:.4f}")
assert rel_l2 < 0.15, rel_l2