Deep Operator Networks: Labs
Learning Objectives¶
After studying this notebook, you should be able to
- say what an operator is, and explain why a network that takes a whole input function removes the need to retrain for every new case;
- build a DeepONet from a branch net and a trunk net, and read the prediction as an inner product $\sum_k b_k(u)\,\tau_k(z)$ of learned coefficients with a learned basis;
- prepare the branch input by sampling the input function at $m$ fixed sensor locations, and explain why the sensors must stay fixed while the trunk coordinate may be anything;
- train an operator on the governing equation alone, with the residual evaluated at coordinates drawn afresh at every step and no solution data anywhere in the loss;
- extend a trunk net from one coordinate to two, so that the output becomes a field rather than a curve;
- add Fourier features to a trunk net when the response is oscillatory, and say what problem they solve;
- combine measured responses with physics residuals when neither alone is enough, and explain why input functions that carry no measurement can still be used by the residual;
- judge a trained operator on held-out input functions, and design a test — such as a constant input the network never saw — that separates learning the physics from learning the training family;
- set up a PI-DeepONet for plane stress with two coupled displacement fields and residual and boundary terms built from
torch.autograd, trained with no solved examples at all; - combine a handful of CFD solutions with the Navier-Stokes residual to learn the map from viscosity to the velocity and pressure fields, and explain why a trained operator acts as a computational accelerator.
Where we are going. The theory behind everything used here — why an operator is the right object to learn, where the branch–trunk form comes from, and what the universal approximation theorem for operators actually promises — is developed in Deep Neural Operators. This notebook is the practice. Five problems are solved by one architecture: a draining tank, a channel flow, a two-degree-of-freedom oscillator, a stretched thin plate, and the flow around a cylinder. Read them as one recipe applied five times; what changes from lab to lab is what the branch net reads, what the trunk net reads, and how many fields come out.
1. From a Parameter to a Function¶
In the previous notebook a PINN was asked to recover one number. The tank drained, nine water-level readings were available, and the network returned a single drainage constant $C$. Everything the problem did not know was contained in that one scalar.
Now the unknown is a whole function. A pump feeds the tank, and its flow rate $Q_{in}(t)$ changes over time. Every different pump schedule gives a different water level history. A PINN trained on one schedule knows nothing about the next one, so solving a hundred schedules means training a hundred networks.
What we want instead is a single network that takes the schedule itself as input and returns the corresponding solution. That is an operator: a map from a function to a function.
1.1. What Changes from a PINN¶
| Level 1 (PINN) | Level 2 (DeepONet) | |
|---|---|---|
| Input | a number, $C$ | a function, $Q_{in}(t)$ |
| Output | one solution $h(t)$ | a solution for every input function |
| A new case costs | a full retraining | one forward pass |
The physics does not change. The tank obeys the same mass balance it did before, and we still teach it to the network as a residual rather than with labelled solutions.
1.2. The DeepONet Architecture¶
A DeepONet splits the job between two networks.
- The branch net reads the input function. A network cannot read a function, so we sample it at $m$ fixed sensor locations and feed the resulting vector.
- The trunk net reads the coordinate at which we want the answer: $t$ for the tank, $(y, t)$ for a field.
Their outputs are two vectors of the same length $p$, and the prediction is their inner product:
$$G_\theta(u)(z) \;=\; \sum_{k=1}^{p} b_k(u)\, \tau_k(z)$$
$b_k$ comes from the branch net and $\tau_k$ from the trunk net. There is no bias term on the product. The trunk net therefore learns a basis over the coordinate, and the branch net learns the coefficients that pick out one member of the solution family.
Physics-informed means the training signal is the governing equation evaluated on the output, not a set of precomputed solutions. No solver is run to produce labels.
1.3. Notation¶
| Symbol | Meaning | Python | Shape |
|---|---|---|---|
| $u$ | input function sampled at the sensors | U |
$(N_b, m)$ |
| $m$ | number of sensors | m |
scalar |
| $p$ | width of the branch and trunk output | p |
scalar |
| $Q_{in}(t)$ | inflow rate into the tank | Q_train |
$(N, m)$ |
| $h(t)$ | water level | h |
$(N_b, 1)$ |
| $C$ | drainage constant of the tank | — | scalar, 1 |
| $G(t)$ | pressure gradient driving the channel | G_train |
$(N, m)$ |
| $u(y,t)$ | streamwise velocity | — | $(N_b, 1)$ |
| $y$ | coordinate across the gap | y |
$(N_b, 1)$ |
$u$ denotes the input function of the DeepONet in Lab 1 and the velocity field in Lab 2. The number of sensors $m$ and the width $p$ are fixed within a lab but change between labs: $m = 20$ and $p = 100$ in Labs 1 and 2, $m = 100$ and $p = 32$ in Lab 3, $m = 250$ and $p = 100$ in Lab 4, and a single sensor with $p = 100$ in Lab 5.
1.4. The Five Labs¶
| Problem | Branch input | Trunk input | Output | |
|---|---|---|---|---|
| Lab 1 | tank with a time-varying inflow | $Q_{in}(t)$ | $t$ | $h(t)$ |
| Lab 2 | unsteady flow between parallel plates | $G(t)$ | $(y, t)$ | $u(y,t)$ |
| Lab 3 | 2-DOF forced vibration | $F(t)$ | $t$ | $x_1(t),\,x_2(t)$ |
| Lab 4 | thin plate under an edge traction | $P(y)$ | $(x, y)$ | $u(x,y),\,v(x,y)$ |
| Lab 5 | flow around a cylinder | $\mu$ | $(x, y)$ | $u, v, p$ |
Lab 1 keeps the physics of the previous notebook and changes only what the input is, so that the operator idea arrives on familiar ground. Lab 2 adds a space coordinate, and the output becomes a field. The same two nets serve both; only the trunk input width changes from 1 to 2.
Lab 3 keeps a single coordinate but predicts two coupled outputs, and it is the first hybrid objective of the notebook: measured responses supply a data term, and the equations of motion and the initial conditions are added on top of them. Its physics terms are evaluated on three hundred further forces that were never solved, which is the practical reason for mixing the two kinds of supervision.
Labs 4 and 5 are the two largest runs of the notebook, and both come with pretrained models. Lab 4 is the physics-informed DeepONet in its purest form: a thin plate under a family of edge tractions, two coupled displacement fields, and nothing but the plane-stress equations and the boundary conditions in the loss. Lab 5 is the laboratory closest to industrial practice: the flow around a cylinder at varying viscosity, where nine CFD solutions and the Navier-Stokes residual together train an operator that answers a new Reynolds number in a single forward pass.
2.1. Problem Setup¶
A tank drains through a small orifice while a pump feeds it at a rate that changes over time. Mass conservation gives
$$A\frac{dh}{dt} = Q_{in}(t) - C\sqrt{h}, \qquad h(0) = h_0$$
- Problem properties (nondimensionalized)
$$A = 1, \qquad C = 1, \qquad h_0 = 1, \qquad t \in [0, 2]$$
The outflow term $C\sqrt{h}$ is Torricelli's law and it is what makes the equation nonlinear, so the residual is not trivially satisfied.
- Operator to be learned
$$Q_{in}(t) \;\longmapsto\; h(t)$$
- Loss
$$\mathcal{L} = \mathcal{L}_{ODE} + \mathcal{L}_{IC}$$
No solution data enters the loss. A reference solution is computed for the test inputs only, and only to measure how well the operator did.
import numpy as np
import torch
import torch.nn as nn
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from IPython.display import HTML
CUDA
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(device)
Generate Data
The input functions are random pump schedules: a smooth Gaussian random field with mean 0.67 and length scale 0.5, sampled at the $m = 20$ sensor times and kept above 0.1 so that the pump never stops. The seed is fixed, so every run of the notebook sees the same family. Each row of Q_train is one pump schedule — that vector, and nothing else, is what the branch net ever sees of the function.
No solution is computed for the 100 training inputs. Only the ten test inputs are solved, and only to grade the operator after training. Between sensors the inflow is interpolated linearly, and the reference solution uses the same interpolation as the training residual, so the two are grading the same problem.
T, m = 2.0, 20 # time horizon, number of sensors
t_sensor = np.linspace(0, T, m) # the fixed sensor times
t_ref = np.linspace(0, T, 401) # the times at which the reference is evaluated
rng = np.random.default_rng(2)
def random_inflows(n, length = 0.5): # smooth random pump schedules at the sensor times, (n, m)
d = t_sensor[:, None] - t_sensor[None, :]
L = np.linalg.cholesky(np.exp(-d**2/(2*length**2)) + 1e-6*np.eye(m)) # squared-exponential covariance
z = rng.standard_normal((n, m)) @ L.T
return np.clip(0.67 + 0.22*z, 0.1, None) # mean inflow 0.67, never below 0.1
def solve_tank(Q): # reference h(t) for one inflow, Q_in interpolated as in sample_at
sol = solve_ivp(lambda t, h: np.interp(t, t_sensor, Q) - np.sqrt(h), [0, T], [1.0],
t_eval = t_ref, rtol = 1e-8, atol = 1e-10, max_step = T/(m - 1))
return sol.y[0]
Q_train = random_inflows(100) # inputs only: no solution is computed for these
Q_test = random_inflows(10)
h_test = np.array([solve_tank(Q) for Q in Q_test])
print('Q_train :', Q_train.shape)
print('Q_test :', Q_test.shape)
print('h_test :', h_test.shape)
plt.figure(figsize = (6, 4))
plt.plot(t_sensor, Q_train[:10].T, 'k', alpha = 0.3)
plt.xlabel('t')
plt.ylabel('Inflow rate, Q_in')
plt.title('Ten of the 100 Training Inflow Functions')
plt.grid(True)
plt.show()
Data Batch Generator
Two things are being sampled at once, and it is worth separating them. The input functions are drawn from the generated family, and the collocation times are drawn fresh from $U(0, T)$ at every step. A DeepONet therefore needs a batch in two dimensions: which function, and where on it.
U = torch.tensor(Q_train, dtype = torch.float32).to(device) # (100, 20): the branch input
U_test = torch.tensor(Q_test, dtype = torch.float32).to(device) # (10, 20)
t_s = torch.tensor(t_sensor, dtype = torch.float32).to(device) # (20,): the sensor times
def sample_at(u, t): # u(t) by linear interpolation between the sensor values, (pairs, 1)
i = torch.clamp(torch.searchsorted(t_s, t[:, 0].detach().contiguous()) - 1, 0, m - 2) # the sensor interval of each t
w = ((t[:, 0] - t_s[i])/(t_s[i + 1] - t_s[i])).reshape(-1, 1)
row = torch.arange(len(u), device = device)
return (1 - w)*u[row, i].reshape(-1, 1) + w*u[row, i + 1].reshape(-1, 1)
def get_batch(batch_size = 64):
i = torch.randint(0, len(U), (batch_size,)) # which inflow
t = T*torch.rand(batch_size, 1, device = device, requires_grad = True) # which time
return U[i], t
Define Network and Hyper-parameters
The branch net reads the 20 sensor values and the trunk net reads $t$; both are seven layers of width 100 and both end at the same width $p = 100$, which is what makes the inner product $\sum_k b_k \tau_k$ possible. There is no bias on the product: whatever the prediction cannot be written as a combination of trunk basis functions, it cannot represent at all.
def new_net(n_in, n_out):
return nn.Sequential(nn.Linear(n_in, 100), nn.Tanh(),
nn.Linear(100, 100), nn.Tanh(),
nn.Linear(100, 100), nn.Tanh(),
nn.Linear(100, 100), nn.Tanh(),
nn.Linear(100, 100), nn.Tanh(),
nn.Linear(100, 100), nn.Tanh(),
nn.Linear(100, n_out)).to(device)
p = 100 # number of basis functions
torch.manual_seed(0)
branch = new_net(m, p) # 20 sensor values -> p coefficients
trunk = new_net(1, p) # t -> p basis values
def deeponet(u, y):
return torch.sum(branch(u)*trunk(y), dim = 1, keepdim = True)
print('number of parameters :', sum(w.numel() for w in branch.parameters()) + sum(w.numel() for w in trunk.parameters()))
Define Functions to Compute Physics Loss
The initial-condition term is written as $h(0) - h_0$ so that every loss term is measured against zero.
Note where the derivative is taken. The quantity differentiated with respect to the trunk input is the output of the whole DeepONet, so what the residual constrains is the operator's prediction rather than one solution curve. It is the same differentiation used in the PINN labs, applied to a larger object.
def derivative(f, x):
return torch.autograd.grad(f, x, torch.ones_like(f), create_graph = True)[0]
def ode_loss():
u, t = get_batch()
h = deeponet(u, t)
r = derivative(h, t) - sample_at(u, t) + torch.sqrt(torch.clamp(h, min = 0) + 1e-8) # dh/dt = Q_in - sqrt(h)
return torch.mean(r**2)
def ic_loss():
u, t = get_batch()
return torch.mean((deeponet(u, torch.zeros_like(t)) - 1)**2) # h(0) = 1
Train
(This process may require a significant amount of time; therefore, we recommend not running it. However, we have provided a pretrained model for your convenience.)
Read the loop once even if you do not run it: nothing in it refers to a solution. A batch of pump schedules and a batch of times go in, the residual of the mass balance and the initial-condition error come out, and their sum is minimized. The network is being taught the equation for a whole family of inputs at once.
optimizer = torch.optim.Adam(list(branch.parameters()) + list(trunk.parameters()), lr = 5e-4)
n_iter = 20000
scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma = 0.02**(1/n_iter)) # 5e-4 -> 1e-5 over the run
# for epoch in range(1, n_iter + 1):
# loss_ode = ode_loss()
# loss_ic = ic_loss()
# loss = loss_ode + loss_ic
# optimizer.zero_grad()
# loss.backward()
# optimizer.step()
# scheduler.step()
# if epoch % 2000 == 0:
# print('epoch %6d | loss %.4e | ode %.4e | ic %.4e' % (epoch, loss.item(), loss_ode.item(), loss_ic.item()))
# torch.save([branch, trunk], 'tank_deeponet_model.pt')
!gdown 1dFr__e1yKq7o8aZ-CYvkknF_M3iZQ5Wb -O /content/tank_deeponet_model.pt
branch, trunk = torch.load('/content/tank_deeponet_model.pt', map_location = device, weights_only = False)
Result
The ten test functions below were drawn from the same family but never shown to the network during training.
"Never shown" is the point of the test. For a PINN, a new pump schedule would mean a new training run; here it is one forward pass through a network whose weights are already fixed.
y_ref = torch.tensor(t_ref.reshape(-1, 1), dtype = torch.float32).to(device) # the reference times
with torch.no_grad():
h_pred = (branch(U_test) @ trunk(y_ref).T).cpu().numpy() # (10, times): every test inflow at every time
rel = np.linalg.norm(h_pred - h_test, axis = 1)/np.linalg.norm(h_test, axis = 1)
print('mean relative L2 error over 10 unseen inflows : %.4f' % rel.mean())
plt.figure(figsize = (9, 3))
for j in range(3):
plt.subplot(1, 3, j + 1)
plt.plot(t_ref, h_test[j], 'k', alpha = 0.3, linewidth = 4, label = 'Reference')
plt.plot(t_ref, h_pred[j], 'r--', label = 'DeepONet')
plt.ylim([0, 1.2])
plt.title('Unseen inflow %d' % (j + 1))
plt.xlabel('t')
if j == 0:
plt.ylabel('Water level, h')
plt.legend()
plt.tight_layout()
plt.show()
One network answered three different pump schedules without being retrained, and it was never shown a single solution while learning. What it was shown is the equation.
Another of the test schedules in motion. The marker on the left rides along $Q_{in}(t)$ and the level on the right is drawn up to that instant, with the whole reference faint behind it. Every turning point of $h$ is an instant where inflow and outflow balance. The dashed curve is produced from the left panel alone, in one forward pass, for a schedule the network never saw.
frames = range(0, len(t_ref), 8)
fig, (ax_q, ax_h) = plt.subplots(1, 2, figsize = (9, 3))
j = 3
ax_q.plot(t_sensor, Q_test[j], 'b')
mark, = ax_q.plot([], [], 'bo')
ax_q.set_xlim([0, T])
ax_q.set_xlabel('t')
ax_q.set_ylabel('Inflow, Q_in')
ax_q.set_title('Input, Q_in(t)')
ax_h.plot(t_ref, h_test[j], c = '0.85', linewidth = 0.8)
line_ref, = ax_h.plot([], [], 'k', alpha = 0.3, linewidth = 4, label = 'Reference')
line_pred, = ax_h.plot([], [], 'r--', label = 'DeepONet')
clock = ax_h.text(0.04, 0.90, '', transform = ax_h.transAxes)
ax_h.set_xlim([0, T])
ax_h.set_ylim([0, 1.2])
ax_h.set_xlabel('t')
ax_h.set_ylabel('Water level, h')
ax_h.set_title('Water level, h(t)')
ax_h.legend(loc = 'lower left')
plt.tight_layout()
plt.close()
def update(k):
mark.set_data([t_ref[k]], [np.interp(t_ref[k], t_sensor, Q_test[j])])
line_ref.set_data(t_ref[:k + 1], h_test[j, :k + 1])
line_pred.set_data(t_ref[:k + 1], h_pred[j, :k + 1])
clock.set_text('t = %.2f' % t_ref[k])
anim = FuncAnimation(fig, update, frames = frames)
HTML(anim.to_jshtml(fps = 12))
2.3. What This Lab Showed¶
One network answered ten pump schedules it had never seen, with a mean relative $L^2$ error of about half a percent, and it was never shown a single solution while learning. What it was shown is the mass balance and the initial condition, evaluated at times drawn afresh at every step.
Three things carried over from the PINN labs unchanged: the residual, the way the derivative of the prediction is taken, and the habit of writing every loss term as a deviation from zero. What changed is the input. A PINN takes a coordinate and returns a number; this network takes a function, sampled at twenty fixed times, and returns the solution that function produces. The cost of a new case fell from a training run to a forward pass, and that is the whole of what the operator formulation buys.
Two limits are worth carrying forward. The sensor vector is the entire extent of what the branch net knows about its input, so a schedule that varies faster than twenty samples can resolve cannot be read correctly however well the network is trained. And the ten test schedules were drawn from the same family as the hundred training ones, which makes this an interpolation test; nothing here says what the operator would do with an input of a different character.
3.1. Problem Setup¶
Two stationary parallel plates a distance apart, with fluid between them driven along by a pressure gradient that varies in time. The flow is fully developed, so it depends on the cross-stream coordinate $y$ and on time, and the convective term is exactly zero. Starting from rest,
$$\frac{\partial u}{\partial t} = G(t) + \frac{\partial^2 u}{\partial y^2}, \qquad y \in [0, 1], \quad t \in [0, 1]$$
$$u(0, t) = u(1, t) = 0 \quad \text{(no-slip)}, \qquad u(y, 0) = 0 \quad \text{(starts from rest)}$$
$u$ is the streamwise velocity, $y$ runs across the gap, and $G(t) = -\frac{1}{\rho}\frac{\partial p}{\partial x}(t)$ is the driving pressure gradient. Constants are set to 1.
Operator to be learned
$$G(t) \;\longmapsto\; u(y, t)$$
- Loss
$$\mathcal{L} = \mathcal{L}_{PDE} + \mathcal{L}_{BC} + \mathcal{L}_{IC}$$
Two things are new relative to Lab 1. The trunk net now reads two coordinates instead of one, so the output is a field rather than a curve. And $G(t)$ carries no sign constraint - a pulsating pressure gradient genuinely reverses, and the flow reverses with it.
If $G$ is held constant the solution settles to the fully developed parabola $u = \frac{G}{2}y(1-y)$, which is the profile the PINN lecture already draws. That limit is used as a check at the end.
3.2. Solve the Operator Problem¶
Generate Data
The pressure gradients are drawn much as the pump schedules of Lab 1 were: a smooth Gaussian random field at the 20 sensor times, each function scaled so that its peak $|G|$ is 1. One thing differs, and it matters for the physics: no sign filter is applied, so most of the gradients reverse at least once and the flow reverses with them.
As in Lab 1, only the ten test inputs are solved. Expanding $u$ in $\sin(n\pi y)$, which already satisfies both walls, turns the PDE into one decoupled ODE per mode,
$$\frac{db_n}{dt} = -(n\pi)^2\, b_n + \frac{4}{n\pi}\,G(t), \qquad n = 1, 3, 5, \dots$$
and the even modes stay at zero because a uniform push has no even component. The reference is therefore exact up to the truncation of the series, and it is computed for the test inputs only.
T, m = 1.0, 20 # time horizon, number of sensors
t_sensor = np.linspace(0, T, m) # the fixed sensor times
t_eval = np.linspace(0, T, 115) # the grid on which the reference is evaluated
y_eval = np.linspace(0, 1, 51)
rng = np.random.default_rng(5)
def random_gradients(n, length = 0.25): # pressure gradients at the sensor times, (n, m)
d = t_sensor[:, None] - t_sensor[None, :]
w, V = np.linalg.eigh(np.exp(-d**2/(2*length**2))) # squared-exponential covariance
modes = V[:, -5:]*np.sqrt(w[-5:]) # its five leading Karhunen-Loeve modes
G = rng.standard_normal((n, 5)) @ modes.T
return G/np.abs(G).max(axis = 1, keepdims = True) # peak |G| = 1, either sign
n_mode = np.arange(1, 200, 2) # odd sine modes; the even ones receive no forcing
lam, c = (n_mode*np.pi)**2, 4/(n_mode*np.pi) # decay rate and forcing weight of each mode
def solve_channel(G): # reference u(y, t) for one gradient, as a sine series in y
sol = solve_ivp(lambda t, b: c*np.interp(t, t_sensor, G) - lam*b, [0, T], np.zeros(len(n_mode)),
t_eval = t_eval, method = 'BDF', jac = np.diag(-lam), rtol = 1e-8, atol = 1e-10)
return sol.y.T @ np.sin(np.pi*np.outer(n_mode, y_eval)) # (times, positions)
G_train = random_gradients(100) # inputs only: no solution is computed for these
G_test = random_gradients(10)
u_test = np.array([solve_channel(G) for G in G_test])
print('G_train :', G_train.shape)
print('G_test :', G_test.shape)
print('u_test :', u_test.shape)
plt.figure(figsize = (6, 4))
plt.plot(t_sensor, G_train[:10].T, 'k', alpha = 0.3)
plt.axhline(0, c = 'b', linewidth = 0.8, label = 'G = 0, flow reverses below')
plt.legend()
plt.xlabel('t')
plt.ylabel('Pressure gradient, G')
plt.title('Ten of the 100 Training Pressure Gradients')
plt.grid(True)
plt.show()
Data Batch Generator
Each term of the loss samples where its own condition lives: the residual at interior points, the no-slip condition on the two walls, and the initial condition at $t = 0$. The pressure gradient is interpolated from its 20 sensor values at whatever interior time is drawn, so the input function is read at the same resolution the reference solution used.
U = torch.tensor(G_train, dtype = torch.float32).to(device) # (100, 20)
U_test = torch.tensor(G_test, dtype = torch.float32).to(device) # (10, 20)
t_s = torch.tensor(t_sensor, dtype = torch.float32).to(device)
def get_batch(batch_size = 64):
i = torch.randint(0, len(U), (batch_size,)) # which pressure gradient
y = torch.rand(batch_size, 1, device = device) # which position
t = T*torch.rand(batch_size, 1, device = device) # which time
return U[i], torch.cat([y, t], dim = 1).requires_grad_(True)
Define Network and Hyper-parameters
The same two nets as Lab 1, with one number changed: the trunk input width goes from 1 to 2. Everything that makes this a field rather than a curve is contained in that number. The trunk net now receives $(y, t)$ and therefore learns basis functions over the plane; the branch net and the inner product are untouched.
torch.manual_seed(0)
branch = new_net(m, p) # 20 sensor values -> p coefficients
trunk = new_net(2, p) # (y, t) -> p basis values
Define Functions to Compute Physics Loss
The residual needs a second derivative in $y$ and a first derivative in $t$, and both come from the same automatic differentiation used for the ordinary derivative of Lab 1 — applied twice for the second derivative, exactly as the beam's higher derivatives were built in the PINN labs. Moving from a curve to a field adds no new machinery to the loss.
def pde_loss():
u, yt = get_batch()
v = deeponet(u, yt)
v_y, v_t = derivative(v, yt).split(1, dim = 1)
v_yy, _ = derivative(v_y, yt).split(1, dim = 1)
r = v_t - sample_at(u, yt[:, 1:2]) - v_yy # u_t = G(t) + u_yy
return torch.mean(r**2)
def bc_loss():
u, yt = get_batch()
t = yt[:, 1:2]
bottom = deeponet(u, torch.cat([torch.zeros_like(t), t], dim = 1)) # u = 0 on y = 0
top = deeponet(u, torch.cat([torch.ones_like(t), t], dim = 1)) # u = 0 on y = 1
return torch.mean(bottom**2) + torch.mean(top**2)
def ic_loss():
u, yt = get_batch()
y = yt[:, 0:1]
return torch.mean(deeponet(u, torch.cat([y, torch.zeros_like(y)], dim = 1))**2) # u = 0 at t = 0
Train
Same optimizer and same schedule as Lab 1.
(Again, this process may require a significant amount of time; therefore, we recommend not running it. However, we have provided a pretrained model for your convenience.)
optimizer = torch.optim.Adam(list(branch.parameters()) + list(trunk.parameters()), lr = 5e-4)
n_iter = 100000
scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma = 0.02**(1/n_iter))
for epoch in range(1, n_iter + 1):
loss_pde = pde_loss()
loss_bc = bc_loss()
loss_ic = ic_loss()
loss = loss_pde + loss_bc + loss_ic
optimizer.zero_grad()
loss.backward()
optimizer.step()
scheduler.step()
if epoch % 10000 == 0:
print('epoch %6d | loss %.4e | pde %.4e | bc %.4e | ic %.4e' % (epoch, loss.item(), loss_pde.item(), loss_bc.item(), loss_ic.item()))
Save and Load
Since running the training process takes a long time, we are providing the pre-trained networks, allowing you to use them for evaluating unseen data directly.
As in Lab 1, everything below this point is evaluation.
# torch.save([branch, trunk], 'unsteadyflow_deeponet_model.pt')
!gdown 1ibMy5tP4NHLxISo3ww9iNX9JLMJO7jV4 -O /content/unsteadyflow_deeponet_model.pt
branch, trunk = torch.load('/content/unsteadyflow_deeponet_model.pt', map_location = device, weights_only = False)
Result
The left panel is the input the two fields answer - one pressure gradient the network never saw. Read it together with them: all three share $t$ as the horizontal axis, so a vertical slice of a field is the velocity across the gap at the instant marked on the left. Where $G$ crosses zero the push reverses, and the field follows with a delay.
yy, tt = np.meshgrid(y_eval, t_eval)
pts = torch.tensor(np.column_stack([yy.ravel(), tt.ravel()]), dtype = torch.float32).to(device) # every (y, t) of the evaluation grid
with torch.no_grad():
u_pred = (branch(U_test) @ trunk(pts).T).cpu().numpy().reshape(-1, len(t_eval), len(y_eval)) # (10, times, positions)
rel = np.linalg.norm(u_pred - u_test, axis = (1, 2))/np.linalg.norm(u_test, axis = (1, 2))
print('mean relative L2 error over 10 unseen pressure gradients : %.4f' % rel.mean())
span = np.abs(u_test[0]).max()
titles = ['Reference u(y, t)', 'DeepONet u(y, t)']
fields = [u_test[0], u_pred[0]]
plt.figure(figsize = (13, 3))
plt.subplot(1, 3, 1)
plt.plot(t_sensor, G_test[0], 'k', label = 'Input, G(t)')
plt.axhline(0, c = 'b', linewidth = 0.8, label = 'G = 0, flow reverses below')
plt.xlim([0, 1])
plt.legend()
plt.title('The Input This Field Answers')
plt.xlabel('t')
plt.ylabel('Pressure gradient, G')
for j in range(2):
plt.subplot(1, 3, j + 2)
plt.imshow(fields[j].T, cmap = 'rainbow', origin = 'lower', aspect = 'auto', extent = [0, 1, 0, 1], vmin = -span, vmax = span)
plt.colorbar(label = 'Velocity, u')
plt.title(titles[j])
plt.xlabel('t')
plt.ylabel('y')
plt.tight_layout()
plt.show()
The field is easier to judge as a set of profiles. Each curve below is the velocity across the gap at one instant.
A contour plot shows the whole field at once but hides the shape of the profile; a set of profiles shows the shape but only at a few instants. Both are worth looking at, which is why both are here.
k_list = [11, 34, 68, 114]
lim = 1.1*np.abs(u_test[0]).max()
plt.figure(figsize = (10, 3))
for j, k in enumerate(k_list):
plt.subplot(1, 4, j + 1)
plt.plot(u_test[0, k], y_eval, 'k', alpha = 0.3, linewidth = 4, label = 'Reference')
plt.plot(u_pred[0, k], y_eval, 'r--', label = 'DeepONet')
plt.xlim([-lim, lim])
plt.title('t = %.2f' % t_eval[k])
plt.xlabel('Velocity, u')
if j == 0:
plt.ylabel('y')
plt.legend()
plt.tight_layout()
plt.show()
The four panels above are four instants out of 115. Below is the same test input in motion: the marker on the left rides along $G(t)$, and the profile on the right is the velocity across the gap at that instant, with every instant drawn faintly behind it.
Watch the profile cross to the other side of zero after $G$ does, not with it. For this input the pressure gradient changes sign at $t = 0.49$ and the velocity at mid-gap follows at $t = 0.58$. That delay is the fluid's inertia: the gradient sets the acceleration and not the velocity, so the profile is always catching up with the forcing, and it is the part a set of still frames does not show.
frames = range(0, len(t_eval), 3)
fig, (ax_g, ax_u) = plt.subplots(1, 2, figsize = (9, 3))
j = 0
ax_g.plot(t_sensor, G_test[j], 'b')
ax_g.axhline(0, c = 'k', linewidth = 0.8)
mark, = ax_g.plot([], [], 'bo')
ax_g.set_xlim([0, 1])
ax_g.set_xlabel('t')
ax_g.set_ylabel('Pressure gradient, G')
ax_g.set_title('Input, G(t)')
for k in frames:
ax_u.plot(u_test[j, k], y_eval, c = '0.85', linewidth = 0.6)
line_ref, = ax_u.plot([], [], 'k', alpha = 0.3, linewidth = 4, label = 'Reference')
line_pred, = ax_u.plot([], [], 'r--', label = 'DeepONet')
clock = ax_u.text(0.04, 0.90, '', transform = ax_u.transAxes)
ax_u.set_xlim([-lim, lim])
ax_u.set_xlabel('Velocity, u')
ax_u.set_ylabel('y')
ax_u.set_title('Velocity profile, u(y, t)')
ax_u.legend(loc = 'lower right')
plt.tight_layout()
plt.close()
def update(k):
mark.set_data([t_eval[k]], [np.interp(t_eval[k], t_sensor, G_test[j])])
line_ref.set_data(u_test[j, k], y_eval)
line_pred.set_data(u_pred[j, k], y_eval)
clock.set_text('t = %.2f' % t_eval[k])
anim = FuncAnimation(fig, update, frames = frames)
HTML(anim.to_jshtml(fps = 12))
Check Against the Fully Developed Limit
Hold $G$ at a constant value and the flow should relax onto the textbook parabola $u = \frac{G}{2}y(1-y)$. The network was never trained on a constant $G$ - every training function varied - so this is an honest test of whether it learned the physics rather than the family.
G_const = torch.ones(len(y_eval), m).to(device) # G = 1 at every sensor, repeated for every y
yt_late = torch.tensor(np.column_stack([y_eval, np.ones(len(y_eval))]), dtype = torch.float32).to(device) # (y, t = 1)
with torch.no_grad():
u_const = deeponet(G_const, yt_late).cpu().numpy().ravel()
plt.figure(figsize = (6, 4))
plt.plot(0.5*y_eval*(1 - y_eval), y_eval, 'k', alpha = 0.3, linewidth = 4, label = 'Textbook parabola, G y(1 - y)/2')
plt.plot(u_const, y_eval, 'r--', label = 'DeepONet')
plt.legend()
plt.xlabel('Velocity, u')
plt.ylabel('y')
plt.title('Fully Developed Profile for Constant G')
plt.grid(True)
plt.show()
3.3. What This Lab Showed¶
The trunk input went from $t$ to $(y, t)$ and the output became a field. Nothing else in the construction moved — the same branch net, the same inner product, the same residual-only loss — and the mean relative $L^2$ error over ten unseen pressure gradients is about two percent. Extending an operator from a curve to a field is a change of one number in the trunk's input width.
Two results should be kept separate. The first is that the operator reproduces the reference field for inputs drawn from the training family, which, as in Lab 1, is an interpolation test. The second is stronger. Held at a constant $G$ — an input unlike every training function, all of which varied — the network returns the fully developed parabola $u = \frac{G}{2}y(1-y)$. That is evidence that what was learned is the equation rather than the family, and it is the kind of test worth designing deliberately, because a model that had merely memorized its training inputs would fail it.
The physics the field shows is a lag. The pressure gradient sets the acceleration and not the velocity, so the profile always trails the forcing: in the case animated above, the pressure gradient changes sign at $t = 0.49$ and the velocity at mid-gap does not reverse until $t = 0.58$. A reader who expected the flow to turn the instant $G$ does would be treating the problem as quasi-steady, which it becomes only in the constant-$G$ limit the last check examines.
4.1. Problem Setup¶
A 2-DOF coupled mass-spring-damper system is subjected to a family of sinusoidal input forces,
$$ F(t)=A\sin(\omega t). $$
The branch network receives the discretized force history and the trunk network receives time. One DeepONet with two outputs predicts $x_1(t)$ and $x_2(t)$.
This lab differs from the first two in two ways. The system has two coupled degrees of freedom, so the operator has two outputs and two residuals; and the loss is hybrid: measured responses supply a data term, and the equations of motion and the initial conditions supply a physics term, both minimized together.
The equation of motion for the 2-DOF coupled mass-spring-damper system is
$$ M\ddot{\mathbf{x}}+C\dot{\mathbf{x}}+K\mathbf{x} = \begin{bmatrix} F(t)\\ 0 \end{bmatrix}, \qquad \mathbf{x} = \begin{bmatrix} x_1\\ x_2 \end{bmatrix}. $$
The initial conditions are
$$ x_1(0)=x_2(0)=\dot{x}_1(0)=\dot{x}_2(0)=0. $$
import numpy as np
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
CUDA
A GPU is worth having here: the DeepONet is trained for five thousand steps, and every step differentiates the outputs twice.
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(device)
m1, m2 = 1.0, 1.0 # masses
c1, c2 = 0.2, 0.1 # dampers
k1, k2 = 4.0, 2.0 # springs
M = np.array([[m1, 0], [0, m2]])
C = np.array([[c1 + c2, -c2], [-c2, c2]])
K = np.array([[k1 + k2, -k2], [-k2, k2]])
T_end = 20.0
time_scale = 10.0 # the trunk reads y = t/time_scale
wn = np.sort(np.sqrt(np.linalg.eigvals(np.linalg.solve(M, K))).real)
print('1st natural frequency : %.4f rad/s' % wn[0])
print('2nd natural frequency : %.4f rad/s' % wn[1])
Generate Training and Test Data
The responses are generated here by integrating the equations of motion from rest. Each loading case is $F(t)=A\sin(\omega t)$ with $A\sim\mathcal{U}(0.5,\,1.5)$ and $\omega\sim\mathcal{U}(0.5,\,3.0)$ rad/s, a band that contains both natural frequencies. One hundred cases are solved for training. A further three hundred forces are drawn from the same family but not solved: they have no measured response and enter only the physics terms of the loss. The test set, drawn with its own random generator, holds eight more random cases plus two resonance cases, $A=1$ with $\omega=\omega_1$ and $\omega=\omega_2$.
Unlike Labs 1 and 2, this lab does use measured responses — that is what supplies the data term. The simulated responses play the role of measurements here, and the equations of motion are added on top of them. Because the residual needs only the force, it can be evaluated on inputs for which no measurement exists, and those extra forces fill the gaps between the measured ones at the cost of nothing but random numbers.
from scipy.integrate import solve_ivp
n_sensor = 100
force_t = np.linspace(0, T_end, n_sensor) # the sensor times of the force (and of the responses)
def simulate(A, w): # F, x1, x2 at the sensor times for F(t) = A sin(w t), starting from rest
def rhs(t, z):
x, v = z[:2], z[2:]
a = np.linalg.solve(M, np.array([A*np.sin(w*t), 0.0]) - C @ v - K @ x)
return np.concatenate([v, a])
sol = solve_ivp(rhs, (0, T_end), np.zeros(4), t_eval = force_t, rtol = 1e-8, atol = 1e-10)
return A*np.sin(w*force_t), sol.y[0], sol.y[1]
def make_dataset(params): # params: (cases, 2) rows of (A, w)
F, x1, x2 = zip(*[simulate(A, w) for A, w in params])
return np.array(F), np.array(x1), np.array(x2)
rng = np.random.default_rng(0)
train_params = np.column_stack([rng.uniform(0.5, 1.5, 100), rng.uniform(0.5, 3.0, 100)])
phys_params = np.column_stack([rng.uniform(0.5, 1.5, 300), rng.uniform(0.5, 3.0, 300)]) # forces only, never solved
rng_test = np.random.default_rng(123)
test_params = np.column_stack([rng_test.uniform(0.5, 1.5, 8), rng_test.uniform(0.5, 3.0, 8)])
test_params = np.vstack([test_params, [[1.0, wn[0]], [1.0, wn[1]]]]) # the last two: 1st and 2nd mode resonance
F_train, x1_train, x2_train = make_dataset(train_params)
F_test, x1_test, x2_test = make_dataset(test_params)
F_phys = phys_params[:, :1]*np.sin(phys_params[:, 1:]*force_t) # (300, 100): forces for the physics terms only
time_train = force_t.reshape(-1, 1) # the times at which the responses are recorded
time_test = force_t.reshape(-1, 1)
print('train forces :', F_train.shape)
print('train responses :', x1_train.shape, x2_train.shape)
print('physics forces :', F_phys.shape)
print('test forces :', F_test.shape)
print('test responses :', x1_test.shape, x2_test.shape)
Define Desired Output Locations
Time is divided by a constant before it reaches the trunk net, so that the coordinate the network sees is of order one rather than of order twenty. This is the same piece of hygiene as feature scaling in ordinary regression, and it matters as much: a tanh network asked about an input of 20 is saturated before it begins.
Y = torch.tensor(time_train/time_scale, dtype = torch.float32).to(device) # (100, 1): normalized times
Y_test = torch.tensor(time_test/time_scale, dtype = torch.float32).to(device)
External Loading Conditions
Five of the 100 training forces are shown below. A hundred functions is a small family, but the network sees each of them at many different times, and it is those condition-time pairs, not the functions alone, that supply the training signal.
plt.figure(figsize = (6, 4))
for i in range(5):
plt.plot(force_t, F_train[i], alpha = 0.7)
plt.xlabel('Time (s)')
plt.ylabel('Force')
plt.title('Example Training Forces')
plt.grid(True)
plt.show()
Data Batch Generator
The same two-dimensional sampling as in Lab 1 — which function, and at what time — except that here the batch carries measured displacements as targets rather than a residual to be driven to zero.
U = torch.tensor(F_train, dtype = torch.float32).to(device) # (100, 100): the force at 100 sensors
GT = torch.tensor(np.stack([x1_train, x2_train], axis = -1), dtype = torch.float32).to(device) # (100, 100, 2): x1 and x2 at the recorded times
U_phys = torch.tensor(np.vstack([F_train, F_phys]), dtype = torch.float32).to(device) # (400, 100): every force the physics terms may use
def get_batch(batch_size = 128):
i = torch.randint(0, len(U), (batch_size,)) # which force
j = torch.randint(0, len(Y), (batch_size,)) # which time
return U[i], Y[j], GT[i, j]
Define DeepONet
The branch net reads the 100-point force history and the trunk net reads normalized time, to which Fourier features are added: $\sin(k\pi y)$ and $\cos(k\pi y)$ for $k = 1, \dots, 10$ alongside the raw coordinate.
The reason is the shape of the answer. A plain multilayer network is biased towards smooth, slowly varying functions and struggles to represent an oscillation of many periods, which is exactly what a forced vibration over twenty seconds is. Handing it the oscillations as inputs removes the difficulty rather than asking the optimizer to overcome it.
n_freq = 10 # number of Fourier features
p = 32 # number of basis functions per output
def fourier(y): # y, sin(k pi y), cos(k pi y) for k = 1, ..., n_freq
k = np.pi*torch.arange(1, n_freq + 1, device = device)
return torch.cat([y, torch.sin(y*k), torch.cos(y*k)], dim = 1)
def new_net(n_in, n_out):
return nn.Sequential(nn.Linear(n_in, 64), nn.Tanh(),
nn.Linear(64, 64), nn.Tanh(),
nn.Linear(64, n_out)).to(device)
class DeepONet(nn.Module):
def __init__(self):
super().__init__()
self.branch = new_net(100, 2*p) # force history -> 2p coefficients
self.trunk = new_net(1 + 2*n_freq, 2*p) # Fourier features of y -> 2p basis values
def forward(self, u, y): # u and y paired row by row -> (pairs, 2): x1 and x2
b1, b2 = self.branch(u).split(p, dim = 1)
t1, t2 = self.trunk(fourier(y)).split(p, dim = 1)
x1 = torch.sum(b1*t1, dim = 1, keepdim = True)
x2 = torch.sum(b2*t2, dim = 1, keepdim = True)
return torch.cat([x1, x2], dim = 1)
Define Functions to Compute Physics Loss
The model is fitted to the displacement measurements and, at the same time, enforces the two equations of motion and the zero initial conditions. The physics terms draw on all four hundred forces, the hundred that were solved together with the three hundred that were not.
The residual of a second-order system needs two derivatives of each output, and the two equations are coupled, so $x_1$ and $x_2$ appear in each other's residuals. The initial conditions constrain both the displacements and the velocities at $t = 0$.
def derivative(f, x):
return torch.autograd.grad(f, x, torch.ones_like(f), create_graph = True)[0]
def sample_at(u, t): # the force at time t by linear interpolation between the sensor values, (pairs, 1)
s = t[:, 0]/T_end*(u.shape[1] - 1) # position of t on the sensor grid
i = torch.clamp(s.long(), max = u.shape[1] - 2)
w = (s - i).reshape(-1, 1)
row = torch.arange(len(u), device = device)
return (1 - w)*u[row, i].reshape(-1, 1) + w*u[row, i + 1].reshape(-1, 1)
def ode_loss(model, n_pairs = 100):
u = U_phys[torch.randint(0, len(U_phys), (n_pairs,))] # which force, measured or not
y = 2*torch.rand(n_pairs, 1, device = device, requires_grad = True) # which time, y = t/time_scale in [0, 2]
x1, x2 = model(u, y).split(1, dim = 1)
x1_t = derivative(x1, y)/time_scale
x2_t = derivative(x2, y)/time_scale
x1_tt = derivative(x1_t, y)/time_scale
x2_tt = derivative(x2_t, y)/time_scale
f = sample_at(u, y*time_scale)
r1 = m1*x1_tt + (c1 + c2)*x1_t - c2*x2_t + (k1 + k2)*x1 - k2*x2 - f
r2 = m2*x2_tt - c2*x1_t + c2*x2_t - k2*x1 + k2*x2
return torch.mean(r1**2) + torch.mean(r2**2)
def ic_loss(model, n_pairs = 100):
u = U_phys[torch.randint(0, len(U_phys), (n_pairs,))]
y0 = torch.zeros(n_pairs, 1, device = device, requires_grad = True)
x1, x2 = model(u, y0).split(1, dim = 1)
x1_t = derivative(x1, y0)/time_scale
x2_t = derivative(x2, y0)/time_scale
return torch.mean(x1**2) + torch.mean(x2**2) + torch.mean(x1_t**2) + torch.mean(x2_t**2) # at rest at t = 0
Train
The loss has three terms: the data term on the measured $x_1$ and $x_2$, the residual of the equations of motion at random times, and the initial conditions. The weights $0.1$ and $10$ balance the three terms; the initial-condition term gets the largest weight because it is enforced at a single instant only, while the residual is enforced across the whole time window.
def train(model, epochs = 5000):
optimizer = torch.optim.Adam(model.parameters(), lr = 2e-3)
scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma = 0.05**(1/epochs)) # 2e-3 -> 1e-4 over the run
history = []
for epoch in range(1, epochs + 1):
u, y, gt = get_batch()
loss_data = torch.mean((model(u, y) - gt)**2, dim = 0).sum() # x1 and x2 errors added
loss_ode = ode_loss(model)
loss_ic = ic_loss(model)
loss = loss_data + 0.1*loss_ode + 10*loss_ic
optimizer.zero_grad()
loss.backward()
optimizer.step()
scheduler.step()
if epoch % 100 == 0:
history.append([loss_data.item(), loss_ode.item(), loss_ic.item()])
if epoch % 1000 == 0:
print('epoch %6d | loss %.4e | data %.4e | ode %.4e | ic %.4e'
% (epoch, loss.item(), loss_data.item(), loss_ode.item(), loss_ic.item()))
return np.array(history)
Physics-informed DeepONet
torch.manual_seed(0)
model = DeepONet()
history = train(model)
Training History
The three terms are plotted separately. A data term that keeps falling while the residual stalls, or the reverse, would mean that one part of the loss is being bought at the expense of another; here all three fall, the residual of the equations of motion the least smoothly.
plt.figure(figsize = (6, 4))
epochs = 100*np.arange(1, len(history) + 1)
plt.semilogy(epochs, history[:, 0], label = 'Data')
plt.semilogy(epochs, history[:, 1], label = 'Equations of motion')
plt.semilogy(epochs, history[:, 2], label = 'Initial conditions')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Training History')
plt.legend()
plt.show()
Evaluation
The model is evaluated on unseen loading conditions, including resonance cases.
Resonance is the interesting case. A forcing frequency near a natural frequency produces the largest responses of the whole family, and the response grows over the whole time window instead of settling into a steady oscillation.
U_test = torch.tensor(F_test, dtype = torch.float32).to(device) # (10, 100)
x_test = np.stack([x1_test, x2_test], axis = -1) # (10, times, 2)
def predict(model, u): # one force at every test time, (times, 2)
with torch.no_grad():
return model(u.repeat(len(Y_test), 1), Y_test).cpu().numpy()
def relative_error(pred, true): # one number per output
return np.linalg.norm(pred - true, axis = 0)/np.linalg.norm(true, axis = 0)
error = np.array([relative_error(predict(model, U_test[i]), x_test[i]) for i in range(10)]) # (10, 2)
print('relative L2 error')
print('8 unseen loading cases (mean) | x1 %.4f | x2 %.4f' % tuple(error[:8].mean(axis = 0)))
print('1st mode resonance | x1 %.4f | x2 %.4f' % tuple(error[-2]))
print('2nd mode resonance | x1 %.4f | x2 %.4f' % tuple(error[-1]))
Test Response
def plot_case(i, title):
pred = predict(model, U_test[i])
plt.figure(figsize = (9, 3))
for k in range(2):
plt.subplot(1, 2, k + 1)
plt.plot(time_test, x_test[i, :, k], 'k', alpha = 0.3, linewidth = 4, label = 'Reference')
plt.plot(time_test, pred[:, k], 'r--', label = 'Physics-informed DeepONet')
plt.xlabel('Time (s)')
plt.ylabel('$x_%d$' % (k + 1))
plt.title('%s, mass %d' % (title, k + 1))
plt.grid(True)
if k == 0:
plt.legend()
plt.tight_layout()
plt.show()
plot_case(0, 'Unseen loading')
plot_case(-2, '1st mode resonance')
plot_case(-1, '2nd mode resonance')
4.3. What This Lab Showed¶
One DeepONet with two outputs learns the map from a sinusoidal force history to the coupled displacements $x_1(t)$ and $x_2(t)$. Its loss combines the measured responses with the coupled equations of motion and the initial conditions, the hybrid recipe that returns in Lab 5. The physics terms are evaluated on three hundred additional forces that were never solved, which is the practical advantage of the hybrid loss: measurements are expensive, input functions are free. With them, the mean error on unseen loads falls to about five percent. Fourier features are kept as part of the trunk representation because they help represent the oscillatory response.
5. Lab 4: Elastic Deformation for Thin Plate (PI-DeepONet)¶
Overview
The purpose of this laboratory is the physics-informed DeepONet in its purest form, trained with no solved examples at all. The data-driven laboratories of Deep Neural Operators had the comfortable property that the ground truth could be written down analytically or computed by a solver. This laboratory changes two things at once, and the section is organised around that pair.
The first change is in the output. The domain is a two-dimensional region, so the trunk network takes a coordinate pair $(x, y)$, and the solution has two components rather than one, since a point in a stretched plate moves both along and across the direction of pull. Nothing conceptual changes — the trunk is still a function of position and the branch is still a function of the sampled input — but the bookkeeping grows.
The second change is the important one. This is a PI-DeepONet, trained with no data whatsoever. There is no reference solution in the objective, no solver output, and no measurement. The entire supervision comes from the governing equations imposed at interior collocation points and from the boundary conditions imposed on the four edges. This is the construction of Physics-informed Neural Networks (PINN) applied at the operator level, exactly as Section 2.4 of Deep Neural Operators described: the residual replaces the label, and because the residual can be evaluated for any input function at any point, the requirement for solved examples disappears. What does not disappear is the need to choose which input functions to train over, and that choice still determines where the trained operator can be trusted.
A numerical solution is used at the end, but only for comparison. It plays no part in training.
5.1. Problem Setup¶
We solve the thin plate equations for the displacement and stress distribution of a square elastic plate under a non-uniform stretching force, using PI-DeepONet with physics only — the PDE and the boundary conditions — and no data.
Plane-stress linear elasticity
The thin plate is modelled as a two-dimensional, in-plane linear elastic solid under plane stress ($\sigma_{zz}=\sigma_{xz}=\sigma_{yz}=0$). The unknowns are the in-plane displacements $u(x,y)$ and $v(x,y)$. The equilibrium equations and plane-stress constitutive relations below determine the displacement and stress fields.
Geometry and loading
A non-uniform stretching force is applied to a square elastic plate. Because both the geometry and the in-plane forces are symmetric about the two centre lines, only one quarter of the plate need be considered — the yellow domain in the figure. This is a standard and worthwhile reduction: it reduces the computational domain and converts two of the four edges into symmetry planes with simple boundary conditions. On the plane $x = 0$ the material cannot move across the symmetry line, so $u = 0$ there, and on $y = 0$ likewise $v = 0$. On both symmetry edges, the shear traction also vanishes ($\sigma_{xy}=0$). These are symmetry conditions, not physical clamps.
Problem properties
The problem in the figure is stated in physical units: $E = 50$ MPa, $\nu = 0.3$, a full width of $20$ mm so that the quarter plate is $a=b=10$ mm square, and a thickness of $1$ mm. The code retains these physical units: coordinates and displacements are in millimetres and tractions are in MPa. The problem is
$$ \Omega = [0,a]\times[0,b],\qquad a=b=10\ \mathrm{mm},\qquad \sigma_{xx}(a,y)=P(y),\qquad \sigma_{xy}(a,y)=0. $$
Here $P(y)$ is a smooth, low-frequency traction profile sampled at $250$ sensors along the loaded edge. The supplied arrays contain $128$ training profiles and $32$ held-out test profiles. Both the shape and magnitude vary across the family, rather than only the amplitude of a single cosine.
Governing equations
The in-plane equilibrium equations for the isotropic elastic plate, in the displacements in physical units $u$ and $v$, are
$$ \begin{align*} &{E \over 1 - \nu^2}\left({\partial^2 u \over \partial x^2} + {1 - \nu \over 2}{\partial^2 u \over \partial y^2} + {1 + \nu \over 2}{\partial^2 v \over \partial x \partial y} \right) = 0\\\\ &{E \over 1 - \nu^2}\left({\partial^2 v \over \partial y^2} + {1 - \nu \over 2}{\partial^2 v \over \partial x^2} + {1 + \nu \over 2}{\partial^2 u \over \partial x \partial y} \right) = 0 \end{align*} $$
These are not arbitrary. They are the two equilibrium equations $\partial\sigma_{xx}/\partial x + \partial\sigma_{xy}/\partial y = 0$ and $\partial\sigma_{xy}/\partial x + \partial\sigma_{yy}/\partial y = 0$ with the plane-stress constitutive relations substituted in, and the derivation is short enough to be worth doing once. Using
$$ \sigma_{xx} = \frac{E}{1-\nu^2}\left(\frac{\partial u}{\partial x} + \nu \frac{\partial v}{\partial y}\right), \qquad \sigma_{xy} = \frac{E}{2(1+\nu)}\left(\frac{\partial v}{\partial x} + \frac{\partial u}{\partial y}\right) $$
the first equilibrium equation becomes
$$ \frac{E}{1-\nu^2}\left(\frac{\partial^2 u}{\partial x^2} + \nu \frac{\partial^2 v}{\partial x \partial y}\right) + \frac{E}{2(1+\nu)}\left(\frac{\partial^2 v}{\partial x \partial y} + \frac{\partial^2 u}{\partial y^2}\right) = 0 $$
Since $1 - \nu^2 = (1-\nu)(1+\nu)$, the second prefactor can be written as $\frac{E}{1-\nu^2}\cdot\frac{1-\nu}{2}$, and factoring $E/(1-\nu^2)$ out of both terms leaves
$$ \frac{\partial^2 u}{\partial x^2} + \frac{1-\nu}{2}\frac{\partial^2 u}{\partial y^2} + \left(\nu + \frac{1-\nu}{2}\right)\frac{\partial^2 v}{\partial x \partial y} = 0 $$
and $\nu + \frac{1-\nu}{2} = \frac{2\nu + 1 - \nu}{2} = \frac{1+\nu}{2}$, which is the first equation above. The second follows by exchanging the roles of $x$ and $y$. Recognising the equations as equilibrium plus constitution is what makes the boundary conditions below intelligible, since they are also statements about stress.
Boundary conditions
Two symmetry conditions on the planes $x = 0$ and $y = 0$ (B.C.① and B.C.②), each fixing the normal displacement and the shear traction:
$$ v(x,y) = 0, \quad \sigma_{xy} = 0 \qquad \text{at} \quad y = 0\\\\ u(x,y) = 0, \quad \sigma_{xy} = 0 \qquad \text{at} \quad x = 0 $$
Two free-surface conditions at $y = b$ (B.C.③), where the plate edge carries no traction:
$$ \sigma_{yy} = 0,\quad \sigma_{yx} = 0 $$
And at $x = a$ (B.C.④), the loaded edge, a prescribed normal traction together with a vanishing shear:
$$ \sigma_{xx} = P(y),\quad \sigma_{xy} = 0 $$
Note the two kinds of condition. A condition on a displacement, such as $u = 0$ on the symmetry plane, constrains the solution value directly and can be enforced by penalising the network output. A condition on a stress, which is a combination of first derivatives of the displacements, requires differentiating the network output with respect to the trunk coordinates; this is exactly the situation Physics-informed Neural Networks (PINN) called a Neumann condition, and it is handled the same way here. The symmetry planes carry one condition of each kind, and the free and loaded edges carry stress conditions only.
What is being learned
The input function is the in-plane traction distribution along the loaded edge $x = a$, and it is presented to the branch network as its values at a fixed set of sensor points running along that edge. The output functions are the two displacement components $u(x,y)$ and $v(x,y)$ over the quarter plate, from which the stress components follow by differentiation. Because there are two output fields, the DeepONet has two outputs, each with its own basis and its own coefficients, as Section 5.2 shows.
The neural network and the loss functions are arranged as follows.
Numerical solution
A finite element solution for the first held-out traction profile is used below, with the displacements $u$ and $v$ in the $x$ and $y$ directions. We solve the same problem with PI-DeepONet and compare.
Be clear about the role of this solution. It is not training data. It never enters the loss, and the operator is trained before it is ever consulted. It exists so that the physics-only prediction can be judged against an independent computation, which is the only way to assess accuracy in the absence of an analytical solution — and, as Section 2.4 of Deep Neural Operators warned, it is a luxury unavailable at the moment when a surrogate is actually being used in anger.
The FEM displacement fields and the matching operator predictions are plotted together in Section 5.2.
Define properties
The material properties and quarter-plate dimensions are kept in physical units.
import numpy as np
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(device)
E, nu = 50.0, 0.3 # Young's modulus (MPa), Poisson's ratio
a, b = 10.0, 10.0 # quarter-plate dimensions (mm)
h = 1.0 # thickness (mm)
Generate collocation points
A physics-informed model needs points at which to evaluate its residuals, and the points fall into two groups with different jobs. The interior points are where the equilibrium equations are enforced; a uniform grid over the quarter plate is used. The edge points are where the four boundary conditions are enforced, and they must be generated separately for each edge because each edge carries a different condition.
This is the same sampling structure as in Physics-informed Neural Networks (PINN), with one difference that is easy to overlook: here collocation points are paired with sampled input functions during training, so the same geometric point serves as a residual location for the whole family of loads. The plot shows a dense grey interior grid with four coloured edges.
Nx, Ny = 50, 50
x, y = np.linspace(0, a, Nx), np.linspace(0, b, Ny)
xx, yy = np.meshgrid(x, y)
xy = np.column_stack([xx.ravel(), yy.ravel()])
left = np.column_stack([np.zeros(Ny), y])
bottom = np.column_stack([x, np.zeros(Nx)])
top = np.column_stack([x, b*np.ones(Nx)])
right = np.column_stack([a*np.ones(Ny), y])
plt.figure(figsize = (5, 5))
plt.scatter(xy[:, 0], xy[:, 1], s = 4, c = 'gray', alpha = 0.3, label = 'Domain')
for points, name in zip([left, bottom, top, right], ['Left', 'Bottom', 'Top', 'Right']):
plt.scatter(points[:, 0], points[:, 1], s = 8, label = name)
plt.xlabel('x (mm)')
plt.ylabel('y (mm)')
plt.title('Collocation Points')
plt.axis('equal')
plt.legend()
plt.show()
Load $P(y)$
The supplied smooth traction profiles define the family of problems: $128$ training loads and $32$ test loads, each sampled at $250$ sensor points along the loaded edge. The files below also include the FEM test solutions, used only for the final comparison.
These arrays are the sensor representation of the input function, and they are the only thing that distinguishes one member of the problem family from another. Everything else — the geometry, the material, the equations, the form of the boundary conditions — is shared.
%pip -q install gdown
!gdown 13FJ5XIYTI6lvVgUGrZC_vbCZOXk5EWHh -O tp_P_train.npy
!gdown 1XsIB6nmnLFjPVxeMw684mjtcXPilSiZG -O tp_P_test.npy
!gdown 1i8MquyDRV8hVjRTkS_2-YB6SZQ4ysiXd -O fem_test.npz
P_train = np.load('tp_P_train.npy')
P_test = np.load('tp_P_test.npy')
m = P_train.shape[1] # 250 traction sensors
sensor_y = np.linspace(0, b, m)
print('Training loads :', P_train.shape)
print('Test loads :', P_test.shape)
Plot
The first ten training profiles are plotted together with the first test profile. Each is the traction along the loaded edge, sampled at a fixed set of points. Both the profile shape and its magnitude vary across the family.
plt.figure(figsize = (6, 4))
plt.plot(sensor_y, P_train[:10].T, 'b', alpha = 0.3)
plt.plot(sensor_y, P_test[0], 'r', label = 'Test case 0')
plt.xlabel('y (mm)')
plt.ylabel('p(y) (MPa)')
plt.title('External Loading Conditions')
plt.legend()
plt.grid(True)
plt.show()
Generate dataset
The dataset here is unlike the data-driven laboratories of Deep Neural Operators in one decisive respect: it contains conditions and locations but no ground truth. There is nothing to compare against, because the supervision will come from the equations.
It is also too large to be used whole. $128$ loads at $2500$ interior points are $320000$ condition-point pairs, each needing second derivatives, so every iteration works on a random batch instead. The unit of training is the pair — one input function together with one point — and a batch mixes loads and positions freely, so that one iteration enforces the equations for many members of the family at once.
U = torch.tensor(P_train, dtype = torch.float32).to(device) # branch inputs only, no FEM labels
U_test = torch.tensor(P_test, dtype = torch.float32).to(device)
Y = torch.tensor(xy, dtype = torch.float32).to(device)
Y_left = torch.tensor(left, dtype = torch.float32).to(device)
Y_bottom = torch.tensor(bottom, dtype = torch.float32).to(device)
Y_top = torch.tensor(top, dtype = torch.float32).to(device)
Y_right = torch.tensor(right, dtype = torch.float32).to(device)
def get_batch(points = Y, batch_size = 2500):
i = torch.randint(len(U), (batch_size,), device = device)
j = torch.randint(len(points), (batch_size,), device = device)
return U[i], points[j].detach().requires_grad_(True)
def sample_at(u, y): # linear interpolation of the traction
s = (y[:, 0]/b*(m - 1)).clamp(0, m - 1)
i = s.long().clamp(max = m - 2)
w = s - i
row = torch.arange(len(u), device = device)
return ((1 - w)*u[row, i] + w*u[row, i + 1]).reshape(-1, 1)
Define the neural networks and hyper-parameters
Two DeepONets produce the displacement components: each has its own branch and trunk, and their inner products give $u$ and $v$. The two fields are trained together because the equilibrium equations couple $u$ and $v$ through the mixed derivative terms and neither can be corrected without regard to the other.
The networks here are considerably larger than in the data-driven laboratories of Deep Neural Operators, in depth, in width, and in the number of learned basis functions. That is a reasonable response to a harder problem — a two-dimensional domain, two coupled fields, and a solution family that must satisfy the equations pointwise rather than merely pass through supplied labels — though as always the capacity is chosen by experiment rather than derived.
def new_net(n_in, n_out):
net = nn.Sequential(nn.Linear(n_in, 100), nn.Tanh(),
nn.Linear(100, 100), nn.Tanh(),
nn.Linear(100, 100), nn.Tanh(),
nn.Linear(100, 100), nn.Tanh(),
nn.Linear(100, 100), nn.Tanh(),
nn.Linear(100, 100), nn.Tanh(),
nn.Linear(100, n_out)).to(device)
for layer in net:
if isinstance(layer, nn.Linear):
nn.init.xavier_normal_(layer.weight)
nn.init.zeros_(layer.bias)
return net
p = 100 # number of learned basis functions
torch.manual_seed(0)
branch_u = new_net(m, p)
trunk_u = new_net(2, p)
branch_v = new_net(m, p)
trunk_v = new_net(2, p)
def deeponet(u, y, branch, trunk): # one load-point pair per row
return torch.sum(branch(u)*trunk(y), dim = 1, keepdim = True)
def displacement(u, y):
u_pred = deeponet(u, y, branch_u, trunk_u)
v_pred = deeponet(u, y, branch_v, trunk_v)
return u_pred, v_pred
parameters = (list(branch_u.parameters()) + list(trunk_u.parameters())
+ list(branch_v.parameters()) + list(trunk_v.parameters()))
optimizer = torch.optim.Adam(parameters, lr = 1e-4)
Define functions to compute the physics loss
Every residual below is built by differentiating the predicted displacements with respect to the trunk coordinates, the same automatic differentiation used throughout Physics-informed Neural Networks (PINN), applied twice where a second derivative is needed.
def derivative(f, x):
return torch.autograd.grad(f, x, torch.ones_like(f), create_graph = True)[0]
Governing equations (PDE)
The two in-plane equilibrium residuals are
$$ \begin{align*} &r_x = {E \over 1 - \nu^2}\left({\partial^2 u \over \partial x^2} + {1 - \nu \over 2}{\partial^2 u \over \partial y^2} + {1 + \nu \over 2}{\partial^2 v \over \partial x \partial y} \right) \\ &r_y = {E \over 1 - \nu^2}\left({\partial^2 v \over \partial y^2} + {1 - \nu \over 2}{\partial^2 v \over \partial x^2} + {1 + \nu \over 2}{\partial^2 u \over \partial x \partial y} \right). \end{align*} $$
The PDE function returns the two residuals without squaring or averaging them; the training loop forms the loss.
def PDE(load, y):
u, v = displacement(load, y)
u_x, u_y = derivative(u, y).split(1, dim = 1)
v_x, v_y = derivative(v, y).split(1, dim = 1)
u_xx, u_xy = derivative(u_x, y).split(1, dim = 1)
_, u_yy = derivative(u_y, y).split(1, dim = 1)
v_xx, v_xy = derivative(v_x, y).split(1, dim = 1)
_, v_yy = derivative(v_y, y).split(1, dim = 1)
r_x = (u_xx + 0.5*(1 - nu)*u_yy + 0.5*(1 + nu)*v_xy)*E/(1 - nu**2)
r_y = (v_yy + 0.5*(1 - nu)*v_xx + 0.5*(1 + nu)*u_xy)*E/(1 - nu**2)
return r_x, r_y
Symmetry boundary conditions (B.C. ① and ②)
At the bottom ($y=0$) and left ($x=0$), the normal displacement and shear traction vanish:
$$ \begin{aligned} \text{B.C. ① (bottom):}\quad &v(x,0)=0,\qquad \sigma_{xy}(x,0)=0,\\ \text{B.C. ② (left):}\quad &u(0,y)=0,\qquad \sigma_{xy}(0,y)=0. \end{aligned} $$
Traction boundary conditions (B.C. ③ and ④)
At the top ($y=b$), the edge is traction-free. At the right ($x=a$), the normal traction is the input function $P(y)$:
$$ \begin{aligned} \text{B.C. ③ (top):}\quad &\sigma_{yy}(x,b)=0,\qquad\sigma_{yx}(x,b)=0,\\ \text{B.C. ④ (right):}\quad &\sigma_{xx}(a,y)=P(y),\qquad\sigma_{xy}(a,y)=0. \end{aligned} $$
The stresses are calculated from the plane-stress constitutive relations:
$$ \begin{aligned} \sigma_{xx} &= \frac{E}{1-\nu^2}\left(\frac{\partial u}{\partial x}+\nu\frac{\partial v}{\partial y}\right),\\ \sigma_{yy} &= \frac{E}{1-\nu^2}\left(\frac{\partial v}{\partial y}+\nu\frac{\partial u}{\partial x}\right),\\ \sigma_{xy} &= \sigma_{yx}=\frac{E}{2(1+\nu)}\left(\frac{\partial v}{\partial x}+\frac{\partial u}{\partial y}\right). \end{aligned} $$
Each boundary function below returns two unsquared residuals.
def BC_bottom(load, y):
# y = 0
u, v = displacement(load, y)
u_x, u_y = derivative(u, y).split(1, dim = 1)
v_x, v_y = derivative(v, y).split(1, dim = 1)
sig_xy = (v_x + u_y)*E/(2*(1 + nu))
return v, sig_xy
def BC_left(load, y):
# x = 0
u, v = displacement(load, y)
u_x, u_y = derivative(u, y).split(1, dim = 1)
v_x, v_y = derivative(v, y).split(1, dim = 1)
sig_xy = (v_x + u_y)*E/(2*(1 + nu))
return u, sig_xy
def BC_top(load, y):
# y = b
u, v = displacement(load, y)
u_x, u_y = derivative(u, y).split(1, dim = 1)
v_x, v_y = derivative(v, y).split(1, dim = 1)
sig_yy = (v_y + nu*u_x)*E/(1 - nu**2)
sig_xy = (v_x + u_y)*E/(2*(1 + nu))
return sig_yy, sig_xy
def BC_right(load, y):
# x = a
u, v = displacement(load, y)
u_x, u_y = derivative(u, y).split(1, dim = 1)
v_x, v_y = derivative(v, y).split(1, dim = 1)
sig_xx = (u_x + nu*v_y)*E/(1 - nu**2)
sig_xy = (v_x + u_y)*E/(2*(1 + nu))
return sig_xx - sample_at(load, y[:, 1:2]), sig_xy
Train
The objective is the sum of two equilibrium residual losses and eight boundary residual losses (two on each edge). Each residual is squared and averaged explicitly in the training loop, as in the PINN laboratory:
$$ \mathcal{L} = \mathcal{L}_{\mathrm{PDE}} + \mathcal{L}_{\mathrm{BC}}. $$
There is no data loss: the 128 training profiles specify applied tractions, not reference displacement fields. The training loop samples load-coordinate pairs and passes them to the PDE and boundary functions. The FEM solutions are used only for evaluation.
The following training loop is commented out so the supplied pretrained models can be evaluated directly. Remove the comments to train from scratch. After training, save each displacement model with its branch and trunk networks using the commented torch.save commands below.
# for epoch in range(1, 150001):
# ## PDE loss
# load, Y_domain_batch = get_batch(Y)
# r_x, r_y = PDE(load, Y_domain_batch)
# loss_pde = torch.mean(r_x**2) + torch.mean(r_y**2)
# ## BC loss
# load_bottom, Y_bottom_batch = get_batch(Y_bottom)
# load_left, Y_left_batch = get_batch(Y_left)
# load_top, Y_top_batch = get_batch(Y_top)
# load_right, Y_right_batch = get_batch(Y_right)
# bottom_v, bottom_shear = BC_bottom(load_bottom, Y_bottom_batch)
# left_u, left_shear = BC_left(load_left, Y_left_batch)
# top_sigyy, top_sigxy = BC_top(load_top, Y_top_batch)
# right_sigxx, right_sigxy = BC_right(load_right, Y_right_batch)
# loss_bc = (torch.mean(bottom_v**2) + torch.mean(bottom_shear**2)
# + torch.mean(left_u**2) + torch.mean(left_shear**2)
# + torch.mean(top_sigyy**2) + torch.mean(top_sigxy**2)
# + torch.mean(right_sigxx**2) + torch.mean(right_sigxy**2))
# loss = loss_pde + loss_bc
# optimizer.zero_grad()
# loss.backward()
# optimizer.step()
# if epoch % 1000 == 0:
# print('epoch %6d | loss %.4e | pde %.4e | bc %.4e' % (epoch, loss.item(), loss_pde.item(), loss_bc.item()))
# torch.save([branch_u, trunk_u], 'tp_model_u.pt')
# torch.save([branch_v, trunk_v], 'tp_model_v.pt')
Test the pretrained model
The supplied pretrained DeepONets can be evaluated directly, without running the training loop above:
The two model files contain the learned branch and trunk networks. The displacement fields are predicted by DeepONet; stresses are derived by differentiating these fields with respect to the spatial coordinates.
!gdown 1daCjgIxFXlu6VIW2rDayFZ0OLMlzxftu -O tp_model_u.pt
!gdown 15MCQyklR-9JCD8Ce0advVtioyppMVWqj -O tp_model_v.pt
branch_u, trunk_u = torch.load("tp_model_u.pt", map_location = device, weights_only = False)
branch_v, trunk_v = torch.load("tp_model_v.pt", map_location = device, weights_only = False)
for net in [branch_u, trunk_u, branch_v, trunk_v]:
net.eval()
The trained operator is evaluated on a test load not used during training, at the FEM comparison coordinates. The trunk network accepts coordinate queries without interpolation of stored field values.
def predict(load, points):
y = torch.tensor(points, dtype = torch.float32).to(device)
with torch.no_grad():
u, v = displacement(load.reshape(1, -1).repeat(len(y), 1), y)
return torch.cat([u, v], dim = 1).cpu().numpy()
with np.load('fem_test.npz') as fem:
fem_xy = fem['xy']
reference = fem['truth'][0, :, :2]
prediction = predict(U_test[0], fem_xy)
error = np.abs(prediction - reference)
relative_error = 100*np.linalg.norm(prediction - reference, axis = 0)/np.linalg.norm(reference, axis = 0)
print('relative L2 error | u %.3f%% | v %.3f%%' % tuple(relative_error))
The panels below compare the two displacement components with the FEM test solution.
What should you expect to see? The plate is being pulled in the $x$ direction, so $u$ should increase from zero at the symmetry plane $x = 0$ to its largest value at the loaded edge. The Poisson effect should contract the plate transversely, so $v$ should be negative and grow in magnitude away from $y = 0$; a positive $v$ field would indicate a sign error somewhere.
Compare the predictions against the FEM solution. Agreement over the whole quarter plate would mean that a model which never saw a solved example has reproduced the output of a finite element analysis, using nothing but the equations and the boundary conditions, for a load it was not trained on. That is the claim of PI-DeepONet, and this is the experiment that supports it.
titles = ['FEM', 'PI-DeepONet', 'Absolute error']
for j, name in enumerate(['u', 'v']):
fields = [reference[:, j], prediction[:, j], error[:, j]]
low = min(fields[0].min(), fields[1].min())
high = max(fields[0].max(), fields[1].max())
limits = [(low, high), (low, high), (0, fields[2].max())]
plt.figure(figsize = (12, 3))
for k in range(3):
plt.subplot(1, 3, k + 1)
plt.scatter(fem_xy[:, 0], fem_xy[:, 1], c = fields[k], cmap = 'rainbow', s = 7,
vmin = limits[k][0], vmax = limits[k][1])
plt.title(titles[k] + ': $' + name + '$')
plt.xlabel('x (mm)')
plt.ylabel('y (mm)')
plt.axis('equal')
plt.colorbar(label = 'mm')
plt.tight_layout()
plt.show()
5.3. What This Lab Showed¶
This laboratory trained an operator with no solved examples at all. The whole objective was the two equilibrium residuals at interior points together with the eight boundary terms on the four edges; no finite element run entered the loss, and the test reference exists only to judge the result afterwards.
What that buys is the removal of the data-generation stage: a data-driven operator for this problem would have needed one finite element solution per training load before training could begin. What it costs is a much longer optimization and a new failure mode. The load enters the problem through one boundary term only, so an objective that satisfies the interior equations while neglecting that term describes a plate under no load at all. The trivial solution satisfies the interior homogeneous equilibrium equations, but not the nonzero applied-traction boundary condition. The boundary loss is therefore essential to exclude that solution.
Two properties of the trunk net appeared here for the first time in this notebook. Stress is not an output of the model; it is recovered by differentiating the predicted displacements, so a quantity that was never supervised is available afterwards. And because the trunk is a function of position rather than a lookup over a grid, the fields can be read at the FEM coordinates after training on collocation points taken from a fifty-by-fifty one, with no resampling step in between.
One limit should be read alongside the success. The test uses a held-out smooth traction profile from the supplied low-frequency family. It assesses changes in load shape as well as magnitude, but does not establish accuracy for arbitrary high-frequency profiles.
6. Lab 5: Flow Around a Cylinder¶
Overview
The purpose of this laboratory is to combine the two: solved data where it exists and the equations everywhere else. It is the laboratory closest to industrial practice, and it differs from the laboratories before it in three ways that are each worth stating before the details arrive.
First, the equations are nonlinear. The steady Navier-Stokes equations contain convective terms in which the velocity multiplies its own derivative, so the solution does not depend linearly on the data and cannot be assembled from a superposition of training cases. Every laboratory of Deep Neural Operators had a linear solution operator, and the honest observation in each of them — that the test case was a linear combination of the training cases and therefore easy — does not apply here. Section 2.3 of Deep Neural Operators predicted where the difficulty would land: with a fixed trunk basis, all nonlinearity of the operator must be carried by the branch network's dependence on the input.
Second, the input is a scalar material property rather than a boundary profile. What varies across the family is the viscosity $\mu$, a single number multiplying the diffusive terms of the momentum equations. Formally this is the degenerate one-sensor case again, as in Lab 1 of Deep Neural Operators, but the physical consequence is far from degenerate: changing the viscosity changes the Reynolds number, and with it the entire character of the flow.
Third, the objective is hybrid on a problem where neither half would do. The data-driven laboratories of Deep Neural Operators used data only and Lab 4 used physics only; Lab 3 already mixed the two, but on a linear system with measurements at every recorded time. Here solved fields from a CFD simulation supply a data term, the momentum and continuity residuals supply a physics term, and the boundary conditions supply a third. This is the arrangement Section 2.4 of Deep Neural Operators recommended as the usual practical choice, and it is here that the weighting between competing loss terms becomes visible.
6.1. Problem Setup¶
We solve the Navier-Stokes equations with a varying viscosity term using PI-DeepONet, trained on both
- data from a numerical solver, and
- physics, namely the PDE and the boundary conditions.
This combines PINNs with DeepONet, so that the model learns from data and from the governing equations at the same time.
Problem properties
$$ \rho = 1\,\mathrm{kg/m^3}, \quad \mu \in [0.002, 0.2]\,\mathrm{N\cdot s/m^2}, \quad D = 0.1\,\mathrm{m}, \quad h = 0.5\,\mathrm{m}, \quad L = 2\,\mathrm{m}, \quad u_{in} = 1\,\mathrm{m/s}, \quad \nu = \frac{\mu}{\rho} $$
The geometry is a channel with an obstacle: a uniform stream enters at the left, passes a circular cylinder, and leaves at the right, with solid walls above and below. The density and the inlet velocity are fixed, and the viscosity is the only quantity that changes from case to case.
That single varying quantity is what makes the family interesting, because it is the one that sets the Reynolds number $\mathrm{Re} = \rho u_{in} D / \mu$. With $\rho$ and $u_{in}$ held fixed, the Reynolds number is inversely proportional to $\mu$, so a viscosity range spanning two orders of magnitude is a Reynolds number range spanning two orders of magnitude. At the viscous end the flow creeps around the cylinder in a nearly symmetric pattern dominated by diffusion; at the other end inertia dominates, the wake lengthens, and the fields become far less smooth. The operator must cover all of it with one set of weights, and this is a much stronger demand than anything asked in the laboratories of Deep Neural Operators.
2D Navier-Stokes equations, steady state
$$ \begin{align*} \rho \left(u{\partial u \over \partial x} + v{\partial u \over \partial y} + {1 \over \rho}{\partial p \over \partial x}\right) - \mu \ \left({\partial^2 u \over {\partial x^2}} + {\partial^2 u \over {\partial y^2}}\right) &= 0\\\\ \rho \left(u{\partial v \over \partial x} + v{\partial v \over \partial y} + {1 \over \rho}{\partial p \over \partial y}\right) - \mu \ \left({\partial^2 v \over {\partial x^2}} + {\partial^2 v \over {\partial y^2}}\right) &= 0\\\\ {\partial u \over \partial x} + {\partial v \over \partial y} &= 0 \end{align*} $$
The first two lines are momentum balance in the $x$ and $y$ directions, and the third is the incompressibility constraint: what flows into any small region must flow out of it. Three residuals therefore have to be driven to zero rather than the two of Lab 4, and there are three unknown fields, $u$, $v$ and $p$, rather than two.
The terms $u\,\partial u/\partial x$ and $v\,\partial u/\partial y$ are the source of the nonlinearity. They are products of the solution with derivatives of the solution, so doubling the inlet velocity does not double the solution and two solutions cannot be added to make a third. Everything comfortable about the data-driven laboratories of Deep Neural Operators came from linearity, and none of it survives here.
Boundary conditions
Two Dirichlet conditions on the channel walls, the no-slip condition:
$$ u(x,y) = 0, \quad v(x,y) = 0 \qquad \text{at} \quad y = \pm h $$
Two Dirichlet conditions at the inlet, a uniform stream entering:
$$u(-0.5,\,y) = u_{\text{in}}, \quad v(-0.5,\,y) = 0$$
Two Dirichlet conditions at the outlet, a reference pressure with no transverse velocity:
$$p(1.5,\,y) = 0, \quad v(1.5,\,y) = 0$$
Two Dirichlet conditions on the cylinder surface, no-slip again:
$$u(\text{cylinder}) = 0, \quad v(\text{cylinder}) = 0$$
All eight are Dirichlet conditions on the solution values, so unlike Lab 4 no derivatives are needed to evaluate them. The pressure condition at the outlet deserves a note: for an incompressible flow the pressure is determined only up to an additive constant, so a reference value must be pinned somewhere or the problem is ill-posed and the optimizer will drift.
What is being learned
The operator maps the viscosity to the whole steady flow field:
$$G: \mu \; \longmapsto \; (u, v, p)$$
The input function is the viscosity, presented to the branch network as a single value, so this is the one-sensor case. The output functions are the two velocity components and the pressure, each a function of position over the channel, and one network carries all three: the branch and trunk outputs are split into three groups, one per field.
The network and the loss functions are arranged as follows: a single input neuron for the branch network, two input neurons for the trunk network.
Define conditions
Nine viscosity values are used for training and two unseen values for testing.
Look at the two arrays with the out-of-distribution warning of Section 2.4 of Deep Neural Operators in mind. The training viscosities run from $0.002$ to $0.2$, a factor of one hundred, and the two test values $0.006$ and $0.04$ each fall between two neighbouring training values — $0.006$ between $0.004$ and $0.008$, and $0.04$ between $0.02$ and $0.06$. Both tests are therefore interpolation tests, as in every laboratory of this notebook. Note also that the training values are spaced roughly geometrically rather than uniformly, which is the sensible choice when the physically meaningful quantity is the Reynolds number and the Reynolds number goes as the reciprocal of the viscosity.
import numpy as np
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(device)
train_vis = np.array([0.002, 0.004, 0.008, 0.01, 0.02, 0.06, 0.08, 0.1, 0.2]).reshape(-1, 1)
test_vis = np.array([0.006, 0.04]).reshape(-1, 1)
Where the training data comes from
To obtain the ground truth for the training data, a CFD simulation has been performed. The ground truth values are stored only at the mesh points of that simulation.
This sentence is the practical heart of the notebook and deserves to be read slowly. Each training viscosity required a complete computational fluid dynamics run: mesh generation, discretisation, iterative solution to convergence. Nine training conditions means nine such runs, performed before any learning could begin, and their cost dwarfs everything else in this laboratory. A data-driven operator network is not a way of avoiding simulation; it is a way of amortising it. The simulations are run once, offline, and the trained model then answers new viscosity queries in a forward pass. Whether that trade is worth making is an arithmetic question about how many queries you expect: for a single answer it is absurd, for a thousand-point design sweep or a real-time controller it is decisive.
The physics term in the loss softens the arithmetic but does not change its shape. Because the residual supplies supervision everywhere, fewer solved cases are needed than a purely data-driven model of comparable accuracy would require — nine, in a problem whose input range spans two decades of Reynolds number.
Download
%pip -q install gdown
!gdown 1641kHe0XHXn7b6EIbnafMhaTY3zRghtW -O viscosity_train_xy_domain.npy
!gdown 1i9QVG1wuGP7ywAlfDojqDYCMKyLpUwSH -O viscosity_train_gt_domain.npy
!gdown 13mUjqmiBSV-k8MMsgFAtqSPpAtvaWE2O -O viscosity_test_xy_domain.npy
!gdown 1uJmDx7Kk861RDwDQbQSjW6_jVTvIrKX6 -O viscosity_test_gt_domain.npy
mesh_pts = np.load('viscosity_train_xy_domain.npy')
train_gt = np.load('viscosity_train_gt_domain.npy')
test_pts = np.load('viscosity_test_xy_domain.npy')
test_gt = np.load('viscosity_test_gt_domain.npy')
print('train XY shape :', mesh_pts.shape)
print('train gt shape :', train_gt.shape)
print('test XY shape :', test_pts.shape)
print('test gt shape :', test_gt.shape)
Note
- 19,692 mesh points
- $(y_1, y_2)$, that is $(x, y)$, are the desired output locations
- for each mesh point the output values include the velocity in the $x$ direction, the velocity in the $y$ direction, and the pressure
The scale of the label array is worth registering. Nine conditions at nearly twenty thousand points each, with three fields per point, is a large supervised dataset by the standards of this notebook — and it came from nine solver runs. That ratio, a handful of expensive simulations yielding hundreds of thousands of training targets, is characteristic of scientific machine learning and is the reason data-driven operator learning is feasible at all.
Collocation points for the PDE (desired output locations)
All the mesh points of the CFD simulation are good candidates for collocation points in PI-DeepONet, and they are used as such here, ensuring the model respects the physics-informed constraints.
Reusing the mesh in this way is convenient and also slightly subtle. The same set of points now plays two roles: it is where the labels live, so it supplies the data term, and it is where the residual is evaluated, so it supplies the physics term. There is no requirement that the two coincide — the residual could be enforced at points where no label exists, and in Lab 4 that was the only option — but when a mesh is already available it is a natural choice, since a CFD mesh is refined exactly where the solution varies fastest, which is also where the residual is hardest to satisfy.
y_pts = mesh_pts[0]
Collocation points for the boundary conditions
The domain boundary has four distinct parts, each carrying its own condition: the two channel walls, the inlet, the outlet, and the cylinder surface. Points are generated on each, with the walls and the cylinder grouped together since both impose no-slip.
x_wall = np.linspace(-0.5, 1.5, 200)
top_pts = np.column_stack([x_wall, 0.5*np.ones(200)])
bottom_pts = np.column_stack([x_wall, -0.5*np.ones(200)])
y_in = np.linspace(-0.5, 0.5, 100)
inlet_pts = np.column_stack([-0.5*np.ones(100), y_in])
outlet_pts = np.column_stack([1.5*np.ones(100), y_in])
theta = np.linspace(0, 2*np.pi, 200)
cylinder_pts = 0.05*np.column_stack([np.cos(theta), np.sin(theta)])
wall_pts = np.concatenate([top_pts, bottom_pts, cylinder_pts])
plt.figure(figsize = (8, 3))
plt.plot(y_pts[:, 0], y_pts[:, 1], 'k.', markersize = 1, label = 'Mesh')
plt.plot(wall_pts[:, 0], wall_pts[:, 1], 'r.', markersize = 3, label = 'Wall')
plt.plot(inlet_pts[:, 0], inlet_pts[:, 1], 'b.', markersize = 3, label = 'Inlet')
plt.plot(outlet_pts[:, 0], outlet_pts[:, 1], 'g.', markersize = 3, label = 'Outlet')
plt.axis('scaled')
plt.xlabel('x')
plt.ylabel('y')
plt.title('Collocation Points')
plt.legend(loc = 'upper right')
plt.show()
The figure shows the mesh points of the simulation together with the boundary point sets. The hole in the middle of the point cloud is the cylinder, and the fact that the mesh is visibly denser near it is the CFD mesh generator doing its job: resolution is placed where the gradients are.
Data batch generator
The data and physics terms draw the same kind of batch, a viscosity together with a mesh location, and the label at that location comes with it. The boundary batches carry no label, because nothing was solved on the boundary: there the condition itself is the target.
All training mesh locations are used for sampling and output scaling. The two unseen test viscosities and their CFD labels are used only for the final comparison.
U = torch.tensor(train_vis, dtype = torch.float32).to(device)
U_test = torch.tensor(test_vis, dtype = torch.float32).to(device)
Y = torch.tensor(y_pts, dtype = torch.float32).to(device)
GT = torch.tensor(train_gt, dtype = torch.float32).to(device)
Y_wall = torch.tensor(wall_pts, dtype = torch.float32).to(device)
Y_inlet = torch.tensor(inlet_pts, dtype = torch.float32).to(device)
Y_outlet = torch.tensor(outlet_pts, dtype = torch.float32).to(device)
mean = GT.mean(dim = (0, 1))
scale = GT.std(dim = (0, 1))
data_scale = torch.sqrt(torch.mean(GT**2, dim = 1))
print('Training labels :', GT.shape)
def get_batch(batch_size = 128):
i = torch.randint(0, len(U), (batch_size,), device = device)
j = torch.randint(0, len(Y), (batch_size,), device = device)
return U[i], Y[j].detach().requires_grad_(True), GT[i, j], data_scale[i]
def get_bc_batch(pts, batch_size = 128):
i = torch.randint(0, len(U), (batch_size,), device = device)
j = torch.randint(0, len(pts), (batch_size,), device = device)
return U[i], pts[j]
Define the DeepONet structure
The nets are those of Lab 4, with a one-dimensional branch input, since the condition is a single viscosity; here one branch and one trunk carry all three fields, their outputs split into three groups.
The branch input is $\mu/0.2$. The three outputs are restored to physical units using the saved training mean and standard deviation before the losses are evaluated.
def new_net(n_in, n_out):
return nn.Sequential(nn.Linear(n_in, 100), nn.Tanh(),
nn.Linear(100, 100), nn.Tanh(),
nn.Linear(100, 100), nn.Tanh(),
nn.Linear(100, 100), nn.Tanh(),
nn.Linear(100, 100), nn.Tanh(),
nn.Linear(100, 100), nn.Tanh(),
nn.Linear(100, n_out)).to(device)
p = 100 # number of basis functions per field
torch.manual_seed(0)
branch = new_net(1, 3*p)
trunk = new_net(2, 3*p)
def deeponet(u, y):
b1, b2, b3 = branch(u/0.2).split(p, dim = 1) # viscosity range: 0.002 to 0.2
t1, t2, t3 = trunk(y).split(p, dim = 1)
v1 = torch.sum(b1*t1, dim = 1, keepdim = True)
v2 = torch.sum(b2*t2, dim = 1, keepdim = True)
pres = torch.sum(b3*t3, dim = 1, keepdim = True)
return torch.cat([v1, v2, pres], dim = 1)*scale + mean
optimizer = torch.optim.Adam(list(branch.parameters()) + list(trunk.parameters()), lr = 5e-4)
scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma = 0.02**(1/150000))
def PLOT(u, Y_domain, gt_domain):
y = torch.tensor(Y_domain, dtype = torch.float32).to(device)
with torch.no_grad():
pred = deeponet(u.repeat(len(y), 1), y).cpu().numpy()
for q, name in enumerate(['u', 'v', 'Pressure']):
plt.figure(figsize = (9, 2))
plt.subplot(1, 3, 1)
plt.title('CFD: ' + name)
plt.scatter(Y_domain[:, 0], Y_domain[:, 1], c = gt_domain[:, q], s = 0.5, cmap = 'jet')
plt.axis('scaled')
plt.subplot(1, 3, 2)
plt.title('Prediction: ' + name)
plt.scatter(Y_domain[:, 0], Y_domain[:, 1], c = pred[:, q], s = 0.5, cmap = 'jet')
plt.axis('scaled')
plt.subplot(1, 3, 3)
plt.title('Error: ' + name)
plt.scatter(Y_domain[:, 0], Y_domain[:, 1], c = np.abs(pred[:, q] - gt_domain[:, q]), s = 0.5, cmap = 'jet')
plt.axis('scaled')
plt.tight_layout()
plt.show()
Since the output includes the velocities in the $x$ and $y$ directions and the pressure, the branch and trunk networks output $3p$ numbers each, split into three groups of $p$, one per field. Three separate DeepONets would have served equally well; the single network keeps the code short and shares the hidden layers across the three fields.
The three fields are optimized jointly under a single loss, and they must be, because the momentum and continuity residuals couple all three fields: the convective terms multiply velocities together, the pressure gradient appears in both momentum equations, and continuity constrains the two velocity components against each other. Separately trained fields would have no mechanism for enforcing those couplings.
Define functions to compute the physics loss
The residual is evaluated at the same conditions and the same mesh points that the data term drew, so every batch constrains the fields twice over: once against the solved values and once against the equations.
Define the PDE
The three residuals of the steady two-dimensional Navier-Stokes system (with $\rho = 1$):
$$ \begin{align*} \rho \left(u{\partial u \over \partial x} + v{\partial u \over \partial y} + {1 \over \rho}{\partial p \over \partial x}\right) - \mu \ \left({\partial^2 u \over {\partial x^2}} + {\partial^2 u \over {\partial y^2}}\right) &= 0\\\\ \rho \left(u{\partial v \over \partial x} + v{\partial v \over \partial y} + {1 \over \rho}{\partial p \over \partial y}\right) - \mu \ \left({\partial^2 v \over {\partial x^2}} + {\partial^2 v \over {\partial y^2}}\right) &= 0\\\\ {\partial u \over \partial x} + {\partial v \over \partial y} &= 0 \end{align*} $$
One point about the residual is specific to the operator setting and has no counterpart in Physics-informed Neural Networks (PINN). The viscosity is not a single fixed value here; it is the scalar branch input, and it therefore appears inside the residual, multiplying the diffusive terms differently for each member of the batch. The physics term is consequently a different equation for each training condition, which is exactly what it means to enforce a PDE across a family rather than for one instance.
def derivative(f, x):
return torch.autograd.grad(f, x, torch.ones_like(f), create_graph = True)[0]
def PDE(u, y):
v1, v2, p = deeponet(u, y).split(1, dim = 1)
v1_x, v1_y = derivative(v1, y).split(1, dim = 1)
v1_xx, _ = derivative(v1_x, y).split(1, dim = 1)
_, v1_yy = derivative(v1_y, y).split(1, dim = 1)
v2_x, v2_y = derivative(v2, y).split(1, dim = 1)
v2_xx, _ = derivative(v2_x, y).split(1, dim = 1)
_, v2_yy = derivative(v2_y, y).split(1, dim = 1)
p_x, p_y = derivative(p, y).split(1, dim = 1)
rho = 1.0
mu = u
r1 = rho*(v1*v1_x + v2*v1_y) + p_x - mu*(v1_xx + v1_yy) # x-momentum
r2 = rho*(v1*v2_x + v2*v2_y) + p_y - mu*(v2_xx + v2_yy) # y-momentum
r3 = v1_x + v2_y # continuity
return r1, r2, r3
Boundary conditions
No-slip on the channel walls,
$$u(x,y) = 0, \quad v(x,y) = 0 \qquad \text{at} \quad y = \pm h$$
no-slip on the cylinder surface,
$$u(\text{cylinder}) = 0, \quad v(\text{cylinder}) = 0$$
a prescribed uniform stream at the inlet,
$$u(-0.5,\,y) = u_{\text{in}}, \quad v(-0.5,\,y) = 0$$
and a reference pressure with no transverse velocity at the outlet,
$$p(1.5,\,y) = 0, \quad v(1.5,\,y) = 0$$
def BC_wall(u, y):
v1_wall, v2_wall, _ = deeponet(u, y).split(1, dim = 1)
return v1_wall, v2_wall
def BC_inlet(u, y):
v1_inlet, v2_inlet, _ = deeponet(u, y).split(1, dim = 1)
return v1_inlet - 1, v2_inlet
def BC_outlet(u, y):
_, v2_outlet, p_outlet = deeponet(u, y).split(1, dim = 1)
return v2_outlet, p_outlet
Train
The objective now has three groups of terms: a data term comparing the three predicted fields against the CFD labels, a physics term comprising the two momentum residuals and the continuity residual, and a boundary term comprising the six conditions above.
The three groups enter with different weights, and the weights are the one set of numbers in this laboratory worth examining rather than skipping. The data term is multiplied by ten, the boundary term by a tenth, and the PDE term by a hundredth. This is not a statement that the physics matters a hundred times less than the data. It is a statement about magnitudes. The residual of a momentum equation has units and a numerical scale entirely unlike those of a velocity error, and a term that is naturally a thousand times larger will dominate the gradient and starve the others regardless of its importance. Scalarizing several objectives into one weighted sum requires the terms to be commensurable, and choosing the weights is the price of that convenience — a point Physics-informed Neural Networks (PINN) made and which becomes more acute here, because there are more terms to balance.
The run is long, and as in Lab 4 the notebook recommends against executing it and provides a pretrained model instead. If you do run it, the printed breakdown of the loss into its data, PDE and boundary components is the diagnostic to watch: what you want is for all three to decrease together, and what you should be suspicious of is any one of them collapsing while another plateaus.
The training loop is commented out by default, and the supplied pretrained networks can be loaded instead. To train the model, uncomment the loop, which runs for 150001 iterations, and save the trained networks. The data error is scaled by the training-field RMS.
# for epoch in range(1, 150001):
# u, y, gt, weight = get_batch()
# ## PDE loss
# r1, r2, r3 = PDE(u, y)
# loss_pde = torch.mean(r1**2) + torch.mean(r2**2) + torch.mean(r3**2)
# ## BC loss
# u_wall, y_wall = get_bc_batch(Y_wall)
# u_inlet, y_inlet = get_bc_batch(Y_inlet)
# u_outlet, y_outlet = get_bc_batch(Y_outlet)
# v1_wall, v2_wall = BC_wall(u_wall, y_wall)
# v1_inlet, v2_inlet = BC_inlet(u_inlet, y_inlet)
# v2_outlet, p_outlet = BC_outlet(u_outlet, y_outlet)
# loss_bc = (torch.mean(v1_wall**2) + torch.mean(v2_wall**2)
# + torch.mean(v1_inlet**2) + torch.mean(v2_inlet**2)
# + torch.mean(v2_outlet**2) + torch.mean(p_outlet**2))
# ## Data loss
# u_obs, v_obs, p_obs = deeponet(u, y).split(1, dim = 1)
# loss_data = torch.mean(((torch.cat([u_obs, v_obs, p_obs], dim = 1) - gt)/weight)**2)
# loss = 0.01 * loss_pde + 0.1 * loss_bc + 10 * loss_data
# optimizer.zero_grad()
# loss.backward()
# optimizer.step()
# scheduler.step()
# if epoch % 1000 == 0:
# print('epoch %6d | data %.4e | pde %.4e | bc %.4e' % (epoch, loss_data.item(), loss_pde.item(), loss_bc.item()))
# # Save the model
# torch.save([branch, trunk, mean, scale], 'viscosity_model.pt')
6.3. Test Pre-trained Model¶
Because training takes a long time, the pretrained networks are provided so that unseen conditions can be evaluated directly.
The comparison below takes one of the two withheld viscosities, evaluates the trained velocity and pressure operators at every mesh point of the simulation, and places the CFD result, the prediction, and their absolute difference side by side.
For each field, read the three panels in order. The first two should be visually indistinguishable at the scale of the plot: a uniform stream that accelerates around the shoulders of the cylinder and slows in its wake. The third panel is the informative one, and the question to ask of it is not how large the error is but where it is. Errors concentrated in the wake and near the cylinder surface would indicate that the model has learned the easy, smooth part of the field and struggles where gradients are steep, which is the expected failure mode and the one a finer sensor or a heavier physics weight would target. Errors spread uniformly over the domain would indicate something else entirely, most likely an underfitted model.
!gdown 1iSqVubw0SbkNrsBapWZr81uJAqGa8vVE -O viscosity_model.pt
# Load the pre-trained model
branch, trunk, mean, scale = torch.load('viscosity_model.pt', map_location = device, weights_only = False)
branch.eval()
trunk.eval()
PLOT(U_test[0], test_pts[0], test_gt[0])
6.4. Fast Computational Accelerator¶
We have demonstrated PI-DeepONet on the flow around a cylinder. For an unseen viscosity, the trained model infers the velocity and pressure fields directly.
This framework serves as a fast computational accelerator, or an AI-based surrogate model, reducing the computational cost dramatically compared with traditional numerical simulation. The animation below shows predicted velocity and pressure fields across varying viscosity in near real time.
It is worth stating the argument carefully, because this is the claim on which the entire field rests and it is easy to overstate. What has been demonstrated is that a model trained once, at considerable expense, can thereafter answer a new query in the time of a forward pass. Nothing has been demonstrated about total cost: the nine CFD runs plus the long training loop certainly cost more than nine CFD runs alone. The surrogate wins only when the number of queries is large, and it wins by exactly the margin that the query count exceeds the break-even point.
Three settings have that property, and they are the settings operator learning is really for. A design sweep evaluates hundreds of configurations to find the promising region of a parameter space. An optimization loop evaluates the same model thousands of times as it descends. A real-time controller must answer within a control interval, which no CFD run can meet at any budget. In all three, an expensive offline stage followed by a negligible online one is the correct shape of cost, and it is the shape that an instance-wise PINN could not offer at all.
And the limits of Section 2.4 of Deep Neural Operators apply here undiminished. The surrogate is trustworthy for viscosities inside the trained range and says nothing reliable outside it. It reports no error bar, and its accuracy on a new case can be checked only against a solver run, which is the thing it was built to avoid. The proper role of such a model is therefore to screen, to explore, and to steer — with a classical solver retained to verify whatever design the exploration finally selects.
Summary¶
Key takeaways
| Previous notebook | Lab 1 | Lab 2 | Lab 3 | Lab 4 | Lab 5 | |
|---|---|---|---|---|---|---|
| Operator | $C \mapsto h(t)$ | $Q_{in}(t) \mapsto h(t)$ | $G(t) \mapsto u(y,t)$ | $F(t) \mapsto x_1, x_2$ | $P(y) \mapsto u, v$ | $\mu \mapsto u, v, p$ |
| Method | PINN | PI-DeepONet | PI-DeepONet | DeepONet, data and PI | PI-DeepONet | DeepONet, data and PI |
| Branch input | — | function | function | function | function | scalar |
| Trunk input | $t$ | $t$ | $(y, t)$ | $t$ | $(x, y)$ | $(x, y)$ |
| Output | curve | curve | field | two curves | two fields | three fields |
- The input became a function and the network barely had to change. The same branch and trunk nets served Labs 1 and 2; the trunk input width went from $1$ to $2$ and nothing else moved. The later labs add outputs rather than machinery — two coupled displacements in Lab 3, two displacement fields in Lab 4, three flow fields in Lab 5.
- A trained operator answers a new input function in one forward pass. That is the whole point of the exercise: a PINN would have to be retrained for every new pump schedule, pressure gradient, forcing history or edge traction.
- The branch–trunk product is a learned basis with learned coefficients. The trunk net supplies functions of the coordinate, the branch net supplies the weights that pick one member of the solution family, and their inner product is the prediction. Nothing else is fitted.
- No solution data was needed in Labs 1, 2 and 4. Those operators were trained on the governing equation alone, and the reference solutions exist in the notebook only to grade the result. Labs 3 and 5 mix the two, because measurements and equations constrain different things.
- Physics can be added to an operator exactly as it was to a PINN. In Lab 3 the equations of motion and the initial conditions are simply further terms in the same loss as the measured responses, and because the residual needs only the force, it can be evaluated on three hundred inputs that were never solved.
- The plate lab is the physics-informed DeepONet with nothing else in the loss. Two displacement fields are coupled through the plane-stress equations and eight boundary terms, the loads form a family of smooth profiles with varying shapes and magnitudes, and the operator is trained without a single solved example and tested against FEM displacements.
- The cylinder lab is the hybrid in its practical form. Nine CFD runs supply labels where a mesh already exists, the Navier-Stokes residual constrains the fields everywhere else, and the trained operator answers a new viscosity in one pass, which is what makes it a computational accelerator; the price is a long, delicately weighted optimization.
- The tests that matter are the ones the training never saw. Every lab is judged on held-out input functions, and Lab 2 goes further by holding the pressure gradient constant — an input drawn from no training family at all — to see whether the network reproduces the textbook parabola.
Looking ahead. Every lab here learned the operator through a branch–trunk product, and the two-dimensional labs paid for it with long training runs. A different architecture takes over next: the Fourier Neural Operator, which learns the operator in the frequency domain rather than through a branch–trunk product.