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,
- explain why a decaying solution weights its own residual loss towards early times, and how dividing out the known decay removes that weighting, 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 of the next section 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.
Time. 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. Time defeats a PINN in a second and unrelated way as well. When the solution decays as it evolves, the loss is dominated by the times at which the solution is still large, and the end of the interval goes unconstrained whatever the order of training. The remedy there is to take the known decay out of what the network represents.
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 of this lab 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 in this lab is measured against it.
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 stresses and the von Mises stress, with $a = 1$ and stresses in units of $\sigma_0$, serve three purposes in this lab: the ground truth, the outer boundary condition, and the error.
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 von Mises stress is largest 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 Cartesian grid over the square, with the points that fall inside the hole removed. What remains is $256$ points spread evenly over the plate, so that the annulus $a<r\le2a$, which carries the whole stress concentration, receives fewer than a third of them.
Focused collocation. The same $256$ points redistributed: a fine grid inside the annulus $a<r\le2a$ and a coarse grid outside it, with the two chosen so that the counts add up to exactly $256$ and the point density inside the annulus is three times the density outside.
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)))
Both sets 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 more than half of its points in the annulus and leaves the rest to 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$. Both models start from the same weights.
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 . $$
These three left-hand sides, evaluated at the collocation points by automatic differentiation, are the physics residuals.
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 loss adds the three physics residuals at the collocation points and the two boundary conditions at the boundary points, with weights,
$$ \mathcal{L} = \mathcal{L}_{\mathrm{pde}} + 10\,\mathcal{L}_{\mathrm{hole}} + 10\,\mathcal{L}_{\mathrm{outer}} . $$
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.
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 and its absolute error against the Kirsch solution, on the colour scales used for every model of this lab.
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 loss ended small. 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, where it amounts to several percent of the peak stress; the far field is nearly black, and the mean error in the annulus is several times the mean over the whole plate.
The pattern is the one the sampling predicts. The naive grid placed most of its points 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 number of iterations; only the 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()))
The two predictions side by side with the ground truth, and their errors on a common scale.
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])))
The higher final loss in the focused run does not indicate a worse solution, but rather reflects that the loss was measured in more challenging areas where the residual is hardest to reduce. The true performance is evidenced by the error relative to the exact solution, where the focused run significantly reduced both the mean error near the hole and the overall maximum error. Visually, the concentrated error "ring" found in the naive run was replaced by a faint, evenly spread residual. These results demonstrate that by redistributing the point budget from the far field to the stress-concentrated annulus, the model achieved higher precision without increasing the total budget or changing the network. Ultimately, this shows that where the residual is evaluated is more critical than the total number of points for achieving a physically accurate solution.
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). In the previous experiment, the stress concentration was known to be located at the hole, allowing the refinement to be placed manually. However, when the location of the difficulty is not known in advance, the residual itself serves as the guide. The residual represents the degree to which the current neural network solution fails to satisfy the governing physical equations; a larger residual indicates a region where the model's prediction is least accurate and the physics are not yet well-captured. By evaluating this residual on a dense candidate set, collocation points can be added or moved to the areas where the residual is largest. This residual-adaptive sampling is the PINN counterpart of adaptive mesh refinement, enabling the model to automatically recover the focused allocation described above without being explicitly told where the hole or other high-difficulty regions are located.
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.
This chapter considers three such cases.
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¶
In the animation a network is 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.
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()
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.
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)
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 the whole run the loss is still far from zero, 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 scheme, the same way of sampling the collocation points, and the same 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 by the same loop as the naive one, with the embedding placed in front of the network and nothing else changed.
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 two orders of magnitude below the naive model's. The loss trace tells the same story from the other side: the naive loss was still high at the end of the run, while the Fourier-feature loss had fallen by three orders of magnitude by the first printed checkpoint. 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 initialization scheme 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.
3.4.3. Network and Physics-Informed Loss¶
Both models use the same network 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 the 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 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 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.
q = 1.0 # distributed load
P = 1.0 # point load
a = 0.5 # location of the point load
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; 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.
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 at one half and cannot go lower, and the prediction shows why: it is a smooth line that splits the difference between the two halves of the exact shear force, wrong on both sides of the load rather than only at it. No continuous function can satisfy the equation and the equilibrium condition at once, and the optimizer has found the best compromise a smooth network allows.
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 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 enriched PINN successfully captures the ground truth, with the coefficient $\alpha$ converging to the mechanics-predicted value of $P=1$.
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()
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 almost exactly. 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} $$
Note: Since the temperature does not vary in the y-direction, this problem is physically equivalent to a one-dimensional heat conduction problem in a rod. However, it is formulated and solved in two dimensions here to facilitate better visualization and intuitive understanding.
The two conductivities and the piecewise linear temperature, evaluated on a $201\times201$ grid that serves as the ground truth throughout, followed by the collocation points: interior points in each material, points on the four edges, and points along the interface.
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()
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.
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 derivative(y, t):
df = torch.autograd.grad(y, t, torch.ones_like(y), create_graph = True)[0]
return df[:, 0:1], df[:, 1:2]
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}$; 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$. Both are written for two networks, one per material, so that the naive model is the special case in which the same network is used on both sides.
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
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 flux term stops falling almost at once and stays where it is while the residual and the boundary terms keep going down: the optimizer has found the compromise a single network allows, satisfying the equation and giving up on the interface.
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.
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()))
3.6.3. Results¶
Both predictions are evaluated on the same 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.
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))
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 one constant 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 flux makes the failure explicit. The naive flux jumps at the interface by the ratio of the two conductivities: 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 one flat line 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 final limitation concerns time-dependent problems, which differ fundamentally from the spatial challenges discussed previously. A transient solution is not merely a field with an added time axis; it is a history. In physics, the state at a later time is produced from the state at an earlier time via the governing equations. Consequently, the solution at $t_2$ cannot be accurately determined until the solution for all times leading up to $t_2$ is established.
Traditional numerical time-stepping schemes respect this principle by construction: they compute the solution sequentially, one step at a time. A PINN, however, treats time as just another input coordinate. Collocation points are scattered across the entire space-time domain, and the optimizer attempts to minimize the residual at $t = 1.5$ with the same priority as at $t = 0.1$. Because the loss function is averaged globally, it lacks an inherent sense of direction. While the underlying physics are causal, the optimization objective is not.
On simple problems, this lack of causality may not be an issue. However, for systems with fast dynamics or trajectories that can be "mimicked" by a nearly stationary field, the optimizer may converge to a solution that satisfies the mathematical equation (low residual) but completely fails to capture the actual physics.
To address this failure, two primary remedies are employed:
- Causal Training: This approach maintains a single network over the entire time interval but introduces a weighting mechanism. The residual at later times is only allowed to contribute to the loss after the solution at earlier times has converged.
- Time Marching: This method decomposes the time interval into smaller windows. The network is trained on these windows sequentially, with each window using the final state of the previous one as its initial condition.
While the first method modifies the loss function, the second modifies the problem structure.
Finally, time-dependent problems present an additional, qualitatively different difficulty: solution decay. When a solution diminishes in magnitude over time, the mean squared residual becomes implicitly weighted by the size of the solution. As a result, the later stages of the interval may stop contributing to the objective function entirely. Solving this requires a different approach—not a change in weights or scheduling, but a change in the actual quantity the network is asked to represent.
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 of the next section 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
$$ \frac{\partial u}{\partial t} - \nu\,\frac{\partial^2 u}{\partial x^2} + 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)
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from IPython.display import HTML
fig, ax = plt.subplots(figsize=(6, 4))
line, = ax.plot([], [], linewidth=2, color='b')
ax.set_xlim(x_ref.min(), x_ref.max())
ax.set_ylim(u_ref.min(), u_ref.max())
ax.set_xlabel('x')
ax.set_ylabel('u(x, t)')
ax.grid(True)
time_text = ax.text(0.05, 0.9, '', transform=ax.transAxes, fontsize=12)
def init():
line.set_data([], [])
time_text.set_text('')
return line, time_text
def update(frame):
t_current = t_ref[frame]
u_current = u_ref[frame, :]
line.set_data(x_ref, u_current)
time_text.set_text(f't = {t_current:.2f}')
return line, time_text
ani = FuncAnimation(fig, update, frames=len(t_ref),
init_func=init, blit=True, interval=50)
plt.close()
HTML(ani.to_jshtml())
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 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 = \frac{\partial u_\theta}{\partial t} + 5u_\theta^3 - 5u_\theta - \nu\,\frac{\partial^2 u_\theta}{\partial x^2} , $$
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 large weight on the initial condition so that the start of the trajectory is pinned down firmly; nothing in this objective says in which order the rest of the interval is to be learned.
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)
# # 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()
The naive prediction drifts away from the reference as time goes on: its error grows steadily from $t = 0$, and by the end of the interval the field bears little resemblance to the plateaus and transition layers of the reference, although the residual there is small, since a nearly stationary field satisfies the equation almost as well as the true transient does. This case illustrates a fundamental challenge in training PINNs for time-dependent problems: the optimizer may satisfy the governing equations without correctly propagating the initial state through time. Because the standard objective function treats all time points equally, the network is not required to resolve the solution at the beginning of the interval before fitting the later stages. To overcome this, the training process must be modified to prioritize the early stages of the time axis, ensuring that the solution is established chronologically, starting from the initial condition.
A causal PINN keeps one network over the whole domain and changes only how much each time slice is allowed to matter. The change is small enough to be worth deriving rather than asserting.
The construction begins by putting an order on time. Introduce $N_t + 1$ nodes,
$$ 0 = t_0 < t_1 < t_2 < \cdots < t_{N_t} = T , $$
which cut the interval into $N_t$ consecutive slabs,
$$ [t_0, t_1) , \;\; [t_1, t_2) , \;\; \ldots , \;\; [t_{N_t - 1}, t_{N_t}] , $$
the last one closed so that the final time belongs to a slab. Each slab carries a strip of space-time with it,
$$ \Omega_i = \bigl\{ (x, t) \;:\; x \in \Omega , \;\; t_{i-1} \le t < t_i \bigr\} , \qquad i = 1, \ldots, N_t , $$
and the collocation points inherit the partition. Writing $S_i$ for the points that fall in $\Omega_i$ and $N_i = |S_i|$, the full collocation set is the disjoint union of the $S_i$, so that $\sum_i N_i = N$. Define the residual loss restricted to one slab,
$$ L_i = \frac{1}{N_i} \sum_{(x_j, t_j) \,\in\, S_i} r_\theta(x_j, t_j)^2 . $$
In Section 4.1 and in the lab below the collocation points are drawn on a grid, $N_t$ time levels with $N_x$ points each. Each slab then contains exactly one level, which we continue to write $t_i$, so that $N_i = N_x$ and the definition reduces to the average over that level,
$$ L_i = \frac{1}{N_x} \sum_{j=1}^{N_x} r_\theta(x_j, t_i)^2 . $$
Nothing that follows depends on which of the two is used. The slabs matter only because they are ordered, and the same partition reappears in Section 4.4, where the windows of time marching are these slabs trained one at a time rather than weighted.
These slab losses are not a new object; they are the old loss taken apart. Collecting the sum over all collocation points slab by slab,
$$ \mathcal{L}_{\mathrm{PDE}} = \frac{1}{N} \sum_{j=1}^{N} r_\theta(x_j, t_j)^2 = \sum_{i=1}^{N_t} \frac{N_i}{N}\, L_i , $$
and when the slabs are sampled uniformly, so that $N_i = N / N_t$,
$$ \mathcal{L}_{\mathrm{PDE}} = \frac{1}{N_t} \sum_{i=1}^{N_t} L_i . $$
Nothing has happened yet; this is the same number written differently. But in this form the objection of Section 4.1 can be stated exactly. The loss is already a sum over time slabs, and every slab enters with the same coefficient $1/N_t$. The coefficient multiplying the last slab does not depend on how well the first has been fitted. That constant, identical for every $i$, is the whole of the problem.
Causal training replaces it by a coefficient that does depend on the earlier slices,
$$ \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) , $$
with a single new parameter $\epsilon > 0$. The sum in the exponent runs over the slices strictly before $i$, so it is the accumulated residual error upstream in time. Written out, the first few weights are
$$ w_1 = e^0 = 1 , \qquad w_2 = e^{-\epsilon L_1} , \qquad w_3 = e^{-\epsilon(L_1 + L_2)} , \qquad w_4 = e^{-\epsilon(L_1 + L_2 + L_3)} , $$
and in general each weight is the previous one multiplied by one more factor,
$$ w_{i+1} = w_i\,e^{-\epsilon L_i} , \qquad e^{-\epsilon L_i} \in (0, 1] . $$
Three properties follow from this recursion alone. The first slice has an empty sum in its exponent, so $w_1 = 1$ for the whole of training: the slice next to the initial condition is never gated. The weights are non-increasing in time,
$$ 1 = w_1 \ge w_2 \ge \cdots \ge w_{N_t} > 0 , $$
a staircase descending into the future. And a weight is close to one only when every slice before it has small loss, since one large $L_k$ contributes a small factor that suppresses all the weights after it.
What the weights do to training is seen in the gradient. Since they are held fixed when the derivative is taken,
$$ \nabla_\theta\,\mathcal{L}_{\mathrm{PDE}}^{\mathrm{causal}} = \frac{1}{N_t}\sum_{i=1}^{N_t} w_i\,\nabla_\theta L_i , $$
so $w_i$ is how loudly slice $i$ speaks in a parameter update. At the first iteration $L_1$ is large, $w_2$ and everything after it are near zero, and the update is driven by the first slice alone: the network is in effect solving the equation next to the initial condition that anchors it. As $L_1$ falls, $w_2$ rises towards one and the second slice begins to contribute; as $L_2$ falls in turn, the third follows. Training 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. It has a scale rather than a universal value, since $L_k$ carries the units of the residual squared and a value tuned for one equation does not transfer to another whose residual is measured differently.
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; otherwise the optimizer would discover that $\sum_i w_i L_i$ can be reduced by shrinking the weights rather than the residuals. 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.
Because the weights descend, the smallest of them is always the last,
$$ \min_i w_i = w_{N_t} , $$
and that single number says whether training ever finished. If $w_{N_t}$ has risen close to one by the end of the run, every gate has opened and the equation has been enforced over the whole interval. If it is still small, the active front never reached the end: the late slices were never fitted, and the prediction there should not be trusted however small the reported total loss is, because the total is small precisely in virtue of having weighted those slices out of it. The quantity that gates the training doubles as the diagnostic that validates it, and is worth recording alongside the loss. It also gives a practical rule for $\epsilon$: take the largest value for which $w_{N_t}$ still reaches one within the iteration budget.
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 only change to the training loop is the causal weighting of the residual, with $\epsilon = 100$, a large value because the slice losses here are small numbers. The smallest weight is printed alongside the losses as training proceeds, so that the advancing front can be watched.
# # 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 grew steadily from the start. 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
$$ \frac{\partial u}{\partial t} + \beta\,\frac{\partial u}{\partial 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 is made at $\beta = 80$.
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from IPython.display import HTML
beta = 80.0
x = np.linspace(0, 2*np.pi, 300)
t_steps = np.linspace(0, 1, 100)
fig, ax = plt.subplots(figsize=(6, 4))
line, = ax.plot(x, np.sin(x - beta * 0), color='gray', lw=2)
ax.set_xlim(0, 2*np.pi)
ax.set_ylim(-1.1, 1.1)
ax.set_title(r'1D Wave Propagation ($\beta = 80$)')
ax.set_xlabel('x')
ax.set_ylabel('u(x, t)')
ax.grid(True)
def update(frame):
t = t_steps[frame]
y = np.sin(x - beta * t)
line.set_ydata(y)
return [line]
ani = FuncAnimation(fig, update, frames=len(t_steps), interval=50, blit=True)
plt.close()
HTML(ani.to_jshtml())
4.5.2. Naive and Time-Marching Formulations¶
Both models use the same network and the same residual $r_\theta = \dfrac{\partial u_\theta}{\partial t} + \beta\,\dfrac{\partial u_\theta}{\partial 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 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¶
The naive model is trained on the whole rectangle at once, with the iteration budget that the time-marching run will share.
# # 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 equal length, each trained for a tenth of the naive run's budget, so that the two runs cost the same. 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.
# # 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 supplying the times of its own window.
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 steadily from zero before saturating at the level of predicting nothing at all. 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.
Section 4.1 showed one way in which time can defeat a PINN. The optimizer worked on the late times before the early times were right, and it settled on a trajectory that the initial condition does not produce. That was a problem of order, and causal weighting fixed it by making the network learn the early times first.
This section is about a second problem, and it is worth keeping the two apart because they look alike from the outside: in both, the prediction is bad at late times. Here, however, the order of training is fine, the collocation points are fine, and the trajectory is the right one. The problem is size. A diffusing solution shrinks as time goes on, and by the end of the interval it can be a hundred times smaller than at the start. The loss is built from squared residuals, and whatever is a hundred times smaller becomes ten thousand times smaller once it is squared. The late times are still in the loss, but they have become too small to be heard.
4.6.1. A Decaying Solution Weights Its Own Loss¶
Start with a picture. Suppose the network is wrong by ten percent at every time. At $t = 0$ the solution has size $1$, so the error is $0.1$ and its square is $0.01$. At $t = T$ the solution has shrunk to $0.01$, so the same ten-percent error is $0.001$ and its square is $10^{-6}$. The optimizer sees the second mistake as ten thousand times less important than the first, even though the two are equally bad in relative terms. Nobody asked for the late times to be discounted; the decay of the solution did it.
Now the same statement made exact. Solutions of the heat equation with fixed ends decay like $e^{-\lambda t}$. So write the solution as the known decay times a field that does not decay,
$$ \begin{align*} \frac{u(x, t)}{e^{-\lambda t}} &= v(x, t)\\\\ u(x, t) &= e^{-\lambda t}\, v(x, t) \end{align*} $$
Here $v$ is nothing more than $u$ magnified by $e^{\lambda t}$, enough to undo the decay, so that $v$ stays of size one over the whole interval. Substitute into the heat equation. By the product rule,
$$ u_t = e^{-\lambda t}\bigl(v_t - \lambda v\bigr) , \qquad u_{xx} = e^{-\lambda t} v_{xx} $$
so the residual of the original equation is a shrinking factor times a residual of $v$:
$$ \boxed{\;r_u = u_t - \alpha u_{xx} = e^{-\lambda t}\underbrace{\bigl(v_t - \alpha v_{xx} - \lambda v\bigr)}_{r_v}\;} $$
Square it and average over collocation points spread uniformly in time, and you have the loss the PINN actually minimizes:
$$ \mathcal{L} = \frac{1}{N_f}\sum_{j=1}^{N_f} r_u(x_j, t_j)^2 = \frac{1}{N_f}\sum_{j=1}^{N_f} e^{-2\lambda t_j}\; r_v(x_j, t_j)^2 $$
Read the right-hand side as a weighted loss on $v$: the point at time $t_j$ carries the weight $e^{-2\lambda t_j}$. This is the ten-thousand-to-one of the picture above, now with its exact value. The plain loss on $u$ was a weighted loss all along, and the weights were chosen by the decay, not by us.
How bad is it? For the heat equation on the unit interval with fixed ends, $\lambda = \alpha\pi^2 \approx 9.9$ when $\alpha = 1$. The table gives, for three interval lengths, how far the solution has decayed by $t = T$, how much weight the loss gives to $t = T$ compared with $t = 0$, and how early half of the whole loss has already been spent.
| $T$ | $u(T)/u(0)$ | weight at $t = T$ relative to $t = 0$ | half of the loss accumulates below |
|---|---|---|---|
| $0.1$ | $0.37$ | $1.4\times10^{-1}$ | $t = 0.029$, the first $28\%$ |
| $0.3$ | $0.052$ | $2.7\times10^{-3}$ | $t = 0.035$, the first $12\%$ |
| $0.5$ | $0.0072$ | $5.2\times10^{-5}$ | $t = 0.035$, the first $7\%$ |
Take the last row. At $T = 0.5$ half of the objective is decided in the first seven percent of the interval. Over the remaining ninety-three percent of the time axis the optimizer is almost free to do as it likes: the network can be wrong there by more than the solution is large, and the loss barely notices.
Remark (causal weighting does not fix this). The causal weights of Section 4.2 make a late slice small only while the earlier slices are still badly fitted. Once those are learned, every weight returns to one and the loss is the plain average again, with the $e^{-2\lambda t}$ imbalance untouched. Causal weighting decides when a time is allowed to count; this section is about how much it counts once it does. Two different problems, two different remedies.
4.6.2. Dividing Out the Known Decay¶
The identity also tells us what to do. If training on $u$ is secretly training on $v$ with bad weights, then train on $v$ directly: let the network represent the magnified field $v$, and recover $u = e^{-\lambda t}v$ afterwards. Think of it as changing the ruler. Instead of measuring the error in absolute units, which shrink along with the solution, measure it relative to the current size of the solution, which does not.
What equation does $v$ satisfy? Set $r_v = 0$ in the boxed identity:
$$ \boxed{\;\frac{\partial v}{\partial t} = \alpha\frac{\partial^2 v}{\partial x^2} + \lambda v\;} $$
It is the heat equation with one extra term, $+\lambda v$, which pushes $v$ up at exactly the rate at which diffusion would pull it down. The slowest part of $v$ therefore stays constant instead of fading, and the faster parts decay only relative to it. The boundary conditions are unchanged, and so is the initial condition, because $e^{-\lambda t} = 1$ at $t = 0$. A mean squared residual of this equation over uniformly drawn times counts every time alike, and no weights appear anywhere.
What has to be known. Only the decay rate $\lambda$, and for a linear diffusion problem it is not an unknown. It is the slowest rate at which anything in the domain can decay, and it depends only on the geometry and the boundary conditions, not on the solution. On an interval of length $L$ with fixed ends it is $\lambda = \alpha\pi^2/L^2$; on an $L_x \times L_y$ rectangle it is $\alpha\pi^2\bigl(1/L_x^2 + 1/L_y^2\bigr)$; on a disc of radius $R$ it is $\alpha(2.405/R)^2$. Every mode of the solution decays at least this fast, so dividing by $e^{-\lambda t}$ never turns a decaying mode into a growing one.
Where this construction has been seen before. It is the same move as the enrichment of Section 3.4 and the masking of a boundary condition, now applied along the time axis. In each case something known before training is taken out of the network's job: added to the output for a singularity, multiplied in for a boundary value, and here multiplied in for a decay. The network is left with the part that is genuinely unknown, and that part is always the easier one to represent.
4.6.3. Lab: A Decaying Heat Problem¶
The demonstration is the plainest transient there is: the heat equation on the unit interval with both ends held at zero and a two-mode initial condition,
$$ \frac{\partial u}{\partial t} = \frac{\partial^2 u}{\partial x^2} , \qquad u(0, t) = u(1, t) = 0 , \qquad u(x, 0) = \sin(\pi x) + \sin(3\pi x) , $$
over $t \in (0, 0.5]$. Two modes give the solution some structure early on, and the exact solution is available without a solver,
$$ u(x, t) = e^{-\pi^2 t}\sin(\pi x) + e^{-9\pi^2 t}\sin(3\pi x) . $$
The fast mode, $\sin(3\pi x)$, is gone by about $t = 0.05$. After that the solution is a single hump that keeps shrinking, and at $t = 0.5$ it stands at $0.72\%$ of its initial height. The magnified field is
$$ v(x, t) = e^{\pi^2 t}u(x, t) = \sin(\pi x) + e^{-8\pi^2 t}\sin(3\pi x) , $$
which, instead of vanishing, settles onto the fixed shape $\sin(\pi x)$.
Two models are compared. They share everything, the architecture, the initial weights, the collocation points and the training schedule, and differ in a single term of the residual:
| Model A | Model B | |
|---|---|---|
| the network represents | $u$ | $v$ |
| residual driven to zero | $u_t - u_{xx}$ | $v_t - v_{xx} - \pi^2 v$ |
| prediction | $u_\theta$ | $e^{-\pi^2 t}v_\theta$ |
Both models impose the boundary and initial conditions exactly, by the hard constraint
$$ w_\theta(x, t) = a(x) + \lambda t\,x(1 - x)\,N_\theta(x, t) , $$
which equals $a(x)$ at $t = 0$ and vanishes at both ends whatever the network returns. Only the equation residual is left in the loss, so the comparison isolates the effect of this section from the loss weighting treated in Regularization, MAP Estimation, and Loss Weighting.
The collocation times are drawn uniformly on $(0, T]$, on purpose. Crowding the points at early times would hide the effect under a second one.
The decay rate $\lambda = \pi^2$ is the smallest Dirichlet eigenvalue of $-d^2/dx^2$ on the unit interval, and the two modes of the initial condition decay at $\lambda$ and $9\lambda$.
import numpy as np
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('Device:', device)
lam = np.pi**2 # smallest Dirichlet eigenvalue of -d2/dx2 on (0, 1)
T = 0.5
N_f = 2000
steps = 20000
def a(x): # initial condition, zero at both ends
return torch.sin(np.pi*x) + torch.sin(3*np.pi*x)
def exact(x, t):
return torch.exp(-lam*t)*torch.sin(np.pi*x) + torch.exp(-9*lam*t)*torch.sin(3*np.pi*x)
def derivative(y, x):
return torch.autograd.grad(y, x, torch.ones_like(y), create_graph = True)[0]
The two models start from identical weights and see identical collocation points at every step. The hard-constrained quantity the network represents is $u$ for Model A and $v$ for Model B, and in both cases it is turned into a prediction of $u$.
def new_pinn():
torch.manual_seed(0)
return nn.Sequential(nn.Linear(2, 64), nn.Tanh(),
nn.Linear(64, 64), nn.Tanh(),
nn.Linear(64, 64), nn.Tanh(),
nn.Linear(64, 1)).to(device)
def field(net, x, t): # the quantity the network represents: u or v
return a(x) + t*x*(1 - x)*lam*net(torch.cat([x, t], dim = 1))
def solution(net, x, t, rescaled):
w = field(net, x, t)
return torch.exp(-lam*t)*w if rescaled else w
One function serves both models. The growth term is the only difference: zero for the residual of $u$, and $\lambda$ for the residual of $v$.
def residual(net, x, t, rescaled):
w = field(net, x, t)
w_t = derivative(w, t)
w_x = derivative(w, x)
w_xx = derivative(w_x, x)
growth = lam if rescaled else 0.0
return w_t - w_xx - growth*w
def train(rescaled):
net = new_pinn()
optimizer = torch.optim.Adam(net.parameters(), lr = 1e-3)
for k in range(steps):
x = torch.rand(N_f, 1, device = device).requires_grad_(True)
t = T*torch.rand(N_f, 1, device = device).requires_grad_(True)
loss = torch.mean(residual(net, x, t, rescaled)**2)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (k + 1) % 5000 == 0:
print(f" step {k+1:6d} loss = {loss.item():.3e}")
return net
The two runs take a few minutes each. The printed losses are not comparable with one another: Model A reports a residual of $u$ and Model B a residual of $v$. The test errors that follow are both relative errors in $u$, and those can be compared.
print('Model A : loss on the u residual')
net_A = train(rescaled = False)
print('Model B : loss on the v residual')
net_B = train(rescaled = True)
The error is reported at each time separately, as a relative $L^2$ error over the interval in $x$,
$$ E(t) = \frac{\bigl\lVert u_\theta(\cdot, t) - u(\cdot, t)\bigr\rVert_2}{\bigl\lVert u(\cdot, t)\bigr\rVert_2} = \left(\frac{\sum_i \bigl(u_\theta(x_i, t) - u(x_i, t)\bigr)^2}{\sum_i u(x_i, t)^2}\right)^{1/2} , $$
on a grid of $201$ equally spaced points, the grid spacing cancelling between numerator and denominator.
Dividing by $\lVert u(\cdot, t)\rVert$ is what makes the measure usable here. An absolute error would shrink along with the solution: at $t = 0.5$ the exact solution stands at $0.72\%$ of its initial height, so a model that has lost it altogether still reports a small absolute error. The relative error asks instead how large the error is compared with the solution it is meant to be, and it therefore reads on the same scale at every time.
Two values are worth fixing in mind. A model returning zero everywhere gives $E = 1$ exactly, since the numerator is then $\lVert u\rVert$; above one, a prediction is worse than no prediction at all. Below about $10^{-2}$ the prediction and the exact solution are hard to tell apart in a plot.
The error is deliberately not pooled over space and time into one number. A relative error taken over the whole rectangle would be dominated, in its numerator and its denominator alike, by the early times where the solution is large, and it would hide precisely the failure this section is about.
x_plot = torch.linspace(0, 1, 201).reshape(-1, 1).to(device)
def relative_error(net, rescaled, t_value):
with torch.no_grad():
t = t_value*torch.ones_like(x_plot)
u = solution(net, x_plot, t, rescaled)
return (torch.norm(u - exact(x_plot, t))/torch.norm(exact(x_plot, t))).item()
print(" t Model A Model B")
for t_value in [0.002, 0.01, 0.05, 0.1, 0.25, 0.5]:
print(f" {t_value:5.3f} {relative_error(net_A, False, t_value):9.4f} "
f"{relative_error(net_B, True, t_value):9.4f}")
The two columns tell the story. Up to $t = 0.05$, while the solution is still large, Model A is slightly better than Model B: its loss is spent almost entirely on these times, and it spends it well. From $t = 0.1$ on, the two part company. Model A's error grows by orders of magnitude over the rest of the interval and ends within a hair of the error of predicting zero. Model B's error stays below one percent at every time in the table.
The same error as a continuous curve over the interval, on logarithmic axes.
t_curve = np.logspace(-3, np.log10(T), 60)
err_A = [relative_error(net_A, False, t) for t in t_curve]
err_B = [relative_error(net_B, True, t) for t in t_curve]
plt.figure(figsize = (6, 4))
plt.loglog(t_curve, err_A, 'r', linewidth = 2, label = 'Model A : loss on u')
plt.loglog(t_curve, err_B, 'b', linewidth = 2, label = 'Model B : loss on v')
plt.xlabel('t')
plt.ylabel('relative L2 error')
plt.legend()
plt.grid(True, which = 'both', alpha = 0.3)
plt.show()
The curves make the same point at a glance. The red curve is flat and low until $t \approx 0.1$ and then rises almost as a straight line on the logarithmic axes, heading for the level of no prediction at all. The blue curve is flat from beginning to end. The crossing of the two near $t = 0.1$ is where the solution has fallen to about a third of its initial size, which is where the weight $e^{-2\lambda t}$ has fallen to about one seventh.
The solution itself at two times, one while both modes are still present and one at the end of the interval, where the exact solution stands at $0.72\%$ of its initial height.
plt.figure(figsize = (10, 4))
for j, t_value in enumerate([0.05, 0.5]):
with torch.no_grad():
t = t_value*torch.ones_like(x_plot)
xs = x_plot.cpu().numpy().ravel()
plt.subplot(1, 2, j + 1)
plt.plot(xs, exact(x_plot, t).cpu().numpy().ravel(), 'k', linewidth = 3, alpha = 0.3)
plt.plot(xs, solution(net_A, x_plot, t, False).cpu().numpy().ravel(), 'r--')
plt.plot(xs, solution(net_B, x_plot, t, True).cpu().numpy().ravel(), 'b--')
plt.legend(['exact', 'Model A', 'Model B'])
plt.title(f't = {t_value}')
plt.xlabel('x')
plt.ylabel('u')
plt.tight_layout()
plt.show()
At $t = 0.05$ both models lie on top of the exact solution. At $t = 0.5$ the exact solution is a small hump, less than one percent of its initial height. Model B reproduces it. Model A returns a curve with a lower peak and a negative lobe over the right half of the interval, a shape the heat equation with these boundary conditions cannot produce. Its loss is nevertheless small, because at this time the residual of $u$ is multiplied by $e^{-2\lambda t}$, a factor of order $10^{-5}$, before it reaches the optimizer.
4.6.4. What the Lab Showed¶
Model A was accurate for the first tenth of the interval and nearly worthless at the end of it, and nothing about its training was wrong. The network had the capacity, the points were where they should be, and the loss went down steadily. What failed was the loss itself: a mean squared residual of $u$ is not a neutral measure of how well the equation is satisfied. It is a measure weighted by the size of the solution, and for a decaying transient that weighting spans orders of magnitude across the interval. The weighting is there whether or not anyone intends it.
Model B fixed the problem by dividing out the known decay, and it did so without tuning anything: no weight was chosen, no schedule introduced, no gradient measured. A function known before training began was moved from the network's job into the model, and the equation the network solves changed by one term. The price was small, a slightly larger error at the earliest times, where Model A had the advantage of concentrating all its effort.
This is the third time the notebook has made the same move: in the input, for a frequency (Section 3.3); in the output, for a singularity (Section 3.4); and here in time, for a decay. Whenever part of the solution is known in advance, take it out, and let the network represent only what is not known.
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.
- Causality in time. 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.
- Scale in time. For a decaying solution the plain residual loss is exactly a weighted loss on the rescaled field, with weights $e^{-2\lambda t}$, so the end of the interval is unconstrained however the training is ordered. Dividing out the known decay, $u = e^{-\lambda t}v$, gives an equation for a field that does not decay, and the weighting disappears with it.
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.