From a single independent variable to a two-dimensional space-time domain — how PINNs handle partial derivatives, boundary conditions, and the full field reconstruction problem.
The first two articles in this series dealt with ordinary differential equations — functions of a single independent variable, typically time. The natural world rarely cooperates with that simplification. Real engineering problems have fields: temperature distributions across a reactor wall, drug concentration profiles in a tissue slab, stress states in a structural component. These are functions of both space and time, and the equations governing them are partial differential equations.
The 1D heat (diffusion) equation is the canonical entry point:
This single equation governs an enormous range of phenomena. In pharmaceutical manufacturing: drug diffusion through membrane barriers, heat generation and dissipation in lyophilization (freeze-drying), and tablet dissolution fronts all satisfy diffusion-type equations in some spatial dimension. In battery engineering, thermal runaway analysis requires solving the heat equation across cell layers with spatially varying conductivity. In chemical engineering, tubular reactor temperature profiles with axial diffusion reduce to the same form.
The classical numerical approach is to discretise space with a finite-difference or finite-element grid and march forward in time. This works well for regular geometries and known boundary conditions. The PINN approach differs fundamentally: rather than discretising, the neural network is the solution, and training enforces the governing equation at scattered collocation points without any mesh. The payoff is flexibility in geometry and the ability to seamlessly blend sparse experimental data with physics — a capability we will exploit heavily starting in Inverse Problems.
The problem from notebook 03_heat_equation.ipynb is deliberately chosen to admit a clean analytical solution, allowing rigorous error quantification rather than just qualitative visual agreement.
with thermal diffusivity $\alpha = 0.4$ and $T = 1.0$. The parameter $\alpha$ sets the timescale of temperature decay: large $\alpha$ means fast diffusion, so the profile flattens quickly.
The initial temperature profile is a single half-sine arch, taking value zero at both ends and peaking at $x = 0.5$. The boundary conditions fix both ends permanently at zero — the rod is in contact with a heat reservoir at zero temperature.
Because the initial condition is exactly the first eigenfunction of the spatial operator $-\partial^2/\partial x^2$ on $[0,1]$ with Dirichlet boundaries, the Fourier series solution collapses to a single mode:
The spatial shape is frozen as $\sin(\pi x)$ forever; only the amplitude decays exponentially with rate $\alpha\pi^2 \approx 3.948$. At $t = 0.8$, the amplitude has decayed to $e^{-3.948 \times 0.8} \approx 0.044$ of its initial value — nearly flat. This rapid decay makes the problem a good test of whether the network correctly captures multi-scale temporal dynamics.
In the ODE notebooks the network mapped $t \mapsto u(t)$: a single scalar input. The moment we step into PDE territory, every collocation point has two coordinates $(x, t)$, and the network must accept both simultaneously. The change is minimal in code but profound in what it represents — the network is now learning a surface in $(x,t)$ space rather than a curve in $t$.
class PINN(nn.Module): def __init__(self, layers): super().__init__() seq = [] for i in range(len(layers)-1): seq.append(nn.Linear(layers[i], layers[i+1])) if i < len(layers)-2: seq.append(nn.Tanh()) self.net = nn.Sequential(*seq) 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, t): # Concatenate x and t along the feature dimension return self.net(torch.cat([x, t], dim=1)) # Architecture: [2, 64, 64, 64, 1] — 2 inputs, 3 hidden layers, 1 output model = PINN([2, 64, 64, 64, 1]).to(device) # Total parameters: 2*64 + 64*64 + 64*64 + 64*1 + biases ≈ 12,737
The key line is torch.cat([x, t], dim=1). Both x and t are tensors of shape (N, 1), and concatenating along the feature dimension produces shape (N, 2) — exactly what the first Linear(2, 64) layer expects. The rest of the forward pass is unchanged from the ODE case.
Why Tanh activations rather than ReLU? Tanh is infinitely differentiable. Since we need $\partial u/\partial t$ and $\partial^2 u/\partial x^2$ from autograd, we require two orders of differentiability everywhere. ReLU is piecewise linear: its second derivative is zero almost everywhere and undefined at the kinks, making it fundamentally unsuitable for PDE residual computation.
The PDE residual requires $\partial u / \partial t$ (first-order in time) and $\partial^2 u / \partial x^2$ (second-order in space). PyTorch's automatic differentiation engine computes these through the computational graph built during the forward pass — no finite differences, no numerical approximation.
There are two critical requirements. First, both x and t must be created with requires_grad=True so that gradients can flow back through them. Second, each autograd.grad call must use create_graph=True to keep the higher-order computation graph alive, enabling the second derivative to be computed from the first.
def pde_residual(model, x, t): u = model(x, t) # First-order time derivative ∂u/∂t u_t = torch.autograd.grad( u, t, grad_outputs=torch.ones_like(u), create_graph=True )[0] # First-order space derivative ∂u/∂x (needed for second order) u_x = torch.autograd.grad( u, x, grad_outputs=torch.ones_like(u), create_graph=True )[0] # Second-order space derivative ∂²u/∂x² u_xx = torch.autograd.grad( u_x, x, grad_outputs=torch.ones_like(u_x), create_graph=True )[0] # Residual: should be zero everywhere inside the domain return u_t - alpha * u_xx
The call to compute u_xx differentiates u_x with respect to x. For this second differentiation to be valid, the computational graph connecting u_x to x must still exist — which is why the first call must set create_graph=True. Without it, the graph is freed after computing u_x, and attempting to differentiate again raises a runtime error.
This is the most common source of confusion for newcomers to PINNs. The rule is simple: any gradient that will itself be differentiated must be computed with create_graph=True.
The PINN total loss is the sum of three mean-squared residuals, each enforcing a different constraint on the solution:
Explicitly:
def loss_fn(model, pts): x_f, t_f, x_ic, t_ic, u_ic, x_b0, x_b1, t_bc = pts # PDE residual at interior collocation points r = pde_residual(model, x_f, t_f) L_pde = torch.mean(r**2) # Initial condition: u(x,0) = sin(πx) L_ic = torch.mean((model(x_ic, t_ic) - u_ic)**2) # Boundary conditions: u(0,t) = 0 and u(1,t) = 0 L_bc = (torch.mean(model(x_b0, t_bc)**2) + torch.mean(model(x_b1, t_bc)**2)) total = L_pde + L_ic + L_bc return total, (L_pde.item(), L_ic.item(), L_bc.item())
Notice that all three terms have equal weight of 1.0 here. In practice, relative scaling matters: the IC loss starts large (the initial profile is $O(1)$) while the BC loss starts near zero if the network already predicts small values at the boundary. Unequal loss magnitudes during early training can cause gradient imbalance. For this clean problem the default equal weighting works well, but Adaptive Sampling covers sophisticated loss-weighting strategies for more difficult cases.
The training domain is now a rectangle $[0,1] \times [0,T]$. Interior points for the PDE residual are sampled uniformly at random. This is the simplest valid strategy, but it distributes effort uniformly across regions that may have very different physics. The notebook uses:
N_f = 4000 # interior collocation points (PDE residual) N_ic = 200 # points on the initial condition line t=0 N_bc = 200 # points on each boundary (x=0 and x=1) x_f = torch.rand(N_f, 1, device=device) # x ~ Uniform[0,1] t_f = torch.rand(N_f, 1, device=device) * T # t ~ Uniform[0,T] x_f.requires_grad_(True); t_f.requires_grad_(True) x_ic = torch.rand(N_ic, 1, device=device) t_ic = torch.zeros(N_ic, 1, device=device) u_ic = torch.sin(np.pi * x_ic) t_bc = torch.rand(N_bc, 1, device=device) * T x_b0 = torch.zeros(N_bc, 1, device=device) # left boundary x=0 x_b1 = torch.ones(N_bc, 1, device=device) # right boundary x=1
The model is trained with Adam for 12,000 epochs at learning rate $10^{-3}$. Collocation points are sampled once and held fixed throughout training (static sampling). The final relative $L_2$ error over the full space-time grid is typically in the range $10^{-3}$ to $10^{-4}$ — three to four orders of magnitude improvement over the initial random-initialization error.
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) EPOCHS = 12000 pts = sample() hist = [] for ep in range(EPOCHS): optimizer.zero_grad() loss, parts = loss_fn(model, pts) loss.backward(); optimizer.step() hist.append(loss.item()) if ep % 2000 == 0: print(f"epoch {ep:5d} | loss {loss.item():.3e} | " f"pde {parts[0]:.2e} ic {parts[1]:.2e} bc {parts[2]:.2e}") # Evaluate on a 100×100 space-time grid nx, nt = 100, 100 xs = np.linspace(0, 1, nx); ts = np.linspace(0, T, nt) X, Tg = np.meshgrid(xs, ts) # ... (flatten, run model, reshape) relL2 = np.linalg.norm(U - Ue) / np.linalg.norm(Ue) print(f"relative L2 error: {relL2:.3e}") # typically ~3e-4
The most intuitive view is temperature $u(x, t)$ sliced at fixed times. The four curves below show the profile at $t = 0, 0.15, 0.4, 0.8$, demonstrating rapid amplitude decay while the sinusoidal shape is preserved exactly.
The three loss terms converge at different rates. The IC loss drops rapidly in the first 4,000 epochs because the initial profile is a smooth, easily representable function. The BC loss converges almost immediately: zero-valued Dirichlet conditions are trivial for the network to enforce. The PDE residual — which must hold across 4,000 interior points and demands accurate second derivatives — converges last and slowest.
The PINN framework extends to PDEs with exactly four structural changes. Everything else — the loss formulation, the Adam training loop, the residual-penalty philosophy — remains identical.
| Aspect | ODE (notebooks 01/02) | PDE (this notebook) |
|---|---|---|
| Network input | t — scalar, shape (N,1) |
(x, t) — via torch.cat, shape (N,2) |
| Derivative types | $du/dt$, $d^2u/dt^2$ (total) | $\partial u/\partial t$, $\partial^2 u/\partial x^2$ (partial, mixed) |
| Constraints | Initial conditions (point values) | Initial field + boundary conditions on two edges |
| Collocation domain | 1D line $t \in [0,T]$ | 2D rectangle $[0,1]\times[0,T]$ |
| Loss terms | 2 (PDE + IC) | 3 (PDE + IC + BC) |
| Autograd depth | Max 2nd order in $t$ | 2nd order in $x$, 1st order in $t$; two separate chains |
The residual pattern is identical in spirit: compute the governing operator applied to the network output, square it, and minimise. The complexity lies only in which variables the derivatives are taken with respect to, not in any fundamentally new concept.
Uniform random sampling across $[0,1]\times[0,T]$ is the simplest valid choice and works well for this problem because the solution is globally smooth. But consider what happens with a problem that develops sharp gradients in a localised region — for example, a material with a sharp thermal conductivity interface, or a reaction front that moves through the domain. Uniform sampling devotes most collocation points to regions where the solution is easy and the residual is already small, wasting computational budget.
Two principled alternatives have been developed in the literature. Residual-Based Adaptive Refinement (RAR), introduced by Lu et al. (2021) in the DeepXDE library, evaluates the PDE residual magnitude on a candidate point cloud and preferentially adds new collocation points where the residual is large. This concentrates training effort at the hardest parts of the domain. Adaptive collocation with self-attention (explored in several 2022–2023 papers) learns a non-uniform sampling distribution end-to-end as part of the training procedure.
For the heat equation with its smooth single-mode solution, uniform sampling is sufficient. Adaptive Sampling in this series covers adaptive strategies in detail, using the Burgers equation (which develops a genuine near-shock) as the motivating example.
Mastering the 1D heat equation establishes the full PINN-for-PDEs vocabulary: multi-input networks, mixed partial derivatives, initial-plus-boundary constraint separation, and 2D collocation domains. The same framework scales to progressively harder problems by changing only the residual definition.
Burgers' equation (Burgers & Reaction PDEs) introduces a nonlinear convective term $u\,\partial u / \partial x$ that the PDE residual must encode. At low viscosity, the solution develops a steep internal layer — a near-shock — that stress-tests the spectral bias of deep networks and motivates the Adam $\to$ L-BFGS two-stage training recipe.
The wave equation $\partial^2 u/\partial t^2 = c^2 \partial^2 u/\partial x^2$ requires a second time derivative in the residual and a second initial condition (initial velocity field). It introduces hyperbolic PDE behavior: information propagates at finite speed, which creates challenges for the PINN if the network does not “know” about the wave speed structure from the collocation distribution.
2D and 3D heat equations on irregular geometries are where the meshless character of PINNs becomes a genuine competitive advantage. A finite-element discretisation of a complex geometry requires mesh generation, which can take hours for industrial parts. A PINN requires only the ability to sample points inside the domain and evaluate the boundary conditions — operations that are trivial with parametric geometry descriptions or even point-cloud representations.