Scientific Machine Learning:
PINN: Derivatives by Automatic Differentiation and Finite Differences
A physics-informed neural network puts the governing equation into the loss, and a governing equation is a statement about derivatives. The network returns a value $u_\theta(x)$; the residual needs $u_{\theta,x}$, $u_{\theta,xx}$ and their companions in the other coordinates. Where those derivatives come from is a choice, and this notebook is about the two ways of making it: automatic differentiation, which differentiates the network through the computation graph that produced its output, and the finite difference method, which never differentiates anything and instead combines the network's values at neighbouring points. Both deliver the same residual to the same loss. They differ in accuracy, in cost, in where they can be evaluated, and in what they need from the collocation points, and the choice between them is part of designing a PINN.
The notebook first takes each method apart on a toy example small enough to follow by hand, then puts the two side by side in one laboratory, steady heat conduction in a plate with a localized heater, where everything except the derivative evaluation is held identical.
Learning Objectives¶
After studying this notebook, you should be able to
- explain why a physics-informed loss needs derivatives of the network output, and name the two ways of obtaining them;
- describe automatic differentiation as the chain rule applied along a computation graph, and compute a first and a second derivative of a small network by hand;
- derive the central finite-difference formulas for $u_x$ and $u_{xx}$ from Taylor expansions, and state their truncation error;
- explain the effect of the step size $h$ on a finite-difference derivative, including the roundoff that appears when $h$ is too small;
- implement the physics residual of a PINN by automatic differentiation and by finite differences, and say what each needs from the collocation points;
- compare the two implementations on a two-dimensional heat conduction problem in accuracy, cost and sensitivity to $h$; and
- choose between automatic differentiation and finite differences for a given problem, and justify the choice.
Setup
NumPy, Matplotlib and PyTorch, and the clock for timing.
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import time
A data-driven network is trained by comparing its output with observations: the loss is $\lVert u_\theta - u_{\mathrm{data}}\rVert^2$, and nothing but the output is needed. A physics-informed network adds a second kind of term. The governing equation, say $-k\,(u_{xx} + u_{yy}) = q$, is evaluated with the network's output in place of $u$, and the amount by which it fails to hold, the residual $r_\theta$, is penalized,
$$ \mathcal{L} = \underbrace{\lVert u_\theta - u_{\mathrm{data}}\rVert^2}_{\text{data}} + \underbrace{\lVert r_\theta\rVert^2}_{\text{physics}} + \underbrace{\lVert u_\theta - u_{\mathrm{BC/IC}}\rVert^2}_{\text{boundary, initial}} . $$
The residual is the new ingredient, and it is where the difficulty sits. The network returns $u_\theta$ at a point; the residual needs $u_{\theta,xx}$ and $u_{\theta,yy}$ at that point. Something has to turn the value into its derivatives before the loss can be formed, and that something is run at every collocation point at every iteration, so its accuracy and its cost both matter.
The PINN workflow. The network predicts $u_\theta$; the physics residual needs derivatives of $u_\theta$ with respect to the input coordinates, which have to be produced by a separate step between the network output and the loss.
There are two ways to fill the derivative box in the figure.
Automatic differentiation (AD) uses the fact that the network is a composition of elementary operations whose derivatives are known. Differentiating the output with respect to an input is the chain rule applied along that composition, and a framework such as PyTorch does it exactly, to machine precision, by walking back through the operations it recorded when the output was computed. This is the standard choice for a PINN and the one used in every earlier notebook.
The finite difference method (FDM) does not differentiate the network at all. It evaluates the network at the point of interest and at neighbours a distance $h$ away, and combines the values with the familiar stencils, $u_{xx} \approx (u_{i-1} - 2u_i + u_{i+1})/h^2$. The derivative is approximate, with an error that depends on $h$, but it costs only a few extra forward passes and needs no record of how the output was computed.
The two are treated in turn in the next two chapters, then compared in a PINN.
2.1. Elementary Operations and the Computation Graph¶
A neural network, however large, is built from a short list of elementary operations: multiplication by a weight, addition of a bias, and the activation function. Each of these has a derivative that is known in closed form. When the network computes its output, a framework records the sequence of operations as a computation graph, a directed graph whose nodes are intermediate values and whose edges are the operations connecting them. Automatic differentiation is the chain rule applied along this graph: the derivative of the output with respect to any earlier value is the product of the local derivatives along the path between them, summed over all paths.
The smallest useful example is a single neuron. With inputs $h_1$, $h_2$, weights $w_1$, $w_2$, bias $b$ and activation $\sigma$,
$$ u_\theta = \sigma(w_1 h_1 + w_2 h_2 + b) $$
decomposes into the elementary operations $z_1 = w_1 h_1$, $z_2 = w_2 h_2$, $z_3 = z_1 + z_2 + b$, $u_\theta = \sigma(z_3)$, with local derivatives $\partial u/\partial z_3 = \sigma'(z_3)$, $\partial z_3/\partial z_1 = \partial z_3/\partial z_2 = 1$, $\partial z_1/\partial h_1 = w_1$ and $\partial z_2/\partial h_1 = 0$.
A neuron as a computation graph. Each node is an intermediate value, each edge an elementary operation with a known local derivative. The graph for a whole network is this picture repeated at every neuron.
The derivative of the output with respect to $h_1$ follows the one path from $h_1$ to $u_\theta$ and multiplies the local derivatives along it,
$$ \frac{\partial u_\theta}{\partial h_1} = \frac{\partial u_\theta}{\partial z_3}\,\frac{\partial z_3}{\partial z_1}\,\frac{\partial z_1}{\partial h_1} = \sigma'(z_3)\cdot 1\cdot w_1 = w_1\,\sigma'(z_3) . $$
In a PINN the inputs of the first layer are the coordinates $x$, $y$ and $t$, so $\partial u_\theta/\partial x$ is this computation with $h_1 = x$ and the path extended through every layer. The framework performs it with one call, and the result is exact: no approximation is made anywhere, and the only error is floating-point roundoff.
The computation above is reproduced in PyTorch for the neuron with $\sigma = \tanh$, $w_1 = 0.8$, $w_2 = -0.5$, $b = 0.1$, at the input $(h_1, h_2) = (0.7, 0.2)$, and checked against the chain rule written out by hand. The input $h_1$ is created with requires_grad = True, which tells the framework to record the graph of every operation that uses it; torch.autograd.grad(u, h1) then walks the graph back from u to h1. Passing create_graph = True makes the derivative itself part of a graph, so that it can be differentiated again, which is how the second derivative is obtained in the last line.
w1 = 0.8
w2 = -0.5
b = 0.1
h1 = torch.tensor([[0.7]], requires_grad = True)
h2 = torch.tensor([[0.2]])
z1 = w1*h1
z2 = w2*h2
z3 = z1 + z2 + b
u = torch.tanh(z3)
du_dh1 = torch.autograd.grad(u, h1, create_graph = True)[0]
d2u_dh1 = torch.autograd.grad(du_dh1, h1)[0]
print('autograd : du/dh1 = %.6f, d2u/dh1^2 = %.6f' % (du_dh1.item(), d2u_dh1.item()))
s = np.tanh(w1*0.7 + w2*0.2 + b)
print('chain rule: du/dh1 = %.6f, d2u/dh1^2 = %.6f' % (w1*(1 - s**2), w1**2*(-2*s*(1 - s**2))))
Autograd and the chain rule agree to every printed digit. For $\sigma = \tanh$ the derivatives are $\sigma' = 1 - \sigma^2$ and $\sigma'' = -2\sigma(1 - \sigma^2)$, so with $z_3 = 0.8\times0.7 - 0.5\times0.2 + 0.1 = 0.56$ the two numbers are $w_1\sigma'(z_3) = 0.593567$ and $w_1^2\sigma''(z_3) = -0.482430$. The second derivative is not a new idea: it is the derivative of a derivative, and the graph that produced $\partial u/\partial h_1$ is differentiated the same way as the graph that produced $u$.
The graph extended for a second derivative. The first derivative $\partial u_\theta/\partial h_1 = w_1\sigma'(z_3)$ is itself a computation with nodes $\sigma'$ and $\times$, so it has a graph of its own, and differentiating along it gives $\partial^2 u_\theta/\partial h_1^2 = w_1^2\sigma''(z_3)$.
The extension to a full network is mechanical. A derivative(u, x) helper wraps the autograd call with create_graph = True, so that derivative(derivative(u, x), x) is $u_{xx}$ and a third call is $u_{xxx}$. The price is that every derivative adds a layer of graph on top of the previous one: the graph for $u_{xx}$ contains the graph for $u_x$, which contains the graph for $u$. Memory grows with the order, and so does time. The cell below measures the forward pass and three successive derivatives of a $1 \to 64 \to 64 \to 64 \to 1$ network at $10000$ points.
def new_pinn():
torch.manual_seed(0)
return nn.Sequential(nn.Linear(1, 64), nn.Tanh(),
nn.Linear(64, 64), nn.Tanh(),
nn.Linear(64, 64), nn.Tanh(),
nn.Linear(64, 1))
def derivative(u, x):
return torch.autograd.grad(u, x, torch.ones_like(u), create_graph = True)[0]
net = new_pinn()
x = torch.linspace(0, 1, 10000).reshape(-1, 1).requires_grad_(True)
u = net(x) # one untimed pass to warm up
u_xxx = derivative(derivative(derivative(u, x), x), x)
t0 = time.time()
for k in range(100):
u = net(x)
t1 = time.time()
for k in range(100):
u_x = derivative(u, x)
t2 = time.time()
for k in range(100):
u_xx = derivative(u_x, x)
t3 = time.time()
for k in range(100):
u_xxx = derivative(u_xx, x)
t4 = time.time()
print('forward pass : %.4f s' % ((t1 - t0)/100))
print('1st derivative : %.4f s' % ((t2 - t1)/100))
print('2nd derivative : %.4f s' % ((t3 - t2)/100))
print('3rd derivative : %.4f s' % ((t4 - t3)/100))
On this machine the forward pass took 0.0071 s and the first derivative 0.0045 s, about the same, which is the usual rule for one reverse pass. From there each order costs more than the one before: the second derivative 0.0119 s, the third 0.0265 s, each roughly two to three times the last, because the graph being differentiated has grown by one layer each time. The ratios vary from machine to machine, but the trend does not, and in a PINN that trains for tens of thousands of iterations with second derivatives at every one of them, it is the dominant cost. This is the main weakness of automatic differentiation, and it is why a cheaper route to the same derivatives is worth having.
Strengths. The derivatives are exact, with no discretization error. They can be evaluated at any point whatever, so the collocation points may be scattered freely over the domain, which suits the coordinate-based networks of the PINN construction. And nothing has to be written for each equation beyond the residual itself.
Weaknesses. The computation graphs must be retained through the derivative computation, which costs memory. Each additional order of derivative adds a graph on top of the last, which costs time. For an equation of fourth order, or for a residual that needs many mixed derivatives, these costs can dominate the training.
AD on scattered points. Because the derivative is computed from the network itself, it can be evaluated at any point, and the collocation points need no structure (left). Each derivative is one more pass through the graph (right).
3.1. Stencils from Taylor Expansion¶
The finite difference method approximates a derivative at a point from the values of the function at that point and its neighbours. The formulas come from Taylor's theorem. For a smooth $u$ and a step $h$,
$$ u(x + h) = u(x) + h\,u'(x) + \frac{h^2}{2}\,u''(x) + \frac{h^3}{6}\,u'''(x) + O(h^4) , \qquad u(x - h) = u(x) - h\,u'(x) + \frac{h^2}{2}\,u''(x) - \frac{h^3}{6}\,u'''(x) + O(h^4) . $$
Subtracting the second from the first cancels the even terms and gives the central difference for the first derivative; adding them cancels the odd terms and gives the second derivative:
$$ u'(x) = \frac{u(x + h) - u(x - h)}{2h} + O(h^2) , \qquad u''(x) = \frac{u(x - h) - 2u(x) + u(x + h)}{h^2} + O(h^2) . $$
Both are second-order accurate: halving $h$ divides the error by four. The error term is called the truncation error, because it is what was truncated from the Taylor series, and it is proportional to $h^2 u''''$ for the second-derivative formula, so it is large where the function bends sharply and small where it is nearly linear.
On a grid with spacing $h$ the formulas are usually written with indices, $u_i = u(x_i)$,
$$ u_x(x_i) \approx \frac{u_{i+1} - u_{i-1}}{2h} , \qquad u_{xx}(x_i) \approx \frac{u_{i-1} - 2u_i + u_{i+1}}{h^2} , $$
and the pattern of points and coefficients, here $(1, -2, 1)/h^2$ over three neighbouring points, is called a stencil.
The three-point stencil. The derivatives at $x_i$ are formed from the values at $x_i$ and its two neighbours a distance $h$ away.
In two dimensions the Laplacian $\nabla^2 u = u_{xx} + u_{yy}$ is the sum of the one-dimensional formula applied in $x$ and in $y$. On a grid with the same spacing $h$ in both directions, with $u_{i,j} = u(x_i, y_j)$,
$$ \nabla^2 u(x_i, y_j) \approx \frac{u_{i+1,j} + u_{i-1,j} + u_{i,j+1} + u_{i,j-1} - 4u_{i,j}}{h^2} , $$
which uses the point and its four neighbours, the five-point stencil. When $u$ is available on the whole grid as an array, the formula is the convolution of that array with the kernel
$$ \frac{1}{h^2}\begin{bmatrix} 0 & 1 & 0 \\ 1 & -4 & 1 \\ 0 & 1 & 0 \end{bmatrix} , $$
which is why finite differences are natural when the network's output is an image or a field rather than a value at a point: one convolution gives the Laplacian everywhere at once. The same stencil, used the other way round, is also how the reference solution of the laboratory is computed: the equation $-\nabla^2 u = q$ at every grid point becomes a linear system in the unknown grid values, which is solved directly.
The accuracy of a finite difference depends on $h$, and in two opposite ways. The test function is $u = \sin(2\pi x)$ on $[0, 1]$, whose second derivative $-(2\pi)^2\sin(2\pi x)$ is known. The first cell compares the three-point formula with the exact $u_{xx}$ for three values of $h$.
def u_fun(x):
return np.sin(2*np.pi*x)
def u_xx_exact(x):
return -(2*np.pi)**2*np.sin(2*np.pi*x)
x = np.linspace(0, 1, 101)
plt.figure(figsize = (6, 6))
plt.subplot(2, 1, 1)
plt.plot(x, u_xx_exact(x), 'k', linewidth = 3, alpha = 0.3, label = 'Exact')
for h in [0.2, 0.1, 0.05]:
u_xx_fdm = (u_fun(x - h) - 2*u_fun(x) + u_fun(x + h))/h**2
plt.plot(x, u_xx_fdm, '--', linewidth = 1.5, label = 'FDM, h = %g' % h)
plt.ylabel('$u_{xx}$')
plt.xlim([0, 1])
plt.legend()
plt.grid(True)
plt.subplot(2, 1, 2)
for h in [0.2, 0.1, 0.05]:
u_xx_fdm = (u_fun(x - h) - 2*u_fun(x) + u_fun(x + h))/h**2
plt.plot(x, u_xx_fdm - u_xx_exact(x), '--', linewidth = 1.5, label = 'h = %g' % h)
plt.xlabel('x')
plt.ylabel('FDM error')
plt.xlim([0, 1])
plt.legend()
plt.grid(True)
plt.show()
With $h = 0.2$, a fifth of the period, the stencil badly underestimates the curvature, because the three points straddle a large part of the wave and the parabola through them is flatter than the sine; the error is largest where $u_{xx}$ is largest. With $h = 0.1$ the error has fallen to a fraction, and with $h = 0.05$ the two curves are hard to tell apart. The error is the truncation error of Section 3.1, and it shrinks as $h^2$: from $h = 0.2$ to $0.05$, a factor of $4$ in $h$, the error falls by about a factor of $16$. The next cell follows the error down to much smaller $h$, in double precision and in the single precision that neural networks usually use.
h_list = [0.2, 0.1, 0.05, 0.02, 0.01, 0.005, 0.002, 0.001]
err_64 = []
err_32 = []
for h in h_list:
x64 = np.linspace(0, 1, 101)
u_xx_fdm = (u_fun(x64 - h) - 2*u_fun(x64) + u_fun(x64 + h))/h**2
err_64.append(np.max(np.abs(u_xx_fdm - u_xx_exact(x64))))
x32 = np.float32(x64)
h32 = np.float32(h)
u_xx_fdm = (u_fun(x32 - h32) - 2*u_fun(x32) + u_fun(x32 + h32))/h32**2
err_32.append(np.max(np.abs(u_xx_fdm - u_xx_exact(x64))))
for i in range(len(h_list)):
print('h = %.3f | error (float64) = %.3e | error (float32) = %.3e' % (h_list[i], err_64[i], err_32[i]))
plt.figure(figsize = (6, 4))
plt.loglog(h_list, err_64, 'ko-', label = 'float64')
plt.loglog(h_list, err_32, 'rs--', label = 'float32')
plt.loglog(h_list, 130*np.array(h_list)**2, 'k:', label = r'$\propto h^2$')
plt.xlabel('h')
plt.ylabel('max error of $u_{xx}$')
plt.legend()
plt.grid(True)
plt.show()
In double precision the error follows the $h^2$ line all the way down to $h = 0.001$, where it is $1.3\times10^{-4}$. In single precision it follows the line only to about $h = 0.01$, where it reaches $1.4\times10^{-2}$, and then it grows as $h$ shrinks further, to $0.48$ at $h = 0.001$. This is roundoff error. The numerator $u(x-h) - 2u(x) + u(x+h)$ is a small difference of numbers of order one; single precision carries about seven digits, so the difference is accurate to about $10^{-7}$, and dividing by $h^2 = 10^{-6}$ turns that into an error of order $0.1$. The two errors pull in opposite directions, truncation falling and roundoff rising as $h$ decreases, and the best $h$ sits between them. For a PINN, which runs in single precision, the practical range is roughly $10^{-3} \lesssim h \lesssim 10^{-1}$, and the laboratory will show where the problem's own grid spacing falls in it.
Remark (why the roundoff scales as $1/h^2$). Each function value carries a relative error $\epsilon$ of about $6\times10^{-8}$ in single precision. The three-term numerator therefore has an absolute error of order $\epsilon\,\lvert u\rvert$, and the division by $h^2$ amplifies it to $\epsilon\lvert u\rvert/h^2$. Adding the truncation error $h^2\lvert u''''\rvert/12$ and minimizing over $h$ gives an optimal step of order $\epsilon^{1/4} \approx 0.016$ for the second derivative, with a best achievable error of order $\epsilon^{1/2} \approx 2.5\times10^{-4}$ relative to $u''''$. Both agree with the table.
Strengths. The derivative is a few additions and one division, with no graph to build or keep, so it is cheap in memory and simple to write. When the output is a field on a grid, the whole Laplacian is one convolution. And the method is the one every numerical analyst already knows, so a finite-difference PINN reads like a finite-difference code.
Weaknesses. The derivative carries a discretization error, $O(h^2)$ in truncation plus roundoff when $h$ is small. Neighbouring points are required: the network must be evaluated at $x \pm h$, $y \pm h$ as well as at $(x, y)$, which is five forward passes for a two-dimensional Laplacian instead of one, and if the collocation points are meant to lie on a grid, the grid must be structured. And the error is not under the optimizer's control: a residual computed with a coarse $h$ is the residual of a slightly different equation, and the network will faithfully solve that one.
A derivative from neighbouring predictions. The network is evaluated at $x$, $x - h$ and $x + h$; the stencil combines the three values. The error of the result is set by $h$, not by the network.
4.1. Same Loss, Different Derivative¶
Inside a PINN the two methods occupy exactly the same place. The loss is the same, the network is the same, the optimizer is the same; what differs is the few lines that turn $u_\theta$ into $u_{\theta,xx}$ before the residual is formed. With automatic differentiation the network is evaluated once at $x_i$, and autograd walks the graph back to the input. With finite differences the network is evaluated at $x_i - h$, $x_i$ and $x_i + h$, and the three values go into the stencil.
Two routes to the same residual. AD differentiates the network output through its computation graph at one point; FDM evaluates the network at neighbouring points and combines the values. Both produce $\partial^2 u/\partial x^2$ for the same PDE loss.
The two residual functions are shown on one network in one dimension, before the laboratory uses them in two. The AD version needs x created with requires_grad_(True) and two calls of derivative. The FDM version needs nothing of x except its value, and three forward passes, at x - h, x and x + h, with $h = 0.01$. The second derivatives of a freshly initialized network are compared at nine points.
net = new_pinn()
x = torch.linspace(0.1, 0.9, 9).reshape(-1, 1).requires_grad_(True)
u = net(x)
u_xx_AD = derivative(derivative(u, x), x)
h = 0.01
u_xx_FDM = (net(x - h) - 2*net(x) + net(x + h))/h**2
for i in range(9):
print('x = %.1f | AD : %9.5f | FDM : %9.5f | difference : %.1e'
% (x[i].item(), u_xx_AD[i].item(), u_xx_FDM[i].item(), abs(u_xx_AD[i] - u_xx_FDM[i]).item()))
The two columns agree to about two or three significant digits, with differences of order $10^{-4}$ to $10^{-3}$, which is the truncation error $h^2 u''''/12$ of the stencil at $h = 0.01$ for a function whose fourth derivative is of order one. AD is the reference here; FDM is close to it, and how close is decided by $h$ alone. Which route is better in a PINN cannot be settled on nine points: it depends on whether the error of the stencil matters for the solution, and on what the cheaper derivative buys in training time. That is the laboratory's question.
5.1. Problem Setup¶
A thin square plate of side $L$ is heated by a small heater at its centre, and all four edges are held at the ambient temperature. Heat spreads from the heater by conduction and leaves through the edges; the steady temperature field is the balance of the two. This is the situation of a chip on a substrate, a heated panel or any thin structural component with a localized heat source. The steady heat equation is
$$ -k\,(T_{xx} + T_{yy}) = q(x, y) , \qquad T = T_{\mathrm{ambient}} \ \text{ on } \partial\Omega , $$
with a Gaussian heater of peak $q_0$ and spread $\sigma$ at the centre,
$$ q(x, y) = q_0\exp\left[-\frac{(x - L/2)^2 + (y - L/2)^2}{2\sigma^2}\right] . $$
| Parameter | Value | Meaning |
|---|---|---|
| $L$ | $1.0\ \mathrm{m}$ | plate width and height |
| $k$ | $10\ \mathrm{W/(m\,K)}$ | thermal conductivity |
| $T_{\mathrm{ambient}}$ | $20^\circ\mathrm{C}$ | edge temperature |
| $q_0$ | $1000\ \mathrm{W/m^3}$ | peak heat generation |
| $\sigma$ | $0.1\,L$ | heater spread |
The network is not asked for $T$ in degrees. Temperatures of order $20$ and a source of order $1000$ are poor numbers to train on, so the problem is made dimensionless first. With the coordinates measured in units of $L$ and the temperature rise above ambient measured in units of $q_0 L^2/k$,
$$ \theta = \frac{T - T_{\mathrm{ambient}}}{q_0 L^2/k} , \qquad -(\theta_{xx} + \theta_{yy}) = \hat q(x, y) = \exp\left[-\frac{(x - 0.5)^2 + (y - 0.5)^2}{2\sigma^2}\right] , \qquad \theta = 0 \ \text{ on } \partial\Omega , $$
every quantity is of order one or smaller, and the physical temperature is recovered at the end from $T = T_{\mathrm{ambient}} + (q_0 L^2/k)\,\theta$. The scale $q_0 L^2/k$ is $100^\circ\mathrm{C}$ for the parameters above.
Plate and heater
The parameters, the temperature scale, and the dimensionless source written once for NumPy and once for PyTorch.
L = 1.0 # plate size (m)
k = 10.0 # conductivity (W/m/K)
T_ambient = 20.0 # edge temperature (deg C)
q0 = 1000.0 # peak heat generation (W/m^3)
sigma = 0.1 # heater spread (in units of L)
scale = q0*L**2/k # temperature scale (deg C): T = T_ambient + scale*theta
def source(x, y):
return np.exp(-((x - 0.5)**2 + (y - 0.5)**2)/(2*sigma**2))
def source_torch(x, y):
return torch.exp(-((x - 0.5)**2 + (y - 0.5)**2)/(2*sigma**2))
print('temperature scale: %.1f deg C' % scale)
The problem has no closed-form solution, so the reference is a conventional finite-difference solution on a fine grid, $201 \times 201$ points, far finer than anything the networks will see. The five-point stencil of Section 3.2 turns $-(\theta_{xx} + \theta_{yy}) = \hat q$ at every interior point into one linear equation in the five surrounding unknowns; the $199 \times 199 = 39601$ equations are assembled as a sparse matrix and solved directly. The one-dimensional second-difference matrix $D$ with $(-1, 2, -1)$ on its diagonals is built first, and the two-dimensional operator is $I \otimes D + D \otimes I$, the Kronecker sum, which is the five-point stencil in matrix form.
from scipy.sparse import diags, eye, kron
from scipy.sparse.linalg import spsolve
n = 199 # interior points per direction
h_ref = 1/(n + 1)
x_ref = np.linspace(0, 1, n + 2)
X_ref, Y_ref = np.meshgrid(x_ref, x_ref)
# 1D second-difference matrix and its 2D (Kronecker) assembly
D = diags([-np.ones(n - 1), 2*np.ones(n), -np.ones(n - 1)], [-1, 0, 1])
A = kron(eye(n), D) + kron(D, eye(n))
q = source(X_ref[1:-1, 1:-1], Y_ref[1:-1, 1:-1])
theta_ref = np.zeros((n + 2, n + 2))
theta_ref[1:-1, 1:-1] = spsolve(A.tocsr(), h_ref**2*q.ravel()).reshape(n, n)
T_ref = T_ambient + scale*theta_ref
print('reference maximum temperature: %.3f deg C' % T_ref.max())
Reference temperature
plt.figure(figsize = (6, 5))
plt.contourf(X_ref, Y_ref, T_ref, levels = np.linspace(20, 21.7, 35), cmap = 'inferno')
plt.colorbar(label = 'T (deg C)')
plt.plot(0.5, 0.5, 'cx', markersize = 10)
plt.title('Reference Temperature')
plt.xlabel('x')
plt.ylabel('y')
plt.axis('square')
plt.show()
The temperature rises from $20^\circ\mathrm{C}$ at the edges to $21.627^\circ\mathrm{C}$ at the centre, a rise of only $1.627^\circ\mathrm{C}$: the heater is weak and the plate conducts well. The field is radially symmetric near the centre, where only the heater matters, and squares off towards the edges, where the boundary condition takes over. In dimensionless terms the peak is $\theta = 0.0163$, a small number, which is one reason the boundary term will need a large weight in the loss.
Both PINNs use the same collocation points, the same network, the same initialization, the same loss and the same optimizer and schedule. Only the function that computes the residual differs.
Collocation points
A $21 \times 21$ grid of interior points with spacing $h = 1/22$, so that $441$ residuals are evaluated, and $80$ points on each edge, $320$ in all, for the boundary condition. A regular grid is used for both models so that the comparison is fair: the FDM-PINN will take its step $h$ equal to the grid spacing, and the AD-PINN could have used any points at all. The interior coordinates are created with requires_grad_(True) for the AD model; the FDM model simply ignores that flag.
h = 1/22 # grid spacing of the collocation grid
x_in = np.linspace(h, 1 - h, 21)
X_in, Y_in = np.meshgrid(x_in, x_in)
x_pde = torch.tensor(X_in.reshape(-1, 1), dtype = torch.float32).requires_grad_(True)
y_pde = torch.tensor(Y_in.reshape(-1, 1), dtype = torch.float32).requires_grad_(True)
s = np.linspace(0, 1, 80)
x_bc = np.concatenate([s, s, np.zeros(80), np.ones(80)]) # bottom, top, left, right
y_bc = np.concatenate([np.zeros(80), np.ones(80), s, s])
x_bc = torch.tensor(x_bc.reshape(-1, 1), dtype = torch.float32)
y_bc = torch.tensor(y_bc.reshape(-1, 1), dtype = torch.float32)
print('interior points: %d, boundary points: %d' % (x_pde.shape[0], x_bc.shape[0]))
plt.figure(figsize = (5, 5))
plt.plot(X_in.ravel(), Y_in.ravel(), 'b.', markersize = 4, label = 'PDE points')
plt.plot(x_bc.numpy(), y_bc.numpy(), 'r.', markersize = 4, label = 'Boundary points')
plt.plot(0.5, 0.5, 'cx', markersize = 10, label = 'Heater')
plt.xlabel('x')
plt.ylabel('y')
plt.axis('square')
plt.legend(loc = 'upper right')
plt.show()
Network
A $\tanh$ network $2 \to 64 \to 64 \to 64 \to 1$; new_pinn() resets the seed so that both models start from identical weights. The network's raw output is the temperature rise in degrees, a number of order one, which is what a freshly initialized network produces naturally; theta_net concatenates the two coordinate columns into the input and divides the output by the scale, so that everything downstream is in terms of $\theta$. derivative is the autograd wrapper of Section 2.3.
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))
def theta_net(net, x, y):
return net(torch.cat([x, y], dim = 1))/scale # the network predicts the temperature rise in deg C
def derivative(u, x):
return torch.autograd.grad(u, x, torch.ones_like(u), create_graph = True)[0]
Loss
The loss is the mean squared residual over the interior points plus the mean squared boundary error with weight $100$,
$$ \mathcal{L} = \frac{1}{441}\sum_{i}\, r_\theta(x_i, y_i)^2 + 100\cdot\frac{1}{320}\sum_{j}\,\theta(x_j, y_j)^2 , $$
trained with Adam at $10^{-3}$ for $20000$ iterations. The weight is needed because $\theta$ is of order $10^{-2}$ while the source is of order one: without it the boundary term would be invisible next to the residual, and the network would be free to shift its whole solution by a constant. The training loop is the same for both models, so it is written out twice with only the residual function changed.
Residual by automatic differentiation
Two calls of derivative in $x$ and two in $y$. The coordinates are passed as separate columns so that each can be differentiated on its own.
def residual_AD(net, x, y):
theta = theta_net(net, x, y)
theta_xx = derivative(derivative(theta, x), x)
theta_yy = derivative(derivative(theta, y), y)
return -(theta_xx + theta_yy) - source_torch(x, y)
Train
net_AD = new_pinn()
optimizer = torch.optim.Adam(net_AD.parameters(), lr = 1e-3)
loss_AD = []
t0 = time.time()
for epoch in range(20001):
r = residual_AD(net_AD, x_pde, y_pde)
loss_pde = torch.mean(r**2)
loss_bc = torch.mean(theta_net(net_AD, x_bc, y_bc)**2)
loss = loss_pde + 100*loss_bc
optimizer.zero_grad()
loss.backward()
optimizer.step()
if epoch % 5000 == 4999: # halve the learning rate every 5000 iterations
optimizer.param_groups[0]['lr'] = optimizer.param_groups[0]['lr']/2
loss_AD.append(loss.item())
if epoch % 2000 == 0:
print('EPOCH : %5d | Loss_Total : %.3e | Loss_PDE : %.3e | Loss_BC : %.3e'
% (epoch, loss.item(), loss_pde.item(), loss_bc.item()))
time_AD = time.time() - t0
print('training time: %.1f s' % time_AD)
The loss falls from $3.4\times10^{-2}$ to $4.9\times10^{-8}$ over the $20000$ iterations, four orders of magnitude in the first 2000 and two more after that. The residual term dominates throughout: the boundary term starts at $8\times10^{-8}$, because the initial network output is small on the scale of $\theta$, and stays three orders below the residual. The learning-rate halvings at $5000$, $10000$ and $15000$ are visible as the moments when the printed loss stops fluctuating and settles. Training took 200 s on this machine.
Residual by finite differences
The network is evaluated at the point and at its four neighbours a distance $h$ away, and the five values go into the stencil of Section 3.2. There is no autograd call: the only derivative taken anywhere in this model is the one the optimizer takes with respect to the weights, which is unaffected. The step $h$ is an argument, so that it can be varied later; for the main comparison it equals the grid spacing $1/22$, so that the neighbours of each interior point are other collocation points or boundary points.
def residual_FDM(net, x, y, h):
theta = theta_net(net, x, y)
theta_xx = (theta_net(net, x + h, y) - 2*theta + theta_net(net, x - h, y))/h**2
theta_yy = (theta_net(net, x, y + h) - 2*theta + theta_net(net, x, y - h))/h**2
return -(theta_xx + theta_yy) - source_torch(x, y)
Train
net_FDM = new_pinn()
optimizer = torch.optim.Adam(net_FDM.parameters(), lr = 1e-3)
loss_FDM = []
t0 = time.time()
for epoch in range(20001):
r = residual_FDM(net_FDM, x_pde, y_pde, h)
loss_pde = torch.mean(r**2)
loss_bc = torch.mean(theta_net(net_FDM, x_bc, y_bc)**2)
loss = loss_pde + 100*loss_bc
optimizer.zero_grad()
loss.backward()
optimizer.step()
if epoch % 5000 == 4999: # halve the learning rate every 5000 iterations
optimizer.param_groups[0]['lr'] = optimizer.param_groups[0]['lr']/2
loss_FDM.append(loss.item())
if epoch % 2000 == 0:
print('EPOCH : %5d | Loss_Total : %.3e | Loss_PDE : %.3e | Loss_BC : %.3e'
% (epoch, loss.item(), loss_pde.item(), loss_bc.item()))
time_FDM = time.time() - t0
print('training time: %.1f s' % time_FDM)
The printed losses track the AD run closely, agreeing to two significant digits at every printed epoch up to $6000$, and at the end the FDM loss is $5.4\times10^{-8}$ against the AD loss of $4.9\times10^{-8}$: by the loss alone the two models are the same model. The difference between them is not in how well they satisfy their residuals but in what their residuals are, and that only the comparison with the reference can show. Training took 132 s, about two thirds of the AD time: five forward passes of a small network are cheaper than one forward pass plus two orders of autograd.
Evaluation on the reference grid
Both networks are evaluated on the full $201\times201$ reference grid, converted back to degrees, and compared with the reference by the relative $L_2$ error of the dimensionless temperature rise and by the maximum temperature.
x_eval = torch.tensor(X_ref.reshape(-1, 1), dtype = torch.float32)
y_eval = torch.tensor(Y_ref.reshape(-1, 1), dtype = torch.float32)
theta_AD = theta_net(net_AD, x_eval, y_eval).detach().numpy().reshape(201, 201)
theta_FDM = theta_net(net_FDM, x_eval, y_eval).detach().numpy().reshape(201, 201)
T_AD = T_ambient + scale*theta_AD
T_FDM = T_ambient + scale*theta_FDM
error_AD = np.linalg.norm(theta_AD - theta_ref)/np.linalg.norm(theta_ref)
error_FDM = np.linalg.norm(theta_FDM - theta_ref)/np.linalg.norm(theta_ref)
print(' reference AD-PINN FDM-PINN')
print('max T (deg C) : %8.3f %8.3f %8.3f' % (T_ref.max(), T_AD.max(), T_FDM.max()))
print('relative L2 : - %.3e %.3e' % (error_AD, error_FDM))
print('max |T error| : - %8.4f %8.4f' % (np.max(np.abs(T_AD - T_ref)), np.max(np.abs(T_FDM - T_ref))))
print('training time : - %6.1f s %6.1f s' % (time_AD, time_FDM))
Temperature fields
On the colour scale of the reference.
levels = np.linspace(20, 21.7, 35)
plt.figure(figsize = (12, 4))
plt.subplot(1, 3, 1)
plt.contourf(X_ref, Y_ref, T_ref, levels = levels, cmap = 'inferno')
plt.title('Reference')
plt.xlabel('x')
plt.ylabel('y')
plt.axis('square')
plt.subplot(1, 3, 2)
plt.contourf(X_ref, Y_ref, T_AD, levels = levels, cmap = 'inferno')
plt.title('AD-PINN')
plt.xlabel('x')
plt.axis('square')
plt.subplot(1, 3, 3)
plt.contourf(X_ref, Y_ref, T_FDM, levels = levels, cmap = 'inferno')
plt.title('FDM-PINN')
plt.xlabel('x')
plt.axis('square')
plt.show()
Error maps
Absolute error in degrees on a fixed scale from $0$ to $0.05^\circ\mathrm{C}$.
levels = np.linspace(0, 0.05, 26)
plt.figure(figsize = (8, 4))
plt.subplot(1, 2, 1)
plt.contourf(X_ref, Y_ref, np.abs(T_AD - T_ref), levels = levels, cmap = 'gray', extend = 'max')
plt.title('AD-PINN Error (deg C)')
plt.xlabel('x')
plt.ylabel('y')
plt.axis('square')
plt.subplot(1, 2, 2)
plt.contourf(X_ref, Y_ref, np.abs(T_FDM - T_ref), levels = levels, cmap = 'gray', extend = 'max')
plt.title('FDM-PINN Error (deg C)')
plt.xlabel('x')
plt.axis('square')
plt.show()
Profile through the centre
plt.figure(figsize = (6, 4))
plt.plot(x_ref, T_ref[100, :], 'k', linewidth = 3, alpha = 0.3, label = 'Reference')
plt.plot(x_ref, T_AD[100, :], 'r--', linewidth = 1.5, label = 'AD-PINN')
plt.plot(x_ref, T_FDM[100, :], 'b--', linewidth = 1.5, label = 'FDM-PINN')
plt.xlabel('x')
plt.ylabel('T (deg C) at y = 0.5')
plt.xlim([0, 1])
plt.legend()
plt.grid(True)
plt.show()
Both PINNs reproduce the field. The three colour maps are the same to the eye, the centre-line profiles lie on the reference, and both maximum temperatures are within $0.013^\circ\mathrm{C}$ of the reference $21.627^\circ\mathrm{C}$. The error maps separate them: the AD-PINN's largest error is $0.0013^\circ\mathrm{C}$, the FDM-PINN's $0.0122^\circ\mathrm{C}$, and the relative $L_2$ errors are $3.6\times10^{-4}$ and $3.6\times10^{-3}$, a factor of $10$ in favour of AD. The FDM error is largest in a round spot centred on the heater, which is where the field bends most and the stencil's $h^2 u''''$ truncation error is largest: the FDM-PINN solved, very accurately, an equation that differs from the true one by that truncation term. Against this, it trained in two thirds of the time, $132$ s against $200$ s.
Training loss
plt.figure(figsize = (6, 4))
plt.semilogy(loss_AD, 'r', linewidth = 1.5, label = 'AD-PINN')
plt.semilogy(loss_FDM, 'b', linewidth = 1.5, label = 'FDM-PINN')
plt.xlabel('Epoch')
plt.ylabel('Total loss')
plt.legend()
plt.grid(True)
plt.show()
The two loss curves lie on top of each other for the whole run, with the three learning-rate halvings showing as steps where the fluctuation collapses. The difference in loss is nil compared with the factor of ten in error, which is a reminder that the loss measures the residual on the collocation points under the method's own definition of the derivative, and the FDM residual can be small for a field that is not the solution of the true equation.
The FDM-PINN's accuracy was set by $h$, so $h$ is now varied with everything else fixed: the same $21 \times 21$ collocation points, the same network and training, and the stencil evaluated with $h = 0.2$, $0.1$, $1/22$ and $0.01$. Note that nothing requires $h$ to equal the grid spacing: the network can be evaluated anywhere, so the neighbours $x \pm h$ need not be collocation points. The four runs take four times the FDM training time.
h_list = [0.2, 0.1, 1/22, 0.01]
error_h = []
for h_try in h_list:
net = new_pinn()
optimizer = torch.optim.Adam(net.parameters(), lr = 1e-3)
for epoch in range(20001):
r = residual_FDM(net, x_pde, y_pde, h_try)
loss = torch.mean(r**2) + 100*torch.mean(theta_net(net, x_bc, y_bc)**2)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if epoch % 5000 == 4999: # halve the learning rate every 5000 iterations
optimizer.param_groups[0]['lr'] = optimizer.param_groups[0]['lr']/2
theta_h = theta_net(net, x_eval, y_eval).detach().numpy().reshape(201, 201)
error_h.append(np.linalg.norm(theta_h - theta_ref)/np.linalg.norm(theta_ref))
print('h = %.4f | final loss = %.3e | relative L2 error = %.3e' % (h_try, loss.item(), error_h[-1]))
Error against step size
The AD-PINN error is drawn as the horizontal line it is, since AD has no $h$.
plt.figure(figsize = (6, 4))
plt.loglog(h_list, error_h, 'bs-', linewidth = 1.5, label = 'FDM-PINN')
plt.loglog([min(h_list), max(h_list)], [error_AD, error_AD], 'r--', linewidth = 1.5, label = 'AD-PINN')
plt.xlabel('h')
plt.ylabel('relative $L_2$ error')
plt.legend()
plt.grid(True)
plt.show()
The FDM-PINN error falls with $h$ as the truncation error says it should: $9.2\times10^{-2}$ at $h = 0.2$, $1.8\times10^{-2}$ at $h = 0.1$, $3.6\times10^{-3}$ at the grid spacing $h = 1/22$ and $4.7\times10^{-4}$ at $h = 0.01$. Each halving of $h$ divides the error by about five, close to the factor of four of a pure $h^2$ law, and at $h = 0.01$ the curve flattens towards the AD-PINN's $3.6\times10^{-4}$, which is the floor set by the network and the optimizer rather than by the stencil. Two things are worth noticing. The final losses of the four runs are all of the same size, $10^{-7}$, so the loss does not reveal which run is accurate; the stencil error is invisible to the optimizer. And for $h = 0.1$ and $0.2$ the neighbours of the points nearest the edges lie outside the plate, where the network is simply extrapolated, which the method tolerates because nothing in it requires the neighbours to be collocation points. Smaller $h$ would run into the single-precision roundoff of Section 3.3, which for this field, with $\theta$ of order $10^{-2}$, sets in around $h \approx 10^{-3}$; between there and $10^{-2}$ the FDM-PINN is as accurate as the AD-PINN, at two thirds of the cost.
Two PINNs with the same network, the same points, the same loss and the same optimizer solved the same heat conduction problem, and the one difference between them, how $\theta_{xx} + \theta_{yy}$ was computed, showed up exactly where the theory of Chapters 2 and 3 said it would. The AD-PINN had the smaller error, because its derivatives were exact. The FDM-PINN was faster, because five forward passes cost less than two orders of autograd, and its error was the truncation error of its stencil, falling with $h$ until roundoff took over. Both reached the same answer to within about a hundredth of a degree on a field whose total rise is under two degrees, so for this problem either method is adequate, and the choice is a trade between the last digit of accuracy and the training time.
The laboratory settles the comparison for one problem; the general guidance follows from what each method needs and what each costs.
| Prefer AD when | Prefer FDM when |
|---|---|
| the collocation points are arbitrary or scattered | the output is a structured field or image |
| the coordinates are explicit network inputs | high-order AD is computationally expensive |
| discretization error matters | memory is the constraint |
| the equation has high-order or mixed derivatives and accuracy is the priority | a cheap, familiar derivative is enough |
Automatic differentiation is the default for a coordinate-based PINN. It is exact, it imposes nothing on the points, and the code for a new equation is the residual and nothing else. Finite differences become attractive when the derivatives are what the training spends its time on, when the network already produces its output on a grid, so that the stencil is a convolution, or when memory is short. Their price is the discretization error, which the optimizer cannot see and will not remove, and the requirement of neighbours, which ties the method to a step $h$ chosen in the window between truncation and roundoff. Whichever is chosen, the residual that reaches the loss is the same physics, and everything else in the PINN, the loss weighting, the sampling, the representation and the training order, applies unchanged.
Key takeaways
- A physics-informed loss needs derivatives of the network output with respect to the coordinates, and the step that produces them is run at every collocation point at every iteration.
- Automatic differentiation is the chain rule applied along the computation graph of the network. It is exact and works at any point; each order of derivative adds a graph on top of the last, which is its cost in memory and time.
- Finite differences combine the network's values at neighbouring points with a stencil derived from Taylor expansion. The derivative is cheap and needs no graph, but it carries a truncation error of order $h^2$, and roundoff makes it worse again when $h$ is too small, so $h$ has a usable window.
- In a PINN the two occupy the same place and feed the same loss. On the heated plate the AD-PINN was more accurate and the FDM-PINN was faster, and the FDM error followed $h$ as the theory predicts.
- Choose AD for scattered points, explicit coordinates and accuracy; choose FDM for structured fields, expensive high-order derivatives and tight memory.
Looking ahead. Both methods train one network to solve one problem. Deep Neural Operators and Fourier Neural Operator learn the map from the data of a problem to its solution, so that a new heater or a new conductivity costs a forward pass rather than a training run; and for operators whose outputs are fields on a grid, the finite-difference residual of this notebook is the natural way to add physics to their training.