A debugging checklist

A PINN with a bug in its equation trains without complaint: it finds the solution of the equation it was given, and its loss falls as nicely as a correct one's. The tests that catch this are cheap, and most of them run before training. This section lists them, then applies them to a short script that contains three bugs of the kinds that happen in practice.

A script with three bugs

The problem is Section 8.1's: $u_t = a\,u_{xx}$ with $a = 0.1$ on $0 < x < 1$, $0 < t < 1$, $u = 0$ at both ends and $u(x, 0) = \sin(\pi x)$. A colleague's script, written with its residuals as functions of the trial function, as step 2 asks:

import math
import torch
import torch.nn as nn

def d(y, x):
    return torch.autograd.grad(y, x, torch.ones_like(y), create_graph=True)[0]

a, T = 0.1, 1.0
exact = lambda x, t: torch.exp(-a * math.pi**2 * t) * torch.sin(math.pi * x)

def equation(u, x, t):                                         # residual of u_t = a u_xx
    v = u(x, t)
    return d(v, t) + a * d(d(v, x), x)

def initial(u, x0, t0):                                        # residual of u(x, 0) = sin(pi x)
    return u(x0, t0) - torch.sin(math.pi * x0)

def boundary(u, xb, tb):                                       # residual of u = 0 at the ends
    return u(xb, tb)

def make_points(seed=0, n=1000):
    g = torch.Generator().manual_seed(seed)
    x = 0.5 * torch.rand(n, 1, generator=g)                    # collocation points
    t = T * torch.rand(n, 1, generator=g)
    x0 = torch.rand(200, 1, generator=g)                       # initial-condition points
    t0 = torch.full_like(x0, T)
    tb = T * torch.rand(200, 1, generator=g)                   # boundary points
    xb = (torch.rand(200, 1, generator=g) < 0.5).float()
    return x.requires_grad_(True), t.requires_grad_(True), x0, t0, xb, tb

def train(equation, initial, boundary, points, seed=0, epochs=3000):
    torch.manual_seed(seed)
    net = nn.Sequential(nn.Linear(2, 32), nn.Tanh(), nn.Linear(32, 32), nn.Tanh(),
                        nn.Linear(32, 32), nn.Tanh(), nn.Linear(32, 1))
    u = lambda x, t: net(torch.cat([2 * x - 1, 2 * t / T - 1], dim=1))
    x, t, x0, t0, xb, tb = points
    opt = torch.optim.Adam(net.parameters(), lr=1e-3)
    for epoch in range(epochs):
        opt.zero_grad()
        terms = [(equation(u, x, t) / (a * math.pi**2)) ** 2, initial(u, x0, t0) ** 2, boundary(u, xb, tb) ** 2]
        loss = sum(term.mean() for term in terms)
        if epoch in (0, epochs - 1):
            print(f"  epoch {epoch:4d}: " + "   ".join(f"{name} {term.mean().item():.1e}"
                  for name, term in zip(("equation", "initial", "boundary"), terms)))
        loss.backward()
        opt.step()
    return u

def error_of(u):
    X, Tt = torch.meshgrid(torch.linspace(0, 1, 101), torch.linspace(0, T, 101), indexing="ij")
    X, Tt = X.reshape(-1, 1), Tt.reshape(-1, 1)
    with torch.no_grad():
        return ((u(X, Tt) - exact(X, Tt)).norm() / exact(X, Tt).norm()).item()

print("the colleague's script:")
u = train(equation, initial, boundary, make_points())
print(f"relative L2 error against the exact solution: {error_of(u):.2f}")

the colleague's script:
  epoch    0: equation 1.2e-03   initial 6.9e-01   boundary 1.6e-02
  epoch 2999: equation 8.1e-06   initial 2.1e-06   boundary 1.4e-06
relative L2 error against the exact solution: 0.56

Every term fell by three orders of magnitude or more, and nothing in the log hints at trouble. The error, which only the exact solution can reveal here, is 56%.

Step 2: the residuals on the exact solution

The exact solution is known here; in a real problem a manufactured one would do (Section 4.4). Every term must be zero for it:

def worst(term):
    return term.detach().abs().max().item()

x, t, x0, t0, xb, tb = make_points()
print(f"equation term on the exact solution:  {worst(equation(exact, x, t)):.1e}")
print(f"initial term on the exact solution:   {worst(initial(exact, x0, t0)):.1e}")
print(f"boundary term on the exact solution:  {worst(boundary(exact, xb, tb)):.1e}")

equation term on the exact solution:  2.0e+00
initial term on the exact solution:   6.3e-01
boundary term on the exact solution:  8.7e-08

The boundary term vanishes, as it should. The other two do not. The equation term is 2.0, twice the size of $u_t = -a\pi^2 u$ for this solution ($a\pi^2 \approx 1$): the two sides of the equation add instead of cancelling, and the line d(v, t) + a * d(d(v, x), x) has the wrong sign, which runs the heat equation backwards (bug 1). The initial term is 0.63, which is $1 - e^{-a\pi^2}$: exactly the difference between the solution at $t = 1$ and at $t = 0$, so the initial-condition points are at the wrong time (bug 2). Neither shows in training; both take a second to see here.

Step 4: the points

for name, values in [("collocation x", x), ("collocation t", t), ("initial t", t0), ("boundary x", xb)]:
    v = values.detach()
    print(f"{name:14s} from {v.min().item():.3f} to {v.max().item():.3f}")

collocation x  from 0.000 to 0.500
collocation t  from 0.000 to 0.999
initial t      from 1.000 to 1.000
boundary x     from 0.000 to 1.000

Two more facts: the collocation points cover only $0 \le x \le 0.5$, a left-over from the half-metre rod of Section 4.5 (bug 3), and the initial-condition points all sit at $t = 1$, which is bug 2 seen from the other side.

The fixed script

def equation_fixed(u, x, t):
    v = u(x, t)
    return d(v, t) - a * d(d(v, x), x)                          # minus: heat flows from hot to cold

def make_points_fixed(seed=0, n=1000):
    g = torch.Generator().manual_seed(seed)
    x = torch.rand(n, 1, generator=g)                           # all of 0 < x < 1
    t = T * torch.rand(n, 1, generator=g)
    x0 = torch.rand(200, 1, generator=g)
    t0 = torch.zeros_like(x0)                                   # t = 0
    tb = T * torch.rand(200, 1, generator=g)
    xb = (torch.rand(200, 1, generator=g) < 0.5).float()
    return x.requires_grad_(True), t.requires_grad_(True), x0, t0, xb, tb

points = make_points_fixed()
x, t, x0, t0, xb, tb = points
print(f"on the exact solution: equation {worst(equation_fixed(exact, x, t)):.1e},"
      f" initial {worst(initial(exact, x0, t0)):.1e}, boundary {worst(boundary(exact, xb, tb)):.1e}")
print("the fixed script:")
u = train(equation_fixed, initial, boundary, points)
print(f"relative L2 error against the exact solution: {error_of(u):.1e}")

on the exact solution: equation 1.8e-07, initial 0.0e+00, boundary 8.7e-08
the fixed script:
  epoch    0: equation 2.4e-03   initial 7.6e-01   boundary 1.6e-02
  epoch 2999: equation 3.3e-06   initial 1.2e-06   boundary 1.1e-06
relative L2 error against the exact solution: 2.0e-03

Every term is now zero on the exact solution, to rounding error, and the fixed script reaches 0.2% with the same network and the same training that gave 56% before.

Habits that make the checklist cheap

• Residuals as functions of the trial function. Then the same code evaluates a network, an exact solution or a manufactured one, and step 2 is three lines.
• One function makes all the points, and it returns them, so step 4 can look at them.
• Print every term, not their sum, at the first and last epochs at least.
• Keep a manufactured solution for every new equation you code, even when the real problem has none.

In the Studio

The Studio writes the residuals for you, which removes most coding bugs but not every mistake: a sign convention (the Studio's Poisson block solves $\nabla^2 u = f$, Section 13.2), a condition on the wrong boundary, a Neumann value with the wrong sign (Section 6.3). The generated script's residual functions take the points as their argument and use the module's field for the trial function, so step 2 is a patch of field, the same patch as Section 6.4's, pointed at the exact solution:

# not run
import math
import torch
import poisson_1d as p

p.field = lambda pts, component=0: torch.sin(math.pi * pts[:, 0:1])     # the exact solution
pts = p.sample_interior()
print("largest equation residual on the exact solution:", p.residual_pde1(pts).abs().max().item())

Exercises

1. Why did the first two bugs not stop the colleague's script from training to a small loss?
2. Step 2 found bugs 1 and 2 but not bug 3. Which step found bug 3, and which later step would have found it if step 4 had been skipped?
3. Write a manufactured solution for $u_t = a\,u_{xx} + s(x, t)$ that you could use for step 2 if no exact solution were known.

Recap

• A bug in the equation is invisible in the loss: the network solves the problem it was given.
• Test every residual on a known or manufactured solution before training, and look at the points.
• After training, plot each term, check the residual at fresh points and run several seeds.