DeepONet Surrogate for 1D Diffusion

Why an operator network

Standard MLPs are universal function approximators: they map numbers to numbers. They struggle with operators, which map a whole function to another function, partly because of how the data is shaped. DeepONet is a universal operator approximator.

Here the operator takes a diffusion coefficient $D(x)$ and returns the concentration $u(x, t)$ everywhere in space and time. Once trained, it answers for a new $D(x)$ in a single forward pass, with no equation to solve.

Part of my research internship with Prof. Saurav Sengupta, IIT (ISM) Dhanbad.

The cell below does both live: it solves the equation with the notebook's finite-difference scheme and asks the DeepONet for the same answer. Move the sliders to change $D(x)$.

The problem

Fick's second law with a diffusion coefficient that varies along a 1D pipe, $x \in [0, 1]$, for $t \in [0, 1]$:

$$\frac{\partial u}{\partial t} = \frac{\partial}{\partial x}\left[D(x)\,\frac{\partial u}{\partial x}\right] = D\,\frac{\partial^{2} u}{\partial x^{2}} + \frac{\partial D}{\partial x}\,\frac{\partial u}{\partial x}$$

The pipe starts with $u(x, 0) = \sin \pi x$. Two kinds of ends:

• Dirichlet (open ends): $u = 0$ at $x = 0$ and $x = 1$, so the concentration drains away.
• Neumann (closed ends): $\partial u/\partial x = 0$ at both ends, so nothing leaves and the profile flattens towards its mean.

The data

1,000 random diffusion coefficients, each linear, $D(x) = ax + b$, kept between $m = 0.10$ and $M = 0.45$ at both ends: $b \sim U[m, M]$, then $a \sim U[m - b, M - b]$. Each is sampled at 100 sensor points.

For each $D_j$, the equation is solved with an explicit finite-difference scheme (FTCS) in flux form, with $D$ taken at the cell interfaces: 100 points in $x$ and 10,000 time steps, of which 150 are kept. That makes each solution a $150 \times 100$ grid, 15,000 values: the ground truth. 800 samples train the network and 200 test it.

The network

A DeepONet has two parts:

• The branch net reads the input function at its 100 sensors and outputs 50 numbers $b_k(D)$ (100 → 128 → 128 → 50, ReLU).
• The trunk net reads a query point $(x, t)$ and outputs 50 numbers $\tau_k(x, t)$ (2 → 128 → 128 → 50, tanh).

They are merged by a dot product, plus a bias:

$$G_\theta(D)(x, t) = \sum_{k=1}^{50} b_k(D)\,\tau_k(x, t) + b_0, \qquad \mathbf{U}_{\text{pred}} = \mathbf{B}\mathbf{T}^{\top} + \text{bias}$$

The trunk net learns 50 basis functions over space and time; the branch net learns how much of each a given $D(x)$ needs. The loss is the mean squared error against the finite-difference solutions, minimised with Adam (learning rate 0.0016, batches of 100, 700 epochs).

Results

On the 200 test coefficients, the relative error $\lVert U_{\text{true}} - U_{\text{pred}} \rVert_2 / \lVert U_{\text{true}} \rVert_2$ is 2.56% with Neumann ends and 3.98% with Dirichlet ends.

The notebook doesn't save the trained weights, so the live cell uses weights from re-running its training code unchanged (same data, split and hyperparameters, on a CPU). On the same 200 test coefficients they reach 1.31% (Neumann) and 1.69% (Dirichlet). The difference from the notebook's numbers comes from the random initialisation and batch shuffling, which the notebook doesn't seed.

In the live cell, slide $D$ outside the training range (0.10 to 0.45) to see the network extrapolate: at $D = 0.05$ everywhere the error grows to over 10%, because it has never seen such a coefficient.