Scientific Machine Learning:
Physics-Informed Neural Networks: Labs
The theory behind everything used here — why the physics enters through the loss function, what the collocation points are for, and how the terms of a composite loss interact — is developed in Physics-informed Neural Networks (PINN). This notebook is the practice. Five problems from five different fields are each carried to the identification of something unknown — a stiffness, a discharge coefficient, two epidemic rates, the density and viscosity of a fluid, and an entire missing boundary condition — with every line of the implementation written out. Read the chapters as variations on one recipe: only the residual function and the list of trainable tensors really change.
Learning Objectives¶
After studying this notebook, you should be able to
- assemble a PINN in PyTorch from its four parts — a network mapping the coordinate to the solution, derivatives obtained with
torch.autograd.grad, a residual function that evaluates the governing equation on those derivatives, and a loss that adds the residual to the initial or boundary conditions; - differentiate a network output as many times as the equation requires, and recognize that the fourth derivative in the beam equation costs four nested
gradcalls; - turn a forward problem into an inverse one by declaring the unknown property a trainable parameter alongside the network weights, initialized far from its true value so that agreement is earned rather than assumed;
- explain, and demonstrate, why the physics and the boundary conditions alone may fail to identify that parameter, and why a handful of observations is enough to fix it;
- decide where to place the sensors, from the sensitivity of the solution to the unknown parameter;
- weight a data term against the physics terms, and say what goes wrong when their scales do not match;
- reparameterize a positive unknown as $\log\theta$, and constrain a bounded output with a sigmoid;
- estimate two unknown properties at once in a non-linear, two-dimensional system, and read their traces rather than their final values to decide whether the estimates can be trusted;
- recover an entire unknown boundary condition instead of a scalar, and write a qualitative piece of prior knowledge, “constant but unknown”, as a variance penalty computed from the network's own output;
- read a PINN result critically — which part of the domain is supported by data, which by the equation alone, and which by neither.
1.1. The Problem¶
This chapter returns to the Euler beam of Physics-informed Neural Networks (PINN), Section 5, and rebuilds it from nothing: no library, every derivative taken by hand with torch.autograd.grad, and the problem carried through to the identification of an unknown stiffness.
1.1.1. Governing Equation¶
For an Euler-Bernoulli beam, the transverse displacement $y(x)$ satisfies
$$ EI\frac{d^4y}{dx^4}+q(x)=0. $$
In this example,
$$ E=1,\qquad I=1,\qquad L=1, $$
and the distributed load is
$$ q(x)=x. $$
Therefore, the governing equation becomes
$$ \frac{d^4y}{dx^4}+x=0, \qquad x\in[0,1]. $$
Each derivative of $y$ carries a physical name — slope, bending moment, shear force, and finally the load — so this single fourth-order equation is four familiar statements stacked on top of one another. It is also what makes the beam a demanding first test for a PINN: before the equation can even be evaluated, the network's output has to be differentiated four times.
1.1.2. Problem Setup¶
- Problem properties
$$ E=1,\qquad I=1,\qquad L=1 $$
- Governing differential equation
$$ y''''(x)+x=0, \qquad x\in[0,1] $$
- Simply supported boundary conditions
$$ y(0)=0,\qquad y''(0)=0 $$
$$ y(1)=0,\qquad y''(1)=0 $$
- Quantities of interest
$$ \theta(x)=y'(x),\qquad M(x)=EIy''(x),\qquad V(x)=EIy'''(x) $$
Simply supported means that each end rests on a support which holds it in place but does not resist rotation. The displacement is therefore zero there, and so is the bending moment. That is two conditions at each end, four in all — exactly the number a fourth-order equation requires.
Import Library
Everything here is plain PyTorch: a network from torch.nn, derivatives from torch.autograd, an optimizer from torch.optim. No PINN library is involved, so every term of the loss stays visible.
import numpy as np
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
Set the Random Seed
The weights are initialized randomly, so a fixed seed is what makes a run repeatable.
np.random.seed(0)
torch.manual_seed(0)
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(device)
Define Collocation Points
The PDE residual is evaluated at 25 evenly spaced points over
$$ x\in[0,1]. $$
The two endpoints are also used to enforce the simply supported boundary conditions.
A collocation point is not a measurement. It is simply a location at which the network will be asked to satisfy the differential equation; the value of $y$ there is unknown and is never supplied. This is the precise sense in which a PINN can be trained without data.
x_pts = np.linspace(0, 1, 25)
plt.figure(figsize = (6, 3))
plt.scatter(x_pts, np.zeros_like(x_pts), s = 10)
plt.scatter(0, 0, label = 'x = 0', s = 10, c = 'red')
plt.scatter(1, 0, label = 'x = 1', s = 10, c = 'red')
plt.xlabel('x')
plt.legend()
plt.title('Collocation Points with Boundaries')
plt.show()
Define Network and Hyper-parameters
The neural network takes the coordinate $x$ as input and predicts the transverse displacement $y(x)$.
The architecture is four hidden layers of $32$ units with $\tanh$ activations. Adam with a learning rate of $10^{-3}$ completes the setup; every loss term will be a mean squared residual.
class PINN(nn.Module):
def __init__(self):
super(PINN, self).__init__()
self.net = nn.Sequential(
nn.Linear(1, 32),
nn.Tanh(),
nn.Linear(32, 32),
nn.Tanh(),
nn.Linear(32, 32),
nn.Tanh(),
nn.Linear(32, 32),
nn.Tanh(),
nn.Linear(32, 1)
)
def forward(self, x):
return self.net(x)
model = PINN().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr = 1e-3)
Define Functions to Compute Physics Loss
derivative is the whole of automatic differentiation for our purposes. It differentiates y with respect to x, and because create_graph = True the result is itself differentiable, so calling it again gives the second derivative, and again the third. Any tensor it differentiates with respect to must have been created with requires_grad = True.
def derivative(y, x):
return torch.autograd.grad(y, x, grad_outputs = torch.ones_like(y), create_graph = True)[0]
Define PDE System
The governing equation is
$$ \frac{d^4y}{dx^4}+x=0, \qquad x\in[0,1]. $$
Therefore, the PDE residual is
$$ r(x)=y''''(x)+x. $$
In code this is four nested calls to derivative, each differentiating the previous result. The function returns the residual; turning it into a single number to minimize is the job of the training loop.
def PDE(model, x):
y = model(x)
y_x = derivative(y, x)
y_xx = derivative(y_x, x)
y_xxx = derivative(y_xx, x)
y_xxxx = derivative(y_xxx, x)
pde = y_xxxx + x
return pde
Define Boundary Condition
For a simply supported beam,
$$ y(0)=0, \qquad y''(0)=0, $$
$$ y(1)=0, \qquad y''(1)=0. $$
Each end contributes two residuals, the displacement and the second derivative, both of which should vanish. They come from the same derivative helper, now applied at the single point $x = 0$ or $x = 1$.
def BC_left(model, x_0):
y_0 = model(x_0)
y_x_0 = derivative(y_0, x_0)
y_xx_0 = derivative(y_x_0, x_0)
bc1 = y_0
bc2 = y_xx_0
return bc1, bc2
def BC_right(model, x_1):
y_1 = model(x_1)
y_x_1 = derivative(y_1, x_1)
y_xx_1 = derivative(y_x_1, x_1)
bc3 = y_1
bc4 = y_xx_1
return bc3, bc4
Train
The total loss combines the governing-equation residual and the boundary-condition errors:
$$ \mathcal{L} = \mathcal{L}_{PDE} + \mathcal{L}_{BC}. $$
There is nothing in the loop that an ordinary training loop does not have. Five quantities are driven towards zero — the PDE residual at the collocation points and the four boundary residuals — and their mean squared values are simply added. Note what is missing: no measured displacement appears anywhere. The network is being fitted to an equation, not to data.
x = torch.tensor(x_pts.reshape(-1, 1), dtype = torch.float32, requires_grad = True).to(device)
x0 = torch.tensor([[0.0]], requires_grad = True).to(device)
x1 = torch.tensor([[1.0]], requires_grad = True).to(device)
for epoch in range(2001):
pde_term = PDE(model, x)
bc1, bc2 = BC_left(model, x0)
bc3, bc4 = BC_right(model, x1)
loss_pde = torch.mean(pde_term**2)
loss_bc = torch.mean(bc1**2) + torch.mean(bc2**2) + torch.mean(bc3**2) + torch.mean(bc4**2)
loss = loss_pde + loss_bc
optimizer.zero_grad()
loss.backward()
optimizer.step()
if epoch % 500 == 0:
print(
'Epoch: {} Loss: {:.6f} PDELoss: {:.6f} BCLoss: {:.6f}'.format(
epoch,
loss.item(),
loss_pde.item(),
loss_bc.item()
)
)
Define Exact Solution
The analytical solution is used only for validation after training.
Integrating $y'''' = -x$ four times and imposing the four support conditions gives a closed form, so the PINN can be judged against the truth rather than against its own loss. The derivatives are written out as well, because the slope, moment and shear are what an engineer actually asks for.
def exact_displacement(x):
return -x ** 5 / 120 + x ** 3 / 36 - 7 * x / 360
def exact_slope(x):
return -x ** 4 / 24 + x ** 2 / 12 - 7 / 360
def exact_moment(x):
return -x ** 3 / 6 + x / 6
def exact_shear(x):
return -x ** 2 / 2 + 1 / 6
Result
The trained PINN is evaluated on 100 points and compared with the analytical solution for displacement, slope, bending moment, and shear force.
Each of the four is one more derivative of the same network output, so read them as four increasingly severe tests of a single fit.
Compute Predicted Quantities
The same derivative calls used during training are reused here, with the results detached to NumPy for plotting.
def predict_quantities(model, x_test):
x_tensor = torch.tensor(x_test.reshape(-1, 1), dtype = torch.float32, requires_grad = True).to(device)
y = model(x_tensor)
y_x = derivative(y, x_tensor)
y_xx = derivative(y_x, x_tensor)
y_xxx = derivative(y_xx, x_tensor)
displacement = y.detach().cpu().numpy().flatten()
slope = y_x.detach().cpu().numpy().flatten()
moment = y_xx.detach().cpu().numpy().flatten() # EI y'' with EI = 1
shear = y_xxx.detach().cpu().numpy().flatten() # EI y''' with EI = 1
return displacement, slope, moment, shear
Define Test Points
The hundred test points are finer than the twenty-five used for training. The network is a continuous function of $x$, so it can be evaluated anywhere, including at points it never saw.
x_test = np.linspace(0, 1, 100)
plt.figure(figsize = (6, 3))
plt.scatter(x_test, np.zeros_like(x_test), s = 10)
plt.scatter(0, 0, label = 'x = 0', s = 10, c = 'red')
plt.scatter(1, 0, label = 'x = 1', s = 10, c = 'red')
plt.xlabel('x')
plt.legend()
plt.title('Test points')
plt.show()
y, slope, moment, shear = predict_quantities(model, x_test)
Displacement
The displacement is the network's own output and the quantity the boundary conditions act on directly, so it is the easiest of the four to get right.
# plot displacement (prediction vs exact solution)
plt.figure(figsize = (6, 4))
plt.plot(x_test, exact_displacement(x_test), 'k--', lw = 3, label = 'Exact Solution', alpha = 0.5)
plt.plot(x_test, y, c = 'r', label = 'PINN Prediction')
plt.xlabel('x')
plt.ylabel(r'Displacement, $y(x)$')
plt.legend()
plt.show()
Slope
One derivative away from the output. Nothing in the loss compares the slope with the exact solution; the agreement comes from the residual alone.
plt.figure(figsize = (6, 4))
plt.plot(x_test, exact_slope(x_test), 'k--', lw = 3, label = 'Exact Solution', alpha = 0.5)
plt.plot(x_test, slope, 'r', label = 'PINN Prediction')
plt.xlabel('x')
plt.ylabel(r'Slope, $\theta(x)$')
plt.legend()
plt.show()
Bending Moment
Two derivatives away. $M = EIy''$ is the quantity a designer checks against an allowable stress, so it matters that it is right in the interior and not only at the supports.
plt.figure(figsize = (6, 4))
plt.plot(x_test, exact_moment(x_test), 'k--', lw = 3, label = 'Exact Solution', alpha = 0.5)
plt.plot(x_test, moment, 'r', label = 'PINN Prediction')
plt.xlabel('x')
plt.ylabel(r'Bending Moment, $M(x)$')
plt.legend()
plt.show()
Shear Force
Three derivatives away, and still on top of the exact solution. That is a stronger statement than the picture suggests: $y'''$ never appears in the loss, so its accuracy is inherited entirely from the equation the network was made to satisfy.
plt.figure(figsize = (6, 4))
plt.plot(x_test, exact_shear(x_test), 'k--', lw = 3, label = 'Exact Solution', alpha = 0.5)
plt.plot(x_test, shear, 'r', label = 'PINN Prediction')
plt.xlabel('x')
plt.ylabel(r'Shear Force, $V(x)$')
plt.legend()
plt.show()
1.3.1. Problem Setup¶
The beam stiffness is now unknown. We keep the second moment of area fixed at
$$I=1$$
and learn Young's modulus,
$$\color{red}{E=\;?}$$
from the beam equation
$$E I y^{(4)}(x)+x=0$$
and the same simply supported boundary conditions. The true value is $E=1$, but no displacement observations are included in this example.
$E$ is initialized at $0.5$ and optimized jointly with the network weights.
1.3.2. Solve the Inverse Problem¶
Define Physics Loss with Unknown Young's Modulus
The residual is the one from Lab 1 with a single change: $E$ arrives as an argument instead of being the constant $1$. Because it is a tensor that requires a gradient, the derivative of the loss with respect to $E$ comes for free, exactly as it does for the network weights.
def inverse_PDE(model, x, E):
y = model(x)
y_x = derivative(y, x)
y_xx = derivative(y_x, x)
y_xxx = derivative(y_xx, x)
y_xxxx = derivative(y_xxx, x)
# E*I times the fourth derivative of y + q(x) = 0
pde = E * y_xxxx + x
return pde
Define Network and Unknown Parameter
E is created as a one-element tensor and handed to Adam in the same parameter list as the weights. From the optimizer's point of view it is simply one more number to adjust.
torch.manual_seed(1)
model_inverse_no_data = PINN().to(device)
E_no_data = torch.tensor([0.5], device = device, requires_grad = True)
x_inverse = torch.tensor(x_pts.reshape(-1, 1), dtype = torch.float32, requires_grad = True).to(device)
x0_inverse = torch.tensor([[0.0]], requires_grad = True).to(device)
x1_inverse = torch.tensor([[1.0]], requires_grad = True).to(device)
optimizer = torch.optim.Adam([
{'params': model_inverse_no_data.parameters(), 'lr': 1e-3},
{'params': E_no_data, 'lr': 1e-3}
])
Train Using Only Physics and Boundary Conditions
$$ \mathcal{L} = \mathcal{L}_{PDE} + \mathcal{L}_{BC} $$
The loss is identical to Lab 1. Only the list of unknowns has grown by one.
E_no_data_list = []
for epoch in range(3001):
pde_term = inverse_PDE(model_inverse_no_data, x_inverse, E_no_data)
bc1, bc2 = BC_left(model_inverse_no_data, x0_inverse)
bc3, bc4 = BC_right(model_inverse_no_data, x1_inverse)
loss_pde = torch.mean(pde_term**2)
loss_bc = torch.mean(bc1**2) + torch.mean(bc2**2) + torch.mean(bc3**2) + torch.mean(bc4**2)
loss = loss_pde + loss_bc
optimizer.zero_grad()
loss.backward()
optimizer.step()
E_no_data_list.append(E_no_data.item())
if epoch % 500 == 0:
print('Epoch: %s Loss: %.6f PDELoss: %.6f BCLoss: %.6f E: %.4f' % (epoch, loss.item(), loss_pde.item(), loss_bc.item(), E_no_data.item()))
Result
The inverse problem is underdetermined: physics and boundary conditions alone do not identify Young's modulus.
x_test_inverse = np.linspace(0, 1, 100)
x_test_tensor = torch.tensor(x_test_inverse.reshape(-1, 1), dtype = torch.float32).to(device)
y_no_data = model_inverse_no_data(x_test_tensor).detach().cpu().numpy().flatten()
print('True E: 1.0000, Estimated E without data: %.4f\n\n' % E_no_data.item())
plt.figure(figsize = (6, 4))
plt.plot(x_test_inverse, exact_displacement(x_test_inverse),
'k--', lw = 3, label = 'Exact Solution ($E=1$)', alpha = 0.5)
plt.plot(x_test_inverse, y_no_data, 'r', label = 'PINN Prediction')
plt.xlabel('x')
plt.ylabel(r'Displacement, $y(x)$')
plt.title('Inverse Problem without Data')
plt.legend()
plt.show()
1.4.1. Define Observed Data¶
We now add nine displacement observations generated from the analytical solution with the true value $E=1$. Interior points are used because the displacement is fixed at zero at both supports and those endpoint values provide no information about the stiffness. The trainable modulus is initialized at $E=0.2$, well away from its true value.
x_obs = np.linspace(0.1, 0.9, 9).reshape(-1, 1)
y_obs = exact_displacement(x_obs) + 0.0001*np.random.randn(*x_obs.shape)
plt.figure(figsize = (6, 4))
plt.plot(x_test_inverse, exact_displacement(x_test_inverse),
'k--', lw = 2, label = 'Exact Solution', alpha = 0.3)
plt.plot(x_obs, y_obs, 'bo', markersize = 3, label = 'Observed')
plt.xlabel('x')
plt.ylabel(r'Displacement, $y(x)$')
plt.title('Observed Data')
plt.legend()
plt.show()
1.4.2. Solve the Inverse Problem¶
Define Data-Driven Loss Term
The displacement values are small for this nondimensional beam, so the data loss is weighted by $1000$ to keep it comparable with the PDE and boundary-condition losses.
$$ \mathcal{L} = \mathcal{L}_{PDE} + \mathcal{L}_{BC} + 1000\mathcal{L}_{DATA} $$
torch.manual_seed(2)
model_inverse_data = PINN().to(device)
E_data = torch.tensor([0.2], device = device, requires_grad = True)
optimizer = torch.optim.Adam([
{'params': model_inverse_data.parameters(), 'lr': 1e-3},
{'params': E_data, 'lr': 1e-3}
])
x_data = torch.tensor(x_obs, dtype = torch.float32).to(device)
y_data = torch.tensor(y_obs, dtype = torch.float32).to(device)
Train
Three terms now, and the data term is the only one that has ever seen a displacement value.
lambda_data = 1000
E_data_list = []
loss_list = []
for epoch in range(3001):
pde_term = inverse_PDE(model_inverse_data, x_inverse, E_data)
bc1, bc2 = BC_left(model_inverse_data, x0_inverse)
bc3, bc4 = BC_right(model_inverse_data, x1_inverse)
loss_pde = torch.mean(pde_term**2)
loss_bc = torch.mean(bc1**2) + torch.mean(bc2**2) + torch.mean(bc3**2) + torch.mean(bc4**2)
loss_data = torch.mean((model_inverse_data(x_data) - y_data)**2)
loss = loss_pde + loss_bc + lambda_data*loss_data
optimizer.zero_grad()
loss.backward()
optimizer.step()
E_data_list.append(E_data.item())
loss_list.append(loss.item())
if epoch % 500 == 0:
print('Epoch: %s Loss: %.6f PDELoss: %.6f BCLoss: %.6f DATALoss: %.6f E: %.4f' % (epoch, loss.item(), loss_pde.item(), loss_bc.item(), loss_data.item(), E_data.item()))
Total Loss History
Worth remembering while reading it: in Lab 2 the loss also fell smoothly, and the answer was still wrong. A small loss reports self-consistency, not correctness.
plt.figure(figsize = (6, 4))
plt.semilogy(loss_list, c = 'r')
plt.xlabel('Epoch')
plt.ylabel('Total Loss')
plt.title('Total Training Loss: Inverse Problem with Data')
plt.show()
Refine with L-BFGS
Adam moves every parameter by a fixed step scaled by its running gradient average. Once the loss is small that step is too large for the sharpest directions, so the loss spikes late in training and $E$ keeps drifting past its true value. L-BFGS uses curvature and a line search, which is why PINNs are commonly trained with Adam first and L-BFGS last. The closure recomputes the loss and its gradient at every trial point. L-BFGS decides each step by comparing loss values along a line search. Near $10^{-5}$ those differences fall below what single precision can resolve, so the refinement runs in double precision.
n_adam = len(loss_list)
model_inverse_data.double()
E_data = E_data.detach().double().requires_grad_(True)
x_inverse = x_inverse.double()
x0_inverse, x1_inverse = x0_inverse.double(), x1_inverse.double()
x_data, y_data = x_data.double(), y_data.double()
optimizer = torch.optim.LBFGS(list(model_inverse_data.parameters()) + [E_data],
lr = 1.0, max_iter = 2000, max_eval = 4000,
history_size = 50, tolerance_change = 0,
tolerance_grad = 1e-9,
line_search_fn = 'strong_wolfe')
def closure():
optimizer.zero_grad()
pde_term = inverse_PDE(model_inverse_data, x_inverse, E_data)
bc1, bc2 = BC_left(model_inverse_data, x0_inverse)
bc3, bc4 = BC_right(model_inverse_data, x1_inverse)
loss_pde = torch.mean(pde_term**2)
loss_bc = torch.mean(bc1**2) + torch.mean(bc2**2) + torch.mean(bc3**2) + torch.mean(bc4**2)
loss_data = torch.mean((model_inverse_data(x_data) - y_data)**2)
loss = loss_pde + loss_bc + lambda_data*loss_data
loss.backward()
E_data_list.append(E_data.item())
loss_list.append(loss.item())
return loss
optimizer.step(closure)
print('L-BFGS evaluations: %s Loss: %.6f E: %.4f' % (len(loss_list) - n_adam, loss_list[-1], E_data.item()))
Total Loss History with L-BFGS
loss_lbfgs = np.array(loss_list[n_adam:])
best = np.minimum.accumulate(loss_lbfgs)
keep = np.concatenate(([True], np.diff(best) < 0))
step = np.arange(n_adam, len(loss_list))
plt.figure(figsize = (6, 4))
plt.semilogy(range(n_adam), loss_list[:n_adam], c = 'r', linewidth = 2, label = 'Adam')
plt.semilogy(step[keep], loss_lbfgs[keep], c = 'b', linewidth = 2, label = 'L-BFGS')
plt.xlabel('Step')
plt.ylabel('Total Loss')
plt.title('Total Training Loss: Adam then L-BFGS')
plt.legend()
plt.grid(alpha = 0.3)
plt.show()
Result
y_with_data = model_inverse_data(x_test_tensor.double()).detach().cpu().numpy().flatten()
print('True E: 1.0000, Estimated E with data: %.4f\n\n' % E_data.item())
plt.figure(figsize = (6, 4))
plt.plot(x_test_inverse, exact_displacement(x_test_inverse),
'k--', lw = 3, label = 'Exact Solution', alpha = 0.3)
plt.plot(x_test_inverse, y_with_data, c = 'r', label = 'PINN Prediction')
plt.plot(x_obs, y_obs, 'bo', markersize = 3, label = 'Observed')
plt.xlabel('x')
plt.ylabel(r'Displacement, $y(x)$')
plt.title('Inverse Problem with Data')
plt.legend()
plt.show()
The displacement observations fix the response scale. The network can no longer compensate for an arbitrary value of $E$ by rescaling $y(x)$, so Young's modulus becomes identifiable and is recovered jointly with the displacement field.
| Lab | Loss terms | Outcome |
|---|---|---|
| 1 Forward | $\mathcal{L}_{PDE} + \mathcal{L}_{BC}$ | matches the exact solution |
| 2 Inverse, no data | $\mathcal{L}_{PDE} + \mathcal{L}_{BC}$ | $E$ is not identifiable |
| 3 Inverse, with data | $\mathcal{L}_{PDE} + \mathcal{L}_{BC} + 1000\,\mathcal{L}_{DATA}$ | $E$ is recovered |
The forward problem is solved by the physics and the boundary conditions alone. Making $E$ trainable costs one line, but without observations the equation $E I y^{(4)} + x = 0$ cannot tell a stiff beam with a small deflection from a soft beam with a large one, and the optimizer stopped at $E \approx 0.57$ with the loss as small as in the forward case. Nine displacement measurements fix the scale of the response, and the same loop then recovers $E \approx 1.007$. An inverse problem needs data, but only as much as it takes to remove the ambiguity the physics leaves.
2.1. The Problem¶
- $h(t)$ : water level at time t
- $A_t$ : tank cross-sectional area
- $A_o$ : outlet area
- $C_d$ : discharge coefficient
- $g$ : acceleration of gravity
An open tank is filled with water to an initial level. A small orifice at the bottom lets the water out. The higher the water level, the faster the water leaves; as water leaves, the level drops; as the level drops, the outflow slows down again.
That feedback is what makes this a differential equation.
2.1.1. Governing Equation¶
Mass conservation for the water in the tank states that the rate of change of the stored volume equals minus the outflow rate. Torricelli's law gives the outflow velocity: the water leaving the orifice moves as fast as if it had fallen freely from the free surface,
$$gh = \frac{1}{2}v^2 \quad \Longrightarrow \quad v = \sqrt{2gh}$$
This form assumes the free surface itself is essentially at rest, which holds when the tank is much wider than the orifice, $A_o \ll A_t$.
In practice the jet contracts (vena contracta) and friction takes a further toll, so less comes out. That loss is lumped into a single discharge coefficient $C_d$ (about $0.6$ for a sharp-edged orifice). Writing $A_t$ for the tank cross-section and $A_o$ for the orifice area,
$$A_t \frac{dh}{dt} = -C_d A_o \sqrt{2gh}$$
We collect the four constants into one,
$$C = \frac{C_d A_o \sqrt{2g}}{A_t}$$
so that the problem carries a single number. The equation is a first-order, nonlinear ODE, and it comes with an initial condition, not a boundary condition.
$$\frac{dh}{dt} = -C\sqrt{h}$$
2.1.2. Problem Setup¶
- Problem properties (nondimensionalized)
$$C = 1$$
- Ordinary differential equation
$$\frac{dh}{dt} = -C\sqrt{h}, \qquad \text{where} \quad t \in [0, 2]$$
- One initial condition
$$h(0) = 1$$
- The exact solution
$$ h(t) = \begin{cases} \left(1 - \frac{Ct}{2}\right)^2 \quad & \text{for } t \le \frac{2}{C} \\ \\ 0 &\text{after that} \end{cases} $$
2.1.3. Three Labs¶
| Lab | Setting | Question |
|---|---|---|
| 1 | $C$ is known | Can a PINN reproduce the exact solution? |
| 2 | $C$ is unknown, no data | Can the physics alone pin down $C$? |
| 3 | $C$ is unknown, a few measurements | Can physics and data together recover $C$? |
Lab 2 and Lab 3 are inverse problems: we solve for the solution and for an unknown property of the system at the same time.
Import Library
The same imports as the beam. Nothing in the method changes when the physics does.
import numpy as np
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
Set the Random Seed
Here the seed fixes both the weights and the collocation points, which are redrawn at every iteration below.
np.random.seed(0)
torch.manual_seed(0)
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(device)
Define Collocation Points
The ODE residual is evaluated at 100 points sampled uniformly from
$$ t\in[0,2]. $$
The points are redrawn at every training iteration, which is a small but real difference from the beam lab: the equation is enforced at a different random set of instants every time, so the network cannot satisfy the residual by learning a fixed grid of points.
t_train = np.random.rand(100, 1) * 2
plt.figure(figsize = (6, 3))
plt.scatter(t_train, np.zeros_like(t_train), s = 10)
plt.xlim([0, 2])
plt.yticks([])
plt.xlabel('t')
plt.title('Collocation Points (One Draw)')
plt.show()
Define Network and Hyper-parameters
The neural network takes time $t$ as input and predicts the water level $h(t)$.
The architecture is unchanged from the beam — four hidden layers of $32$ $\tanh$ units — which is itself worth noticing. It was not retuned for a different equation in a different field; only the residual function was rewritten.
class PINN(nn.Module):
def __init__(self):
super(PINN, self).__init__()
self.net = nn.Sequential(
nn.Linear(1, 32),
nn.Tanh(),
nn.Linear(32, 32),
nn.Tanh(),
nn.Linear(32, 32),
nn.Tanh(),
nn.Linear(32, 32),
nn.Tanh(),
nn.Linear(32, 1)
)
def forward(self, t):
return self.net(t)
model = PINN().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr = 1e-3)
Define Functions to Compute Physics Loss
Identical to the beam's helper, and this time only one derivative is ever needed.
def derivative(y, t):
return torch.autograd.grad(y, t, grad_outputs = torch.ones_like(y), create_graph = True)[0]
Define ODE System
The governing equation is
$$ \frac{dh}{dt}+C\sqrt{h}=0. $$
The clamp and the small offset inside the square root are a practical detail rather than part of the mathematics: $\sqrt{h}$ is undefined for negative $h$, and an untrained network will cheerfully predict one in its first few iterations, which would turn the residual into NaN.
def ODE(model, t, C):
h = model(t)
h_t = derivative(h, t)
ode = h_t + C * torch.sqrt(torch.clamp(h, min = 0.0) + 1e-8)
return ode
Define Initial Condition
The initial condition is
$$ h(0)=1. $$
A first-order equation needs exactly one condition, and unlike the beam it is imposed at a single instant rather than at two ends.
def IC(model, t_0):
h_0 = model(t_0)
return h_0
Train
The total loss combines the ODE residual and the initial-condition error:
$$ \mathcal{L} = \mathcal{L}_{ODE} + \mathcal{L}_{IC}. $$
Both terms are evaluated at the freshly drawn points. The initial-condition residual is pushed towards $1$ rather than towards $0$, and that line is the only place where the value $h(0) = 1$ enters the computation at all.
C = 1.0
t_0 = torch.zeros((1, 1), requires_grad = True).to(device)
for epoch in range(5001):
t = torch.tensor(np.random.rand(100, 1)*2, dtype = torch.float32, requires_grad = True).to(device)
ode_term = ODE(model, t, C)
ic_term = IC(model, t_0)
loss_ode = torch.mean(ode_term**2)
loss_ic = torch.mean((ic_term - 1)**2)
loss = loss_ode + loss_ic
optimizer.zero_grad()
loss.backward()
optimizer.step()
if epoch % 500 == 0:
print(
'Epoch: {} Loss: {:.6f} ODELoss: {:.6f} ICLoss: {:.6f}'.format(
epoch,
loss.item(),
loss_ode.item(),
loss_ic.item()
)
)
Define Exact Solution
The analytical solution is used to validate the trained PINN.
Separating variables gives a parabola that reaches zero at $t = 2/C$, after which the tank is empty and the formula no longer applies.
def exact_solution(t):
return (1 - t / 2) ** 2
Result
The prediction should follow the parabola down to the moment the tank empties; the printed error puts a number on how closely.
t_test = np.linspace(0, 2, 100).reshape(-1, 1)
t_test_tensor = torch.tensor(t_test, dtype = torch.float32).to(device)
h_pred = model(t_test_tensor).detach().cpu().numpy()
plt.figure(figsize = (6, 4))
plt.plot(t_test, exact_solution(t_test), 'k--', lw = 3, label = 'Exact Solution', alpha = 0.3)
plt.plot(t_test, h_pred, c = 'r', label = 'PINN Prediction')
plt.xlabel('t')
plt.ylabel('Water Level, h')
plt.title('Forward Problem')
plt.legend()
plt.show()
2.3.1. Problem Setup¶
Now the drainage constant is not given.
- Unknown property
$$\color{red}{C = \;?}$$
- Ordinary differential equation
$$\frac{dh}{dt} = -\color{red}{C}\sqrt{h}, \qquad h(0) = 1$$
$C$ becomes a trainable parameter, optimized jointly with the network weights. It is initialized at $0.5$, away from the true value, so that any agreement at the end is earned rather than assumed.
Nothing else changes — the loss still has only the two physics terms.
2.3.2. Solve the Inverse Problem¶
Define Network and Unknown Parameter
A fresh network and a fresh $C$, initialized at $0.5$ — half the true value, so that any agreement at the end has to be earned.
torch.manual_seed(0)
model_inverse_no_data = PINN().to(device)
C_no_data = torch.tensor([0.5], device = device, requires_grad = True)
optimizer = torch.optim.Adam([
{'params': model_inverse_no_data.parameters(), 'lr': 1e-3},
{'params': C_no_data, 'lr': 1e-3}
])
Train Using Only Physics and Initial Condition
$$ \mathcal{L} = \mathcal{L}_{ODE} + \mathcal{L}_{IC}. $$
Identical to Lab 1 except that $C$ is now optimized alongside the weights. Nothing tells the loop what $C$ ought to be.
C_no_data_list = []
for epoch in range(5001):
t = torch.tensor(np.random.rand(100, 1)*2, dtype = torch.float32, requires_grad = True).to(device)
ode_term = ODE(model_inverse_no_data, t, C_no_data)
ic_term = IC(model_inverse_no_data, t_0)
loss_ODE = torch.mean(ode_term**2)
loss_IC = torch.mean((ic_term - 1)**2)
loss = loss_ODE + loss_IC
optimizer.zero_grad()
loss.backward()
optimizer.step()
C_no_data_list.append(C_no_data.item())
if epoch % 500 == 0:
print(
'Epoch: {} Loss: {:.6f} ODELoss: {:.6f} ICLoss: {:.6f} C: {:.4f}'.format(
epoch,
loss.item(),
loss_ODE.item(),
loss_IC.item(),
C_no_data.item()
)
)
Result
Read the trajectory of $C$ rather than its final value: where it goes, and whether it settles anywhere near $1$.
h_no_data = model_inverse_no_data(t_test_tensor).detach().cpu().numpy()
plt.figure(figsize = (6, 4))
plt.plot(t_test, exact_solution(t_test), 'k--', lw = 3, label = 'Exact Solution', alpha = 0.3)
plt.plot(t_test, h_no_data, c = 'r', label = 'PINN Prediction')
plt.xlim([0, 2])
plt.xlabel('t')
plt.ylabel('Water Level, h')
plt.title('Inverse Problem without Data')
plt.legend()
plt.show()
print('True C: 1.0000, Predicted C: %.4f\n\n' % C_no_data.item())
plt.figure(figsize = (6, 4))
plt.plot(C_no_data_list, c = 'r', linewidth = 2, label = 'Predicted C')
plt.axhline(1.0, c = 'k', alpha = 0.5, label = 'True C')
plt.xlabel('Epoch')
plt.ylabel('C')
plt.xlim([0, 5000])
plt.ylim([0.2, 1.1])
plt.title('Predicted C without Data')
plt.legend()
plt.show()
2.4.1. Define Observed Data¶
Suppose the water level was read nine times during the experiment. The readings are generated from the exact solution with $C = 1$ and corrupted with a small amount of noise; $C$ is of course assumed to be unknown throughout this lab.
The readings are taken between $t = 0.2$ and $t = 1.6$, away from both ends of the domain. How much a reading tells us about $C$ is the sensitivity of the level to $C$,
$$\frac{\partial h}{\partial C} = -t\left(1 - \frac{Ct}{2}\right)$$
Evaluated at the true value $C = 1$, this vanishes at $t = 0$, where the level is pinned by the initial condition, and again at $t = 2$, where the tank has just emptied. It is largest in magnitude at $t = 1$. A sensor placed at either end would carry no information about $C$, which is why the readings sit in the middle.
t_obs = np.linspace(0.2, 1.6, 9).reshape(-1, 1)
h_obs = exact_solution(t_obs) + 0.01*np.random.randn(9, 1)
plt.figure(figsize = (6, 4))
plt.plot(t_test, exact_solution(t_test), 'k--', lw = 2, label = 'Exact Solution', alpha = 0.3)
plt.plot(t_obs, h_obs, 'bo', markersize = 3, label = 'Observed')
plt.xlim([0, 2])
plt.legend()
plt.xlabel('t')
plt.ylabel('Water level, h')
plt.title('Observed Data')
plt.show()
2.4.2. Solve the Inverse Problem¶
Define Data-Driven Loss Term
$$\mathcal{L} = \mathcal{L}_{ODE} + \mathcal{L}_{IC} + \mathcal{L}_{DATA}$$
The three terms are added with equal weight. The water level is of order one, so the data term is already comparable with the others and needs no scaling — unlike the beam, whose displacements were a few thousandths and whose data term needed a factor of $1000$.
torch.manual_seed(0)
model_inverse_data = PINN().to(device)
C_data = torch.tensor([0.5], device = device, requires_grad = True)
optimizer = torch.optim.Adam([
{'params': model_inverse_data.parameters(), 'lr': 1e-3},
{'params': C_data, 'lr': 1e-3}
])
t_data = torch.tensor(t_obs, dtype = torch.float32).to(device)
h_data = torch.tensor(h_obs, dtype = torch.float32).to(device)
Train Using Physics, Initial Condition, and Observed Data
$$ \mathcal{L} = \mathcal{L}_{ODE} + \mathcal{L}_{IC} + \mathcal{L}_{DATA}. $$
One more term, nine more points, and no other change to the loop.
C_data_list = []
for epoch in range(5001):
t = torch.tensor(np.random.rand(100, 1)*2, dtype = torch.float32, requires_grad = True).to(device)
ode_term = ODE(model_inverse_data, t, C_data)
ic_term = IC(model_inverse_data, t_0)
loss_ode = torch.mean(ode_term**2)
loss_ic = torch.mean((ic_term - 1)**2)
loss_data = torch.mean((model_inverse_data(t_data) - h_data)**2)
loss = loss_ode + loss_ic + loss_data
optimizer.zero_grad()
loss.backward()
optimizer.step()
C_data_list.append(C_data.item())
if epoch % 500 == 0:
print(
'Epoch: {} Loss: {:.6f} ODELoss: {:.6f} ICLoss: {:.6f} '
'DATALoss: {:.6f} C: {:.4f}'.format(
epoch,
loss.item(),
loss_ode.item(),
loss_ic.item(),
loss_data.item(),
C_data.item()
)
)
Refine with L-BFGS
Adam moves every parameter by a fixed step scaled by its running gradient average. Once the loss is small that step is too large for the sharpest directions, so the loss spikes late in training and $C$ stops short of its true value. L-BFGS uses curvature and a line search, which is why PINNs are commonly trained with Adam first and L-BFGS last. The closure recomputes the loss and its gradient at every trial point, and L-BFGS chooses each step by comparing those values along the line search. Two things follow. The collocation points are drawn once and held fixed, because a line search cannot compare losses measured on a different draw each time. And the refinement runs in double precision, because near $10^{-5}$ the differences the line search is comparing fall below what single precision can resolve.
n_adam = len(C_data_list)
model_inverse_data.double()
C_data = C_data.detach().double().requires_grad_(True)
t_col = torch.tensor(np.random.rand(100, 1)*2, dtype = torch.float64, requires_grad = True).to(device)
t_0 = torch.zeros((1, 1), dtype = torch.float64, requires_grad = True).to(device)
t_data, h_data = t_data.double(), h_data.double()
optimizer = torch.optim.LBFGS(list(model_inverse_data.parameters()) + [C_data],
lr = 1.0, max_iter = 2000, max_eval = 4000,
history_size = 50, tolerance_change = 0,
tolerance_grad = 1e-9,
line_search_fn = 'strong_wolfe')
loss_list = []
def closure():
optimizer.zero_grad()
ode_term = ODE(model_inverse_data, t_col, C_data)
ic_term = IC(model_inverse_data, t_0)
loss_ode = torch.mean(ode_term**2)
loss_ic = torch.mean((ic_term - 1)**2)
loss_data = torch.mean((model_inverse_data(t_data) - h_data)**2)
loss = loss_ode + loss_ic + loss_data
loss.backward()
C_data_list.append(C_data.item())
loss_list.append(loss.item())
return loss
optimizer.step(closure)
print('L-BFGS evaluations: %s Loss: %.6f C: %.4f' % (len(loss_list), loss_list[-1], C_data.item()))
Result
Compare this curve of $C$ with the one from Lab 2. Same network, same equation, same optimizer; the only difference is nine numbers.
h_with_data = model_inverse_data(t_test_tensor.double()).detach().cpu().numpy().flatten()
plt.figure(figsize = (6, 4))
plt.plot(t_test, exact_solution(t_test), 'k--', lw = 2, label = 'Exact Solution', alpha = 0.3)
plt.plot(t_test, h_with_data, c = 'r', label = 'PINN Prediction')
plt.plot(t_obs, h_obs, 'bo', markersize = 3, label = 'Observed')
plt.xlim([0, 2])
plt.xlabel('t')
plt.ylabel('Water Level, h')
plt.title('Inverse Problem with Data')
plt.legend()
plt.show()
best = np.minimum.accumulate(np.array(loss_list))
keep = np.concatenate(([True], np.diff(best) < 0))
step = np.arange(n_adam, len(C_data_list))
C_lbfgs = np.array(C_data_list[n_adam:])
plt.figure(figsize = (6, 4))
plt.plot(range(n_adam), C_data_list[:n_adam], c = 'r', linewidth = 2, label = 'Adam')
plt.plot(step[keep], C_lbfgs[keep], c = 'b', linewidth = 2, label = 'L-BFGS')
plt.axhline(1.0, c = 'k', alpha = 0.5, label = 'True C')
plt.xlabel('Step')
plt.ylabel('C')
plt.ylim([0.2, 1.1])
plt.title('Predicted C with Data')
plt.legend()
plt.show()
Nine noisy readings are enough to close the system. The same network, the same equation and one extra loss term now recover the drainage constant.
Note where the curve is least accurate. After the last reading at $t = 1.6$ there is no data left, and the solution there is held up by the equation alone — which is exactly the part of the domain the sensitivity argument told us carries no information about $C$.
Converting back to the physical coefficient is a single line,
$$C_d = \frac{C A_t}{A_o \sqrt{2g}}$$
so a value near $0.6$ would tell us the orifice is sharp-edged.
| Lab | Loss terms | Outcome |
|---|---|---|
| 1 Forward | $\mathcal{L}_{ODE} + \mathcal{L}_{IC}$ | matches the exact solution |
| 2 Inverse, no data | $\mathcal{L}_{ODE} + \mathcal{L}_{IC}$ | $C$ is not identifiable |
| 3 Inverse, with data | $\mathcal{L}_{ODE} + \mathcal{L}_{IC} + \mathcal{L}_{DATA}$ | $C$ is recovered |
A PINN solving a forward problem competes with numerical methods that are already good at it. The inverse problem is where it earns its place: the unknown property sits in the loss function as one more trainable number, and a handful of measurements is enough to find it.
3.1. The Problem¶
The SIR model is a foundational mathematical framework in epidemiology that models the spread of an infectious disease by dividing a fixed population ($N$) into three distinct compartments over time $t$:
- $S(t)$ (Susceptible): Individuals who are healthy and can contract the disease.
- $I(t)$ (Infected): Individuals who are currently infected and capable of transmitting the disease.
- $R(t)$ (Recovered / Removed): Individuals who have recovered with immunity or have been removed from the transmission chain.
3.1.1. Governing Equations¶
The SIR model partitions a population into susceptible, infected, and recovered fractions,
$$ S(t),\qquad I(t),\qquad R(t). $$
Their evolution is governed by
$$ \frac{dS}{dt}=-\beta SI, $$
$$ \frac{dI}{dt}=\beta SI-\gamma I, $$
$$ \frac{dR}{dt}=\gamma I, $$
where $\beta$ is the transmission rate and $\gamma$ is the recovery rate.
3.1.2. Inverse Problem Setup¶
In this example, the reference parameters and initial conditions are
$$ \beta=0.30,\qquad \gamma=0.10, $$
$$ S(0)=0.99,\qquad I(0)=0.01,\qquad R(0)=0. $$
A physics-informed neural network learns $S(t)$, $I(t)$, and $R(t)$ while jointly estimating the unknown parameters $\beta$ and $\gamma$ from sparse observations of $I(t)$.
Import Libraries
This lab has both a solver and a network in it. scipy.integrate.solve_ivp generates the reference epidemic, and the PINN then has to recover it from twenty noisy glimpses.
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
from scipy.integrate import solve_ivp
Set the Random Seed
np.random.seed(42)
torch.manual_seed(42)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('Device:', device)
Define and Solve the Reference SIR System
The reference solution is generated over
$$ t\in[0,100] $$
using a high-accuracy numerical ODE solver.
This plays the part of the real epidemic. The network never sees it — only the handful of noisy samples drawn from it below — and it is kept aside so that the recovered trajectory and parameters can be checked at the end.
# True SIR parameters
beta_true = 0.30
gamma_true = 0.10
# Initial conditions
S0 = 0.99
I0 = 0.01
R0 = 0.00
# Time domain
t_start = 0.0
t_end = 100.0
def sir_rhs(t, y):
S, I, R = y
dSdt = -beta_true * S * I
dIdt = beta_true * S * I - gamma_true * I
dRdt = gamma_true * I
return [dSdt, dIdt, dRdt]
# Dense time grid for ground truth
t_true = np.linspace(t_start, t_end, 500)
solution = solve_ivp(sir_rhs, [t_start, t_end], [S0, I0, R0], t_eval = t_true)
S_true = solution.y[0]
I_true = solution.y[1]
R_true = solution.y[2]
Plot the Reference Solution
plt.figure(figsize = (6, 4))
plt.plot(t_true, S_true, label = 'Susceptible $S(t)$')
plt.plot(t_true, I_true, label = 'Infected $I(t)$')
plt.plot(t_true, R_true, label = 'Recovered $R(t)$')
plt.xlabel('Time')
plt.ylabel('Population fraction')
plt.title('SIR Epidemic Model')
plt.legend()
plt.grid(alpha = 0.3)
plt.xlim([0, 100])
plt.show()
Create Sparse Epidemiological Observations
Twenty infected-population observations are sampled across the time interval. Small Gaussian measurement noise is added to emulate observed epidemiological data.
Only the infected fraction is observed, which is realistic: case counts are reported, while the susceptible and recovered fractions are not measured. $S(t)$ and $R(t)$ must therefore be inferred from the equations alone, and that is the interesting part of this lab.
n_obs = 20
obs_idx = np.linspace(0, len(t_true) - 1, n_obs, dtype = int)
t_obs = t_true[obs_idx]
I_obs = I_true[obs_idx]
# Add small measurement noise
noise_std = 0.002
I_obs_noisy = I_obs + noise_std * np.random.randn(n_obs)
I_obs_noisy = np.clip(I_obs_noisy, 0.0, 1.0)
Plot the Sparse Observations
plt.figure(figsize = (6, 4))
plt.plot(t_true, I_true, 'k--', label = 'True $I(t)$', alpha = 0.3)
plt.scatter(t_obs, I_obs_noisy, s = 20, label = 'Sparse observations')
plt.xlabel('Time')
plt.ylabel('Infected fraction')
plt.title('Sparse Epidemiological Observations')
plt.legend()
plt.grid(alpha = 0.3)
plt.xlim([0, 100])
plt.ylim([-0.1, 1.1])
plt.show()
Define Network and Hyper-parameters
The network takes time $t$ as input and predicts all three state variables:
$$ t\longmapsto \bigl(\hat S(t),\hat I(t),\hat R(t)\bigr). $$
A sigmoid output constrains each predicted population fraction to $[0,1]$. The logarithms of $\beta$ and $\gamma$ are optimized so that both learned parameters remain positive.
class SIR_PINN(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Linear(1, 64),
nn.Tanh(),
nn.Linear(64, 64),
nn.Tanh(),
nn.Linear(64, 64),
nn.Tanh(),
nn.Linear(64, 3)
)
def forward(self, t):
# Scale time approximately to [-1, 1]
t_scaled = 2.0 * (t - t_start) / (t_end - t_start) - 1.0
raw = self.net(t_scaled)
# Keep S, I, R positive
y = torch.sigmoid(raw)
S = y[:, 0:1]
I = y[:, 1:2]
R = y[:, 2:3]
return S, I, R
model = SIR_PINN().to(device)
log_beta = nn.Parameter(torch.tensor(np.log(0.2), dtype = torch.float32, device = device))
log_gamma = nn.Parameter(torch.tensor(np.log(0.2), dtype = torch.float32, device = device))
Define Functions to Compute Physics Loss
The helper below is the same derivative as in the earlier chapters.
def derivative(y, x):
return torch.autograd.grad(y, x, grad_outputs = torch.ones_like(y), create_graph = True)[0]
Define SIR Physics Residuals
Automatic differentiation provides the time derivatives of the network outputs. The SIR residuals are
$$ f_S=\frac{d\hat S}{dt}+\beta\hat S\hat I, $$
$$ f_I=\frac{d\hat I}{dt}-\beta\hat S\hat I+\gamma\hat I, $$
$$ f_R=\frac{d\hat R}{dt}-\gamma\hat I. $$
An ideal physics-consistent solution satisfies
$$ f_S=f_I=f_R=0. $$
def physics_residual(model, t, beta, gamma):
S, I, R = model(t)
dSdt = derivative(S, t)
dIdt = derivative(I, t)
dRdt = derivative(R, t)
# SIR residuals
f_S = dSdt + beta * S * I
f_I = dIdt - beta * S * I + gamma * I
f_R = dRdt - gamma * I
return f_S, f_I, f_R
Define Collocation Points
The physics residuals are evaluated at 200 evenly spaced collocation points over the full time interval.
Two hundred points for a hundred-day interval, against twenty observations. The equations are enforced everywhere; the data speak only at twenty instants.
n_collocation = 200
t_f = torch.tensor(np.linspace(t_start, t_end, n_collocation).reshape(-1, 1),
dtype = torch.float32, requires_grad = True).to(device)
# Observation and initial-condition tensors
t_obs_t = torch.tensor(t_obs[:, None], dtype = torch.float32, device = device)
I_obs_t = torch.tensor(I_obs_noisy[:, None], dtype = torch.float32, device = device)
t0_t = torch.tensor([[0.0]], dtype = torch.float32, device = device)
initial_state = torch.tensor([[S0, I0, R0]], dtype = torch.float32, device = device)
Train
The network weights and the two epidemiological parameters are optimized jointly with Adam. The total loss combines sparse-data agreement, the governing equations, the initial conditions, and population conservation:
$$ \mathcal{L} =10\mathcal{L}_{DATA} +\mathcal{L}_{PHYSICS} +10\mathcal{L}_{IC} +\mathcal{L}_{MASS}. $$
The conservation term encourages
$$ \hat S(t)+\hat I(t)+\hat R(t)=1. $$
optimizer = torch.optim.Adam(list(model.parameters()) + [log_beta, log_gamma], lr = 1e-3)
lambda_data = 10.0
lambda_ic = 10.0
epochs = 10000
loss_history = []
for epoch in range(epochs):
beta = torch.exp(log_beta)
gamma = torch.exp(log_gamma)
# Data loss: sparse infected observations
_, I_pred_obs, _ = model(t_obs_t)
loss_data = torch.mean((I_pred_obs - I_obs_t)**2)
# Physics loss: SIR residuals at the collocation points
f_S, f_I, f_R = physics_residual(model, t_f, beta, gamma)
loss_physics = torch.mean(f_S**2) + torch.mean(f_I**2) + torch.mean(f_R**2)
# Initial-condition loss
S_pred0, I_pred0, R_pred0 = model(t0_t)
loss_ic = torch.mean((S_pred0 - S0)**2) + torch.mean((I_pred0 - I0)**2) + torch.mean((R_pred0 - R0)**2)
# Population-conservation loss
S_f, I_f, R_f = model(t_f)
loss_mass = torch.mean((S_f + I_f + R_f - 1)**2)
loss = (lambda_data * loss_data + loss_physics + lambda_ic * loss_ic + loss_mass)
optimizer.zero_grad()
loss.backward()
optimizer.step()
loss_history.append(loss.item())
if epoch % 1000 == 0:
beta_est = torch.exp(log_beta).item()
gamma_est = torch.exp(log_gamma).item()
print(
'Epoch %5d | Loss: %.3e | Data: %.3e | Physics: %.3e | beta: %.4f | gamma: %.4f' % (epoch, loss.item(), loss_data.item(), loss_physics.item(), beta_est, gamma_est)
)
Refine with L-BFGS
The loss history above spikes repeatedly after the first few thousand epochs: Adam's fixed step is too large for the sharpest directions once the loss is small. As in the beam problem, the network and the two parameters are refined with L-BFGS, which uses curvature and a line search. The closure recomputes the full loss and its gradient at every trial point. L-BFGS decides each step by comparing loss values along a line search. Near $10^{-5}$ those differences fall below what single precision can resolve, so the refinement runs in double precision.
n_adam = len(loss_history)
model.double()
log_beta = nn.Parameter(log_beta.detach().double())
log_gamma = nn.Parameter(log_gamma.detach().double())
t_obs_t, I_obs_t, t0_t, t_f = t_obs_t.double(), I_obs_t.double(), t0_t.double(), t_f.double()
optimizer = torch.optim.LBFGS(list(model.parameters()) + [log_beta, log_gamma],
lr = 1.0, max_iter = 2000, max_eval = 4000, history_size = 50,
tolerance_change = 0, tolerance_grad = 1e-9,
line_search_fn = 'strong_wolfe')
def closure():
optimizer.zero_grad()
beta = torch.exp(log_beta)
gamma = torch.exp(log_gamma)
# Data loss: sparse infected observations
_, I_pred_obs, _ = model(t_obs_t)
loss_data = torch.mean((I_pred_obs - I_obs_t)**2)
# Physics loss: SIR residuals at the collocation points
f_S, f_I, f_R = physics_residual(model, t_f, beta, gamma)
loss_physics = torch.mean(f_S**2) + torch.mean(f_I**2) + torch.mean(f_R**2)
# Initial-condition loss
S_pred0, I_pred0, R_pred0 = model(t0_t)
loss_ic = torch.mean((S_pred0 - S0)**2) + torch.mean((I_pred0 - I0)**2) + torch.mean((R_pred0 - R0)**2)
# Population-conservation loss
S_f, I_f, R_f = model(t_f)
loss_mass = torch.mean((S_f + I_f + R_f - 1)**2)
loss = (lambda_data * loss_data + loss_physics
+ lambda_ic * loss_ic + loss_mass)
loss.backward()
loss_history.append(loss.item())
return loss
optimizer.step(closure)
beta_est = torch.exp(log_beta).item()
gamma_est = torch.exp(log_gamma).item()
print(
'L-BFGS evaluations: %s | Loss: %.3e | beta: %.4f | gamma: %.4f' % (len(loss_history) - n_adam, loss_history[-1], beta_est, gamma_est)
)
Total Loss History
loss_lbfgs = np.array(loss_history[n_adam:])
best = np.minimum.accumulate(loss_lbfgs)
keep = np.concatenate(([True], np.diff(best) < 0))
step = np.arange(n_adam, len(loss_history))
plt.figure(figsize = (7, 4))
plt.semilogy(range(n_adam), loss_history[:n_adam], label = 'Adam')
plt.semilogy(step[keep], loss_lbfgs[keep], label = 'L-BFGS')
plt.xlabel('Step')
plt.ylabel('Loss')
plt.title('PINN Training Loss')
plt.legend()
plt.grid(alpha = 0.3)
plt.show()
Results
The trained network is evaluated at the reference time points, and the learned positive parameters are recovered from their logarithmic representation.
t_test_t = torch.tensor(t_true[:, None], dtype = torch.float64, device = device)
with torch.no_grad():
S_pred, I_pred, R_pred = model(t_test_t)
S_pred = S_pred.cpu().numpy().flatten()
I_pred = I_pred.cpu().numpy().flatten()
R_pred = R_pred.cpu().numpy().flatten()
beta_est = torch.exp(log_beta).item()
gamma_est = torch.exp(log_gamma).item()
Compare the True and Estimated Parameters
The basic reproduction number is
$$ R_0=\frac{\beta}{\gamma}. $$
$R_0$ is the number of new infections caused by one infected individual in a fully susceptible population: above one the epidemic grows, below one it dies out. It is the quantity an epidemiologist actually wants, and it follows from the two learned parameters.
print('True beta : %.4f' % beta_true)
print('Estimated beta : %.4f' % beta_est)
print()
print('True gamma : %.4f' % gamma_true)
print('Estimated gamma : %.4f' % gamma_est)
print()
print('True R0 : %.3f' % (beta_true/gamma_true))
print('Estimated R0 : %.3f' % (beta_est/gamma_est))
Compare the Predicted and Reference Solutions
The learned trajectories are compared with the numerical reference solution, while the sparse infected-population observations show the data available during training.
Check the two unobserved curves in particular. Nothing in the data constrained $S$ or $R$ at any time; whatever agreement they show was supplied by the residuals and the conservation term.
plt.figure(figsize = (8, 5))
# Ground truth
plt.plot(t_true, S_true, '--', label = 'True $S(t)$')
plt.plot(t_true, I_true, '--', label = 'True $I(t)$')
plt.plot(t_true, R_true, '--', label = 'True $R(t)$')
# PINN predictions
plt.plot(t_true, S_pred, label = 'PINN $S(t)$')
plt.plot(t_true, I_pred, label = 'PINN $I(t)$')
plt.plot(t_true, R_pred, label = 'PINN $R(t)$')
# Sparse observations
plt.scatter(t_obs, I_obs_noisy, s = 20, label = 'Observed $I(t)$')
plt.xlabel('Time')
plt.ylabel('Population fraction')
plt.title('PINN Reconstruction of SIR Dynamics')
plt.legend(ncol = 2)
plt.grid(alpha = 0.3)
plt.xlim([0, 100])
plt.ylim([-0.1, 1.1])
plt.show()
| Component | Role |
|---|---|
| Sparse infected observations | constrain $\hat I(t)$ at measured times |
| SIR residuals | enforce the governing equations throughout the domain |
| Initial-condition loss | anchors $\hat S(0)$, $\hat I(0)$, and $\hat R(0)$ |
| Population-conservation loss | encourages $\hat S+\hat I+\hat R=1$ |
| Trainable $\beta$ and $\gamma$ | identify the transmission and recovery rates |
The inverse PINN combines limited observations with the SIR equations to reconstruct the full epidemic trajectory and estimate the governing parameters.
Overview
The three chapters so far identified one unknown constant in a one-dimensional problem. This chapter keeps the arrangement and raises the difficulty on every axis: the domain is two-dimensional with a hole in it, the equations are the non-linear Navier-Stokes system acting on three coupled fields, and two properties are unknown instead of one. We are given a flow around a cylinder, we are given the form of the equations governing it, and we are told that the density and the viscosity of the fluid are unknown. Seventeen sensor readings are available. Find the fluid properties, and the flow field with them.
Section 3 in Physics-informed Neural Networks (PINN) argued that this requires no new algorithm, only a longer decision vector, and this laboratory is where the claim is cashed. Two scalars are created as trainable tensors and handed to the optimizer alongside the network parameters; nothing else in the training loop changes. The two scalars are then estimated to within about one percent from measurements at seventeen points, which is less than a tenth of one percent of the domain.
It is worth appreciating in advance how much harder this problem would be by classical means. A conventional identification would wrap an optimization around a CFD solver: guess $\rho$ and $\mu$, solve the flow, compare against the sensors, adjust, and repeat, paying for a full flow solution at every iteration and needing derivatives of the solution with respect to the parameters, which must be obtained by finite differencing or by a separately derived adjoint. Here the flow field and the parameters descend together in one loop.
The geometry is a channel with a circular cylinder in it, and the flow through it is steady and laminar. Three fields are unknown, the two velocity components and the pressure, and so are two properties of the fluid.
Unknown properties
$$\rho = ? \quad \mu = ?$$
The density and the dynamic viscosity are the quantities to be identified. The data used here was generated by a CFD computation with $\rho = 1$ and $\mu = 0.01$, and those values are recorded in the code so that the estimates can be scored afterwards. Throughout the training they are treated as completely unknown, and both are initialized at $1$, so the viscosity begins a factor of one hundred away from its true value.
Since the inlet speed is $1$ and the cylinder has radius $0.05$ and hence diameter $D = 0.1$, the Reynolds number of the true flow is
$$ \mathrm{Re} = \frac{\rho U D}{\mu} = \frac{(1)(1)(0.1)}{0.01} = 10, $$
or $100$ if the channel height of $1$ is used as the length scale instead. Either way the flow is firmly laminar and there is no vortex shedding, which is what makes it legitimate to look for a steady solution with no time variable at all.
Governing equations and boundary conditions
$$ \color{red}\rho\left(u\frac{\partial u}{\partial x} + v\frac{\partial u}{\partial y} + \frac{1}{\color{red}\rho}\frac{\partial p}{\partial x} \right) - \color{red}\mu\left( \frac{\partial^2u}{\partial^2 x} + \frac{\partial^2u}{\partial^2 y}\right) = 0$$
$$ \color{red}\rho\left(u\frac{\partial v}{\partial x} + v\frac{\partial v}{\partial y} + \frac{1}{\color{red}\rho}\frac{\partial p}{\partial y} \right) - \color{red}\mu\left( \frac{\partial^2v}{\partial^2 x} + \frac{\partial^2v}{\partial^2 y}\right) = 0$$
These are the steady momentum equations, with the unknown parameters marked in red. Reading them term by term: $\rho(u\,\partial u/\partial x + v\,\partial u/\partial y)$ is the convective acceleration, the rate at which a fluid particle's momentum changes because it is being carried into a region of different velocity; $\partial p/\partial x$ is the pressure gradient driving the flow; and $\mu \nabla^2 u$ is the viscous diffusion of momentum. The equations say that these three balance.
Notice the algebraic detail that the factor $1/\rho$ inside the bracket cancels the $\rho$ outside it, so the pressure gradient enters with coefficient exactly $1$. The unknown $\rho$ therefore multiplies the convective term only, and $\mu$ the viscous term only. This is not a cosmetic point and Section 4.2 returns to it, because it determines what can and cannot be identified from the measurements.
- Inlet boundary condition ($x=-0.5$):
$$u|_{x = -0.5} = 1$$
- Outlet boundary condition ($x=1.5$):
$$v|_{x = 1.5} = 0$$
$$p|_{x = 1.5} = 0$$
- Non-slip boundary condition ($y=\pm 0.5$):
$$u|_{y = \pm 0.5} = 0$$
$$v|_{y = \pm 0.5} = 0$$
The physical readings are worth naming. At the inlet the fluid arrives with a uniform speed of $1$ in the $x$-direction and none in $y$. At the outlet the flow is required to leave straight, $v=0$, and the pressure is fixed at zero, which is a datum rather than a physical statement: only pressure differences drive an incompressible flow, so the pressure field is determined by the equations only up to an additive constant, and one boundary value must be nominated to remove that freedom. At the channel walls and on the cylinder surface the no-slip condition applies: a viscous fluid does not slide along a solid boundary, so both velocity components vanish there.
Make a neural network and loss functions like below :
Import Library
The imports are as before, with the addition of the two true parameter values, retained purely for scoring. They are never used in the loss.
import numpy as np
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
# true values
C1true = 1.0
C2true = 0.01
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(device)
Load the Data
Two files are loaded: the coordinates of the CFD mesh points and the corresponding flow variables. (The files are not distributed with this material; the cells of this lab that depend on them are shown with the outputs of the original run.)
The printed shapes are $(19692, 2)$ and $(19692, 3)$: nearly twenty thousand points, each with two coordinates and three field values, namely the two velocity components and the pressure. The coordinates will serve as collocation points for the equation residual, which is a legitimate use of the file since only the geometry is being taken from it. The field values are the ground truth, and all but seventeen of them are withheld from training.
!gdown 1V6wUNUirlaW7aZ4Hx4SMmxZJpqFeFA_B -O /content/lab3_XY_domain.npy
!gdown 1r5sFfiRRVQgVXliwa6KB6rBBZlxoZeDw -O /content/lab3_gt_domain.npy
Y_domain = np.load('/content/lab3_XY_domain.npy')
gt_domain = np.load('/content/lab3_gt_domain.npy')
print(Y_domain.shape)
print(gt_domain.shape)
Define Collocation Points
The boundary point sets are constructed by hand, one per physical surface: two hundred points along each of the top and bottom walls, one hundred at the inlet and one hundred at the outlet, and two hundred distributed around the cylinder by sampling the angle $\theta$ uniformly on $[0, 2\pi]$ and mapping to $(r\cos\theta, r\sin\theta)$ with $r = 0.05$.
The cylinder illustrates a real advantage of the mesh-free formulation. A curved boundary in a finite volume code requires a body-fitted mesh, or a cut-cell treatment, or an immersed boundary method, each of which is a substantial piece of engineering. Here the curved boundary is a list of parametric coordinates, and the no-slip condition is imposed at those coordinates by the same mean squared error used everywhere else. Geometry that is difficult to mesh is not difficult to sample.
The top wall, the bottom wall, and the cylinder are concatenated into a single array, since all three carry the identical no-slip condition and can therefore share one loss term. The scatter plot then shows the domain points together with the four boundary sets; as in Section 6 in Physics-informed Neural Networks (PINN), the value of the plot is that it catches construction errors that would otherwise appear only as an inexplicable failure to converge, and the geometry here is complicated enough that the check is worth its cell.
# Boundary Conditions
bc_top_x = np.linspace(-0.5, 1.5, 200).reshape(-1, 1)
bc_top_y = 0.5 * np.ones_like(bc_top_x).reshape(-1, 1)
bc_bottom_x = np.linspace(-0.5, 1.5, 200).reshape(-1, 1)
bc_bottom_y = -0.5 * np.ones_like(bc_bottom_x).reshape(-1, 1)
bc_inlet_y = np.linspace(-0.5, 0.5, 100).reshape(-1, 1)
bc_inlet_x = -0.5 * np.ones_like(bc_inlet_y).reshape(-1, 1)
bc_outlet_y = np.linspace(-0.5, 0.5, 100).reshape(-1, 1)
bc_outlet_x = 1.5 * np.ones_like(bc_outlet_y).reshape(-1, 1)
Y_top = np.concatenate((bc_top_x, bc_top_y), 1)
Y_bottom = np.concatenate((bc_bottom_x, bc_bottom_y), 1)
Y_inlet = np.concatenate((bc_inlet_x, bc_inlet_y), 1)
Y_outlet = np.concatenate((bc_outlet_x, bc_outlet_y), 1)
radius = 0.05
theta = np.linspace(0, 2 * np.pi, 200)
bc_cylinder_x = (0 + radius * np.cos(theta)).reshape(-1, 1)
bc_cylinder_y = (0 + radius * np.sin(theta)).reshape(-1, 1)
Y_cylinder = np.concatenate((bc_cylinder_x, bc_cylinder_y), 1)
Y_wall = np.concatenate((Y_top, Y_bottom, Y_cylinder))
plt.scatter(Y_domain[:, 0], Y_domain[:, 1], s = 1, label = 'domain')
plt.scatter(Y_wall[:, 0], Y_wall[:, 1], s = 2, label = 'wall')
plt.scatter(Y_inlet[:, 0], Y_inlet[:, 1], s = 2, label = 'inlet')
plt.scatter(Y_outlet[:, 0], Y_outlet[:, 1], s = 2, label = 'outlet')
plt.axis('scaled')
plt.legend(loc = 'upper right')
plt.show()
Define Observation Points
This CFD data is where the density $\rho$ is 1 and the viscosity $\mu$ is 0.01.
Of course, these two values are assumed to be unknown throughout this example.
Also, assume that 17 sensor locations are used.
Assume that x-velocity, y-velocity, and pressure are available at these locations.
The cell that places the sensors repays reading, because the placement is deliberate rather than random. Twelve locations are laid out on a rectangular lattice spanning the channel, four values of $x$ from $-0.3$ to $1.3$ crossed with three values of $y$ from $-0.45$ to $0.45$. Five more are placed evenly around a circle of radius $0.12$ centred on the cylinder, that is, just outside the cylinder's own radius of $0.05$. None of these seventeen points is guaranteed to be a node of the CFD mesh, so the cell measures the distance from each to every one of the $19692$ domain points and keeps the nearest.
The arrangement is worth a moment's thought, because it is a small piece of experimental design. The lattice samples the free stream and the wake; the ring concentrates five of the seventeen sensors, nearly a third of the budget, in the annulus immediately around the cylinder where the velocity gradients are largest and therefore where the viscous term $\mu(u_{xx}+u_{yy})$ is largest. Since $\mu$ can only be identified from places where the viscous term is not negligible, spending a third of the sensors there is exactly the right instinct. PINN: Limitations and Extensions returns to this idea in its general form.
x = np.linspace(-0.3, 1.3, 4)
y = np.linspace(-0.45, 0.45, 3)
xy = np.meshgrid(x, y)
XY = np.concatenate([xy[0].reshape(-1, 1), xy[1].reshape(-1, 1)], axis = 1)
theta = np.linspace(0, 2 * np.pi, 6)[:-1]
x_e = (radius + 0.07) * np.cos(theta)
y_e = (radius + 0.07) * np.sin(theta)
XY_e = np.concatenate([x_e.reshape(-1, 1), y_e.reshape(-1, 1)], axis = 1)
XY = np.concatenate([XY, XY_e], axis = 0)
diff = np.expand_dims(Y_domain, 0) - np.expand_dims(XY, 1)
dists = np.linalg.norm(diff, axis = 2)
idx_obs = np.argmin(dists, axis = 1)
Y_obs = Y_domain[idx_obs]
gt_obs = gt_domain[idx_obs]
print('Y obs: {}'.format(Y_obs.shape))
print('GT obs: {}\n\n'.format(gt_obs.shape))
plt.scatter(Y_domain[:, 0], Y_domain[:, 1], c = gt_domain[:, 0], s = 1, cmap = 'jet')
plt.scatter(Y_obs[:, 0], Y_obs[:, 1], marker = 'x', linewidths = 1, s = 20, c = 'k')
plt.axis('scaled')
plt.show()
The printed shapes confirm the selection: seventeen sensor coordinates and seventeen triples of ground truth. Consider what that means as a quantity of information. The unknowns are a network with a few hundred thousand parameters plus two physical constants, and the entire observational input is seventeen triples, fewer than one point in a thousand of the domain. No purely data-driven method could do anything with this. The reason the problem is nevertheless tractable is that the equation residual constrains the network at all $19692$ collocation points, and it is the sensors' job only to select among the fields the equations permit, and to pin down the two constants.
The plot overlays the sensor locations as black crosses on the true $x$-velocity field, and both the lattice and the ring around the cylinder are visible in it.
Define PINN Network
The network maps the two coordinates to three outputs, $u$, $v$ and $p$, through four hidden layers of two hundred and fifty-six units. It is considerably larger than the plate network of Section 6 in Physics-informed Neural Networks (PINN), reflecting a harder problem: a non-linear system of three coupled equations on a domain with a hole in it.
Predicting the pressure with the same network that predicts the velocities is the natural choice here, since the momentum equations couple all three fields at every point. Weights are again Xavier-normal initialized with zero biases.
The plotting helper that follows displays three panels for the $x$-velocity, namely the CFD reference, the current prediction, and the absolute error between them, followed by two panels tracing the estimated density and viscosity against iteration.
Those two parameter traces are the most informative diagnostic in this laboratory, and it is worth knowing how to read them. A parameter estimate that is still moving steadily has not converged, whatever the loss is doing; one that has flattened onto a horizontal line has converged, and the value it has flattened onto is the estimate. A trace that oscillates without settling indicates that the learning rate is too large for that parameter, or that the parameter is not identifiable from the available information, which is a possibility examined directly below.
class PINN(nn.Module):
def __init__(self):
super(PINN, self).__init__()
self.net = nn.Sequential(
nn.Linear(2, 256),
nn.Tanh(),
nn.Linear(256, 256),
nn.Tanh(),
nn.Linear(256, 256),
nn.Tanh(),
nn.Linear(256, 256),
nn.Tanh(),
nn.Linear(256, 3)
)
for m in self.net:
if isinstance(m, nn.Linear):
nn.init.xavier_normal_(m.weight)
nn.init.zeros_(m.bias)
def forward(self, x):
output = self.net(x)
v1, v2, p = output[:, 0:1], output[:, 1:2], output[:, 2:3]
return v1, v2, p
from IPython.display import clear_output
def PLOT(model, Y_domain, rho_list, vis_list):
v1, _, _ = model(Y_domain)
xy = Y_domain.detach().cpu().numpy()
v1 = v1.detach().cpu().numpy()
plt.figure(figsize = (9, 2))
plt.subplot(1, 3, 1)
plt.title('CFD')
plt.scatter(xy[:, 0], xy[:, 1], c = gt_domain[:, 0], s = 0.1, cmap = 'jet')
plt.axis('scaled')
plt.subplot(1, 3, 2)
plt.title('Prediction')
plt.scatter(xy[:, 0], xy[:, 1], c = v1, s = 0.1, cmap = 'jet')
plt.axis('scaled')
plt.subplot(1, 3, 3)
plt.title('Error')
plt.scatter(xy[:, 0], xy[:, 1], c = np.abs(v1 - gt_domain[:, 0:1]), s = 0.1, cmap = 'jet')
plt.axis('scaled')
plt.tight_layout()
plt.show()
plt.figure(figsize = (6, 2))
plt.subplot(1, 2, 1)
plt.title('Rho Prediction')
plt.plot(rho_list)
plt.subplot(1, 2, 2)
plt.title('Vis Prediction')
plt.plot(vis_list)
plt.tight_layout()
plt.show()
clear_output(wait = True)
Define Functions to Compute Physics Loss
The helper is identical to that of Section 6 in Physics-informed Neural Networks (PINN): the derivative function returns both partial derivatives of a scalar field.
def derivative(y, t):
df = torch.autograd.grad(y, t, grad_outputs = torch.ones_like(y), create_graph = True)[0]
df_x = df[:, 0:1]
df_y = df[:, 1:2]
return df_x, df_y
Define PDE
The two-dimensional steady Navier-Stokes system, in the form the code implements:
$$ \begin{align*} \rho \left(u{\partial u \over \partial x} + v{\partial u \over \partial y} + {1 \over \rho}{\partial p \over \partial x}\right) - \mu \ \left({\partial^2 u \over {\partial x^2}} + {\partial^2 u \over {\partial y^2}}\right) &= 0\\\\ \rho \left(u{\partial v \over \partial x} + v{\partial v \over \partial y} + {1 \over \rho}{\partial p \over \partial y}\right) - \mu \ \left({\partial^2 v \over {\partial x^2}} + {\partial^2 v \over {\partial y^2}}\right) &= 0\\\\ {\partial u \over \partial x} + {\partial v \over \partial y} &= 0 \end{align*} $$
The third equation is new and is the reason the system has three equations for three unknowns. It is the continuity equation for an incompressible fluid: the divergence of the velocity field vanishes, so that no volume is created or destroyed anywhere. It contains neither $\rho$ nor $\mu$, and therefore carries no information about them, but it is indispensable, since without it there would be two equations for three fields.
The residual function computes seven derivatives, two first and two second for each velocity component and one first derivative for the pressure, and assembles the three residuals. The two unknown parameters appear in it as ordinary multiplicative factors. This is the mechanical heart of the inverse construction: because they are ordinary factors in an expression built from tensors that are being tracked, the residual is a differentiable function of them, and $\partial\mathcal{L}/\partial\rho$ costs nothing beyond the backward pass that is happening anyway.
def PDE(model, Y_domain):
v1, v2, p = model(Y_domain)
dv1_x, dv1_y = derivative(v1, Y_domain)
dv1_xx, _ = derivative(dv1_x, Y_domain)
_, dv1_yy = derivative(dv1_y, Y_domain)
dv2_x, dv2_y = derivative(v2, Y_domain)
dv2_xx, _ = derivative(dv2_x, Y_domain)
_, dv2_yy = derivative(dv2_y, Y_domain)
dp_x, dp_y = derivative(p, Y_domain)
pde_v1 = rho * (v1 * dv1_x + v2 * dv1_y) + dp_x - vis * (dv1_xx + dv1_yy)
pde_v2 = rho * (v1 * dv2_x + v2 * dv2_y) + dp_y - vis * (dv2_xx + dv2_yy)
pde_cont = dv1_x + dv2_y
return pde_v1, pde_v2, pde_cont
Boundary Conditions
- Two Dirichlet boundary conditions on the plate boundary (no-slip condition),
$$u(x,y) = 0, \quad v(x,y) = 0 \qquad \text{at} \quad y = \frac{D}{2} \ \; \text{or} \; -\frac{D}{2}$$
- Two Dirichlet boundary conditions at the cylinder boundary
$$u(\text{cylinder}) = 0, \quad v(\text{cylinder}) = 0$$
- Two Dirichlet boundary conditions at the inlet boundary
$$u(-1,y) = u_{\text{in}}, \quad v(-1,y) = 0$$
- Two Dirichlet boundary conditions at the outlet boundary
$$p(1,y) = 0, \quad v(1,y) = 0$$
The inlet and outlet coordinates are written schematically here; the code places them at $x=-0.5$ and $x=1.5$, matching the geometry of Section 4.1. All eight conditions are Dirichlet, prescribing values rather than derivatives, so the three residual functions below are simply network evaluations, with a subtraction where the prescribed value is non-zero. The inlet function subtracts a tensor of ones from the $x$-velocity, implementing $u = u_{\text{in}} = 1$; every other condition is homogeneous and the residual is the raw output.
The wall function takes the concatenated array of top wall, bottom wall and cylinder points and returns the two velocity components, which is why those three surfaces were merged earlier: one function and one loss term serve all of them.
def BC_wall(model, Y_wall):
v1_wall, v2_wall, _ = model(Y_wall)
return v1_wall, v2_wall
def BC_inlet(model, Y_inlet):
v1_inlet, v2_inlet, _ = model(Y_inlet)
return v1_inlet - 1, v2_inlet
def BC_outlet(model, Y_outlet):
_, v2_outlet, p_outlet = model(Y_outlet)
return v2_outlet, p_outlet
Parameters to be Optimized
All we need to do is to define the unknown parameters, specifically $\rho$ (density) and $\mu$ (viscosity), as trainable variables within the PINN. The next cell is where the entire inverse machinery of Section 3 in Physics-informed Neural Networks (PINN) is implemented, and it is two lines long.
Both parameters start at $1$. For the density that is the true value by luck; for the viscosity it is a factor of one hundred too large, so the optimizer must transport that parameter from $1$ down to $0.01$.
The optimizer is constructed with three parameter groups, one for the network and one for each scalar, all at a learning rate of $10^{-4}$. Groups exist so that different learning rates can be assigned, and here the facility is set up but not used, all three rates being equal. It is worth doing an order-of-magnitude check on whether that rate suffices. Adam normalizes each update by a running estimate of the gradient magnitude, so the step taken in any one parameter is at most of order the learning rate. Moving the viscosity from $1$ to $0.01$ therefore requires at least about $0.99/10^{-4} \approx 10^4$ iterations even at full speed, which is why the loop below runs for fifty thousand and not for the thousand that sufficed in Section 5 in Physics-informed Neural Networks (PINN).
# Initialize rho and mu as trainable tensors with the value 1.
rho = torch.tensor([1.0], device = device, requires_grad = True)
vis = torch.tensor([1.0], device = device, requires_grad = True)
model = PINN().to(device)
optimizer = torch.optim.Adam([{'params': model.parameters(), 'lr': 1e-4},
{'params': rho, 'lr': 1e-4},
{'params': vis, 'lr': 1e-4}])
Train
$$ \begin{align*} \mathcal{L} &= \omega_{\text{data}} \mathcal{L}_{\text{data}} + \omega_{\text{PDE}} \mathcal{L}_{\text{PDE}} + \omega_{\text{BC}} \mathcal{L}_{\text{BC}}\\\\\\ \mathcal{L}_{\text{data}} &= \frac{1}{N_{\text{data}}} \sum \left( \left(u - \hat u \right)^2 + \left(v - \hat v \right)^2 + \left(p - \hat p \right)^2 \right) \\ \end{align*} $$
The first cell initializes the bookkeeping lists and the second runs the loop, in which the loss is
$$ \mathcal{L} = \mathcal{L}_{\text{PDE}} + \mathcal{L}_{\text{BC}} + 10\,\mathcal{L}_{\text{data}}, $$
and the factor of ten is the first non-unit weight in this notebook. It is deliberate and the reason is structural rather than empirical. In the forward problems of Sections 4 to 6, every term of the loss was pushing towards the same unique answer, so the weighting affected only the rate of convergence. Here the equation and boundary terms have a whole family of solutions available to them, one for each pair $(\rho,\mu)$, and the data term is the only term that distinguishes among them. Under-weighting it would let the network settle into some equation-satisfying flow with wrong properties and a small total loss. The data term is carrying the entire identification, and seventeen points against nearly twenty thousand collocation points would otherwise be a whisper against a shout.
Everything else in the loop is familiar: three equation residuals over the domain, six boundary residuals, three data residuals for the three observed fields, all as mean squared errors, summed and backpropagated. A single optimizer step updates the network weights, the density, and the viscosity together, which is Section 3 in Physics-informed Neural Networks (PINN)'s central claim reduced to one line of code. The estimates are appended to the trace lists at every iteration, so the two parameter plots are drawn from a record of all fifty thousand steps.
What is identifiable, and why the pressure observations matter
Before reading the answer, here is an argument worth following, because it decides whether this problem can be solved at all.
Suppose $(u, v, p, \rho, \mu)$ satisfies all three equations. Consider the scaled candidate $(u, v, cp, c\rho, c\mu)$ for any constant $c > 0$: the velocities are untouched, the pressure and both parameters are multiplied by $c$. Substituting into the first momentum equation, every term acquires exactly one factor of $c$, since the convective term carries $\rho$, the pressure term carries $p$, and the viscous term carries $\mu$. The residual is therefore multiplied by $c$ and remains zero. Continuity is untouched because it involves only the velocities. The boundary conditions are untouched as well: the inlet condition constrains $u$, and the outlet condition sets $p=0$, which is preserved by scaling. So the scaled candidate satisfies the physics exactly as well as the original.
The consequence is that the equations and boundary conditions alone determine only the ratio $\rho/\mu$, which is to say the Reynolds number, and not the individual values. If the sensors reported velocities only, the two parameters would be individually unidentifiable, the traces would wander along the one-parameter family, and no amount of training would fix it. It is the observed pressure that breaks the degeneracy, because scaling $p$ by $c$ changes the data residual. Note that this is a property of the problem, not of the method: a classical identification from velocity data alone would face exactly the same obstruction. Recognizing such degeneracies before running anything is a large part of the practical skill in inverse problems.
Y_domain = torch.tensor(Y_domain, dtype = torch.float32, requires_grad = True).to(device)
Y_wall = torch.tensor(Y_wall, dtype = torch.float32).to(device)
Y_inlet = torch.tensor(Y_inlet, dtype = torch.float32).to(device)
Y_outlet = torch.tensor(Y_outlet, dtype = torch.float32).to(device)
Y_obs = torch.tensor(Y_obs, dtype = torch.float32).to(device)
gt_obs = torch.tensor(gt_obs, dtype = torch.float32).to(device)
epoch = 0
loss_list = []
rho_list, vis_list = [], []
while epoch < 50001:
## PDE loss
PDE_v1, PDE_v2, PDE_cont = PDE(model, Y_domain)
loss_pde = torch.mean(PDE_v1**2) + torch.mean(PDE_v2**2) + torch.mean(PDE_cont**2)
## BC loss
v1_wall, v2_wall = BC_wall(model, Y_wall)
v1_inlet, v2_inlet = BC_inlet(model, Y_inlet)
v2_outlet, p_outlet = BC_outlet(model, Y_outlet)
loss_bc = (torch.mean(v1_wall**2) + torch.mean(v2_wall**2)
+ torch.mean(v1_inlet**2) + torch.mean(v2_inlet**2)
+ torch.mean(v2_outlet**2) + torch.mean(p_outlet**2))
## Data loss
u_obs, v_obs, p_obs = model(Y_obs)
loss_data = (torch.mean((u_obs - gt_obs[:, 0:1])**2)
+ torch.mean((v_obs - gt_obs[:, 1:2])**2)
+ torch.mean((p_obs - gt_obs[:, 2:3])**2))
loss = loss_pde + loss_bc + 10 * loss_data
optimizer.zero_grad()
loss.backward()
optimizer.step()
loss_list.append(loss.item())
rho_list.append(rho.item())
vis_list.append(vis.item())
if epoch % 500 == 0:
print('EPOCH : %6d/%6d | Loss_PDE : %5.4f | Loss_BC : %5.4f | Loss_DATA : %5.4f | RHO : %.4f | VIS : %.4f' \
%(epoch, 50000, loss_pde.item(), loss_bc.item(), loss_data.item(), rho.item(), vis.item()))
PLOT(model, Y_domain, rho_list, vis_list)
epoch += 1
# Save the model and the parameter histories
torch.save(model, 'lab3_model.pt')
np.save('rho_list.npy', np.array(rho_list))
np.save('vis_list.npy', np.array(vis_list))
print('Model saved as lab3_model.pt')
print('Rho list saved as rho_list.npy')
print('Viscosity list saved as vis_list.npy')
Result
The trained model and the two parameter traces are provided so that the result can be inspected without a fifty-thousand-iteration run.
!gdown 1ebNl9iZphjsvmWw_vCSFwnA8forR8woj -O /content/lab3_model.pt
!gdown 1KWDXGCsooxXvH5rmoQWRS_143eGspTrD -O /content/rho_list.npy
!gdown 18CIt0fsRLbaQjQQrFzXq2i3M8sfwIyz5 -O /content/vis_list.npy
# Load the estimated parameters and the pre-trained model
rho_list = np.load('/content/rho_list.npy').tolist()
vis_list = np.load('/content/vis_list.npy').tolist()
model = torch.load('/content/lab3_model.pt', map_location = device, weights_only = False)
print('Predicted RHO: %.4f, Predicted VIS: %.4f\n' % (rho_list[-1], vis_list[-1]))
Y_domain_tensor = torch.tensor(np.load('/content/lab3_XY_domain.npy'), dtype = torch.float32).to(device)
PLOT(model, Y_domain_tensor, rho_list, vis_list)
The cell prints a predicted density of $0.9945$ and a predicted viscosity of $0.0099$, and redraws the diagnostic panels.
Compare against the true values $\rho = 1$ and $\mu = 0.01$. The density is recovered with a relative error of $0.55\%$ and the viscosity with a relative error of about $1\%$, the latter having travelled from its initial value of $1$ down by two orders of magnitude to reach $0.0099$. Both estimates were obtained from seventeen measurements, fewer than one domain point in a thousand, together with the knowledge that the flow obeys the Navier-Stokes equations.
That result deserves a moment's reflection rather than a nod. The network was never told what a density is. It was given a functional form containing two unknown numbers, a set of coordinates, and seventeen readings, and it recovered both numbers by gradient descent, at the same time and in the same loop as it recovered the flow field. The estimated field, shown in the first row of panels against the CFD reference and the pointwise error, is a by-product; in an identification problem it is the two numbers that were wanted.
Read the parameter traces as well as the printed values. Both should have flattened well before the end of the run, and it is the flatness rather than the final value that tells you the estimate is trustworthy. A trace that is still descending at the last iteration reports where the optimizer happened to be when the loop ended, which is not an estimate of anything.
One caveat completes the picture. The sensor readings here were taken from the CFD solution without added noise, so the data is exactly consistent with the equations, and the estimates are correspondingly clean. With noisy measurements the identified parameters would carry uncertainty, and quantifying that uncertainty, rather than reporting a single point estimate, is where serious inverse-problem practice begins. The framework of this section gives a point estimate and no error bar.
Overview
Another common type of inverse problem involves unknown boundary conditions. It arises constantly in practice: a component is instrumented on its accessible faces, and the condition on the inaccessible one, buried, submerged, or in contact with something unmeasured, is exactly what one wants to know. Here the setting is two-dimensional steady heat conduction on a square, with three of the four edges at known temperatures and the fourth entirely unknown. By leveraging a PINN we can infer the missing boundary condition from observed temperature data while remaining consistent with the governing heat equation.
This section is arranged as a controlled experiment in three stages, and the arrangement is the point. Section 5.2 solves the problem with physics alone and shows the failure. Section 5.3 adds nine thermometers. Section 5.4 adds one further piece of knowledge, namely that the unknown boundary temperature is uniform along the edge, without saying what value it takes. Section 5.5 compares the three recovered boundary profiles against the truth on a single plot. Each stage adds information rather than tuning, and the improvements are attributable to the information added.
The setting is steady two-dimensional heat conduction on a square plate. Three edges are held at known temperatures; the fourth is inaccessible, and its temperature profile is what we want.
Problem properties
$$ \kappa = 1 $$
Governing equation and boundary conditions
$$ \kappa \left(\frac{\partial^2 T}{\partial x^2} + \frac{\partial^2 T}{\partial y^2} \right) = 0$$
This is Laplace's equation for the temperature, the steady-state form of the heat equation. It follows from conservation of energy together with Fourier's law: the heat flux is $\mathbf{q} = -\kappa \nabla T$, and in the steady state with no internal generation the flux is divergence-free, so $\nabla\cdot(\kappa\nabla T)=0$, which for constant $\kappa$ reduces to $\nabla^2 T = 0$. PINN: Limitations and Extensions takes up the case where $\kappa$ is not constant, and the distinction between the two forms of the equation matters a great deal there.
Since $\kappa$ is a positive constant, multiplying the Laplacian by it changes nothing about which fields satisfy the equation. It is retained here because the code writes it, and because in a two-material problem the corresponding factor is not constant.
The domain is the square $[-1,1]\times[-1,1]$ and three of its edges carry prescribed temperatures.
- Left boundary condition ($x=-1$):
$$T|_{x = -1} = 0.5$$
- Right boundary condition ($x=1$):
$$T|_{x = 1} = 1$$
- Bottom boundary condition ($y = -1$):
$$T|_{y = -1} = 0$$
- The top boundary condition $(y = 1)$ is unknown.
$$T|_{y = 1} = \; ?$$
Why this is a genuinely different situation
Laplace's equation on a bounded domain is well posed when a condition is given on the entire boundary, and it is not well posed when part of the boundary is left free. The consequence is not that the answer is slightly less accurate. It is that there is no unique answer: any smooth profile assigned to the top edge produces a different harmonic function inside, and all of them satisfy the equation and the three known edges exactly. The solution set is an infinite-dimensional family.
A conventional finite element or finite difference code cannot even be started in this situation, since assembling its system of equations requires a value or a flux at every boundary node. A PINN, by contrast, will run happily: it will minimize a loss that contains no information about the top edge, and it will return one member of the family. Whether that member is the right one is another question, and Section 5.2 answers it.
Make a neural network and loss functions like below:
Import Library
import numpy as np
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(device)
Load the Data
The reference solution comes from a finite difference computation on a $100\times100$ grid over the square. (The file is not distributed with this material; the cells that depend on it are shown with the outputs of the original run.)
As in Section 6 in Physics-informed Neural Networks (PINN), this reference plays two roles that should be kept distinct: it supplies the nine sensor readings used for training in Sections 8.3 and 8.4, and it supplies the ground truth against which all three models are judged. Only nine of its ten thousand values ever reach the optimizer.
!gdown 1PHr2D76v31bQ58vrw9O1B5yqyJUovaC4 -O /content/lab4_gt_domain.npy
gt_domain = np.load('/content/lab4_gt_domain.npy')
Define Observation Points
Assume that nine thermometers are positioned at specific locations as shown below. These thermometers provide observed temperature data, which will be used to infer the unknown boundary condition in the 2D heat transfer problem.
The cell selects the grid indices $15$, $50$ and $85$ in each direction and takes all nine combinations, giving a regular $3\times3$ array of sensors at coordinates approximately $-0.697$, $0.010$ and $0.717$ in each axis. They therefore sit well inside the domain, none of them close to the unknown top edge, which is realistic: an inaccessible boundary is inaccessible to instrumentation as well.
Nine numbers is a thin diet even compared with the seventeen of Section 4, and the unknown here is not a pair of scalars but an entire function along an edge. That mismatch is why Section 5.3 improves the answer without perfecting it, and why Section 5.4 needs to supply structure rather than more data.
The plot shows the finite difference field with the nine sensor locations marked in red. Read the field before continuing: the right edge is hot at $T=1$, the left edge is at $T=0.5$, the bottom edge is cold at $T=0$, and the top edge is where the interesting question lies.
plt.figure(figsize = (5, 4))
x = np.linspace(-1, 1, 100)
y = np.linspace(-1, 1, 100)
x, y = np.meshgrid(x, y)
idx_xs = np.array([15, 50, 85])
idx_ys = np.array([15, 50, 85])
Y_obs = []
gt_obs = []
for idx_x in idx_xs:
for idx_y in idx_ys:
Y_obs.append(np.concatenate((x[idx_x, idx_y].reshape(-1, 1),
y[idx_x, idx_y].reshape(-1, 1)), 1)[0, :])
gt_obs.append(gt_domain[idx_x, idx_y])
Y_obs = np.array(Y_obs)
gt_obs = np.array(gt_obs).reshape(-1, 1)
plt.pcolormesh(x, y, gt_domain, cmap = 'jet')
plt.colorbar()
plt.scatter(Y_obs[:, 0], Y_obs[:, 1], s = 30, c = 'r')
plt.title('FDM')
plt.xlabel('x')
plt.ylabel('y')
plt.axis('square')
plt.show()
Define Collocation Points
A $100\times100$ lattice gives ten thousand interior collocation points, and each of the four edges is sampled with a hundred points; the last lines turn every point set into a tensor on the selected device, with gradients tracked only for the interior points, where the equation is differentiated. Note that the top edge is included in the point sets even though no condition is imposed on it in this section; it will be needed in Section 5.4, where the prior knowledge term is evaluated there, and in Section 5.5, where the recovered profile is plotted along it.
# 'Domain Points'
x = np.linspace(-1, 1, 100)
y = np.linspace(-1, 1, 100)
x, y = np.meshgrid(x, y)
Y_domain = np.concatenate((x.reshape(-1, 1), y.reshape(-1, 1)), 1)
# 'Boundary Conditions'
bc_top_x = np.linspace(-1, 1, 100).reshape(-1, 1)
bc_top_y = np.ones_like(bc_top_x)
bc_bottom_x = np.linspace(-1, 1, 100).reshape(-1, 1)
bc_bottom_y = -1 * np.ones_like(bc_bottom_x)
bc_right_y = np.linspace(-1, 1, 100).reshape(-1, 1)
bc_right_x = np.ones_like(bc_right_y)
bc_left_y = np.linspace(-1, 1, 100).reshape(-1, 1)
bc_left_x = -1 * np.ones_like(bc_left_y)
Y_top = np.concatenate((bc_top_x, bc_top_y), 1)
Y_bottom = np.concatenate((bc_bottom_x, bc_bottom_y), 1)
Y_right = np.concatenate((bc_right_x, bc_right_y), 1)
Y_left = np.concatenate((bc_left_x, bc_left_y), 1)
plt.scatter(Y_domain[:, 0], Y_domain[:, 1], s = 1, alpha = 0.5)
plt.scatter(Y_top[:, 0], Y_top[:, 1], s = 5)
plt.scatter(Y_bottom[:, 0], Y_bottom[:, 1], s = 5)
plt.scatter(Y_right[:, 0], Y_right[:, 1], s = 5)
plt.scatter(Y_left[:, 0], Y_left[:, 1], s = 5)
plt.title('Collocation Points')
plt.axis('scaled')
plt.show()
Y_domain = torch.tensor(Y_domain, dtype = torch.float32, requires_grad = True).to(device)
Y_bottom = torch.tensor(Y_bottom, dtype = torch.float32).to(device)
Y_right = torch.tensor(Y_right, dtype = torch.float32).to(device)
Y_left = torch.tensor(Y_left, dtype = torch.float32).to(device)
Y_top = torch.tensor(Y_top, dtype = torch.float32).to(device)
Y_obs = torch.tensor(Y_obs, dtype = torch.float32).to(device)
gt_obs = torch.tensor(gt_obs, dtype = torch.float32).to(device)
Define PINN Network
Two inputs, one output, four hidden layers of two hundred and fifty-six units, with Adam at a learning rate of $10^{-4}$. The network is nearly as large as the Navier-Stokes network of Section 4, which is generous for a linear equation in a single field. Since the whole section is a comparison between three training regimes, a fixed and generous architecture is the right choice: it removes capacity from the list of possible explanations for the differences observed.
The plotting helper draws the reference field, the current prediction, and the absolute error, with the error panel's colour scale fixed to $[0, 0.2]$ so that the three regimes can be compared on identical terms. Fixing that scale by hand is what makes the comparison in Sections 8.3 and 8.4 legible.
class PINN(nn.Module):
def __init__(self):
super(PINN, self).__init__()
self.net = nn.Sequential(
nn.Linear(2, 256),
nn.Tanh(),
nn.Linear(256, 256),
nn.Tanh(),
nn.Linear(256, 256),
nn.Tanh(),
nn.Linear(256, 256),
nn.Tanh(),
nn.Linear(256, 1)
)
for m in self.net:
if isinstance(m, nn.Linear):
nn.init.xavier_normal_(m.weight)
nn.init.zeros_(m.bias)
def forward(self, x):
return self.net(x)
model = PINN().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr = 1e-4)
from IPython.display import clear_output
def PLOT(model_test, Y_domain):
pred = model_test(Y_domain).detach().cpu().numpy()
xy = Y_domain.detach().cpu().numpy()
gt = gt_domain.reshape(-1, 1)
plt.figure(figsize = (9, 2.3))
plt.subplot(1, 3, 1)
plt.title('CFD')
plt.scatter(xy[:, 0], xy[:, 1], c = gt, vmin = 0, vmax = 1, s = 1, cmap = 'jet')
plt.colorbar()
plt.axis('scaled')
plt.subplot(1, 3, 2)
plt.title('Prediction')
plt.scatter(xy[:, 0], xy[:, 1], c = pred, vmin = 0, vmax = 1, s = 1, cmap = 'jet')
plt.colorbar()
plt.axis('scaled')
plt.subplot(1, 3, 3)
plt.title('Error')
plt.scatter(xy[:, 0], xy[:, 1], c = np.abs(pred - gt), vmin = 0, vmax = 0.2, s = 1, cmap = 'jet')
plt.colorbar()
plt.axis('scaled')
plt.tight_layout()
plt.show()
clear_output(wait = True)
Define Functions to Compute Physics Loss
The same helper as before.
def derivative(y, t):
df = torch.autograd.grad(y, t, grad_outputs = torch.ones_like(y), create_graph = True)[0]
df_x = df[:, 0:1]
df_y = df[:, 1:2]
return df_x, df_y
Define Governing Equation and Three Known Boundary Conditions
The residual function takes two derivatives in each direction and returns $\kappa(T_{xx} + T_{yy})$ with $\kappa = 1$ written explicitly as a factor. The three boundary functions evaluate the network on the bottom, right and left edges and subtract the prescribed values $0$, $1$ and $0.5$ respectively.
There is no function for the top edge. That absence is the entire content of this section, and it is worth pausing on the fact that nothing in the code will complain about it. The loss will be assembled from four terms, will descend, and will produce a field. The missing physics announces itself nowhere except in the answer.
def PDE(model, Y_domain):
T = model(Y_domain)
dT_x, dT_y = derivative(T, Y_domain)
dT_xx, _ = derivative(dT_x, Y_domain)
_, dT_yy = derivative(dT_y, Y_domain)
pde = 1 * (dT_xx + dT_yy)
return pde
def BC_bottom(model, Y_bottom):
T_bottom = model(Y_bottom)
return T_bottom
def BC_right(model, Y_right):
T_right = model(Y_right)
return T_right - 1
def BC_left(model, Y_left):
T_left = model(Y_left)
return T_left - 0.5
Train
Twenty thousand epochs with
$$ \mathcal{L} = \mathcal{L}_{\text{PDE}} + \mathcal{L}_{\text{bottom}} + \mathcal{L}_{\text{right}} + \mathcal{L}_{\text{left}}, $$
and no term for the top edge and no data. This is the pure physics baseline, and the diagnostic figure is redrawn every thousand epochs so that the field can be watched as it forms.
epochs = 20001
for epoch in range(epochs):
## PDE loss
pde_term = PDE(model, Y_domain)
loss_pde = torch.mean(pde_term**2)
## BC loss
T_bottom = BC_bottom(model, Y_bottom)
T_right = BC_right(model, Y_right)
T_left = BC_left(model, Y_left)
loss_bc = torch.mean(T_bottom**2) + torch.mean(T_right**2) + torch.mean(T_left**2)
loss = loss_pde + loss_bc
optimizer.zero_grad()
loss.backward()
optimizer.step()
if epoch % 1000 == 0:
print('EPOCH : %6d/%6d| Loss : %5.4f | Loss_PDE : %5.4f | Loss_BC : %5.4f' \
%(epoch, epochs-1, loss.item(), loss_pde.item(), loss_bc.item()))
PLOT(model, Y_domain)
# import torch
# torch.save(model, 'lab4_wo_data.pt')
# print("Model successfully saved as lab4_wo_data.pt")
Result
As shown, the prediction deviates significantly from the ground truth due to the absence of one boundary condition, and the reason is the absent condition rather than any failure of training. It is worth being precise about the two contrasting behaviours.
Most conventional computational methods cannot produce any result at all unless every boundary condition is specified, because a missing condition leaves their system of equations singular or under-determined and the solve fails outright. A PINN never assembles such a system. It minimizes a penalty, and a penalty that says nothing about the top edge simply leaves the network free there; the optimizer settles on whichever admissible field its initialization and its trajectory happen to lead it to.
This is directly analogous to an under-determined system of linear equations, which has infinitely many solutions and whose standard algorithms return one particular member of the family, typically the one of smallest norm, without any indication that others exist. A PINN behaves the same way and offers the same false comfort: the loss will look small, because the terms in it really are being satisfied, and the answer will still be wrong.
Load the model parameters
!gdown 1SuQvz8T2SsKo_5pVLvAMXySqKG9jfxiy -O /content/lab4_wo_data.pt
model = torch.load('/content/lab4_wo_data.pt', map_location = device, weights_only = False)
PLOT(model, Y_domain)
We now incorporate the readings from the nine sensors, giving the PINN more information with which to infer the missing boundary condition.
The analogy with systems of equations carries over exactly. Adding nine observations adds nine conditions on the same unknowns, pushing the problem from under-determined towards over-determined, and an over-determined system is handled by least squares, returning the solution that best satisfies all constraints simultaneously. That is precisely what minimizing a sum of squared residuals does, and it is why the sum-of-squares form of the PINN loss is the right one rather than merely a convenient one.
Nine numbers cannot determine an entire boundary profile, so we should expect improvement rather than resolution. The interesting question, which Section 5.5 answers, is how much of the profile nine interior temperatures can recover.
Train with data
The model is rebuilt from scratch so that the comparison is clean, then trained for the same twenty thousand epochs with the identical optimizer and learning rate. The loss now carries one additional term,
$$ \mathcal{L} = \mathcal{L}_{\text{PDE}} + \mathcal{L}_{\text{BC}} + \mathcal{L}_{\text{data}}, $$
with $\mathcal{L}_{\text{data}}$ the mean squared error between the network's prediction and the nine measured temperatures. Its weight is $1$, in contrast to the $10$ used in Section 4. The difference in circumstances justifies the difference in weight: there the data term was the sole means of identifying two scalars against nineteen thousand collocation points, while here it competes against a merely under-determined physics loss and is not being asked to fix a scale.
model = PINN().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr = 1e-4)
for epoch in range(epochs):
## PDE loss
pde_term = PDE(model, Y_domain)
loss_pde = torch.mean(pde_term**2)
## BC loss
T_bottom = BC_bottom(model, Y_bottom)
T_right = BC_right(model, Y_right)
T_left = BC_left(model, Y_left)
loss_bc = torch.mean(T_bottom**2) + torch.mean(T_right**2) + torch.mean(T_left**2)
## Data loss
T_obs = model(Y_obs)
loss_data = torch.mean((T_obs - gt_obs)**2)
loss = loss_pde + loss_bc + loss_data
optimizer.zero_grad()
loss.backward()
optimizer.step()
if epoch % 1000 == 0:
print('EPOCH : %6d/%6d| Loss : %5.4f | Loss_PDE : %5.4f | Loss_BC : %5.4f | Loss_DATA : %5.4f' \
%(epoch, epochs-1, loss.item(), loss_pde.item(), loss_bc.item(), loss_data.item()))
PLOT(model, Y_domain)
# import torch
# torch.save(model, 'lab4_w_data.pt')
# print("Model successfully saved as lab4_w_data.pt")
Result
The results demonstrate improved predictions, although they still do not perfectly match the ground truth. The obvious remedy is more sensors, and it would work; nine measurements are simply too few to determine a function on an edge. But there is a second kind of remedy available, and it is more interesting because it costs nothing.
Load the model parameters
!gdown 1rhxBpNgUoi0dUNFUCbEZfYL_hjuh3aiB -O /content/lab4_w_data.pt
model_w_data = torch.load('/content/lab4_w_data.pt', map_location = device, weights_only = False)
PLOT(model_w_data, Y_domain)
To make this example more interesting, let us assume that we have prior knowledge that the temperature at the missing boundary is unknown but constant. This is a very ordinary kind of engineering knowledge. An edge in contact with a large well-stirred reservoir, or with a boiling liquid, or with a block of high-conductivity metal, is at a uniform temperature whether or not anyone has measured it.
The information is qualitative, and it is not a boundary condition in the usual sense, since it names no value. It is nevertheless enormously informative. The space of possible top-edge profiles collapses from an infinite-dimensional function space to a single real number, and nine sensors are ample to determine one number.
Expressing "constant but unknown" as a loss
The implementation is three lines and is the most elegant construction in this notebook. Evaluate the network along the top edge to get $T_{\text{top}}$, compute the mean of those values, and penalize the mean squared deviation of the values from their own mean:
$$ \mathcal{L}_{\text{prior}} = \frac{1}{N}\sum_{i=1}^{N}\left(T_\theta(x_i, 1) - \bar{T}\right)^2, \qquad \bar{T} = \frac{1}{N}\sum_{j=1}^{N} T_\theta(x_j, 1). $$
This is exactly the sample variance of the predicted top-edge temperature. It is zero when and only when the profile is constant, and it says nothing whatever about which constant. The target of the mean squared error is not a fixed number supplied by the modeller but a quantity computed from the network's own output, and because it is computed with ordinary differentiable operations, gradients flow through it: the penalty pushes each edge value towards the current mean while the mean itself is free to drift to wherever the rest of the loss wants it.
The general lesson is worth extracting from the special case. Anything you know about the solution that can be written as a differentiable function of the network output, vanishing when the knowledge is respected, can be added to the loss as a term. Monotonicity, symmetry, non-negativity, a bound, an integral constraint, a known asymptotic slope: these are all expressible in the same way. The PINN loss is not restricted to differential equations. It is a place to put everything you know, which is the theme Embedding Prior Knowledge into Neural Networks develops at length.
Train again using both data and prior knowledge
The model is rebuilt once more and trained for the same twenty thousand epochs with
$$ \mathcal{L} = \mathcal{L}_{\text{PDE}} + \mathcal{L}_{\text{BC}} + \mathcal{L}_{\text{data}} + \mathcal{L}_{\text{prior}}, $$
all weights equal to one. Predict the outcome before running it. The physics fixes the field given the boundary, the prior reduces the unknown boundary to a single number, and nine sensors overdetermine that number comfortably. This model should be much better than the previous two, and it should be better mostly near the top edge, where the previous models had nothing to go on.
model = PINN().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr = 1e-4)
for epoch in range(epochs):
## PDE loss
pde_term = PDE(model, Y_domain)
loss_pde = torch.mean(pde_term**2)
## BC loss
T_bottom = BC_bottom(model, Y_bottom)
T_right = BC_right(model, Y_right)
T_left = BC_left(model, Y_left)
loss_bc = torch.mean(T_bottom**2) + torch.mean(T_right**2) + torch.mean(T_left**2)
## Data loss
T_obs = model(Y_obs)
loss_data = torch.mean((T_obs - gt_obs)**2)
## Prior Knowledge
T_top = model(Y_top)
loss_prior = torch.mean((T_top - torch.mean(T_top))**2)
loss = loss_pde + loss_bc + loss_data + loss_prior
optimizer.zero_grad()
loss.backward()
optimizer.step()
if epoch % 1000 == 0:
print('EPOCH : %6d/%6d| Loss : %5.4f | Loss_PDE : %5.4f | Loss_BC : %5.4f | Loss_DATA : %5.4f | Loss_PRIOR : %5.4f' \
%(epoch, epochs - 1, loss.item(), loss_pde.item(), loss_bc.item(), loss_data.item(), loss_prior.item()))
PLOT(model, Y_domain)
torch.save(model, 'lab4_w_prior.pt')
print('Model successfully saved as lab4_w_prior.pt')
Load the model parameters
!gdown 1OQPMZEBvz3d_MeodYeiYJBVCrmwrgLms -O /content/lab4_w_prior.pt
model_w_prior = torch.load('/content/lab4_w_prior.pt', map_location = device, weights_only = False)
PLOT(model_w_prior, Y_domain)
Result
It is clearly observed that the final prediction is the most accurate among the three different scenarios, which was predictable from an information argument alone and did not require running anything. Two conclusions follow.
The first concerns the PINN framework itself. We have now seen it absorb three different kinds of knowledge into a single objective function: a governing equation, a set of measurements, and a qualitative structural fact. All three entered in the same way, as a non-negative term that vanishes when the knowledge is respected, and all three were minimized by the same optimizer in the same loop. That uniformity is the framework's real strength, more than any individual result in this notebook.
The second is a strategy for engineering problems generally. When addressing a complex engineering challenge, begin by asking what data is available, and then, rather than stopping there, ask what else is known: governing equations, boundary conditions, conservation laws, symmetries, monotonicity, bounds, results from a simplified analysis, or the engineer's judgement that a certain surface must be at uniform temperature. Each such item is a constraint, each constraint is a term, and each term reduces the size of the space in which the answer is being sought. Data is one source of information among several, and in engineering it is frequently the scarcest.
To evaluate the models' performance more precisely, we plot the predicted temperature along the top boundary itself. This gives a direct comparison against the ground truth and shows how well each model captured the missing boundary condition.
This is the right diagnostic and the reason is worth stating. The colour maps of the interior compare fields that are dominated by the three known edges, so all three models look broadly similar in them; the disagreement is concentrated at the top, where a small region of a colour map conveys very little. Plotting the recovered profile along the unknown edge isolates the quantity that was actually being inferred and puts all three models on one set of axes against the ground truth.
The ground truth appears in the plot as a line of zeros: the true temperature along the top edge is uniformly zero. Note that this is consistent with the prior knowledge used in Section 5.4, which asserted uniformity without naming the value, and that the value itself was never supplied to any of the three models.
Predict the three curves before looking. The model without data has no information about the edge at all, so its profile can be anything and should be furthest from zero. The model with data has nine interior constraints and should be closer but uneven along the edge, since nothing forces it to be flat. The model with the prior should be nearly flat by construction, at whichever constant the sensors imply.
model_wo_data = torch.load('/content/lab4_wo_data.pt', map_location = device, weights_only = False)
model_w_data = torch.load('/content/lab4_w_data.pt', map_location = device, weights_only = False)
model_w_prior = torch.load('/content/lab4_w_prior.pt', map_location = device, weights_only = False)
x_top = Y_top[:, 0].cpu().numpy()
T_wo_data = model_wo_data(Y_top).detach().cpu().numpy().flatten()
T_w_data = model_w_data(Y_top).detach().cpu().numpy().flatten()
T_w_prior = model_w_prior(Y_top).detach().cpu().numpy().flatten()
plt.figure(figsize = (6, 4))
plt.plot(x_top, np.zeros_like(x_top), c = 'k', linestyle = '--', linewidth = 2, label = 'GT')
plt.plot(x_top, T_wo_data, c = 'b', linestyle = '--', linewidth = 2, label = 'w/o data')
plt.plot(x_top, T_w_data, c = 'g', linestyle = '--', linewidth = 2, label = 'w/ data')
plt.plot(x_top, T_w_prior, c = 'r', linewidth = 2, label = 'w/ prior')
plt.xlabel('x')
plt.ylabel('Temperature, T')
plt.title('Predicted Temperature along the Top Boundary')
plt.legend()
plt.grid(True)
plt.show()
The plot confirms the ordering, and the shapes of the curves are as informative as their distances from the truth. The unconstrained model wanders; the data-only model is pulled towards the truth in the middle of the edge, where the sensors have most influence, and drifts near the corners where they have least; the model with the prior is flat, because flatness is what the variance penalty enforces, and sits close to zero, because that is where nine interior temperatures place the constant.
One last observation to close the section. The improvement from the second model to the third came from no new measurements whatsoever. It came from writing down something the engineer already knew. That is a good note on which to leave inverse problems.
Key takeaways
- The same four-part recipe solved a fourth-order boundary value problem in solid mechanics, a first-order nonlinear initial value problem in fluid mechanics, and a coupled three-state system in epidemiology. Between them, only the residual function changed.
- The forward problem is not where a PINN earns its place. On all three problems it competes with numerical methods that are faster and more accurate. The inverse problem is different: the unknown property is one more trainable number in the same loss, and nothing else about the method has to change to accommodate it.
- Physics alone may not identify a parameter. In the beam, any $E$ satisfies the equation and all four boundary conditions if the displacement is rescaled by $1/E$; in the tank, the residual is equally content with a family of values of $C$. In both cases the loss is small and the answer is wrong — a failure that no amount of further training reveals.
- A few observations close the system, because they fix the scale that the physics leaves free. Nine readings sufficed for the beam and for the tank, and twenty noisy observations of the infected fraction alone recovered both $\beta$ and $\gamma$ in the SIR model.
- Where the observations are placed matters. The sensitivity $\partial h/\partial C$ vanishes at both ends of the drainage problem, so a sensor there would carry no information about $C$; the readings belong in the middle. The same argument puts the beam's sensors away from the supports, where the displacement is pinned to zero whatever the stiffness.
- Scales are balanced by hand. The beam's displacements are small enough that the data term needs a factor of $1000$ before it can compete with the residual, and the SIR loss weights the data and the initial conditions by $10$. These weights are chosen, not derived, and they are the most fragile part of a PINN.
- The unknown does not have to be a scalar. In the heated plate it is a function along an edge, which nine interior readings can improve but not resolve. What resolves it is a piece of qualitative knowledge, that the edge is at a uniform temperature, written as the variance of the network's own prediction along that edge — a penalty that names no value and collapses an infinite-dimensional unknown to one number.
- Two properties can be identified at once, and in a non-linear system: the density and the viscosity of the flow past a cylinder are recovered to about one percent from seventeen velocity readings, the viscosity travelling two orders of magnitude from its initial guess. The estimates are trustworthy because their traces flatten, not because the final numbers look plausible.
- Two small devices recur: learn $\log\theta$ when $\theta$ must stay positive, and add a conservation term such as $\hat S + \hat I + \hat R = 1$ when the system has an invariant that the equations only imply.
Looking ahead. Every loss in this notebook was a sum of terms whose weights were chosen by hand, and the results depended on those choices. How to set them — and how to avoid having to set them at all — is the subject of Regularization, MAP Estimation, and Loss Weighting.