Scientific Machine Learning:
PINN: Limitations and Overcoming Strategies
Learning Objectives¶
After studying this lecture, you should be able to:
- explain why a PINN can minimize a correct loss to a wrong solution, and name the four failure modes that account for most such cases,
- describe how collocation-point placement affects the resolution of localized solution features,
- explain why high-frequency solutions, singular / sharp-gradient structures, and gradient discontinuities are difficult for a standard global PINN,
- describe how Fourier feature embedding, enrichment, and domain decomposition modify the representation of difficult solution structures,
- explain why full-domain transient training does not automatically enforce the physical order of information propagation,
- distinguish causal weighting from sequential time marching, and
- interpret the demonstrations as examples of how PINN limitations can be mitigated through changes in sampling, representation, and training order.
Physics-informed neural networks (PINNs) incorporate governing equations, boundary conditions, initial conditions, and available observations into neural-network training. This construction is attractive because the network can be optimized without requiring a labeled solution at every point in the domain. However, encoding physics in the loss function does not by itself guarantee efficient or reliable learning.
Four recurring limitations account for most of what goes wrong: loss imbalance, sampling strategy, complex solution structures, and temporal causality. The first was treated in Regularization, MAP Estimation, and Loss Weighting, together with its remedies, adaptive weighting and hard constraints, and is only recalled here. This lecture takes up the remaining three. For complex solution structures, three representative cases are considered: high-frequency behavior, singularity / sharp gradients, and gradient discontinuity across an interface.
The corresponding extensions are introduced through compact demonstrations: focused collocation for localized resolution, Fourier feature embedding, enrichment, and domain decomposition for difficult solution structures, and causal PINNs and time marching for transient problems.
A PINN is attractive because it needs no mesh and no solution data. The equation itself supplies the training signal: sample points in the domain, evaluate the residual of the differential equation there by automatic differentiation, add the residual of each boundary and initial condition, and hand the sum to an optimizer. Where data are available they can be added as one more term. Nothing about the geometry has to be discretized, and the result is a closed-form function of the coordinates that can be evaluated anywhere.
That much is the promise. In practice a PINN trained on a problem that a classical solver handles without comment can return a solution that is smooth, plausible, and wrong, and the loss curve will not say so.
Write $N_\theta$ for the network and collect the conditions of the problem. A typical steady problem contributes two, and a transient one three:
$$ \mathcal{L}(\theta) = \mathcal{L}_{\mathrm{PDE}}(\theta) + \mathcal{L}_{\mathrm{BC}}(\theta) + \mathcal{L}_{\mathrm{IC}}(\theta) , $$
each term a mean squared residual over its own set of sample points,
$$ \mathcal{L}_{\mathrm{PDE}} = \frac{1}{N_f}\sum_{j=1}^{N_f}\bigl|\mathcal{N}[N_\theta](x_j)\bigr|^2 , \qquad \mathcal{L}_{\mathrm{BC}} = \frac{1}{N_b}\sum_{j=1}^{N_b}\bigl|N_\theta(x_j) - g(x_j)\bigr|^2 . $$
The optimizer sees one number. It does not know that the first term is a statement about the interior and the second a statement about the boundary, that one is exact physics and the other exact data, or that a solution satisfying the first but not the second is not a solution at all. It sees a scalar and descends.
Every failure mode below is a consequence of that reduction, or of a choice that the reduction leaves unconstrained: where to sample, what function class to use, in what order to train.
Loss imbalance. The terms of the sum compete for the same parameters, and one of them can dominate the update so completely that another is never satisfied, while the total loss falls convincingly the whole time. Remedies: weights on the terms, adaptively chosen; or removing a condition from the loss entirely by building it into the output.
Sampling strategy. The residual is only ever evaluated where the collocation points are. A uniform sample spends its budget evenly over a domain whose difficulty is not even, and a feature confined to a small region is resolved by whatever few points happen to land in it. Remedy: put the points where the solution varies.
Complex solution structures. A plain network trained by gradient descent is biased towards smooth, low-frequency functions. Solutions that oscillate rapidly, that blow up at a point, or whose derivative jumps across an interface are all shapes it approaches slowly or not at all. Remedies differ by shape: Fourier features, an enrichment basis, a decomposition of the domain.
Temporal causality. A transient problem has an order, the initial condition determines what comes next, but nothing in the loss respects it. All time levels are trained at once, and the residual at late times can be driven down around a trajectory that never came from the right initial state. Remedies: weight the late residuals by how well the early ones are satisfied, or train the time interval in windows.
The four failure modes. Top left: gradients from different loss terms pulling the same parameter vector in different directions. Top right: collocation points spread uniformly over a domain with a localized stress concentration. Bottom left: three solution shapes that resist a smooth network — a high-frequency oscillation, a sharp peak, and a kink at an interface. Bottom right: a transient solution, whose information must travel from left to right in time.
One observation applies to all four. In each case the symptom is not a large loss but a small loss attached to a wrong solution, so the diagnosis never comes from the loss curve alone. It comes from comparing the terms against each other, from looking at where the error lives in the domain, and from knowing what shape the solution is supposed to have.
The first of the four, loss imbalance, was treated in full in Regularization, MAP Estimation, and Loss Weighting, where the weighting of the terms and the hard-constraint alternative are developed and demonstrated; it is not repeated here. The remaining three are the subject of this notebook, one to a chapter, each with the mechanism, the remedy, and a worked demonstration.
A PINN has no data. Its training set is the set of points at which the residual of the governing equation is evaluated, the collocation points, and the physics loss is an average over them. Two consequences follow, and both are easy to overlook because the equation itself holds everywhere. The residual is constrained only at the sampled points; between them the network is held only by the smoothness of its own function class. And the loss is an average, so a region that occupies a small fraction of the domain contributes a small fraction of the loss, however large the residual is there and however much the answer depends on getting that region right.
The default choice is a uniform sample. It requires knowing nothing about the solution, which is exactly its weakness, because the difficulty of a problem is rarely spread evenly over the domain. A stress concentration, a boundary layer, a crack tip, a shock: each confines the fast variation of the solution to a small region, and a uniform sample spends its points everywhere else. Computational fluid dynamics learned this fifty years ago and answers it with a graded mesh, fine where the gradients are and coarse where they are not. A collocation set is the PINN's counterpart of that mesh, and the same principle applies, at a lower price: the points carry no connectivity, so adding one costs one row in a tensor.
This chapter makes the point with one problem, one point budget, and two ways of spending it.
2.1.1. The Kirsch Problem¶
The test problem is the classical one for stress concentration. Take a large thin plate and pull it in one direction with a uniform stress $\sigma_0$, so that far from any disturbance every cross-section carries the same load. Now cut a circular hole of radius $a$ through it. The material that used to carry the load across the hole is gone, and the load must go around it; the stress is no longer uniform, and next to the hole it is larger than $\sigma_0$. How much larger, and over how wide a region, is the question Kirsch answered in closed form in 1898.
It helps to fix the geometry before writing the formulas. Place the centre of the hole at the origin, let the load act along the $x$-axis, and describe a point by its distance $r$ from the centre and its angle $\theta$ measured from the direction of the load. The plate is taken to be infinite, so that the only length in the problem is $a$, and every distance below is naturally measured in units of it. The stresses are described in polar components: $\sigma_{rr}$ is the normal stress acting on a surface that faces away from the hole, $\sigma_{\theta\theta}$ is the normal stress acting around the hole, the so-called hoop stress, and $\sigma_{r\theta}$ is the shear stress between the two. The Kirsch solution is
$$ \begin{align*} \frac{\sigma_{rr}}{\sigma_0} &= \frac{1}{2}\Bigl(1-\frac{a^2}{r^2}\Bigr) + \frac{1}{2}\Bigl(1 + 3\frac{a^4}{r^4} - 4\frac{a^2}{r^2}\Bigr)\cos 2\theta , \\[6pt] \frac{\sigma_{\theta\theta}}{\sigma_0} &= \frac{1}{2}\Bigl(1+\frac{a^2}{r^2}\Bigr) - \frac{1}{2}\Bigl(1 + 3\frac{a^4}{r^4}\Bigr)\cos 2\theta , \\[6pt] \frac{\sigma_{r\theta}}{\sigma_0} &= -\frac{1}{2}\Bigl(1 - 3\frac{a^4}{r^4} + 2\frac{a^2}{r^2}\Bigr)\sin 2\theta . \end{align*} $$
The formulas look heavy, and they need not be memorized; what matters is what they say in three places.
Far from the hole, $r \gg a$, every term carrying $a/r$ disappears, leaving $\sigma_{rr}/\sigma_0 = \tfrac12(1 + \cos 2\theta) = \cos^2\theta$, $\sigma_{\theta\theta}/\sigma_0 = \sin^2\theta$ and $\sigma_{r\theta}/\sigma_0 = -\sin\theta\cos\theta$. These are nothing but the uniform stress $\sigma_0$ along $x$, rewritten in polar components: the hole is not felt there, as it should not be.
At the edge of the hole, $r = a$, the first and third formulas give $\sigma_{rr} = 0$ and $\sigma_{r\theta} = 0$. This is the statement that the hole is empty: nothing inside it pushes or drags on the edge, so a surface facing the hole carries no normal stress and no shear stress. The condition is called traction free, and it is the boundary condition the PINN will have to impose there. The hoop stress, on the other hand, does not vanish at the edge; the second formula gives
$$ \frac{\sigma_{\theta\theta}}{\sigma_0}\Big|_{r=a} = 1 - 2\cos 2\theta . $$
At $\theta = 0$ and $\theta = \pi$, the two points of the edge that face the load, this is $-1$: the material there is in compression, even though the plate as a whole is being pulled. At $\theta = \pm\pi/2$, the top and bottom of the hole, it is $3$. The load that would have crossed the hole has been squeezed into the material beside it, and the stress there is three times the applied stress. That factor of three is the stress concentration factor of a circular hole, and it does not depend on the size of the hole or on the material; a small hole in a large plate triples the local stress just as a large one does.
Between the two, the disturbance dies away as $a^2/r^2$: at $r = 2a$ the extra stress has fallen to about a quarter of its value at the edge, and at $r = 3a$ to about a ninth. The region in which the hole matters is therefore an annulus about one radius wide around it, and everything outside is nearly uniform. This is the fact the whole lab turns on. The hard part of the solution is confined to a small fraction of the plate, and a sample of collocation points that does not know this will spend most of its points where nothing is happening.
A stress has three components, and to compare two solutions it is convenient to collapse them into one number. The usual choice is the von Mises stress,
$$ \sigma_{\mathrm{vM}} = \sqrt{\sigma_{rr}^2 - \sigma_{rr}\sigma_{\theta\theta} + \sigma_{\theta\theta}^2 + 3\sigma_{r\theta}^2} , $$
which is the scalar an engineer compares against the yield strength to decide whether the material will fail. Under the uniform far-field stress it equals $\sigma_0$; at the top and bottom of the hole it is close to $3\sigma_0$. Every figure below shows this quantity.
The PINN cannot work on an infinite plate, so the problem is posed on the finite square $\Omega = \{(x,y)\in[-3a,3a]^2 : r \ge a\}$ with $a = 1$. Two boundary conditions make the finite problem equivalent to the infinite one. On the hole the traction-free condition $\sigma_{rr} = \sigma_{r\theta} = 0$ is imposed, as in the original problem. On the outer square the Kirsch stresses themselves are prescribed, so that the square carries exactly the stress the infinite plate would carry there. With these two conditions the exact solution of the finite problem is the Kirsch solution, and every error reported below is measured against it.
NumPy, Matplotlib and PyTorch. The two training runs take a few minutes each on a CPU and well under a minute on a GPU.
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('Device:', device)
The Kirsch formulas and the von Mises stress, written once as functions of $(r,\theta)$, with the hole radius $a=1$ and stresses in units of $\sigma_0$. They serve three purposes below: the ground truth, the outer boundary condition, and the error. column turns a NumPy array into the one-column tensor that the network expects; it is used every time points are handed to the network.
def kirsch(r, theta): # hole of radius 1, stresses in units of sigma_0
s_rr = 0.5*(1 - 1/r**2) + 0.5*(1 + 3/r**4 - 4/r**2)*torch.cos(2*theta)
s_tt = 0.5*(1 + 1/r**2) - 0.5*(1 + 3/r**4)*torch.cos(2*theta)
s_rt = -0.5*(1 - 3/r**4 + 2/r**2)*torch.sin(2*theta)
return s_rr, s_tt, s_rt
def von_mises(s_rr, s_tt, s_rt):
return torch.sqrt(s_rr**2 - s_rr*s_tt + s_tt**2 + 3*s_rt**2)
def column(v):
return torch.tensor(v, dtype = torch.float32, device = device).reshape(-1, 1)
All predictions are evaluated on the same $201\times201$ grid over the square, with the hole masked out.
x_grid = np.linspace(-3, 3, 201)
X, Y = np.meshgrid(x_grid, x_grid)
R = np.sqrt(X**2 + Y**2)
r_eval = column(np.maximum(R, 1))
theta_eval = column(np.arctan2(Y, X))
with torch.no_grad():
s_rr, s_tt, s_rt = kirsch(r_eval, theta_eval)
VM_true = von_mises(s_rr, s_tt, s_rt).cpu().numpy().reshape(201, 201)
VM_true[R < 1] = np.nan # the hole
plt.figure(figsize = (6, 5))
plt.contourf(X, Y, VM_true, levels = np.linspace(0, 3, 101), cmap = 'jet')
plt.colorbar()
plt.axis('equal')
plt.xlabel('$x/a$')
plt.ylabel('$y/a$')
plt.title('Ground Truth')
plt.show()
The largest von Mises stress is $2.93\,\sigma_0$, at the top and bottom of the hole where the hoop stress is three times the applied load. The figure shows how localized the feature is: the red region hugs the hole, the blue low-stress lobes on its left and right are only slightly wider, and beyond about $r=2a$ the field is the nearly uniform green of the far-field stress. A network has to represent the whole of this with one set of parameters, and the question is where it should be told to try hardest.
2.1.2. Two Ways to Spend 256 Points¶
Naive collocation. One cell-centred Cartesian grid of $20\times14$ cells over the square, with the points that fall inside the hole removed. What remains is $256$ points spread evenly over the plate, at a density of about seven points per unit area, so that the annulus $a<r\le2a$, which carries the whole stress concentration, receives only $76$ of them.
Focused collocation. The same $256$ points redistributed. A fine $23\times23$ grid is kept only inside the annulus $a<r\le2a$, and a coarse $14\times13$ grid is kept only outside it. The two grid sizes are chosen so that the counts add up to exactly $256$ and the point density inside the annulus is three times the density outside.
The cell below builds both sets and reports the counts.
def grid(nx, ny): # cell centres of an nx by ny grid over the square
hx = 6/nx
hy = 6/ny
x = -3 + hx/2 + hx*np.arange(nx)
y = -3 + hy/2 + hy*np.arange(ny)
X, Y = np.meshgrid(x, y)
return X.ravel(), Y.ravel()
# naive: one 20 x 14 grid, the hole removed
x, y = grid(20, 14)
r = np.sqrt(x**2 + y**2)
x_naive = x[r >= 1]
y_naive = y[r >= 1]
# focused: a fine grid for 1 <= r <= 2, a coarse grid for r > 2
x, y = grid(23, 23)
r = np.sqrt(x**2 + y**2)
x_fine = x[(r >= 1) & (r <= 2)]
y_fine = y[(r >= 1) & (r <= 2)]
x, y = grid(14, 13)
r = np.sqrt(x**2 + y**2)
x_coarse = x[r > 2]
y_coarse = y[r > 2]
x_focused = np.concatenate([x_fine, x_coarse])
y_focused = np.concatenate([y_fine, y_coarse])
print('naive : %d points' % len(x_naive))
print('focused : %d points (%d near the hole + %d far away)' % (len(x_focused), len(x_fine), len(x_coarse)))
Drawn over the exact von Mises field, so that the placement can be judged against the feature it is meant to resolve.
plt.figure(figsize = (6, 10))
plt.subplot(2, 1, 1)
plt.contourf(X, Y, VM_true, levels = np.linspace(0, 3, 101), cmap = 'jet')
plt.plot(x_naive, y_naive, 'k.')
plt.axis('equal')
plt.title('Naive Collocation')
plt.subplot(2, 1, 2)
plt.contourf(X, Y, VM_true, levels = np.linspace(0, 3, 101), cmap = 'jet')
plt.plot(x_focused, y_focused, 'k.')
plt.axis('equal')
plt.title('Focused Collocation')
plt.show()
The naive grid treats the plate as if every part of it were equally hard. The focused grid puts $140$ of its $256$ points in the annulus and the remaining $116$ in the far field, where the solution is nearly constant and a handful of points is enough to pin it. Nothing else differs between the two sets: same number of points, same boundary points, same everything that follows.
2.1.3. The PINN¶
The network takes the position in polar form, $\bigl(r/3,\cos\theta,\sin\theta\bigr)$, and returns the three stress components $\bigl(\sigma_{rr},\sigma_{\theta\theta},\sigma_{r\theta}\bigr)/\sigma_0$ directly. Polar inputs suit the problem, because the solution is naturally a function of $r$ and $\theta$; dividing $r$ by $3$ keeps the input of order one, and feeding $\cos\theta$ and $\sin\theta$ instead of $\theta$ removes the artificial jump at $\theta=\pm\pi$. Four hidden layers of $64$ units with $\tanh$; new_pinn() resets the seed so that both models start from the same weights. The function stresses assembles the input and splits the output into three columns, and derivative wraps the one call to automatic differentiation.
def new_pinn():
torch.manual_seed(1) # both models start from the same weights
return nn.Sequential(nn.Linear(3, 64), nn.Tanh(),
nn.Linear(64, 64), nn.Tanh(),
nn.Linear(64, 64), nn.Tanh(),
nn.Linear(64, 64), nn.Tanh(),
nn.Linear(64, 3)).to(device)
def stresses(net, r, theta):
out = net(torch.cat([r/3, torch.cos(theta), torch.sin(theta)], dim = 1))
return out[:, 0:1], out[:, 1:2], out[:, 2:3]
def derivative(y, x):
return torch.autograd.grad(y, x, torch.ones_like(y), create_graph = True)[0]
In polar coordinates the two equilibrium equations of plane elasticity are
$$ \begin{align*} \frac{\partial \sigma_{rr}}{\partial r} + \frac{1}{r}\frac{\partial \sigma_{r\theta}}{\partial \theta} + \frac{\sigma_{rr}-\sigma_{\theta\theta}}{r} &= 0 \\ \\ \frac{\partial \sigma_{r\theta}}{\partial r} + \frac{1}{r}\frac{\partial \sigma_{\theta\theta}}{\partial \theta} + \frac{2\sigma_{r\theta}}{r} &= 0 \end{align*} $$
and, because the network predicts stresses rather than displacements, a third equation is needed to make the stress field derivable from a displacement: the compatibility condition, which for plane problems without body forces says that the trace $S=\sigma_{rr}+\sigma_{\theta\theta}$ is harmonic,
$$ \frac{\partial^2 S}{\partial r^2} + \frac{1}{r}\frac{\partial S}{\partial r} + \frac{1}{r^2}\frac{\partial^2 S}{\partial\theta^2} = 0 . $$
The function below evaluates the three left-hand sides at the collocation points by automatic differentiation, the compatibility one by differentiating twice.
def residuals(net, r, theta):
s_rr, s_tt, s_rt = stresses(net, r, theta)
eq_r = derivative(s_rr, r) + derivative(s_rt, theta)/r + (s_rr - s_tt)/r # equilibrium
eq_t = derivative(s_rt, r) + derivative(s_tt, theta)/r + 2*s_rt/r
S = s_rr + s_tt # compatibility
S_r = derivative(S, r)
S_theta = derivative(S, theta)
comp = derivative(S_r, r) + S_r/r + derivative(S_theta, theta)/r**2
return eq_r, eq_t, comp
The hole is traction free, $\sigma_{rr}=\sigma_{r\theta}=0$ at $r=a$, sampled at $256$ equally spaced angles. On the outer square the Kirsch stresses are prescribed, sampled at $64$ points per side with the corners left out. These points are the same for both models.
theta_hole = column(np.linspace(0, 2*np.pi, 256, endpoint = False))
r_hole = torch.ones_like(theta_hole)
s = np.linspace(-3, 3, 64)
x_outer = np.concatenate([-3*np.ones(64), 3*np.ones(64), s, s])
y_outer = np.concatenate([s, s, -3*np.ones(64), 3*np.ones(64)])
r_outer = column(np.sqrt(x_outer**2 + y_outer**2))
theta_outer = column(np.arctan2(y_outer, x_outer))
with torch.no_grad():
s_rr_outer, s_tt_outer, s_rt_outer = kirsch(r_outer, theta_outer)
2.1.4. Naive PINN¶
The training loop is written out in full. Each iteration evaluates the three physics residuals at the collocation points and the two boundary conditions at the boundary points, adds the three terms with weights,
$$ \mathcal{L} = \mathcal{L}_{\mathrm{pde}} + 10\,\mathcal{L}_{\mathrm{hole}} + 10\,\mathcal{L}_{\mathrm{outer}} , $$
and takes an Adam step. The boundary terms carry a weight of $10$, because the boundary values are what fix the level of the solution and, as Regularization, MAP Estimation, and Loss Weighting showed, a boundary term with too little weight is simply ignored. The collocation points are converted to tensors once, before the loop, with requires_grad_ so that the residuals can be differentiated with respect to them. The loss is printed every $5000$ of the $30000$ iterations.
r_pde = column(np.sqrt(x_naive**2 + y_naive**2))
theta_pde = column(np.arctan2(y_naive, x_naive))
r_pde.requires_grad_(True)
theta_pde.requires_grad_(True)
net_naive = new_pinn()
optimizer = torch.optim.Adam(net_naive.parameters(), lr = 1e-3)
for k in range(30000):
optimizer.zero_grad()
eq_r, eq_t, comp = residuals(net_naive, r_pde, theta_pde)
loss_pde = torch.mean(eq_r**2) + torch.mean(eq_t**2) + torch.mean(comp**2)
s_rr, s_tt, s_rt = stresses(net_naive, r_hole, theta_hole)
loss_hole = torch.mean(s_rr**2) + torch.mean(s_rt**2)
s_rr, s_tt, s_rt = stresses(net_naive, r_outer, theta_outer)
loss_outer = torch.mean((s_rr - s_rr_outer)**2) + torch.mean((s_tt - s_tt_outer)**2) + torch.mean((s_rt - s_rt_outer)**2)
loss = loss_pde + 10*loss_hole + 10*loss_outer
loss.backward()
optimizer.step()
if (k + 1) % 5000 == 0:
print('Naive %6d | loss %.3e | pde %.3e | hole %.3e | outer %.3e' % (k + 1, loss.item(), loss_pde.item(), loss_hole.item(), loss_outer.item()))
The predicted von Mises field on the same colour scale as the ground truth, $0$ to $3\,\sigma_0$, and the absolute error on a grey scale from $0$ to $0.15\,\sigma_0$; errors above that are shown white.
with torch.no_grad():
s_rr, s_tt, s_rt = stresses(net_naive, r_eval, theta_eval)
VM_naive = von_mises(s_rr, s_tt, s_rt).cpu().numpy().reshape(201, 201)
VM_naive[R < 1] = np.nan
ERR_naive = np.abs(VM_naive - VM_true)
plt.figure(figsize = (6, 10))
plt.subplot(2, 1, 1)
plt.contourf(X, Y, VM_naive, levels = np.linspace(0, 3, 101), cmap = 'jet')
plt.axis('equal')
plt.title('Naive PINN')
plt.subplot(2, 1, 2)
plt.contourf(X, Y, ERR_naive, levels = np.linspace(0, 0.15, 101), cmap = 'gray', extend = 'max')
plt.axis('equal')
plt.title('Naive PINN, Error')
plt.show()
The predicted field looks right. On the colour scale of the ground truth the two panels are hard to tell apart: the red concentration at the top and bottom of the hole, the blue lobes at its sides and the uniform far field are all there, and the total loss ends at $7\times10^{-5}$. The error map says where the agreement ends. Almost all of the error sits in a ring around the hole, with bright spots at the top and bottom of the hole's edge, and the far field is nearly black. The largest error anywhere is $0.154\,\sigma_0$, at the edge of the hole, which is five percent of the peak stress; the mean error in the annulus $a\le r\le 2a$ is $2.1\times10^{-2}$, nearly three times the mean over the whole plate, $8.1\times10^{-3}$.
The pattern is the one the sampling predicts. The naive grid placed $76$ of its $256$ points in the annulus and $180$ in the far field, so the averaged residual the optimizer minimized was dominated by the region where the solution is easy. The hole, where the solution is hard, was resolved by whatever points happened to fall there, and the error is largest exactly where the stress gradient is steepest.
2.1.5. Focused PINN¶
The same loop, the same starting weights, the same boundary points and the same $30000$ iterations; only the $256$ collocation points are the focused set.
r_pde = column(np.sqrt(x_focused**2 + y_focused**2))
theta_pde = column(np.arctan2(y_focused, x_focused))
r_pde.requires_grad_(True)
theta_pde.requires_grad_(True)
net_focused = new_pinn()
optimizer = torch.optim.Adam(net_focused.parameters(), lr = 1e-3)
for k in range(30000):
optimizer.zero_grad()
eq_r, eq_t, comp = residuals(net_focused, r_pde, theta_pde)
loss_pde = torch.mean(eq_r**2) + torch.mean(eq_t**2) + torch.mean(comp**2)
s_rr, s_tt, s_rt = stresses(net_focused, r_hole, theta_hole)
loss_hole = torch.mean(s_rr**2) + torch.mean(s_rt**2)
s_rr, s_tt, s_rt = stresses(net_focused, r_outer, theta_outer)
loss_outer = torch.mean((s_rr - s_rr_outer)**2) + torch.mean((s_tt - s_tt_outer)**2) + torch.mean((s_rt - s_rt_outer)**2)
loss = loss_pde + 10*loss_hole + 10*loss_outer
loss.backward()
optimizer.step()
if (k + 1) % 5000 == 0:
print('Focused %6d | loss %.3e | pde %.3e | hole %.3e | outer %.3e' % (k + 1, loss.item(), loss_pde.item(), loss_hole.item(), loss_outer.item()))
Predicted fields in the top row on the colour scale of the ground truth, errors in the bottom row on the grey scale of the previous figure.
with torch.no_grad():
s_rr, s_tt, s_rt = stresses(net_focused, r_eval, theta_eval)
VM_focused = von_mises(s_rr, s_tt, s_rt).cpu().numpy().reshape(201, 201)
VM_focused[R < 1] = np.nan
ERR_focused = np.abs(VM_focused - VM_true)
plt.figure(figsize = (12, 8))
plt.subplot(2, 3, 1)
plt.contourf(X, Y, VM_true, levels = np.linspace(0, 3, 101), cmap = 'jet')
plt.axis('equal')
plt.title('Ground Truth')
plt.subplot(2, 3, 2)
plt.contourf(X, Y, VM_naive, levels = np.linspace(0, 3, 101), cmap = 'jet')
plt.axis('equal')
plt.title('Naive PINN')
plt.subplot(2, 3, 3)
plt.contourf(X, Y, VM_focused, levels = np.linspace(0, 3, 101), cmap = 'jet')
plt.axis('equal')
plt.title('Focused PINN')
plt.subplot(2, 3, 5)
plt.contourf(X, Y, ERR_naive, levels = np.linspace(0, 0.15, 101), cmap = 'gray', extend = 'max')
plt.axis('equal')
plt.title('Naive PINN, Error')
plt.subplot(2, 3, 6)
plt.contourf(X, Y, ERR_focused, levels = np.linspace(0, 0.15, 101), cmap = 'gray', extend = 'max')
plt.axis('equal')
plt.title('Focused PINN, Error')
plt.show()
The error is averaged over the whole plate and, separately, over the annulus $a\le r\le 2a$ where the points were concentrated; the largest error anywhere is reported as well.
plate = R >= 1
near = (R >= 1) & (R <= 2)
print('Mean error, whole plate : naive %.3e focused %.3e' % (np.mean(ERR_naive[plate]), np.mean(ERR_focused[plate])))
print('Mean error, near the hole : naive %.3e focused %.3e' % (np.mean(ERR_naive[near]), np.mean(ERR_focused[near])))
print('Largest error : naive %.3e focused %.3e' % (np.max(ERR_naive[plate]), np.max(ERR_focused[plate])))
One number should not mislead: the final loss of the focused run, $2.6\times10^{-3}$, is larger than that of the naive run, $6.6\times10^{-5}$. The loss is an average of the residual over the collocation points, and the focused set evaluates it where the residual is hardest to drive down, so a larger loss here is not a worse solution; it is the same quantity measured in a harder place. The error against the exact solution is what decides, and the figure and the numbers below do that.
The upper row shows three fields that agree to the eye; the lower row is where the comparison is made. The naive error map has its bright ring around the hole; the focused error map has no ring. Its residual error is faint and spread over the plate rather than concentrated at the hole.
The numbers confirm it. Near the hole the mean error falls from $2.1\times10^{-2}$ to $8.4\times10^{-3}$, a factor of $2.6$, and the largest error anywhere falls from $0.154$ to $0.047$, a factor of three. Over the whole plate the mean error falls from $8.1\times10^{-3}$ to $5.9\times10^{-3}$, which deserves a second look: the focused grid has fewer points in the far field than the naive grid, $116$ against $180$, and the far field did not get worse. The solution there is nearly uniform, a coarse grid pins it as well as a fine one, and the points taken away from it were not missed. The points added to the annulus, on the other hand, were exactly what the stress concentration needed. Same budget, same network, same training; the only thing that changed was where the residual was looked at.
2.1.6. What the Lab Showed¶
Two models with identical capacity, identical boundary information, identical training and identical numbers of collocation points reached different accuracies, and the difference was decided by where the points were. The naive sample spent most of its budget on the far field, where the solution is nearly constant and almost any placement would do, and left the stress concentration to the few points that happened to fall near the hole. The focused sample put its effort where the solution varies, and it paid for this with nothing: the far field is so smooth that a coarse grid resolves it as well as a fine one.
Two conclusions generalize beyond the plate. A collocation set should be graded like a mesh, dense where the solution is hard and sparse where it is easy, and since the points carry no connectivity there is less excuse for a uniform sample than there is for a uniform mesh. And the loss a PINN reports is a statement about its points, never about the domain: a small averaged residual over a uniform sample says little about the one region that matters if that region holds a small fraction of the points.
Remark (when the hard region is not known in advance). Here the stress concentration was known to sit at the hole, so the refinement could be placed by hand. When the location of the difficulty is not known, the residual itself is the guide: evaluate it on a dense candidate set, and add or move collocation points to where it is largest. This residual-adaptive sampling is the PINN counterpart of adaptive mesh refinement, and it recovers the focused allocation above without being told where the hole is.
A physics-informed neural network asks one network to represent the solution over the whole domain. For a smooth field that varies on a single scale this is a reasonable thing to ask, and the earlier treatments showed it working. Loss imbalance, in Regularization, MAP Estimation, and Loss Weighting, and sampling, in the previous chapter, were difficulties of the training: the terms of the loss competing for the parameters, and the collocation points being in the wrong places. This chapter deals with a difficulty of the representation. Some solutions are hard for a single smooth network to express at all, however well it is trained and however the points are placed, and the remedy is then to change what the network is asked to represent.
Three such structures recur in engineering problems. A solution may oscillate rapidly, so that it contains many wavelengths across the domain; it may vary violently in one small region while being smooth everywhere else, as a stress field does near a crack tip or a corner; or it may be continuous but kinked, with a gradient that jumps across a line where the material changes. Each defeats a plain network for a different reason, and each has a remedy that changes the representation in a different place: in the input, by feeding the network features that already oscillate; in the output, by adding a known singular function to it; or in the domain, by giving each region a network of its own and tying them together at the interface.
The three laboratories of this chapter show one failure and one remedy each. In every case the governing equation, the boundary conditions and the optimizer are the same for the failing model and the repaired one; only the representation differs, and that is the point.
A standard PINN uses a single fully connected network, with smooth activations, as the solution over the entire computational domain. Such a network is a very flexible function, but it is a particular kind of flexible function: it is smooth to all orders, it changes slowly unless its weights are large, and it treats every part of the domain with the same set of parameters. These are the right properties for a field that varies gently and on one scale. They are the wrong properties when the important features of the solution live on a very different scale from the domain, or are confined to a small part of it, or are not smooth at all.
The figure below shows the three cases this chapter considers.
Three solution structures that defeat a single smooth network. In each panel the exact solution is drawn against the prediction of a standard PINN: ten oscillations that the network replaces with a slow drift, a sharp peak that it rounds into a low bump, and a kink at a material interface that it smooths into a curve.
High frequency: the solution oscillates many times across the domain. The network must represent every oscillation, and gradient descent reaches the slow components long before the fast ones.
Singularity / sharp gradient: the solution, or its derivative, grows without bound or changes steeply at one point, while being smooth elsewhere. A smooth network spends its whole capacity trying to bend around that point.
Gradient discontinuity / interface: the solution is continuous but its slope jumps across a line, as it must where two materials meet. A differentiable network has no kink to offer and rounds the corner off.
The first of these is treated at some length, because the reason behind it, the spectral bias of neural networks, is a general fact about how networks learn and will come up again whenever a PINN is asked for fine detail.
Train a network to fit a target that contains both a slow component and a fast one, and watch what it learns first. It is not the whole target at once, and it is not the component with the larger amplitude. It is the slow component. The fast one appears later, often much later, and if training is stopped early it may not appear at all. This ordering is not an accident of one experiment; it is a property of gradient descent on a network, and it is called spectral bias: low-frequency components of a target are learned before high-frequency components, and the gap in learning speed widens as the frequency rises.
The reason can be stated without the theory that makes it precise. A network with smooth activations and weights of ordinary size is itself a slowly varying function, and a small change in its weights changes its output slowly across the input. The gradient of the loss with respect to the weights therefore points, for the most part, in directions that change the output smoothly, and a step of gradient descent corrects the smooth part of the error first. To produce a rapid oscillation the network needs large weights in its first layer, and large weights are reached only by many small steps. The fast component of the target is not forbidden; it is merely expensive, and the optimizer pays for the cheap things first.
A consequence worth stating now, because it will be the lesson of the lab in Section 3.3: a small training loss, or a prediction that looks reasonable at a glance, does not mean that the fine structure of the solution has been recovered. The coarse structure accounts for most of the loss, and it is learned first.
3.2.1. Progressive Frequency Learning¶
The animation below shows a network being trained on a target with several frequency components. The left panel is the prediction against the target in the spatial domain; the right panel is the Fourier spectrum of the prediction, which shows directly which frequencies have been recovered.
Progressive frequency learning. Left: the prediction against a multi-frequency target as training proceeds. Right: the Fourier amplitude of the prediction against that of the target. The low-frequency peak is reached early; the higher peaks rise later and more slowly.
At the start the prediction is nearly flat. The first thing to appear is the slow envelope of the target; the oscillations within it come later, and the fastest of them last. In the spectrum the same story is told by the peaks: the one at the lowest frequency rises first and reaches its target height early, while the peaks at higher frequencies rise later and more slowly. The network is recovering the target one frequency band at a time, from the bottom up.
3.2.2. Evolution of the Predicted Solution¶
The same run, frozen at a few epochs, makes the ordering easier to inspect.
The same run at three epochs. At epoch $0$ the prediction is almost a straight line, at epoch $2000$ it follows the slow envelope of the target, and only by epoch $4000$ does it carry the fast oscillations as well.
At the early epochs the prediction is a smooth curve through the middle of the target, a low-frequency average of it. As training proceeds, finer oscillations are added on top of that curve, and only at the end does the prediction follow the target through every wiggle. Note that the prediction at an intermediate epoch is not a bad approximation of the target in the usual sense; its error is spread evenly and its loss is modest. It is simply missing everything above a certain frequency.
3.2.3. Frequency-wise Convergence¶
The cleanest way to see the bias is to track each frequency component of the target separately and plot the error in that component against the epoch.
Frequency-wise convergence. The ratio of the predicted amplitude to the exact one, for a low-frequency and a high-frequency component of the same target. The low component is recovered within about a hundred epochs; the high one stays near zero until epoch $2000$ and reaches one only at the end of the run.
The lowest-frequency component converges within the first few hundred epochs.
Each higher component converges later, and the error in the highest one is still large when the lowest has long been exact.
The gap in convergence time grows with the frequency; it is not a constant delay but an increasingly steep price.
For a PINN this matters more than it does for an ordinary regression, because the loss is built from derivatives of the network. Differentiating multiplies each frequency component by its frequency, so the residual of a high-frequency solution is dominated by exactly the components the network learns last, and the equation in the loss cannot by itself change the order in which the optimizer learns them. The lab of Section 3.3 shows this on a one-dimensional boundary value problem whose solution has ten oscillations across the domain.
The first laboratory takes a boundary value problem whose exact solution is a pure oscillation, trains a standard PINN on it and watches it fail, and then repairs the failure by changing nothing but the input of the network. The repair, Fourier feature embedding, hands the network sines and cosines of the coordinate at a range of frequencies, so that a rapidly oscillating solution becomes a slowly varying combination of inputs that already oscillate.
3.3.1. Problem Setup¶
Consider the one-dimensional boundary value problem
$$ \frac{d^2u}{dx^2} + \omega^2\sin(\omega x) = 0, \qquad x \in [0, 1], \qquad u(0) = u(1) = 0 , $$
with $\omega = 2 \pi f = 20\pi $. Its exact solution is
$$ u(x) = \sin(\omega x) , $$
which completes ten full oscillations on the unit interval. Nothing about the problem is difficult in the classical sense: it is linear, one-dimensional, and solvable by inspection. Its only feature is the frequency, and that is enough.
The constants of the problem, the evaluation grid, and the two boundary points. The residual will be evaluated on $N_f = 256$ collocation points drawn afresh at every iteration, and every run takes $10000$ Adam steps at learning rate $10^{-3}$.
cycles = 10
omega = 2*np.pi*cycles
N_f = 256
epochs = 10001
lr = 1e-3
bc_weight = 1e-2
x_eval = torch.linspace(0, 1, 2001).reshape(-1, 1).to(device)
u_exact = torch.sin(omega*x_eval)
x_bc = torch.tensor([[0.0], [1.0]]).to(device)
x_plot = x_eval.cpu().numpy().ravel()
u_gt = u_exact.cpu().numpy().ravel()
plt.figure(figsize = (6, 4))
plt.plot(x_plot, u_gt, 'k', linewidth = 2, alpha = 0.3)
plt.title('Ground Truth')
plt.xlabel('x')
plt.ylabel('u(x)')
plt.grid(True)
plt.show()
The target changes sign twenty times on the interval. It is smooth everywhere, and nowhere is it steep in the sense of the lab in Section 3.4; its difficulty is entirely that there is a great deal of it, and a network must reproduce all of it from the equation and two boundary values.
3.3.2. Naive PINN¶
The baseline is the standard construction. The coordinate $x$ is the input, the network output is the solution,
$$ u_\theta(x) = N_\theta(x) , $$
and the loss is the mean squared residual of the equation at the collocation points plus the boundary term. One detail of the residual deserves a word. Written as it stands, the residual $u_{\theta,xx} + \omega^2\sin(\omega x)$ carries the factor $\omega^2 \approx 4000$, and squaring it puts the loss at the scale of $10^{7}$ before training begins. Dividing through by $\omega^2$,
$$ r(x) = \frac{1}{\omega^2}\,\frac{d^2u_\theta}{dx^2} + \sin(\omega x) , $$
changes nothing about which functions satisfy the equation and brings the residual to order one. The boundary term is then weighted by $\lambda_{BC} = 10^{-2}$, a small value that reflects how easy the two boundary values are compared with the interior: the network can satisfy $u(0) = u(1) = 0$ at almost no cost, and a large weight would only distract it.
Four hidden layers of $64$ units with $\tanh$, Xavier initialization. The network is built by a small function that takes the input dimension, because the Fourier-feature model of Section 3.3.3 will use the same architecture with a wider input; the seed is reset inside so that both models start from the same weights.
def new_pinn(input_dim):
torch.manual_seed(0)
net = nn.Sequential(nn.Linear(input_dim, 64), nn.Tanh(),
nn.Linear(64, 64), nn.Tanh(),
nn.Linear(64, 64), nn.Tanh(),
nn.Linear(64, 64), nn.Tanh(),
nn.Linear(64, 1))
for m in net: # Xavier initialization of every linear layer
if isinstance(m, nn.Linear):
nn.init.xavier_normal_(m.weight)
nn.init.zeros_(m.bias)
return net.to(device)
The collocation points are redrawn at every iteration, so the equation is enforced at a different random set of $256$ locations each time. The prediction on the evaluation grid is stored at a few epochs for the comparison of Section 3.3.6.
naive_model = new_pinn(1)
optimizer = torch.optim.Adam(naive_model.parameters(), lr = lr)
snapshot_epochs = [0, 500, 2000, 5000, 10000]
naive_snapshots = np.zeros((5, 2001)) # one row per snapshot epoch
k = 0
for epoch in range(epochs):
x_f = torch.rand(N_f, 1).to(device).requires_grad_(True)
u = naive_model(x_f)
u_x = derivative(u, x_f)
u_xx = derivative(u_x, x_f)
loss_pde = torch.mean((u_xx/omega**2 + torch.sin(omega*x_f))**2)
loss_bc = torch.mean(naive_model(x_bc)**2)
loss = loss_pde + bc_weight*loss_bc
optimizer.zero_grad()
loss.backward()
optimizer.step()
if epoch in snapshot_epochs:
naive_snapshots[k, :] = naive_model(x_eval).detach().cpu().numpy().ravel()
k = k + 1
if epoch % 2000 == 0:
print('EPOCH : %6d/%6d| Loss_Total : %7.6f | Loss_PDE : %7.6f | Loss_BC : %7.6f'
% (epoch, epochs - 1, loss.item(), loss_pde.item(), loss_bc.item()))
The prediction is compared with the exact solution on the fine grid, and the relative $L_2$ error $\lVert u_\theta - u\rVert / \lVert u\rVert$ is reported.
pred_naive = naive_model(x_eval).detach().cpu().numpy().ravel()
error_naive = np.linalg.norm(pred_naive - u_gt)/np.linalg.norm(u_gt)
print('Relative L2 Error : %.6f\n\n' % error_naive)
plt.figure(figsize = (6, 4))
plt.plot(x_plot, u_gt, 'k', linewidth = 2, label = 'Ground Truth', alpha = 0.3)
plt.plot(x_plot, pred_naive, 'r--', linewidth = 1.5, label = 'Naive PINN')
plt.title('Naive PINN')
plt.xlabel('x')
plt.ylabel('u(x)')
plt.legend()
plt.grid(True)
plt.show()
The naive PINN does not recover the oscillation. After ten thousand iterations the loss has fallen only to a few tenths, and the prediction is a slow curve with a few bumps where the target has ten full waves; the relative error is larger than one, which means the prediction is further from the exact solution than zero would be. The equation is in the loss and the loss did fall, but it fell by fitting the smooth part of the residual, exactly as Section 3.2 predicts. The high-frequency solution is there to be found, and the optimizer has not found it.
3.3.3. Fourier Feature Embedding¶
The remedy changes the input and nothing else. Instead of handing the network the coordinate $x$, hand it the coordinate together with sines and cosines of it at a set of frequencies,
$$ \gamma(x) = \bigl[\,x,\ \sin(2\pi x),\ \cos(2\pi x),\ \ldots,\ \sin(2\pi Kx),\ \cos(2\pi Kx)\,\bigr] , $$
and let the network take $\gamma(x)$ as its input. This is Fourier feature embedding. The network no longer has to manufacture an oscillation from scratch; it is given a library of oscillations at the input and has only to combine them, and combining is the kind of smooth, low-frequency operation that a network learns quickly. A target like $\sin(20\pi x)$ is then not a high-frequency function of the input at all, but a simple function of one of the features.
Fourier feature embedding. The coordinate $x$ is replaced at the input by $x$ together with $\sin(2\pi k x)$ and $\cos(2\pi k x)$ for $k=1,\ldots,K$. The network behind the embedding, the differentiation of its output and the physics loss are unchanged.
Everything else stays as it was: the same network architecture behind the embedding, the same initialization, the same collocation points, boundary condition, loss, optimizer and number of iterations. Whatever changes in the result is due to the input representation alone.
The choice of $K$ matters. The features must reach at least the frequency of the solution, here $10$ cycles, and here $K = 16$ is used, so that the target frequency lies comfortably inside the range. A much larger $K$ would hand the network frequencies it does not need and make the fit noisier.
K = 16
frequencies = torch.arange(1, K + 1).reshape(1, -1).to(device)
def fourier_features(x):
phase = 2*np.pi*x*frequencies
return torch.cat([x, torch.sin(phase), torch.cos(phase)], dim = 1)
fourier_model = new_pinn(1 + 2*K)
optimizer = torch.optim.Adam(fourier_model.parameters(), lr = lr)
3.3.4. Naive PINN vs. Fourier-feature PINN¶
The Fourier-feature model is trained with the loop of Section 3.3.2; the only difference is that every call to the network goes through fourier_features first.
fourier_snapshots = np.zeros((5, 2001))
k = 0
for epoch in range(epochs):
x_f = torch.rand(N_f, 1).to(device).requires_grad_(True)
u = fourier_model(fourier_features(x_f))
u_x = derivative(u, x_f)
u_xx = derivative(u_x, x_f)
loss_pde = torch.mean((u_xx/omega**2 + torch.sin(omega*x_f))**2)
loss_bc = torch.mean(fourier_model(fourier_features(x_bc))**2)
loss = loss_pde + bc_weight*loss_bc
optimizer.zero_grad()
loss.backward()
optimizer.step()
if epoch in snapshot_epochs:
fourier_snapshots[k, :] = fourier_model(fourier_features(x_eval)).detach().cpu().numpy().ravel()
k = k + 1
if epoch % 2000 == 0:
print('EPOCH : %6d/%6d| Loss_Total : %7.6f | Loss_PDE : %7.6f | Loss_BC : %7.6f'
% (epoch, epochs - 1, loss.item(), loss_pde.item(), loss_bc.item()))
pred_fourier = fourier_model(fourier_features(x_eval)).detach().cpu().numpy().ravel()
error_fourier = np.linalg.norm(pred_fourier - u_gt)/np.linalg.norm(u_gt)
print('Naive PINN Relative L2 Error : %.6f' % error_naive)
print('Fourier PINN Relative L2 Error : %.6f\n\n' % error_fourier)
plt.figure(figsize = (6, 4))
plt.plot(x_plot, u_gt, 'k', linewidth = 3, label = 'Ground Truth', alpha = 0.3)
plt.plot(x_plot, pred_naive, 'r--', linewidth = 1.5, label = 'Naive PINN')
plt.plot(x_plot, pred_fourier, 'b--', linewidth = 1.5, label = 'Fourier-feature PINN')
plt.title('Naive PINN vs. Fourier-feature PINN')
plt.xlabel('x')
plt.ylabel('u(x)')
plt.legend()
plt.grid(True)
plt.show()
The Fourier-feature PINN lies on the exact solution through all ten oscillations, with a relative error of a few percent against the naive model's error of more than one. The loss trace tells the same story from the other side: the naive run was still near $0.3$ after ten thousand iterations, while the Fourier-feature run was below $10^{-4}$ after five hundred. No data were added and the equation was not changed. The network was given a representation in which the solution is easy to express, and the optimizer found it almost at once.
3.3.5. Frequency-Domain Analysis¶
A plot in the spatial domain shows whether the prediction follows the waveform; a plot in the frequency domain shows whether it contains the right frequency at all, which for an oscillatory solution is the more direct question. The discrete Fourier transform of each prediction on the evaluation grid gives its amplitude at every spatial frequency, in cycles per unit length. The exact solution has all of its amplitude at $f = 10$.
dx = x_plot[1] - x_plot[0]
frequency = np.fft.rfftfreq(len(x_plot), d = dx)
fft_gt = np.abs(np.fft.rfft(u_gt - u_gt.mean()))/len(x_plot)
fft_naive = np.abs(np.fft.rfft(pred_naive - pred_naive.mean()))/len(x_plot)
fft_fourier = np.abs(np.fft.rfft(pred_fourier - pred_fourier.mean()))/len(x_plot)
mask = frequency <= 25
plt.figure(figsize = (6, 4))
plt.plot(frequency[mask], fft_gt[mask], 'k', linewidth = 2, label = 'Ground Truth', alpha = 0.3)
plt.plot(frequency[mask], fft_naive[mask], 'r--', linewidth = 1.5, label = 'Naive PINN')
plt.plot(frequency[mask], fft_fourier[mask], 'b--', linewidth = 1.5, label = 'Fourier-feature PINN')
plt.title('Frequency-Domain Comparison')
plt.xlabel('Spatial frequency')
plt.ylabel('FFT amplitude')
plt.legend()
plt.grid(True)
plt.show()
The exact solution is a single peak at $f = 10$. The Fourier-feature prediction has the same peak at nearly the same height and almost nothing elsewhere. The naive prediction has almost nothing at $f = 10$; what amplitude it has sits at low frequencies, which is the smooth curve seen in Section 3.3.2 described in different terms. The spectrum confirms that the naive model did not partly recover the target frequency and lose the rest; it did not reach that frequency at all.
3.3.6. Training Process¶
The snapshots stored during training show when each model found the oscillation, which is the question spectral bias is about.
plt.figure(figsize = (12, 5))
for k in range(5):
plt.subplot(2, 5, k + 1)
plt.plot(x_plot, u_gt, 'k', linewidth = 1.5, alpha = 0.3)
plt.plot(x_plot, naive_snapshots[k, :], 'r--', linewidth = 1.5)
plt.title('Epoch %d' % snapshot_epochs[k])
plt.ylim(-3.1, 3.6)
plt.grid(True)
if k == 0:
plt.ylabel('Naive PINN')
plt.subplot(2, 5, k + 6)
plt.plot(x_plot, u_gt, 'k', linewidth = 1.5, alpha = 0.3)
plt.plot(x_plot, fourier_snapshots[k, :], 'b--', linewidth = 1.5)
plt.xlabel('x')
plt.ylim(-3.1, 3.6)
plt.grid(True)
if k == 0:
plt.ylabel('Fourier-feature PINN')
plt.tight_layout()
plt.show()
In the top row the naive model is a smooth curve at epoch $500$, a smooth curve with a few bumps at $2000$, and not much more at $10000$; the oscillation is arriving, but at a rate that would need far more iterations than anyone would give it. In the bottom row the Fourier-feature model already has the ten waves at epoch $500$ and spends the rest of the run sharpening them. The two models have the same architecture, the same starting weights and the same loss; the difference is that one of them was handed the frequency it needed and the other had to build it.
3.3.7. What the Lab Showed¶
A PINN inherits the spectral bias of the network inside it, and the equation in the loss does not undo it: the residual of a high-frequency solution is dominated by the components the optimizer learns last, so the loss descends slowly and the prediction stays smooth. Fourier feature embedding puts the needed frequencies into the input, where the network can use them with ordinary, low-frequency weights, and the same optimizer then converges in a few hundred iterations. The remedy costs one line in the forward pass and one choice, the range of frequencies $K$, which should cover the frequencies the solution is expected to contain.
The second difficult structure is a solution that is smooth almost everywhere and violent in one place. Stress near the tip of a crack, the temperature next to a point source, the velocity in a boundary layer, the field where one boundary condition gives way to another: in each case most of the domain is easy and a small region is not, and the small region is usually the one the engineer cares about. A single smooth network must represent both with one set of parameters, and the smooth part wins, because it is most of the domain and most of the loss.
Three ways a solution can be singular. A boundary layer, where the field varies rapidly next to a wall; a crack tip or re-entrant corner, where the stress concentrates; and a concentrated source or load, where the response is sharp around one point.
Boundary singularity: rapid variation next to a boundary, as in a boundary layer.
Geometry singularity: stress concentration at a corner or a crack tip.
Localized source singularity: a strong response around a concentrated load or source.
The remedy of this laboratory is enrichment. When the form of the singular behaviour is known, as it often is from the theory of the problem, the network is not asked to represent it. The known singular function is added to the network output with a trainable coefficient, and the network is left to represent only the smooth remainder.
3.4.1. A Beam with a Point Load¶
The demonstration uses the simplest structure in which a singular load produces a discontinuous response: a simply supported beam carrying a uniform distributed load together with a concentrated load at midspan. Everything is nondimensional.
- Beam span: $0 \le x \le 1$, simply supported at both ends
- Distributed load: $q = 1$
- Concentrated load: $P = 1$ at midspan, $a = 1/2$
The unknown is the shear force $V(x)$, and the two loads act on it in entirely different ways. This is the reason the problem was chosen, and it is worth separating the two before any network is built.
The test problem. A simply supported beam under two loads of different character: a uniform distributed load $q = 1$ over the whole span and a concentrated load $P = 1$ at midspan. The distributed load makes the shear force vary smoothly along the span; the concentrated load makes it jump.
The distributed load produces a smooth variation. Away from the point load the shear force obeys
$$ \frac{dV}{dx} = -q = -1 , $$
so $q$ contributes a straight line of slope $-1$, which is as easy a shape as a network could be asked for.
The concentrated load produces a jump. A point force is not a function but a Dirac delta in the load distribution, so that the full equation reads $dV/dx = -q - P\,\delta(x - a)$. Integrating across $x = a$ gives $V(a^+) - V(a^-) = -P$: the shear force does not merely steepen at the load, it steps, by exactly the magnitude of the force. This is the third of the three singular cases of the figure above, a localized source, and it produces the sharpest response a one-dimensional problem can have.
What the supports say. The two loads are given, but the shear force at the supports is not, and it is worth being precise about what a simple support does and does not tell the model. A simple support supplies a vertical reaction and nothing else: it can push up on the beam with whatever force is needed, but it cannot resist rotation, so it carries no bending moment. Two statements follow, one about forces and one about moments, and each turns into a condition on $V$.
Force equilibrium. The two reactions together must carry the whole load,
$$ R_A + R_B = q\,L + P = 2 . $$
The shear force at a section is the resultant of the forces on one side of it, so at the two ends it is nothing but the reaction: $V(0) = R_A$ and $V(1) = -R_B$, the sign at the right end following from the convention that $V$ is positive when the part of the beam to the left of the section is pushed up. Force equilibrium is therefore a statement about $V$ alone,
$$ V(0) - V(1) = q\,L + P = 2 . $$
This is also what the equation $dV/dx = -q - P\,\delta(x - a)$ gives when it is integrated across the whole span: the drop in $V$ from one end to the other is the total load, the distributed part $qL = 1$ through the smooth slope and the concentrated part $P = 1$ through the jump. It is the only place where $P$ enters the problem that the model will see.
Moment equilibrium. A simple support carries no bending moment, $M(0) = M(1) = 0$. The bending moment and the shear force are related by $dM/dx = V$, so integrating from one support to the other gives
$$ M(1) - M(0) = \int_0^1 V\,dx = 0 . $$
The shear force must have zero mean over the span: the positive part near the left support and the negative part near the right support must cancel. This condition fixes the level of $V$, which force equilibrium alone leaves open, because force equilibrium constrains only the difference $V(0) - V(1)$ and not the two values separately.
What the two conditions determine. Neither condition prescribes $V$ at a point, which is why there is no boundary value in this problem. Together they are enough. For the symmetric loading here the two conditions give $R_A = R_B = 1$, so that $V(0) = 1$ and $V(1) = -1$; in general they are two equations for the two unknown reactions, which is the statically determinate case. The model is never told these reactions. It is given the equation, the force equilibrium and the zero-mean condition, and the reactions are part of what it has to find.
The exact shear force is therefore a smooth part and a singular part added together,
$$ V(x) = \underbrace{1 - x}_{\text{from } q} \;-\; \underbrace{\alpha\,H(x - a)}_{\text{from } P} , $$
with $H$ the Heaviside step and $\alpha$ the height of the jump. Mechanics says what $\alpha$ is: integrating $dV/dx = -q - P\,\delta(x - a)$ across the load gives $\alpha = P = 1$. In what follows this relation is deliberately withheld from the model. The model will be told that a concentrated load produces a step, and where, but not how large the step is; the height $\alpha$ is left as an unknown, and the only thing that can fix it is the equilibrium of the whole beam. Keep that split in view: the enriched model of Section 3.4.2 will mirror it exactly, one term for each, and $\alpha$ will be what it has to find.
Notice what the equation and the equilibrium condition do to a smooth model. Integrating $dV/dx = -1$ across the span gives $V(0) - V(1) = 1$, while equilibrium demands $2$. No continuous function can satisfy both, and the gap between them is exactly the point load. The contradiction is not an accident of this problem; it is the physics telling the model that something it cannot represent is present.
3.4.2. Two Representations¶
The naive model is the standard one: a smooth network is asked for the shear force directly,
$$ V_{\mathrm{Naive}}(x) = N_\theta(x) . $$
The enriched model adds to it a function that already has the shape of the difficulty, with a coefficient to be learned,
$$ V_{\mathrm{Enriched}}(x) = N_\theta(x) + \alpha\,\phi(x) , \qquad \phi(x) = -H(x - a) . $$
The two terms correspond one to one with the two loads of Section 3.4.1. The enrichment term $\alpha\phi$ carries the jump that the point load produces, which is outside the reach of a smooth network entirely; the network $N_\theta$ carries everything else, the slope that the distributed load produces and the level that the supports set. The model is not being asked to learn a discontinuity; it is being handed one.
What is known and what is learned should be kept apart. The basis $\phi$ is known: that a concentrated load produces a step, and that it sits at $x = a$, are part of the problem statement, so the shape and the position of the jump are fixed before training begins. The coefficient $\alpha$ is unknown: it is a single trainable scalar, optimized alongside the network weights, and it is what sets how far the solution jumps. The model is not told the relation $\alpha = P$; it is told that the beam as a whole must be in equilibrium under $q$ and $P$, and the height of the step has to follow from that. Prior knowledge supplies the form; the optimizer supplies the magnitude.
In general the enrichment is a sum, and the construction is the same whatever the shape being supplied,
$$ u_{\mathrm{enriched}}(x) = N_{\theta}(x) + \sum_{i=1}^{m} a_i\,\phi_i(x) , $$
with $\phi_i$ known from the theory of the problem and the coefficients $a_i$ trained together with $\theta$.
Enrichment. The prediction is the sum of the network output and a basis built from prior knowledge, $u_{\mathrm{pred}} = N_\theta + \sum_i a_i\phi_i$. The basis carries the shape of the difficult part of the solution, the coefficients $a_i$ are trained together with the network parameters, and the network is left with the smooth remainder.
The same construction covers the other two singular cases of the figure above. For a crack tip the leading term of the elastic solution behaves as $r^{-1/2}$ in the stress, so $\phi_{\mathrm{sing}}$ is that known field and $a$ is the stress intensity factor; for a point source the basis is the fundamental solution of the operator. In every case the basis carries the form of the singularity and the coefficient carries its strength, and the coefficient is what the optimizer recovers from the equation and the boundary conditions, with no data.
Remark (enrichment and masking). This is the construction of the hard constraint in Regularization, MAP Estimation, and Loss Weighting, used for the opposite purpose. There a known function was multiplied in, so that a boundary condition held identically and left the loss; here a known function is added in, so that a shape the network cannot produce is supplied from outside. In both cases the network is asked to represent only what is left once the known part has been taken out.
The constants of the problem.
q = 1.0 # distributed load
P = 1.0 # point load
a = 0.5 # location of the point load
3.4.3. Network and Physics-Informed Loss¶
Both models use the same network, a $1 \to 32 \to 32 \to 32 \to 1$ multilayer perceptron with $\tanh$, and the same loss,
$$ \mathcal{L} = \frac{1}{N}\sum_{i}\Bigl(\frac{dV}{dx}(x_i) + q\Bigr)^{2} + \bigl(V(0) - V(1) - q - P\bigr)^{2} + \Bigl(\int_0^1 V\,dx\Bigr)^{2} , $$
the residual at $100$ collocation points, the force equilibrium of the whole beam, and the zero end moments of a simple support, with the integral evaluated as the mean of $V$ on a fine grid. The collocation points $x_i = 0.005, 0.015, \ldots, 0.995$ are placed so that none falls on $x = a$, because the equation $dV/dx = -q$ holds only away from the point load. No value of $V$ is prescribed anywhere; the level of the solution has to come out of the two global conditions.
The loss is written as a function of the model $V(x)$, so that the same function serves the naive and the enriched models. The enriched model starts from $\alpha = 0$: at the first iteration it is the naive model exactly, and the only thing that has changed is that $\alpha$ is allowed to move.
def new_pinn():
torch.manual_seed(0)
return nn.Sequential(nn.Linear(1, 32), nn.Tanh(),
nn.Linear(32, 32), nn.Tanh(),
nn.Linear(32, 32), nn.Tanh(),
nn.Linear(32, 1)).to(device)
def derivative(y, x):
return torch.autograd.grad(y, x, torch.ones_like(y), create_graph = True)[0]
x_pde = torch.linspace(0.005, 0.995, 100).reshape(-1, 1).to(device).requires_grad_(True) # no point at x = a
x_end = torch.tensor([[0.0], [1.0]]).to(device)
x_int = torch.linspace(0, 1, 1001).reshape(-1, 1).to(device)
def physics_loss(V):
loss_pde = torch.mean((derivative(V(x_pde), x_pde) + q)**2)
V_end = V(x_end)
loss_force = (V_end[0] - V_end[1] - (q + P))**2 # R_A + R_B = q + P
loss_moment = torch.mean(V(x_int))**2 # M(0) = M(1) = 0
return loss_pde + loss_force + loss_moment
3.4.4. Naive PINN¶
The baseline is trained first, for $3000$ Adam iterations. new_pinn() resets the seed, so that the enriched model will start from exactly the same weights.
naive_net = new_pinn()
optimizer = torch.optim.Adam(naive_net.parameters(), lr = 1e-3)
def V_naive(x):
return naive_net(x)
for epoch in range(1, 3001):
loss = physics_loss(V_naive)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if epoch % 500 == 0:
print('Naive | epoch %4d | loss %.4e' % (epoch, loss.item()))
The prediction is compared with the exact shear force, with the absolute error beside it.
x_plot = np.linspace(0, 1, 1001)
V_gt = 1 - x_plot - P*(x_plot >= a)
x_plot_t = torch.tensor(x_plot, dtype = torch.float32).reshape(-1, 1).to(device)
V_naive_pred = V_naive(x_plot_t).detach().cpu().numpy().ravel()
plt.figure(figsize = (6, 4))
plt.plot(x_plot, V_gt, 'k', linewidth = 3, alpha = 0.3, label = 'Ground Truth')
plt.plot(x_plot, V_naive_pred, 'r--', linewidth = 1.5, label = 'Naive PINN')
plt.axvline(a, color = 'gray', linestyle = ':')
plt.title('Ground Truth vs. Naive PINN')
plt.xlabel('x')
plt.ylabel('V(x)')
plt.legend()
plt.grid(True)
plt.show()
The loss stops falling after a few hundred iterations and stays at $0.5$ for the rest of the run. That plateau is not slow convergence; it is the contradiction of Section 3.4.1 made numerical. A continuous function cannot carry the equation and the equilibrium condition at once, so the optimizer settles on the compromise that costs least: the straight line $V = 0.75 - 1.5x$, whose slope splits the difference between the $-1$ the equation asks for and the $-2$ the equilibrium asks for. The residual term contributes $0.5^2$ and the equilibrium term another $0.5^2$, which is the $0.5$ on the trace, and the error is $0.25$ at the supports and $0.5$ at midspan. Nothing in the loss is satisfied, the error is spread over the whole span rather than concentrated at the jump, and running ten thousand iterations instead of three changes none of this.
3.4.5. Enriched PINN¶
The discontinuous basis $\phi(x)$ is multiplied by one learnable scalar $\alpha$. The network weights and $\alpha$ are optimized together, from the same starting weights, with the same loss and the same optimizer as before, and the height of the jump is never supplied.
enriched_net = new_pinn() # same initial weights as the naive network
alpha = torch.tensor(0.0, requires_grad = True, device = device) # height of the jump, to be learned
optimizer = torch.optim.Adam(list(enriched_net.parameters()) + [alpha], lr = 1e-3)
def phi(x):
return -(x >= a).float() # -H(x - a)
def V_enriched(x):
return enriched_net(x) + alpha*phi(x)
for epoch in range(1, 3001):
loss = physics_loss(V_enriched)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if epoch % 500 == 0:
print('Enriched | epoch %4d | loss %.4e | alpha %.4f' % (epoch, loss.item(), alpha.item()))
print('Learned alpha : %.4f' % alpha.item())
The two predictions and their errors are drawn together, and the learned coefficient is compared with the value mechanics assigns to the jump.
V_enriched_pred = V_enriched(x_plot_t).detach().cpu().numpy().ravel()
plt.figure(figsize = (6, 4))
plt.plot(x_plot, V_gt, 'k', linewidth = 4, alpha = 0.3, label = 'Ground Truth')
plt.plot(x_plot, V_naive_pred, 'r--', linewidth = 1.5, label = 'Naive PINN')
plt.plot(x_plot, V_enriched_pred, 'b--', linewidth = 1.5, label = 'Enriched PINN')
plt.axvline(a, color = 'gray', linestyle = ':')
plt.title('Shear Force')
plt.xlabel('x')
plt.ylabel('V(x)')
plt.legend()
plt.grid(True)
plt.show()
print('Learned alpha : %.4f (mechanics: alpha = P = %.1f)' % (alpha.item(), P))
The coefficient converges to the physics. Starting from $\alpha = 0$ it passes $0.63$ after a thousand iterations, $0.94$ after two thousand, and ends at $0.998$, against the $P = 1$ that mechanics assigns to the jump. The loss falls to $4\times10^{-6}$, five orders of magnitude below the naive plateau, and the relative $L_2$ error falls from $50\%$ to $0.1\%$. Nothing told the model that the step equals the point load; the equilibrium of the whole beam, which the naive model could only approach by tilting a straight line, is satisfied exactly once the step is there, and the optimizer used it to measure the step.
One scalar made the difference. The network, the starting weights, the loss, the optimizer and the number of iterations are the same. What changed is that the model was given a function with a step in it, and was therefore able to satisfy a set of conditions that a continuous model could not reach.
3.4.6. Decomposition of the Enriched Solution¶
The enriched prediction is a sum of two parts, and separating them shows what each one did,
$$ V_{\mathrm{Enriched}}(x) = V_{\mathrm{NN}}(x) + V_{\mathrm{Prior}}(x) , \qquad V_{\mathrm{Prior}}(x) = \alpha\,\phi(x) , $$
where the prior component is drawn with the learned $\alpha$, not with the reference value.
V_nn = enriched_net(x_plot_t).detach().cpu().numpy().ravel()
V_prior = alpha.item()*phi(x_plot_t).cpu().numpy().ravel()
plt.figure(figsize = (9, 3))
plt.subplot(1, 3, 1)
plt.plot(x_plot, V_nn, 'b', linewidth = 2)
plt.title('NN Component')
plt.xlabel('x')
plt.ylabel('V(x)')
plt.ylim(-1.2, 1.2)
plt.grid(True)
plt.subplot(1, 3, 2)
plt.plot(x_plot, V_prior, 'r', linewidth = 2)
plt.title('Prior Component')
plt.xlabel('x')
plt.ylim(-1.2, 1.2)
plt.grid(True)
plt.subplot(1, 3, 3)
plt.plot(x_plot, V_nn + V_prior, 'k', linewidth = 2)
plt.title('Enriched Solution')
plt.xlabel('x')
plt.ylim(-1.2, 1.2)
plt.grid(True)
plt.tight_layout()
plt.show()
The three panels show the division of labour. The network component is the straight line $1 - x$, the only shape the network was ever asked for and one it fits without difficulty; note that its level, the reaction at the left support, was learned. The prior component is the step, scaled by the coefficient the optimizer found. Their sum is the shear-force diagram: linear on each side of midspan and discontinuous across it.
Neither part could have produced the answer alone. The network has no step to offer, and the step knows nothing of the distributed load or the supports. Enrichment is the statement that these are two different jobs, and that only one of them needs a neural network.
3.4.7. What the Lab Showed¶
A smooth network cannot represent a jump, and when the physics demands one it does not fail locally but globally: the optimizer converged, to a loss it could not lower further, and the compromise it found was wrong across the whole span rather than only at midspan. Nothing in the training was at fault, and no amount of further iteration would have helped.
Enrichment supplied the shape the network lacked and asked the optimizer for one number. That number was not given, and no condition in the loss named it; it was implied by the equilibrium of the whole beam, which no continuous function could satisfy together with the equation, and the enriched model recovered it to three decimal places. The requirement is the one that runs through this chapter: the form of the difficulty must be known in advance. In structural mechanics it usually is, because the position of a concentrated load is part of the problem statement, and the exponent of a crack-tip field is fixed by the local geometry before any computation begins.
The third difficult structure arises when the domain is made of parts that obey different local physics, most simply two materials with different properties joined along an interface. The solution is then continuous across the interface, because a temperature or a displacement cannot jump, but its gradient is not, because the flux must be continuous and a continuous flux through materials of different conductivity requires different slopes on the two sides. The solution has a kink.
A single network with smooth activations has no kink to offer. It can approximate one, by bending sharply over a short distance, but the bend costs large weights and is learned slowly, and until it is learned the network smooths the corner off and misplaces the solution on both sides of it. Typical examples are multilayer thermal systems, composite laminates, dielectrics with different permittivities and layered porous media.
Interfaces in engineering problems. An insulated pipe, a multilayer dielectric and a layered aquifer. In each the field is continuous across the interface while its gradient is not, because the flux through materials of different conductivity requires different slopes on the two sides.
The remedy, shown in the lab of Section 3.6, is domain decomposition: one network per material, each responsible only for its own region, with the two solutions joined at the interface by the physical conditions, continuity of the field and continuity of the flux, imposed as loss terms. Each network then represents a smooth function on its own side, which it does well, and the kink appears where the two meet.
Consider steady heat conduction in the square plate $[-1, 1]^2$, made of two materials that meet along the vertical line $x = 0$: the left half with conductivity $k_1 = 0.3$ and the right half with $k_2 = 0.6$. The left edge is held at temperature $1$, the right edge at $0.2$, and the top and bottom edges are insulated. In each material the temperature satisfies
$$ k_i\left(\frac{\partial^2 T_i}{\partial x^2} + \frac{\partial^2 T_i}{\partial y^2}\right) = 0 , \qquad i = 1, 2 , $$
and on the interface two conditions hold, continuity of temperature and continuity of heat flux,
$$ T_1 = T_2 , \qquad k_1\frac{\partial T_1}{\partial x} = k_2\frac{\partial T_2}{\partial x} \qquad \text{at } x = 0 . $$
The second condition is where the kink comes from. With $k_1 \neq k_2$ the two slopes must differ by the factor $k_2/k_1 = 2$ for the fluxes to match. Because the edge temperatures are uniform in $y$ and the top and bottom are insulated, nothing in the problem depends on $y$, and the exact solution is two planes meeting along the interface,
$$ T(x, y) = \begin{cases} \dfrac{7}{15} - \dfrac{8}{15}\,x , & x < 0 , \\[10pt] \dfrac{7}{15} - \dfrac{4}{15}\,x , & x > 0 , \end{cases} $$
with slope $-8/15$ in the poorly conducting left material and $-4/15$ in the right, and the flux $k\,\partial T/\partial x = -0.16$ the same constant everywhere in the plate. Nothing about this problem is hard: no high frequency, no singularity. The only feature is the corner along $x = 0$.
The two-material plate. Left: the geometry, with conductivity $k_1=0.3$ on the left, $k_2=0.6$ on the right, fixed temperatures on the vertical edges and insulated horizontal edges. Right: the exact temperature, which varies only with $x$ and changes slope at the dashed interface.
The two conductivities and the piecewise linear temperature, evaluated on a $201\times201$ grid that serves as the ground truth throughout.
k1 = 0.3
k2 = 0.6
T_left = 1.0
T_right = 0.2
x_grid = np.linspace(-1, 1, 201)
X, Y = np.meshgrid(x_grid, x_grid)
A = (T_left - T_right)*k2/(k1 + k2) # slope magnitude in material 1
B = (T_left - T_right)*k1/(k1 + k2) # slope magnitude in material 2
T_ref = np.where(X < 0, T_left - A*(X + 1), T_right + B*(1 - X))
plt.figure(figsize = (6, 4))
plt.imshow(T_ref, origin = 'lower', extent = [-1, 1, -1, 1], cmap = 'jet', vmin = 0.2, vmax = 1)
plt.colorbar()
plt.axvline(0, color = 'k', linestyle = '--')
plt.title('Exact Temperature')
plt.xlabel('x')
plt.ylabel('y')
plt.show()
The temperature falls from $1$ to $0.2$ across the plate, steeply in the left material and half as steeply in the right, with no variation in $y$. The colour map shows a field that is nearly featureless; the change of slope along the dashed line is what a network has to represent.
Four hundred random points in each material for the residual, $120$ points on each edge for the boundary conditions, with the top and bottom edges split at the interface so that each half belongs to one material, and $200$ points along $x = 0$ for the interface conditions. The same points are used by every model below.
np.random.seed(0)
Y1_domain = np.column_stack([-np.random.rand(400), 2*np.random.rand(400) - 1]) # material 1
Y2_domain = np.column_stack([np.random.rand(400), 2*np.random.rand(400) - 1]) # material 2
Y_left = np.column_stack([-np.ones(120), 2*np.random.rand(120) - 1])
Y_right = np.column_stack([np.ones(120), 2*np.random.rand(120) - 1])
Y1_bottom = np.column_stack([-np.random.rand(120), -np.ones(120)])
Y2_bottom = np.column_stack([np.random.rand(120), -np.ones(120)])
Y1_top = np.column_stack([-np.random.rand(120), np.ones(120)])
Y2_top = np.column_stack([np.random.rand(120), np.ones(120)])
Y_int = np.column_stack([np.zeros(200), np.linspace(-1, 1, 200)]) # interface
plt.figure(figsize = (5, 5))
plt.plot(Y1_domain[:, 0], Y1_domain[:, 1], '.', alpha = 0.5, label = 'Material 1')
plt.plot(Y2_domain[:, 0], Y2_domain[:, 1], '.', alpha = 0.5, label = 'Material 2')
plt.plot(Y_int[:, 0], Y_int[:, 1], 'k.', label = 'Interface')
plt.title('Collocation Points')
plt.xlabel('x')
plt.ylabel('y')
plt.axis('scaled')
plt.legend()
plt.show()
3.6.1. Naive PINN¶
The baseline is a single network $T_\theta(x, y)$ over the whole plate. It is given every piece of physics the problem has: the residual $T_{xx} + T_{yy} = 0$ on both halves, the two edge temperatures, the insulated top and bottom, and the flux condition at the interface, all as loss terms,
$$ \mathcal{L} = \mathcal{L}_{\mathrm{PDE}} + \mathcal{L}_{\mathrm{BC}} + \mathcal{L}_{\mathrm{flux}} , \qquad \mathcal{L}_{\mathrm{flux}} = \frac{1}{N_{\mathrm{int}}}\sum_i\bigl(k_1\,T_{\theta, x}(0, y_i) - k_2\,T_{\theta, x}(0, y_i)\bigr)^2 . $$
The flux term is written exactly as the physics states it, but for one network the two derivatives in it are the same number, since a smooth function has only one slope at a point. The term therefore reduces to $(k_1 - k_2)^2\,T_{\theta, x}^2$ along the interface, and the only way to make it zero is a flat tangent there, which the residual and the edge temperatures do not allow. The loss asks for a kink from a function that cannot have one.
A small $\tanh$ network, $2 \to 32 \to 32 \to 32 \to 1$; new_pinn() resets the seed, so that every model in this lab starts from the same weights. The points are converted to tensors once, with requires_grad_ so that the network output can be differentiated with respect to them. derivative, defined with the training cell below, returns both partial derivatives of a scalar field.
def new_pinn():
torch.manual_seed(0)
return nn.Sequential(nn.Linear(2, 32), nn.Tanh(),
nn.Linear(32, 32), nn.Tanh(),
nn.Linear(32, 32), nn.Tanh(),
nn.Linear(32, 1)).to(device)
def tensor(Y):
return torch.tensor(Y, dtype = torch.float32).to(device).requires_grad_(True)
Y1_domain = tensor(Y1_domain)
Y2_domain = tensor(Y2_domain)
Y_left = tensor(Y_left)
Y_right = tensor(Y_right)
Y1_bottom = tensor(Y1_bottom)
Y2_bottom = tensor(Y2_bottom)
Y1_top = tensor(Y1_top)
Y2_top = tensor(Y2_top)
Y_int = tensor(Y_int)
The residual is $T_{xx} + T_{yy}$, evaluated by differentiating twice; the conductivity does not appear in it, because $k_i\,\nabla^2 T = 0$ and $\nabla^2 T = 0$ are the same equation inside a material. The boundary loss collects the two edge temperatures and the four insulated half-edges, where the condition is $\partial T/\partial y = 0$. It is written for two networks, one per material, so that the same function serves the naive model, with the same network passed twice, and the decomposed model.
def laplacian(net, Y):
T = net(Y)
T_x, T_y = derivative(T, Y)
T_xx, _ = derivative(T_x, Y)
_, T_yy = derivative(T_y, Y)
return T_xx + T_yy
def insulated(net, Y):
_, T_y = derivative(net(Y), Y)
return torch.mean(T_y**2)
def bc_loss(net1, net2):
loss_edge = torch.mean((net1(Y_left) - T_left)**2) + torch.mean((net2(Y_right) - T_right)**2)
loss_insulated = insulated(net1, Y1_bottom) + insulated(net1, Y1_top) + insulated(net2, Y2_bottom) + insulated(net2, Y2_top)
return loss_edge + loss_insulated
Adam at $10^{-3}$ for $5000$ iterations. The three losses are printed every $1000$ iterations.
def derivative(y, t):
df = torch.autograd.grad(y, t, torch.ones_like(y), create_graph = True)[0]
return df[:, 0:1], df[:, 1:2]
model = new_pinn()
optimizer = torch.optim.Adam(model.parameters(), lr = 1e-3)
for epoch in range(5001):
loss_pde = torch.mean(laplacian(model, Y1_domain)**2) + torch.mean(laplacian(model, Y2_domain)**2)
loss_bc = bc_loss(model, model)
T_x, _ = derivative(model(Y_int), Y_int)
loss_flux = torch.mean((k1*T_x - k2*T_x)**2) # one network: same T_x on both sides
loss = loss_pde + loss_bc + loss_flux
optimizer.zero_grad()
loss.backward()
optimizer.step()
if epoch % 1000 == 0:
print('EPOCH : %5d | Loss_Total : %.3e | PDE : %.3e | BC : %.3e | Flux : %.3e'
% (epoch, loss.item(), loss_pde.item(), loss_bc.item(), loss_flux.item()))
The residual and boundary losses fall quickly, to about $10^{-4}$ and $6\times10^{-4}$, and then nothing more happens: the flux loss stops at $1.3\times10^{-2}$ by the first thousand iterations and stays there to the end. The optimizer has found the compromise described above. It satisfies the equation and the edge temperatures with one plane from $1$ to $0.2$, whose slope $-0.4$ makes the flux term $(0.3 - 0.6)^2 \times 0.4^2 \approx 0.014$, and it cannot reduce that term without breaking the other two.
3.6.2. Domain Decomposition¶
Domain decomposition splits the plate at the interface and gives each material its own network,
$$ T_1(x, y) = N_{\theta_1}(x, y) \ \text{ on } [-1, 0]\times[-1, 1] , \qquad T_2(x, y) = N_{\theta_2}(x, y) \ \text{ on } [0, 1]\times[-1, 1] . $$
Each network has its own residual on its own half and its own edge conditions, and the two share no parameters. What joins them is the interface, where the two conditions of Section 3.5 become two more loss terms,
$$ \mathcal{L}_T = \frac{1}{N_{\mathrm{int}}}\sum_i\bigl(T_1(0, y_i) - T_2(0, y_i)\bigr)^2 , \qquad \mathcal{L}_{\mathrm{flux}} = \frac{1}{N_{\mathrm{int}}}\sum_i\bigl(k_1\,T_{1, x}(0, y_i) - k_2\,T_{2, x}(0, y_i)\bigr)^2 , $$
so that the total loss is
$$ \mathcal{L} = \mathcal{L}_{\mathrm{PDE}} + \mathcal{L}_{\mathrm{BC}} + \mathcal{L}_T + \mathcal{L}_{\mathrm{flux}} . $$
The flux term now contains two different derivatives, one from each network, and it can be driven to zero with both slopes nonzero. Neither network is asked for a kink: each represents a plane, and the kink is produced by the two planes meeting along $x = 0$ with the slopes the flux condition requires.
The same points, the same number of iterations and the same learning rate; one optimizer holds the parameters of both networks.
model1 = new_pinn()
model2 = new_pinn()
optimizer = torch.optim.Adam(list(model1.parameters()) + list(model2.parameters()), lr = 1e-3)
for epoch in range(5001):
loss_pde = torch.mean(laplacian(model1, Y1_domain)**2) + torch.mean(laplacian(model2, Y2_domain)**2)
loss_bc = bc_loss(model1, model2)
T1_int = model1(Y_int)
T2_int = model2(Y_int)
T1_x, _ = derivative(T1_int, Y_int)
T2_x, _ = derivative(T2_int, Y_int)
loss_T = torch.mean((T1_int - T2_int)**2) # continuity of T
loss_flux = torch.mean((k1*T1_x - k2*T2_x)**2) # continuity of flux
loss = loss_pde + loss_bc + loss_T + loss_flux
optimizer.zero_grad()
loss.backward()
optimizer.step()
if epoch % 1000 == 0:
print('EPOCH : %5d | Loss_Total : %.3e | PDE : %.3e | BC : %.3e | T : %.3e | Flux : %.3e'
% (epoch, loss.item(), loss_pde.item(), loss_bc.item(), loss_T.item(), loss_flux.item()))
All five terms fall together. By $1000$ iterations the total loss is $4\times10^{-5}$, and at the end it is $10^{-6}$, with the flux term at $5\times10^{-8}$: the condition that the single network could not meet at all is met by the pair to high precision. The prediction is assembled from the two networks, each evaluated on its own half.
3.6.3. Results¶
Both predictions are evaluated on the $201\times201$ grid; the decomposed prediction takes $T_1$ on the left half and $T_2$ on the right. The gradient $\partial T/\partial x$ is evaluated on the same grid for the flux comparison below.
Y_eval = tensor(np.column_stack([X.ravel(), Y.ravel()]))
T_naive_t = model(Y_eval)
T1_t = model1(Y_eval)
T2_t = model2(Y_eval)
T_naive = T_naive_t.detach().cpu().numpy().reshape(201, 201)
T1 = T1_t.detach().cpu().numpy().reshape(201, 201)
T2 = T2_t.detach().cpu().numpy().reshape(201, 201)
T_dd = np.where(X < 0, T1, T2)
Tx_naive = derivative(T_naive_t, Y_eval)[0].detach().cpu().numpy().reshape(201, 201)
T1x = derivative(T1_t, Y_eval)[0].detach().cpu().numpy().reshape(201, 201)
T2x = derivative(T2_t, Y_eval)[0].detach().cpu().numpy().reshape(201, 201)
Tx_dd = np.where(X < 0, T1x, T2x)
error_naive = np.abs(T_naive - T_ref)
error_dd = np.abs(T_dd - T_ref)
plt.figure(figsize = (9, 2.4))
plt.subplot(1, 3, 1)
plt.imshow(T_ref, origin = 'lower', extent = [-1, 1, -1, 1], cmap = 'jet', vmin = 0.2, vmax = 1)
plt.colorbar()
plt.axvline(0, color = 'k', linestyle = '--')
plt.title('Ground Truth')
plt.xlabel('x')
plt.ylabel('y')
plt.subplot(1, 3, 2)
plt.imshow(T_naive, origin = 'lower', extent = [-1, 1, -1, 1], cmap = 'jet', vmin = 0.2, vmax = 1)
plt.colorbar()
plt.axvline(0, color = 'k', linestyle = '--')
plt.title('Naive PINN')
plt.xlabel('x')
plt.ylabel('y')
plt.subplot(1, 3, 3)
plt.imshow(T_dd, origin = 'lower', extent = [-1, 1, -1, 1], cmap = 'jet', vmin = 0.2, vmax = 1)
plt.colorbar()
plt.axvline(0, color = 'k', linestyle = '--')
plt.title('Domain Decomposition')
plt.xlabel('x')
plt.ylabel('y')
plt.tight_layout()
plt.show()
plt.figure(figsize = (8, 3))
plt.subplot(1, 2, 1)
plt.imshow(error_naive, origin = 'lower', extent = [-1, 1, -1, 1], cmap = 'gray', vmin = 0, vmax = 0.15)
plt.colorbar()
plt.axvline(0, color = 'gray', linestyle = '--')
plt.title('Naive PINN Error')
plt.xlabel('x')
plt.ylabel('y')
plt.subplot(1, 2, 2)
plt.imshow(error_dd, origin = 'lower', extent = [-1, 1, -1, 1], cmap = 'gray', vmin = 0, vmax = 0.15)
plt.colorbar()
plt.axvline(0, color = 'gray', linestyle = '--')
plt.title('Domain Decomposition Error')
plt.xlabel('x')
plt.ylabel('y')
plt.tight_layout()
plt.show()
print('max error naive : %.4f' % np.max(error_naive))
print('max error DD : %.4f' % np.max(error_dd))
On the colour maps the three fields look alike, which is the usual situation: an error at an interface is a thin feature, and a map of the whole plate does not show it well. The error maps are more telling. The naive error is a band along $x = 0$, brightest on the interface itself, where it reaches $0.13$, and falling off towards the edges where the temperatures are pinned; it does not depend on $y$, because nothing in the problem does. The decomposed error is dark over the whole plate, with a largest value of $3\times10^{-4}$.
A horizontal cut through the middle of the plate shows the kink itself. The flux $k\,\partial T/\partial x$ is the sharper test, because for the exact solution it is the constant $-0.16$ across the whole plate, interface included.
j = 100 # the row y = 0
k_line = np.where(x_grid < 0, k1, k2)
flux_ref = -k1*k2/(k1 + k2)*(T_left - T_right)*np.ones_like(x_grid)
plt.figure(figsize = (8, 4))
plt.subplot(1, 2, 1)
plt.plot(x_grid, T_ref[j], 'k', linewidth = 4, alpha = 0.3, label = 'Exact')
plt.plot(x_grid, T_naive[j], 'r--', linewidth = 1.5, label = 'Naive PINN')
plt.plot(x_grid, T_dd[j], 'b--', linewidth = 1.5, label = 'Domain Decomposition')
plt.axvline(0, color = 'gray', linestyle = '--')
plt.title('Temperature at y = 0')
plt.xlabel('x')
plt.ylabel('T')
plt.legend()
plt.grid(True)
plt.subplot(1, 2, 2)
plt.plot(x_grid, flux_ref, 'k', linewidth = 4, alpha = 0.3, label = 'Exact')
plt.plot(x_grid, k_line*Tx_naive[j], 'r--', linewidth = 1.5, label = 'Naive PINN')
plt.plot(x_grid, k_line*Tx_dd[j], 'b--', linewidth = 1.5, label = 'Domain Decomposition')
plt.axvline(0, color = 'gray', linestyle = '--')
plt.title('Flux k(x) T_x at y = 0')
plt.xlabel('x')
plt.ylabel('flux')
# plt.legend()
plt.grid(True)
plt.tight_layout()
plt.show()
The naive prediction is the straight line from $(-1, 1)$ to $(1, 0.2)$. It passes through $T = 0.6$ at the interface, where the exact value is $7/15 \approx 0.47$, and it is too high over the whole plate. The decomposed prediction lies on the exact solution: two lines of different slope meeting at $x = 0$.
The flux makes the failure explicit. The naive flux is $-0.12$ on the left and $-0.23$ on the right, a jump at the interface by a factor of two: the single network has one slope, and multiplying it by two different conductivities gives two different fluxes, so heat is not conserved across $x = 0$. The decomposed flux is the flat line at $-0.16$ on both sides, with no step at the interface, because the flux condition was imposed between two slopes that could differ.
3.6.4. What the Lab Showed¶
A continuous solution with a jump in its gradient is beyond a single smooth network, and the failure is not a matter of training: the network was given the interface condition explicitly, and the optimizer converged, to a loss it could not lower further. With one slope at $x = 0$ the flux condition and the governing equation contradict each other, and the compromise the optimizer found satisfies the equation and ignores the interface. Domain decomposition assigns one network to each material, so that every network represents a smooth function, and imposes the physics of the interface, continuity of the field and of the flux, as loss terms that join them. The kink is then not something a network has to learn but something two networks produce by meeting, and it appears where the physics puts it.
Remark (where the interface is). Decomposition requires knowing where to cut. For a material interface that is part of the problem statement, as here; the same construction applies to any line across which the physics changes, a phase boundary or a contact surface, as long as its location is known. When the location is itself unknown, as in a moving-boundary problem, the interface has to be represented and solved for as well, which is a different and harder problem.
Three Structures, Three Remedies
The three laboratories can be read as one argument. In each, a standard PINN with a good optimizer and enough points failed, not because the loss was badly weighted or the points badly placed, but because the function it was asked to represent did not suit the network. And in each, the remedy was not a better optimizer but a different representation, changed at a different place.
| Structure | Why a single network fails | Remedy | Where the representation changes |
|---|---|---|---|
| High frequency | spectral bias: fast components learned last | Fourier feature embedding | the input |
| Singularity / sharp gradient | a smooth function cannot step or blow up | enrichment with a known singular basis | the output |
| Gradient discontinuity | a differentiable function has no kink | domain decomposition with interface conditions | the domain |
All three require knowing something about the solution in advance: its frequency range, the form of its singularity, or the location of the interface. That knowledge is physics, not data, and supplying it to the network is in the spirit of the whole subject: the governing equation told the network what to satisfy, and the representation tells it where to look.
The last limitation belongs to time-dependent problems, and it is of a different kind from the others. A transient solution is not a field that happens to have a time axis; it is a history. The state at a later time is produced from the state at an earlier time by the governing equation, and nothing at $t_2$ can be known until everything up to $t_2$ is. A numerical time-stepping scheme respects this by construction, because it has no choice: it computes the solution one step at a time, in order. A PINN has a choice, and by default it makes the wrong one.
The reason is that time enters a PINN as one more input coordinate. Collocation points are scattered over the whole space-time domain, the residual is averaged over all of them, and from the first iteration the optimizer is pulling on the solution at $t = 1.5$ as hard as on the solution at $t = 0.1$. Nothing in the loss says that the late residual is meaningless until the early solution is right. The equation is causal; the objective is not. On easy problems this does no harm. On problems whose dynamics are fast, or whose trajectory passes through states that a nearly stationary field can imitate at small residual, it leads the optimizer to a solution that satisfies the equation well and the physics not at all.
Two remedies are shown. Causal training keeps one network over the whole interval but weights the residual so that later times count only after earlier times have converged. Time marching splits the interval into windows and trains them one after another, each starting from the end state of the one before. The first changes the loss, the second changes the problem, and the chapter ends by noting that they can be combined.
For a transient system the solution has a direction. If $t_0 < t_1 < t_2$, the trajectory is generated in the order
$$ u(\cdot, t_0) \;\longrightarrow\; u(\cdot, t_1) \;\longrightarrow\; u(\cdot, t_2) , $$
and the state at $t_2$ is not an independent target but the result of applying the dynamics to the state at $t_1$, which is itself the result of applying them to the state at $t_0$. Get $t_1$ wrong and $t_2$ is wrong for that reason alone.
A standard PINN does not see this. With collocation points $\{(x_j, t_i)\}$ laid over the whole domain, the residual loss is
$$ \mathcal{L}_{\mathrm{PDE}} = \frac{1}{N_t N_x}\sum_{i=1}^{N_t}\sum_{j=1}^{N_x} r_\theta(x_j, t_i)^2 , $$
and every time slice contributes to the same update from the first iteration. The optimizer is free to reduce the residual at a late time while the solution near the initial condition is still wrong, and because one set of parameters represents the whole field, a step taken to fix a late residual moves the early solution as well. What the optimizer finds in the end is a field with a small residual everywhere, but a residual can be small on the wrong trajectory: the governing equation has many solutions, and only the initial condition, propagated forward in time, picks out the right one. If the propagation is never enforced in order, the optimizer is free to settle on some other solution of the equation that happens to be easy to represent.
The difficulty is worst when the dynamics are fast, when information has to travel a long way over the interval, or when a nearly stationary state satisfies the equation almost as well as the true transient does. The Allen–Cahn equation below has the last property, and is used throughout Sections 4.1 to 4.3 to show the failure and the first remedy.
Consider the one-dimensional Allen–Cahn equation
$$ u_t - \nu\,u_{xx} + 5u^3 - 5u = 0 , \qquad \nu = 10^{-4} , $$
on $x \in [-1, 1]$, $t \in [0, 1.5]$, with periodic boundaries and the initial condition
$$ u(x, 0) = x^2\cos(\pi x) . $$
The reaction term $5u - 5u^3$ pushes $u$ away from zero and towards $\pm 1$, so that an initial profile of modest amplitude grows into plateaus at $+1$ and $-1$ separated by thin transition layers, which the small diffusion keeps from becoming discontinuous. The important point for what follows is that $u \equiv 0$ is also a solution of the equation: it has zero residual everywhere, it is simply not the solution that starts from this initial condition.
A reference trajectory is computed by the method of lines, a finite difference in $x$ and an adaptive integrator in $t$. It is used only to draw the pictures and to measure the error; the networks never see it.
from scipy.integrate import solve_ivp
nu = 1e-4
x_ref = np.linspace(-1, 1, 512, endpoint = False)
t_ref = np.linspace(0, 1.5, 201)
dx = x_ref[1] - x_ref[0]
u0 = x_ref**2*np.cos(np.pi*x_ref)
def allen_cahn(t, u):
u_xx = (np.roll(u, -1) - 2*u + np.roll(u, 1))/dx**2
return nu*u_xx - 5*u**3 + 5*u
u_ref = solve_ivp(allen_cahn, (0, 1.5), u0, t_eval = t_ref, rtol = 1e-6, atol = 1e-8).y.T # (201, 512)
Three profiles of the reference solution, at $t = 0$, $0.75$ and $1.5$, show what the equation does to the initial condition.
plt.figure(figsize = (6, 4))
for t in [0, 0.75, 1.5]:
plt.plot(x_ref, u_ref[np.argmin(np.abs(t_ref - t))], linewidth = 2, label = 't = %.2f' % t)
plt.xlabel('x')
plt.ylabel('u(x, t)')
plt.legend()
plt.grid(True)
plt.show()
At $t = 0$ the profile is the small bump $x^2\cos(\pi x)$. By $t = 0.75$ the reaction has amplified it: where $u$ was positive it has risen towards $+1$, where negative towards $-1$, and the transitions between the plateaus have sharpened. By $t = 1.5$ the plateaus are flat and the transitions are nearly vertical. Each profile is produced from the one before it, and the final shape is the result of that history, not a fact that can be read off the equation alone.
A standard PINN does not train the three profiles in order. Its collocation points cover the whole $(t, x)$ rectangle at once, and the residual at every one of them enters the same loss at every iteration.
T_col, X_col = np.meshgrid(np.linspace(0, 1.5, 10), np.linspace(-1, 1, 10))
plt.figure(figsize = (6, 4))
plt.imshow(u_ref.T, extent = [0, 1.5, -1, 1], origin = 'lower', aspect = 'auto', cmap = 'jet', vmin = -1, vmax = 1)
plt.plot(T_col.ravel(), X_col.ravel(), 'k.')
plt.title('Collocation Points over the Reference Solution')
plt.xlabel('t')
plt.ylabel('x')
plt.show()
The dots are the collocation points, drawn over the reference solution so that the two can be seen together. The points at $t = 1.5$ are present from the first optimization step, and the equation will be enforced there before the network has any idea what the solution at $t = 0.5$ looks like. In the picture that is a harmless-looking grid; in the optimization it is the whole problem.
The network takes $(x, t)$ and returns $u_\theta(x, t)$. The periodic boundary condition in $x$ is built into the input rather than penalized: the coordinate $x$ is replaced by the Fourier features $\cos(k\pi x)$ and $\sin(k\pi x)$ for $k = 1, \ldots, 10$, all of which have period $2$, so that whatever the network computes from them is automatically periodic on $[-1, 1]$. This is the embedding of Section 3.3 used for a different purpose, periodicity rather than frequency, and it removes the boundary term from the loss altogether.
The residual of the Allen–Cahn equation is
$$ r_\theta = u_{\theta, t} + 5u_\theta^3 - 5u_\theta - \nu\,u_{\theta, xx} , $$
and the initial condition is a penalty,
$$ \mathcal{L}_{\mathrm{IC}} = \frac{1}{N_{\mathrm{IC}}}\sum_i\bigl[u_\theta(x_i, 0) - x_i^2\cos(\pi x_i)\bigr]^2 , $$
so that the naive objective is
$$ \mathcal{L}_{\mathrm{Naive}} = \mathcal{L}_{\mathrm{PDE}} + 100\,\mathcal{L}_{\mathrm{IC}} , $$
with the residual averaged over a $100\times256$ grid of collocation points in $(t, x)$. The weight $100$ on the initial condition is a fixed weight of the kind discussed in Regularization, MAP Estimation, and Loss Weighting; the initial condition is the only thing that distinguishes the wanted solution from every other solution of the equation, and it must not be drowned out by $25600$ residual terms.
k = torch.arange(1, 11).reshape(1, -1).to(device)
def features(x, t):
return torch.cat([t, torch.ones_like(t), torch.cos(k*np.pi*x), torch.sin(k*np.pi*x)], dim = 1)
def new_pinn():
net = nn.Sequential(nn.Linear(22, 128), nn.Tanh(),
nn.Linear(128, 128), nn.Tanh(),
nn.Linear(128, 128), nn.Tanh(),
nn.Linear(128, 128), nn.Tanh(),
nn.Linear(128, 1))
for m in net: # Xavier initialization of every linear layer
if isinstance(m, nn.Linear):
nn.init.xavier_normal_(m.weight)
nn.init.zeros_(m.bias)
return net.to(device)
def derivative(y, x):
return torch.autograd.grad(y, x, torch.ones_like(y), create_graph = True)[0]
def residual(net, x, t):
u = net(features(x, t))
u_t = derivative(u, t)
u_x = derivative(u, x)
u_xx = derivative(u_x, x)
return u_t + 5*u**3 - 5*u - nu*u_xx
T_pde, X_pde = torch.meshgrid(torch.linspace(0, 1.5, 100), torch.linspace(-1, 1, 256), indexing = 'ij')
x_pde = X_pde.reshape(-1, 1).to(device).requires_grad_(True)
t_pde = T_pde.reshape(-1, 1).to(device).requires_grad_(True)
x_ic = torch.tensor(x_ref, dtype = torch.float32).reshape(-1, 1).to(device)
t_ic = torch.zeros_like(x_ic)
u_ic = x_ic**2*torch.cos(np.pi*x_ic)
T_eval, X_eval = np.meshgrid(t_ref, x_ref, indexing = 'ij')
x_eval = torch.tensor(X_eval.reshape(-1, 1), dtype = torch.float32).to(device)
t_eval = torch.tensor(T_eval.reshape(-1, 1), dtype = torch.float32).to(device)
def predict(net):
return net(features(x_eval, t_eval)).detach().cpu().numpy().reshape(201, 512)
Adam at $10^{-3}$ with an exponential decay of the learning rate, a factor of $0.9$ every $5000$ iterations, for $150000$ iterations. Every residual on the grid is averaged with the same weight at every iteration.
# # Training takes a long time. The pre-trained model is loaded in the cells below.
# naive_net = new_pinn()
# optimizer = torch.optim.Adam(naive_net.parameters(), lr = 1e-3)
# scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma = 0.9**(1/5000))
# n_iter = 150000
# for it in range(1, n_iter + 1):
# r = residual(naive_net, x_pde, t_pde)
# loss_pde = torch.mean(r**2)
# loss_ic = torch.mean((naive_net(features(x_ic, t_ic)) - u_ic)**2)
# loss = loss_pde + 100*loss_ic
# optimizer.zero_grad()
# loss.backward()
# optimizer.step()
# scheduler.step()
# if it == 1 or it % 10000 == 0:
# print('Naive | %6d | Total %.3e | PDE %.3e | IC %.3e' % (it, loss.item(), loss_pde.item(), loss_ic.item()))
# torch.save(naive_net, 'allen_cahn_naive.pt')
Training takes a long time, so the pre-trained model is downloaded and loaded instead. The file holds the trained network.
!gdown 1-jAHsy9CpxmAkxbS9k_yOtch63OJFa4e -O allen_cahn_naive.pt
naive_net = torch.load('allen_cahn_naive.pt', map_location = device, weights_only = False)
The prediction is compared with the reference on the full grid: the two fields and their absolute difference, then the mean absolute error at each time,
$$ \mathrm{MAE}(t) = \frac{1}{N_x}\sum_j\bigl|u_\theta(x_j, t) - u^*(x_j, t)\bigr| , $$
which collapses the error map into one curve over $t$.
u_naive = predict(naive_net)
error_naive = np.abs(u_naive - u_ref)
plt.figure(figsize = (9, 3))
plt.subplot(1, 3, 1)
plt.imshow(u_ref.T, extent = [0, 1.5, -1, 1], origin = 'lower', aspect = 'auto', cmap = 'jet', vmin = -1, vmax = 1)
plt.colorbar()
plt.title('Ground Truth')
plt.xlabel('t')
plt.ylabel('x')
plt.subplot(1, 3, 2)
plt.imshow(u_naive.T, extent = [0, 1.5, -1, 1], origin = 'lower', aspect = 'auto', cmap = 'jet', vmin = -1, vmax = 1)
plt.colorbar()
plt.title('Naive PINN')
plt.xlabel('t')
# plt.ylabel('x')
plt.subplot(1, 3, 3)
plt.imshow(error_naive.T, extent = [0, 1.5, -1, 1], origin = 'lower', aspect = 'auto', cmap = 'gray', vmin = 0, vmax = 0.4)
plt.colorbar()
plt.title('Naive PINN Error')
plt.xlabel('t')
plt.ylabel('x')
plt.tight_layout()
plt.show()
plt.figure(figsize = (6, 4))
plt.plot(t_ref, np.mean(error_naive, axis = 1), 'k', linewidth = 2)
plt.xlabel('t')
plt.ylabel('MAE(t)')
plt.grid(True)
plt.show()
Judged by the loss, training is a success: the total loss falls by six orders of magnitude and nothing in the numbers suggests a problem. The pictures say otherwise. The naive prediction matches the reference for a short time after $t = 0$, where the initial condition pins it, and then flattens out: instead of growing into the $\pm 1$ plateaus it decays towards the trivial state $u \approx 0$, which has zero residual and is far easier for the network to represent than two plateaus with sharp transitions. The error map is dark near $t = 0$ and bright everywhere else, and $\mathrm{MAE}(t)$ climbs steadily from zero to about $0.5$ across the interval.
This is the failure described at the start of the chapter, seen in a concrete case. The optimizer satisfied the equation, but it did so with a solution that does not follow from the initial condition, because nothing in the objective required the early solution to be right before the late one was fitted. The residual is small along the wrong trajectory; a decreasing loss is not evidence that the initial state has propagated correctly through the domain.
A causal PINN keeps one network over the whole domain and changes only how much each time slice is allowed to matter. Write the residual loss of the $i$th slice as
$$ L_i = \frac{1}{N_x}\sum_{j=1}^{N_x} r_\theta(x_j, t_i)^2 , $$
so that the ordinary loss is the plain average $\mathcal{L}_{\mathrm{PDE}} = \frac{1}{N_t}\sum_i L_i$. Causal training replaces the average by a weighted one,
$$ \mathcal{L}_{\mathrm{PDE}}^{\mathrm{causal}} = \frac{1}{N_t}\sum_{i=1}^{N_t} w_i\,L_i , \qquad w_i = \exp\Bigl(-\epsilon\sum_{k<i} L_k\Bigr) . $$
The weight of a slice is an exponential of the accumulated loss of all the slices before it. The first slice has nothing before it, so $w_1 = 1$ and it is trained at full strength from the start. A later slice has weight close to zero as long as the earlier slices are still poorly fitted, and its weight rises towards one only as their losses fall. Training therefore proceeds from left to right in time, not because the collocation points are ordered but because the loss is: the later residuals are present at every iteration, and they are ignored until the earlier ones have been dealt with.
The parameter $\epsilon$ sets how strictly the order is enforced. For $\epsilon \to 0$ every weight is one and the ordinary loss is recovered; for large $\epsilon$ even a small early loss suppresses everything after it. Two details of the implementation matter. The accumulated losses inside the exponential are detached from the graph, so that the weights act as a schedule and are not themselves optimized, and the weights are recomputed at every iteration from the current residuals, so that the active region moves forward on its own as the solution improves. Nothing else changes: no new equation, no extra network, no partition of the domain.
The mechanism in miniature, for six time slices. In the left panel all six losses are equal to $1$ and the weights are shown for three values of $\epsilon$; in the middle panel the first loss has dropped to $0.2$; in the right panel the first two have.
def causal_weights(L, eps):
L = np.array(L)
return np.exp(-eps*(np.cumsum(L) - L))
L1 = [1, 1, 1, 1, 1, 1]
L2 = [0.2, 1, 1, 1, 1, 1]
L3 = [0.2, 0.2, 1, 1, 1, 1]
plt.figure(figsize = (9, 3))
plt.subplot(1, 3, 1)
for eps in [0.1, 1.0, 2.0]:
plt.plot(causal_weights(L1, eps), 'o-', label = r'$\epsilon$ = %.1f' % eps)
plt.title('all losses = 1')
plt.xlabel('time slice i')
plt.ylabel(r'weight $w_i$')
plt.ylim(0, 1.05)
plt.legend()
plt.grid(True)
plt.subplot(1, 3, 2)
plt.plot(causal_weights(L2, 1.0), 'o-', color = 'C1')
plt.title(r'$L_0$ decreased, $\epsilon$ = 1')
plt.xlabel('time slice i')
plt.ylim(0, 1.05)
plt.grid(True)
plt.subplot(1, 3, 3)
plt.plot(causal_weights(L3, 1.0), 'o-', color = 'C1')
plt.title(r'$L_0, L_1$ decreased, $\epsilon$ = 1')
plt.xlabel('time slice i')
plt.ylim(0, 1.05)
plt.grid(True)
plt.tight_layout()
plt.show()
In the left panel the weights fall off geometrically with the slice index, slowly for $\epsilon = 0.1$ and steeply for $\epsilon = 2$; at the start of training only the first slice or two are active. When the first loss drops, the weights of all later slices rise together, because the sum inside the exponential has shrunk for every one of them; when the second drops as well, the active front moves one slice further. The region in which the equation is actually enforced advances in time as the solution behind it converges, which is the behaviour of a time-stepping scheme obtained from a single loss.
The causal model is the naive model of Section 4.1 with one change, the weighting of the residual in time. Same network, same Fourier features, same collocation grid, same optimizer and learning-rate schedule, same initial-condition weight and the same $150000$ iterations. Whatever differs in the result is due to the order in which the residuals were allowed to act.
4.3.1. Causal Training¶
The residual on the $100\times256$ grid is reshaped into $100$ rows, one per time, and averaged along each row to give $L_i$. The weights are built from the detached cumulative sum with $\epsilon = 100$, a large value chosen because the losses here are small numbers, and the weighted average replaces the plain one. The smallest weight is printed alongside the losses as training proceeds.
# # Training takes a long time. The pre-trained model is loaded in the cells below.
# causal_net = new_pinn()
# optimizer = torch.optim.Adam(causal_net.parameters(), lr = 1e-3)
# scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma = 0.9**(1/5000))
# eps = 100.0
# n_iter = 150000
# for it in range(1, n_iter + 1):
# r = residual(causal_net, x_pde, t_pde)
# L = torch.mean(r.reshape(100, 256)**2, dim = 1) # one loss per time slice
# w = torch.exp(-eps*(torch.cumsum(L, 0) - L)).detach() # causal weights
# loss_pde = torch.mean(w*L)
# loss_ic = torch.mean((causal_net(features(x_ic, t_ic)) - u_ic)**2)
# loss = loss_pde + 100*loss_ic
# optimizer.zero_grad()
# loss.backward()
# optimizer.step()
# scheduler.step()
# if it == 1 or it % 10000 == 0:
# print('Causal | %6d | Total %.3e | PDE %.3e | IC %.3e | W_min %.2e' % (it, loss.item(), loss_pde.item(), loss_ic.item(), w.min().item()))
# torch.save(causal_net, 'allen_cahn_causal.pt')
The trained network is downloaded and loaded instead.
!gdown 1qT3Fehd64l0dwP5wzJWB7yEIyLXNTJl3 -O allen_cahn_causal.pt
causal_net = torch.load('allen_cahn_causal.pt', map_location = device, weights_only = False)
4.3.2. Results¶
The two predictions side by side with the reference, their error maps, and the time-wise error.
u_causal = predict(causal_net)
error_causal = np.abs(u_causal - u_ref)
plt.figure(figsize = (9, 3))
plt.subplot(1, 3, 1)
plt.imshow(u_ref.T, extent = [0, 1.5, -1, 1], origin = 'lower', aspect = 'auto', cmap = 'jet', vmin = -1, vmax = 1)
plt.colorbar()
plt.title('Ground Truth')
plt.xlabel('t')
plt.ylabel('x')
plt.subplot(1, 3, 2)
plt.imshow(u_naive.T, extent = [0, 1.5, -1, 1], origin = 'lower', aspect = 'auto', cmap = 'jet', vmin = -1, vmax = 1)
plt.colorbar()
plt.title('Naive PINN')
plt.xlabel('t')
plt.ylabel('x')
plt.subplot(1, 3, 3)
plt.imshow(u_causal.T, extent = [0, 1.5, -1, 1], origin = 'lower', aspect = 'auto', cmap = 'jet', vmin = -1, vmax = 1)
plt.colorbar()
plt.title('Causal PINN')
plt.xlabel('t')
plt.ylabel('x')
plt.tight_layout()
plt.show()
plt.figure(figsize = (6, 3))
plt.subplot(1, 2, 1)
plt.imshow(error_naive.T, extent = [0, 1.5, -1, 1], origin = 'lower', aspect = 'auto', cmap = 'gray', vmin = 0, vmax = 0.4)
plt.colorbar()
plt.title('Naive PINN Error')
plt.xlabel('t')
plt.ylabel('x')
plt.subplot(1, 2, 2)
plt.imshow(error_causal.T, extent = [0, 1.5, -1, 1], origin = 'lower', aspect = 'auto', cmap = 'gray', vmin = 0, vmax = 0.4)
plt.colorbar()
plt.title('Causal PINN Error')
plt.xlabel('t')
plt.ylabel('x')
plt.tight_layout()
plt.show()
plt.figure(figsize = (6, 4))
plt.plot(t_ref, np.mean(error_naive, axis = 1), 'r', linewidth = 2, label = 'Naive PINN')
plt.plot(t_ref, np.mean(error_causal, axis = 1), 'b', linewidth = 2, label = 'Causal PINN')
plt.xlabel('t')
plt.ylabel('MAE(t)')
plt.legend()
plt.grid(True)
plt.show()
The causal prediction is the reference solution: the bump grows into the two plateaus, the transition layers sharpen at the right times, and the error map is dark over the whole rectangle except for faint lines along the transitions, where a small misplacement of a steep front produces a visible error. Its $\mathrm{MAE}(t)$ stays near zero across the interval, where the naive model's climbed to $0.5$. Both runs end with a loss of the same small size; the difference between the models is not in how small the loss became but in which solution of the equation it was small on.
During training the printed smallest weight shows the front advancing: it stays at zero while the early part of the interval is being fitted, meaning the last slices are not being trained, and rises to order one only once the early part is right. The late plateaus are then fitted as the consequence of the early growth rather than as an independent target.
4.3.3. What the Lab Showed¶
The naive and causal models solve the same equation with the same initial condition, the same network and the same budget; what the causal model changes is when each residual is permitted to influence the parameters. That a better late-time prediction follows is evidence that order matters in the training of a transient PINN, not evidence about the equation or the boundary treatment, which were untouched. The lesson is that a trajectory has to be built from its initial state forward, and that a loss which does not say so will be minimized in whatever order is cheapest.
Causal weighting keeps one network over the whole interval and arranges the loss so that the network learns the interval in order. Time marching is the more literal approach: it changes the problem rather than the loss, by cutting the interval into windows
$$ [t_0, t_1],\ [t_1, t_2],\ \ldots,\ [t_{M-1}, t_M] $$
and solving them one after another. The first window is an ordinary PINN problem on a short interval, with the true initial condition. When it has converged, its solution at $t_1$ becomes the initial condition of the second window,
$$ u_k(x, t_k) \;\longrightarrow\; u_{k+1}(x, t_k) , $$
imposed by a penalty of the usual form,
$$ \mathcal{L}_{\mathrm{interface}} = \frac{1}{N_x}\sum_i\bigl[u_{k+1}(x_i, t_k) - u_k(x_i, t_k)\bigr]^2 , $$
and the second window is trained; and so on to the end. The network of each window is usually initialized from the trained network of the one before, so that both the state and a good starting point for the weights are passed forward.
What this buys is a short horizon. While the first window is being trained, the far future is not in the problem at all, so the optimizer cannot spend effort on it and cannot be led astray by it; the trajectory is built forward in the same direction as the physics, one window at a time. What it costs is sequential dependence. A later window cannot start until the earlier one is done, so the stages cannot be run in parallel, and an error at an interface is handed to every window after it. The window length sets the trade-off: short windows keep each problem easy but multiply the number of stages and interfaces, long windows do the opposite, and in the limit of one window the method is the naive PINN again.
Time marching. The interval is cut into windows. The shaded window is the one being trained; the windows to its right are not in the problem yet. When a window converges, its solution at the right edge is predicted and handed to the next window as an initial condition.
The figure shows the arrangement. The interval is divided into windows, the one being trained is shaded, and the ones before it have already been solved and frozen; the state at the right edge of a solved window is predicted and handed to the next as its initial condition. Causal weighting and time marching impose the same order at different levels, one inside a single objective and the other across a sequence of problems, and the difference will be visible in the next demonstration: in the time-marching prediction the future is simply absent until its window has been trained.
4.5.1. Problem Setup¶
Time marching is demonstrated on the linear convection equation
$$ u_t + \beta\,u_x = 0 , \qquad x \in [0, 2\pi], \quad t \in [0, 1] , $$
with the periodic boundary condition $u(0, t) = u(2\pi, t)$ and the initial condition $u(x, 0) = \sin x$. The exact solution is
$$ u(x, t) = \sin(x - \beta t) : $$
the initial profile travels to the right at speed $\beta$ without changing shape. For $\beta = 1$ it moves a sixth of the way around the periodic domain in the unit interval. For $\beta = 80$ it goes around nearly thirteen times, and the solution in the $(t, x)$ plane is a dense set of diagonal stripes. The equation is as simple as a PDE can be; the only difficulty is that information has to be carried a long way from the initial condition, which is exactly the situation in which a full-domain PINN loses its way.
The exact solution in the $(t, x)$ plane for the two speeds.
beta = 80.0
T_ex, X_ex = np.meshgrid(np.linspace(0, 1, 300), np.linspace(0, 2*np.pi, 300), indexing = 'ij')
plt.figure(figsize = (6, 3))
plt.subplot(1, 2, 1)
plt.imshow(np.sin(X_ex - 1.0*T_ex).T, extent = [0, 1, 0, 2*np.pi], origin = 'lower', aspect = 'auto', cmap = 'jet', vmin = -1, vmax = 1)
plt.colorbar()
plt.title(r'$\beta$ = 1')
plt.xlabel('t')
plt.ylabel('x')
plt.subplot(1, 2, 2)
plt.imshow(np.sin(X_ex - beta*T_ex).T, extent = [0, 1, 0, 2*np.pi], origin = 'lower', aspect = 'auto', cmap = 'jet', vmin = -1, vmax = 1)
plt.colorbar()
plt.title(r'$\beta$ = 80')
plt.xlabel('t')
plt.ylabel('x')
plt.tight_layout()
plt.show()
With $\beta = 1$ the stripes are nearly vertical and a network has to represent one gentle drift. With $\beta = 80$ the same profile crosses the picture many times, and the network has to represent a field that oscillates rapidly in $t$ as well as in $x$, with every stripe in the right place relative to the initial condition at $t = 0$. The comparison below uses $\beta = 80$.
4.5.2. Naive and Time-Marching Formulations¶
Both models use the same network, a plain $2 \to 64 \to 64 \to 64 \to 64 \to 1$ MLP with $\tanh$ and the inputs scaled to $[-1, 1]$, and the same residual $r_\theta = u_{\theta, t} + \beta\,u_{\theta, x}$. For the naive model the loss is
$$ \mathcal{L} = \mathcal{L}_{\mathrm{PDE}} + \mathcal{L}_{\mathrm{IC}} + \mathcal{L}_{\mathrm{BC}} , $$
with the initial condition $u_\theta(x, 0) = \sin x$ and the periodic condition $u_\theta(0, t) = u_\theta(2\pi, t)$ both as penalties, and the residual on a $256\times64$ grid over the whole rectangle. For time marching the same three terms are used inside each window, with the initial-condition target replaced, from the second window on, by the previous window's prediction at the interface.
def features_conv(x, t):
return torch.cat([x/np.pi - 1, 2*t - 1], dim = 1) # inputs scaled to [-1, 1]
def new_conv_pinn():
net = nn.Sequential(nn.Linear(2, 64), nn.Tanh(),
nn.Linear(64, 64), nn.Tanh(),
nn.Linear(64, 64), nn.Tanh(),
nn.Linear(64, 64), nn.Tanh(),
nn.Linear(64, 1))
return net.to(device)
def pde_loss(net, x, t):
u = net(features_conv(x, t))
r = derivative(u, t) + beta*derivative(u, x)
return torch.mean(r**2)
def bc_loss(net, t):
x0 = torch.zeros_like(t)
x1 = torch.full_like(t, 2*np.pi)
return torch.mean((net(features_conv(x0, t)) - net(features_conv(x1, t)))**2)
x_ic = torch.linspace(0, 2*np.pi, 128).reshape(-1, 1).to(device)
t_ic = torch.zeros_like(x_ic)
u_ic = torch.sin(x_ic)
x_1d = torch.linspace(0, 2*np.pi, 66)[1:-1]
t_1d = torch.linspace(0, 1, 257)[1:]
T_pde, X_pde = torch.meshgrid(t_1d, x_1d, indexing = 'ij')
x_pde = X_pde.reshape(-1, 1).to(device).requires_grad_(True)
t_pde = T_pde.reshape(-1, 1).to(device).requires_grad_(True)
t_bc = torch.linspace(0, 1, 128).reshape(-1, 1).to(device)
t_eval_1d = np.linspace(0, 1, 300)
x_eval_1d = np.linspace(0, 2*np.pi, 300)
T_eval, X_eval = np.meshgrid(t_eval_1d, x_eval_1d, indexing = 'ij')
x_eval = torch.tensor(X_eval.reshape(-1, 1), dtype = torch.float32).to(device)
t_eval = torch.tensor(T_eval.reshape(-1, 1), dtype = torch.float32).to(device)
u_exact = np.sin(X_eval - beta*T_eval)
def predict_conv(net):
return net(features_conv(x_eval, t_eval)).detach().cpu().numpy().reshape(300, 300)
4.5.3. Naive PINN¶
Adam at $10^{-4}$ for $200000$ iterations on the whole rectangle.
# # Training takes a long time. The pre-trained model is loaded in the cells below.
# naive_net = new_conv_pinn()
# optimizer = torch.optim.Adam(naive_net.parameters(), lr = 1e-4)
# n_iter = 200000
# for it in range(1, n_iter + 1):
# loss_pde = pde_loss(naive_net, x_pde, t_pde)
# loss_ic = torch.mean((naive_net(features_conv(x_ic, t_ic)) - u_ic)**2)
# loss_bc = bc_loss(naive_net, t_bc)
# loss = loss_pde + loss_ic + loss_bc
# optimizer.zero_grad()
# loss.backward()
# optimizer.step()
# if it == 1 or it % 20000 == 0:
# print('Naive | %6d | Total %.3e | PDE %.3e | IC %.3e | BC %.3e' % (it, loss.item(), loss_pde.item(), loss_ic.item(), loss_bc.item()))
# torch.save(naive_net, 'convection_naive.pt')
The trained network is downloaded and loaded instead.
!gdown 11w-9bDpi1JLfzy7e6lwJgq0aiZilODl5 -O convection_naive.pt
naive_net = torch.load('convection_naive.pt', map_location = device, weights_only = False)
u_naive = predict_conv(naive_net)
error_naive = np.abs(u_naive - u_exact)
plt.figure(figsize = (9, 3))
plt.subplot(1, 3, 1)
plt.imshow(u_exact.T, extent = [0, 1, 0, 2*np.pi], origin = 'lower', aspect = 'auto', cmap = 'jet', vmin = -1, vmax = 1)
plt.colorbar()
plt.title('Ground Truth')
plt.xlabel('t')
plt.ylabel('x')
plt.subplot(1, 3, 2)
plt.imshow(u_naive.T, extent = [0, 1, 0, 2*np.pi], origin = 'lower', aspect = 'auto', cmap = 'jet', vmin = -1, vmax = 1)
plt.colorbar()
plt.title('Naive PINN')
plt.xlabel('t')
plt.ylabel('x')
plt.subplot(1, 3, 3)
plt.imshow(error_naive.T, extent = [0, 1, 0, 2*np.pi], origin = 'lower', aspect = 'auto', cmap = 'gray', vmin = 0, vmax = 0.4)
plt.colorbar()
plt.title('Naive PINN Error')
plt.xlabel('t')
plt.ylabel('x')
plt.tight_layout()
plt.show()
4.5.4. Time Marching¶
The interval is divided into ten windows of length $0.1$, each trained for $20000$ iterations, so that the total budget equals the naive run's. The first window starts from $\sin x$; each later window starts from the previous network's prediction at the shared time, computed once and held fixed, and its network is initialized from the previous window's weights. The trained networks are kept in a list, one per window, and the full prediction is assembled from them afterwards.
# # Training takes a long time. The pre-trained model is loaded in the cells below.
# n_windows = 10
# n_iter = 200000
# tm_nets = []
# it = 0
# net = new_conv_pinn()
# for w in range(n_windows):
# t_start = w/n_windows
# t_end = (w + 1)/n_windows
# if w > 0:
# u_start = net(features_conv(x_ic, torch.full_like(x_ic, t_start))).detach() # end state of the previous window
# new_net = new_conv_pinn()
# new_net.load_state_dict(net.state_dict())
# net = new_net
# else:
# u_start = u_ic
# optimizer = torch.optim.Adam(net.parameters(), lr = 1e-4)
# T_w, X_w = torch.meshgrid(torch.linspace(t_start, t_end, 257)[1:], x_1d, indexing = 'ij')
# x_w = X_w.reshape(-1, 1).to(device).requires_grad_(True)
# t_w = T_w.reshape(-1, 1).to(device).requires_grad_(True)
# t_bc_w = torch.linspace(t_start, t_end, 128).reshape(-1, 1).to(device)
# for _ in range(n_iter//n_windows):
# it += 1
# loss_pde = pde_loss(net, x_w, t_w)
# loss_ic = torch.mean((net(features_conv(x_ic, torch.full_like(x_ic, t_start))) - u_start)**2)
# loss_bc = bc_loss(net, t_bc_w)
# loss = loss_pde + loss_ic + loss_bc
# optimizer.zero_grad()
# loss.backward()
# optimizer.step()
# if it == 1 or it % 20000 == 0:
# print('Time marching | %6d | Window %2d | Total %.3e' % (it, w + 1, loss.item()))
# tm_nets.append(net)
# torch.save(tm_nets, 'convection_time_marching.pt')
One network per window, stored together in one file.
!gdown 1Kk5iJKANsDRT1i5ESgyDbCRqzGAhllCi -O convection_time_marching.pt
n_windows = 10
tm_nets = torch.load('convection_time_marching.pt', map_location = device, weights_only = False)
4.5.5. Results¶
The time-marching prediction is assembled window by window: each network is evaluated on the whole grid and only the times of its own window are kept.
u_naive = predict_conv(naive_net)
u_tm = np.zeros((300, 300))
for w in range(n_windows):
u_w = predict_conv(tm_nets[w])
rows = (t_eval_1d >= w/n_windows) & (t_eval_1d <= (w + 1)/n_windows)
u_tm[rows] = u_w[rows]
error_naive = np.abs(u_naive - u_exact)
error_tm = np.abs(u_tm - u_exact)
plt.figure(figsize = (9, 3))
plt.subplot(1, 3, 1)
plt.imshow(u_exact.T, extent = [0, 1, 0, 2*np.pi], origin = 'lower', aspect = 'auto', cmap = 'jet', vmin = -1, vmax = 1)
plt.colorbar()
plt.title('Ground Truth')
plt.xlabel('t')
plt.ylabel('x')
plt.subplot(1, 3, 2)
plt.imshow(u_naive.T, extent = [0, 1, 0, 2*np.pi], origin = 'lower', aspect = 'auto', cmap = 'jet', vmin = -1, vmax = 1)
plt.colorbar()
plt.title('Naive PINN')
plt.xlabel('t')
plt.ylabel('x')
plt.subplot(1, 3, 3)
plt.imshow(u_tm.T, extent = [0, 1, 0, 2*np.pi], origin = 'lower', aspect = 'auto', cmap = 'jet', vmin = -1, vmax = 1)
plt.colorbar()
plt.title('Time Marching')
plt.xlabel('t')
plt.ylabel('x')
plt.tight_layout()
plt.show()
plt.figure(figsize = (6, 3))
plt.subplot(1, 2, 1)
plt.imshow(error_naive.T, extent = [0, 1, 0, 2*np.pi], origin = 'lower', aspect = 'auto', cmap = 'gray', vmin = 0, vmax = 0.4)
plt.colorbar()
plt.title('Naive PINN Error')
plt.xlabel('t')
plt.ylabel('x')
plt.subplot(1, 2, 2)
plt.imshow(error_tm.T, extent = [0, 1, 0, 2*np.pi], origin = 'lower', aspect = 'auto', cmap = 'gray', vmin = 0, vmax = 0.4)
plt.colorbar()
plt.title('Time Marching Error')
plt.xlabel('t')
plt.ylabel('x')
plt.tight_layout()
plt.show()
plt.figure(figsize = (6, 4))
plt.plot(t_eval_1d, np.mean(error_naive, axis = 1), 'r', linewidth = 2, label = 'Naive PINN')
plt.plot(t_eval_1d, np.mean(error_tm, axis = 1), 'b', linewidth = 2, label = 'Time Marching')
plt.xlabel('t')
plt.ylabel('MAE(t)')
plt.legend()
plt.grid(True)
plt.show()
The naive prediction reproduces the stripes for about the first third of the interval, then loses amplitude and finally collapses to a flat field near zero, which is once again a zero-residual solution of the equation that the initial condition does not permit. Its error map is dark on the left and bright on the right, and $\mathrm{MAE}(t)$ rises linearly from zero before saturating at about $0.6$, the error of predicting nothing where the solution oscillates between $\pm 1$. The time-marching prediction has the stripes all the way across, with an error map that is nearly uniformly dark and a $\mathrm{MAE}(t)$ that stays small to $t = 1$.
4.5.6. What the Lab Showed¶
Convection at $\beta = 80$ is an easy equation and a hard transport problem: the solution at $t = 1$ is the initial condition carried thirteen times around the domain, and a network asked to fit the whole rectangle at once gives up on the far end and settles for a trivial solution there. Time marching never asks for the far end. Each window is a short problem with its initial condition one window away, the state is handed forward explicitly, and the stripes are carried across the interval the way the equation carries them, from left to right.
The two remedies impose order at different scales, and they combine. Time marching decides the global order, which window is solved when; causal weighting can decide the order inside a window, which of its time slices are active when. The result is a hierarchy,
$$ \text{initial state} \longrightarrow \text{window } 1 \xrightarrow{\ \text{causal weighting}\ } \text{terminal state} \longrightarrow \text{window } 2 \xrightarrow{\ \text{causal weighting}\ } \cdots $$
which is useful when the whole horizon is too long for one network and a single window is still long enough to contain dynamics that a network does not learn uniformly from its start to its end. The window length controls the decomposition and $\epsilon$ controls the ordering within each piece.
The two mechanisms also fail differently, which is a reason to use them together rather than to choose one. Time marching accumulates interface error: whatever a window gets wrong at its right edge is the initial condition of the next. Causal weighting with too large an $\epsilon$ can stall, holding the later part of a window at zero weight long after it could have been trained. A short window limits the first, a moderate $\epsilon$ the second, and the two settings should be chosen together.
Causal weighting inside a time-marching window. Time marching sets the global order, one window after another; within the active window the causal weights set the local order, so that the later time slices gain weight only as the earlier ones converge.
Key takeaways
- A PINN puts the physics into the loss, but putting it there does not make the optimization easy, the sampling adequate, the representation suitable or the training order physical. Each of the limitations treated here, and the loss imbalance treated in Regularization, MAP Estimation, and Loss Weighting, is a way in which a correct loss can still be minimized to a wrong answer, and each remedy changes something other than the equation.
- Sampling. The loss is an average over the collocation points, so a region that holds few points holds little of the loss however large its residual. A collocation set should be graded like a mesh, dense where the solution is hard, and the same point budget redistributed in this way resolved the stress concentration of the Kirsch problem that a uniform grid missed.
- Complex solution structures. Spectral bias makes a network learn high frequencies last; a smooth bounded network cannot represent a singularity; a differentiable one cannot represent a kink. Fourier features change the input, enrichment adds a known singular function to the output, and domain decomposition gives each region its own network joined by interface conditions.
- Temporal causality. A full-domain loss enforces the equation at late times before the early solution is right, and the optimizer settles on a zero-residual solution that the initial condition does not permit. Causal weighting lets later time slices count only after earlier ones have converged; time marching trains the interval window by window, each starting from the end state of the last.
Looking ahead. Every method in this notebook makes one trained network solve one problem: one set of coefficients, one boundary and initial data, one geometry. Change any of them and the training starts again. Deep Neural Operators and Fourier Neural Operator learn the map from the data of a problem to its solution, so that a new instance costs a forward pass rather than a training run.