Systems of equations

Many problems have more than one unknown field: the velocity and pressure of a fluid, the displacements and stresses in a solid, the concentrations in a reaction. A PINN treats them with a network that has one output per field (or with one network per field). Each equation contributes its residual, which may involve any of the outputs and their derivatives; each condition applies to the output it constrains. Nothing else changes.

There is also a choice to be made when there is only one unknown. An equation of high order can usually be rewritten as a system of equations of lower order by naming some of its derivatives as new unknowns. Section 5.4 showed that every extra order of differentiation multiplies the cost; this section measures what the rewriting buys.

A beam, two ways

A beam of length $L$ and bending stiffness $EI$, simply supported at both ends, carries a uniform load $q$. Its deflection $w(x)$ obeys the Euler–Bernoulli equation (Section 4.1),

$$EI\,w'''' = q, \qquad w = 0 \text{ and } M = 0 \text{ at both ends},$$

where $M = -EI\,w''$ is the bending moment: zero at a simple support, because the support cannot resist rotation. In units with $EI = q = L = 1$ the solution is

$$w(x) = \frac{x - 2x^3 + x^4}{24}, \qquad M(x) = \frac{x(1 - x)}{2}.$$

The moment is what an engineer needs most, since the bending stress is proportional to it. The equation can be written in two ways:

The equilibrium equation is the same in both. What differs is where $M$ comes from: the second derivative of an output, or an output in its own right, tied to $w$ by a second equation. The two quantities are of very different sizes, a largest deflection of $5/384$ and a largest moment of $1/8$, so the network's outputs are scaled by these values and the bending residual by $1/8$ (Section 4.5).

import time
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]

w_exact = lambda x: (x - 2 * x**3 + x**4) / 24          # deflection, with EI = q = L = 1
M_exact = lambda x: (x - x**2) / 2                      # bending moment
W0, M0 = 5 / 384, 1 / 8                                 # their largest values: the natural scales

def train_beam(form, seed=0, epochs=5000):
    torch.manual_seed(seed)
    net = nn.Sequential(nn.Linear(1, 32), nn.Tanh(), nn.Linear(32, 32), nn.Tanh(),
                        nn.Linear(32, 32), nn.Tanh(), nn.Linear(32, 2 if form == "two outputs" else 1))
    x = torch.rand(200, 1).requires_grad_(True)
    x_end = torch.tensor([[0.0], [1.0]], requires_grad=True)

    def fields(p):                                      # w and M, from outputs of order one
        out = net(2 * p - 1)
        if form == "two outputs":
            return W0 * out[:, 0:1], M0 * out[:, 1:2]
        w = W0 * out
        return w, -d(d(w, p), p)                        # one output: M = -EI w''

    def loss_fn():
        w, M = fields(x)
        w_end, M_end = fields(x_end)
        residual = (d(d(M, x), x) + 1) ** 2             # equilibrium, M'' + q = 0, in both forms
        if form == "two outputs":
            residual = residual + ((d(d(w, x), x) + M) / M0) ** 2   # bending, EI w'' + M = 0
        return residual.mean() + ((w_end / W0) ** 2).sum() + ((M_end / M0) ** 2).sum()

    opt = torch.optim.Adam(net.parameters(), lr=1e-3)
    start = time.time()
    for epoch in range(epochs):
        opt.zero_grad()
        loss = loss_fn()
        loss.backward()
        opt.step()
    seconds = time.time() - start
    xt = torch.linspace(0, 1, 1001).reshape(-1, 1).requires_grad_(True)
    w, M = (f.detach() for f in fields(xt))
    xt = xt.detach()
    return ((w - w_exact(xt)).norm() / w_exact(xt).norm()).item(), \
           ((M - M_exact(xt)).norm() / M_exact(xt).norm()).item(), seconds

print("form          relative error in w (seeds 0, 1, 2)     relative error in M            seconds per run")
for form in ("one output", "two outputs"):
    runs = [train_beam(form, seed) for seed in range(3)]
    print(f"{form:12s}  " + "  ".join(f"{r[0]:.1e}" for r in runs) + "             "
          + "  ".join(f"{r[1]:.1e}" for r in runs) + f"        {sum(r[2] for r in runs) / 3:.0f}")

form          relative error in w (seeds 0, 1, 2)     relative error in M            seconds per run
one output    6.5e-03  4.9e-04  1.1e-03             2.0e-03  1.7e-04  3.3e-04        44
two outputs   2.3e-03  5.6e-05  1.1e-04             2.7e-03  1.4e-04  3.9e-05        17

With two outputs, each run took about 40% of the time, 17 seconds against 44 on the machine that produced these numbers, because the deepest chain of derivatives is two calls long instead of four (Section 5.4). The deflection was more accurate in every run, by a factor of three to ten. The moment, which the one-output form has to obtain as a second derivative, came out about as accurate either way: better with two outputs in two runs of the three, worse in one. Note also the conditions: with two outputs every condition is a plain value of an output, while the one-output form imposes $M = 0$ through a second derivative at the ends.

One network or several

A single network with several outputs shares its hidden layers between the fields, so it is cheaper, and the fields are represented by the same features, which suits fields as closely linked as $w$ and $M$. Separate networks, one per field, are more flexible when the fields have different character, for instance a smooth displacement and a stress that varies sharply. Both are common in solid mechanics: Haghighat et al. (2021) use a network for each displacement and stress component, and Rao et al. (2021) train displacement and stress together in a mixed formulation like the one above, imposing conditions on the stresses directly as values.

Systems need their residuals balanced like any loss with several terms. Here each residual was divided by its natural size, and each output scaled to be of order one. When the fields have no obvious scales, Chapter 10's methods take over.

In the Studio

The Network block's Network outputs sets the number of outputs, and every PDE, BC and IC block has an Applies to network output setting, so several fields can be solved side by side. A Custom residual, however, can only refer to its own output, $u$, and its derivatives; equations that couple two outputs, like the bending equation above, need the generated residual function edited, where the second output is field(pts, 1):

# not run
def residual_pde2(pts):                                  # bending: w'' + M = 0, coupling outputs 1 and 2
    w = p.field(pts, 0)
    M = p.field(pts, 1)
    w_xx = p.grad_of(p.grad_of(w, pts)[:, 0:1], pts)[:, 0:1]
    return w_xx + M
p.residual_pde2 = residual_pde2

Exercises

1. A cantilever is clamped at $x = 0$ and free at $x = 1$. At a clamp the deflection and slope vanish; at a free end the moment and the shear force, $V = M'$, vanish. Write the four conditions for the two-output formulation.
2. How many nested calls to torch.autograd.grad does each formulation need per training step for its residuals? Does that explain the timing?
3. For a beam on an elastic foundation, $EI\,w'''' + k\,w = q$, write a two-output system.

Recap

• A network with several outputs solves coupled equations; each residual may use any outputs, and each condition applies to the output it constrains.
• Rewriting a fourth-order equation as two second-order ones made the beam cheaper to train and its deflection more accurate, and gave the bending moment as an output, with its conditions imposed as plain values.
• Scale each output and each residual by its natural size. In the Studio, outputs and their conditions are set per block; coupled residuals need the generated script edited.