Conditions as loss terms

Section 6.2 imposed two initial conditions as two squared errors added to the loss. Every condition of Section 4.2 can be imposed the same way, as a soft constraint: a penalty that is zero when the condition holds. This section writes the residual for each kind, tries them on one problem whose answer is known in every case, and then removes a condition to see what a PINN does with a problem that has no unique solution.

One residual per kind

At a boundary point with outward unit normal $\mathbf{n}$, the condition residuals are:

A Neumann or Robin condition needs the network's gradient on the boundary, from the same autograd call as the equation (Section 5.3), and the outward normal: on $[0, 1]$, $\mathbf{n} = -1$ at $x = 0$ and $+1$ at $x = 1$. The Studio's generated code does exactly this, with dudn = torch.sum(g * normals, dim=1, keepdim=True). Getting the normal's sign wrong is the classic Neumann bug: the condition then says the opposite of what was meant, and nothing fails.

One problem, four sets of conditions

Take $-u'' = \pi^2\sin(\pi x)$ on $(0, 1)$. Its solutions are $\sin(\pi x) + A + Bx$ for any $A$ and $B$. We choose conditions that all single out the same member, $u^\ast = \sin(\pi x) + 1 + x$, so that every run can be graded against it. The training function takes a list of conditions, each a point, its outward normal, a kind and a value:

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]

f = lambda x: math.pi**2 * torch.sin(math.pi * x)                 # -u'' = f
exact = lambda x: torch.sin(math.pi * x) + 1 + x                  # the member of the family we aim for

def train(conditions, seed=0, epochs=4000):
    torch.manual_seed(seed)
    net = nn.Sequential(nn.Linear(1, 32), nn.Tanh(), nn.Linear(32, 32), nn.Tanh(), nn.Linear(32, 1))
    u = lambda x: net(2 * x - 1)
    x_col = torch.rand(200, 1).requires_grad_(True)
    opt = torch.optim.Adam(net.parameters(), lr=1e-3)
    for epoch in range(epochs):
        opt.zero_grad()
        loss = ((-d(d(u(x_col), x_col), x_col) - f(x_col)) ** 2).mean()        # the equation
        for x_b, normal, kind, value in conditions:                            # one term per condition
            xb = torch.tensor([[x_b]], requires_grad=True)
            ub = u(xb)
            dudn = normal * d(ub, xb)                                          # outward normal derivative
            if kind == "value":
                res = ub - value
            elif kind == "slope":
                res = dudn - value
            else:                                                              # Robin: a u + b du/dn = g
                a, b, g = value
                res = a * ub + b * dudn - g
            loss = loss + (res ** 2).sum()
        loss.backward()
        opt.step()
    x_t = torch.linspace(0, 1, 1001).reshape(-1, 1).requires_grad_(True)
    u_t = u(x_t)
    residual = (-d(d(u_t, x_t), x_t) - f(x_t)).detach()
    return u_t.detach(), residual, x_t.detach()

pi = math.pi
cases = {                     # point, outward normal, kind, value: all consistent with u* = sin(pi x) + 1 + x
    "value / value": [(0.0, -1, "value", 1.0), (1.0, +1, "value", 2.0)],
    "value / slope": [(0.0, -1, "value", 1.0), (1.0, +1, "slope", 1 - pi)],
    "value / Robin": [(0.0, -1, "value", 1.0), (1.0, +1, "robin", (1.0, 1.0, 3 - pi))],
}
for name, conditions in cases.items():
    u_t, residual, x_t = train(conditions)
    err = u_t - exact(x_t)
    print(f"{name}:  relative L2 error {err.norm() / exact(x_t).norm():.1e}   rms residual {residual.pow(2).mean().sqrt():.1e}")

value / value:  relative L2 error 3.4e-05   rms residual 9.6e-03
value / slope:  relative L2 error 9.1e-05   rms residual 1.0e-02
value / Robin:  relative L2 error 8.9e-05   rms residual 1.1e-02

The values come from $u^\ast$: $u^\ast(0) = 1$ and $u^\ast(1) = 2$; the outward slope at $x = 1$ is $u^{\ast\prime}(1) = \pi\cos\pi + 1 = 1 - \pi$; and the Robin condition $u + \partial u/\partial n = g$ at $x = 1$ has $g = 2 + 1 - \pi$. All three runs land within about $10^{-4}$ of $u^\ast$. A slope condition is as easy to impose as a value, provided its normal is right.

Too few conditions: the loss is satisfied, the answer is not

Now replace the value at $x = 0$ by a slope there too, again consistent with $u^\ast$: $\partial u/\partial n = -u'(0) = -(\pi + 1)$. Section 4.2 predicted the trouble: with slopes at both ends, adding a constant to a solution gives another solution, so the problem has infinitely many. Here is what four training runs, differing only in their seed, make of it:

both_slopes = [(0.0, -1, "slope", -(pi + 1)), (1.0, +1, "slope", 1 - pi)]
for seed in range(4):
    u_t, residual, x_t = train(both_slopes, seed=seed)
    err = u_t - exact(x_t)
    print(f"seed {seed}:  error = a constant {err.mean():+.3f} (varies by {err.std():.0e})   rms residual {residual.pow(2).mean().sqrt():.1e}")

seed 0:  error = a constant -0.083 (varies by 1e-04)   rms residual 1.0e-02
seed 1:  error = a constant -0.152 (varies by 7e-05)   rms residual 1.2e-02
seed 2:  error = a constant +0.251 (varies by 3e-05)   rms residual 1.4e-02
seed 3:  error = a constant +0.316 (varies by 5e-05)   rms residual 1.9e-02

Every run satisfies the equation and both conditions as well as the well-posed runs did, residual for residual. Every run is wrong by a different constant, and the error is almost exactly that constant: the network has found a correct shape and an arbitrary level, chosen by the accident of its starting weights (Figure 6.3). No term of the loss can see the difference, because the problem does not define one. A PINN will not warn you; it is the problem that was incomplete.

The fix is to supply the missing fact. Here one value is enough, for instance $u(0) = 1$ on top of the two slopes:

pinned = both_slopes + [(0.0, -1, "value", 1.0)]
u_t, residual, x_t = train(pinned)
err = u_t - exact(x_t)
print(f"two slopes and one value:  relative L2 error {err.norm() / exact(x_t).norm():.1e}")

two slopes and one value:  relative L2 error 8.7e-05

Physically, a bar insulated at both ends has a temperature fixed only up to a constant; the constant is set by how much heat it holds, so a condition on the average temperature, $\int_0^1 u\,dx = \bar u$, is the natural alternative to pinning a point. Either is one more term in the loss.

Periodic conditions, and the Studio's version

A periodic condition ties two edges together, so its residual compares the network with itself: values at matching points on opposite edges, and slopes. On a rectangle the Studio pairs BC1 with BC2 and BC3 with BC4 (Set the conditions). Periodicity alone, like Neumann alone, leaves a constant undetermined for Poisson's equation, for the same reason.

In the Studio each condition is one BC or IC block, and each produces one term, named after its block (BC1 dirichlet, BC2 neumann) in the generated script's log and loss history. The weights in the Loss block are the $w$ of Section 6.1, one per block. One IC block with an initial velocity produces two terms that share its weight.

Exercises

1. Which sign of normal does the Studio use on the bottom edge of a rectangle, $y = y_\text{min}$? Write the Neumann residual there in terms of $u_y$.
2. In the well-posed runs, the rms residual was about $10^{-2}$ while the error was about $10^{-4}$. Why can a residual of that size coexist with such a small error? (Section 6.5 returns to this.)
3. For the steady heat equation $-u'' = 0$ on $(0, 1)$ with $u'(0) = 0$ and $u'(1) = 1$, what goes wrong, and what would a PINN do?

Recap

• Each condition is one mean-square term: values for Dirichlet, outward slopes for Neumann (mind the normal's sign), a mix for Robin, matched pairs for periodic.
• With conditions that define a unique solution, slopes are as easy to impose as values.
• With too few conditions, a PINN reaches a small loss and an arbitrary member of the family of solutions. Supply the missing fact; the Graph check cannot tell you it is missing.