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
1.1. The Physics Loss Needs Derivatives¶
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. Timing the forward pass and three successive derivatives of a small network at many points shows how the cost grows.
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))
The forward pass and the first derivative take about the same time, which is the usual rule for one reverse pass. From there each order costs more than the one before, 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.
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$,
$$ \begin{align*} u(x + h) &= u(x) + h\,u'(x) + \frac{h^2}{2}\,u''(x) + \frac{h^3}{6}\,u'''(x) + O(h^4) , \\ u(x - h) &= u(x) - h\,u'(x) + \frac{h^2}{2}\,u''(x) - \frac{h^3}{6}\,u'''(x) + O(h^4) . \end{align*} $$
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:
$$ \begin{align*} u'(x) &= \frac{u(x + h) - u(x - h)}{2h} + O(h^2) , \\ u''(x) &= \frac{u(x - h) - 2u(x) + u(x + h)}{h^2} + O(h^2) . \end{align*} $$
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)$,
$$ \begin{align*} u_x(x_i) &\approx \frac{u_{i+1} - u_{i-1}}{2h} , \\ u_{xx}(x_i) &\approx \frac{u_{i-1} - 2u_i + u_{i+1}}{h^2} , \end{align*} $$
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$. The test function is $u = \sin(2\pi x)$ on $[0, 1]$, whose second derivative $-(2\pi)^2\sin(2\pi x)$ is known, and the three-point formula is compared with it 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()
The three approximations lie on the exact curve at the scale of the upper panel, and the lower panel shows what the upper one hides: an error that follows the shape of $u''''$, largest where the curve bends most, and that shrinks by about a factor of four each time $h$ is halved. The stencil is doing what Taylor's theorem says it should, and its accuracy is set by $h$ and by nothing else.
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 a few significant digits, and the small differences between them are the truncation error of the stencil at this $h$. 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 a handful of 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 \left( \frac{\partial^2 T}{\partial x^2} + \frac{\partial^2 T}{\partial y^2} \right) = 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 grid 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. Numbering the interior points row by row, each equation becomes one row of a sparse matrix, with $4$ on the diagonal and $-1$ for each neighbour that lies inside the plate, and the system is solved directly.
from scipy.sparse import lil_matrix
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)
q = source(X_ref[1:-1, 1:-1], Y_ref[1:-1, 1:-1])
# five-point stencil at every interior point: 4*theta_ij - (four neighbours) = h^2 q_ij
A = lil_matrix((n*n, n*n)) # LIL format for assembly, CSR for the solve
b = np.zeros(n*n)
for i in range(n):
for j in range(n):
k = i*n + j # 2D index (i, j) -> 1D index k
A[k, k] = 4
if i > 0:
A[k, (i - 1)*n + j] = -1
if i < n - 1:
A[k, (i + 1)*n + j] = -1
if j > 0:
A[k, i*n + (j - 1)] = -1
if j < n - 1:
A[k, i*n + (j + 1)] = -1
b[k] = h_ref**2*q[i, j]
theta_ref = np.zeros((n + 2, n + 2)) # boundary stays at theta = 0
theta_ref[1:-1, 1:-1] = spsolve(A.tocsr(), b).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 by less than two degrees from the edges to the centre: 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 of order $10^{-2}$, 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 regular grid of interior points with spacing $h$ for the residual, and points along each edge for the boundary condition. A regular grid is used for both models so that the comparison is fair: the FDM-PINN takes its step equal to the grid spacing, and the AD-PINN could have used any points at all.
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
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; it is divided by the scale so that everything downstream is in terms of $\theta$.
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}{N_f}\sum_{i}\, r_\theta(x_i, y_i)^2 + 100\cdot\frac{1}{N_b}\sum_{j}\,\theta(x_j, y_j)^2 . $$
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.
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 by six orders of magnitude, most of it in the first few thousand iterations. The residual term dominates throughout: the boundary term starts small, because the initial network output is small on the scale of $\theta$, and stays far below the residual. The learning-rate halvings are visible as the loss settling in stages, with an occasional spike in between. The training time is recorded for the comparison with the FDM-PINN.
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)
Apart from a few transient spikes, the printed losses track the AD run and end at the same size: 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 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 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 PINN models produce visually similar results, with color maps and maximum temperatures that closely align with the reference data. However, the AD-PINN is more accurate, with errors roughly ten times smaller than those of the FDM-PINN in both maximum and relative $L_2$ norms. The FDM-PINN's inaccuracy is concentrated around the heater, where the field's sharp curvature increases the $h^2 u''''$ truncation error, meaning the model solved a slightly deviated version of the original equation. Despite this, the FDM-PINN offers a clear advantage in speed, requiring only two-thirds of the training time compared to the AD-PINN.
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 error associated with the FDM-PINN decreases as a function of the grid spacing $h$, consistent with theoretical expectations for truncation error. Specifically, each halving of $h$ reduces the error by approximately a factor of four, adhering to an $h^2$ convergence law. As $h$ reaches its minimum value, the error curve plateaus, converging toward the error floor of the AD-PINN.
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 it met the floor set by the network itself, the error of the AD-PINN. 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.