A PINN for the heat equation

This tutorial trains a physics-informed network (PINN) for the 1-D heat equation, using omnibias's closed-form derivatives for the PDE residual.

The heat equation is:

u_t = α · u_xx

The PINN minimizes the residual u_t − α·u_xx at collocation points, plus initial / boundary conditions.

Prerequisites

pip install omnibias-pinn[torch]

Step 1 — the field and residual

omnibias-pinn ships prebuilt PDE residuals. The derivatives u_t and u_xx come from the closed-form tower, not from torch.autograd.grad through the network.

import torch
from omnibias.pinn.torch import PINNHeat, equations

model = PINNHeat(hidden=64, base="tanh", alpha=0.1)

# Collocation points in (x, t)
coords = torch.rand(4096, 2, requires_grad=False)

residual = equations.Heat(alpha=0.1)(model, coords)   # u_t - α·u_xx
loss_pde = (residual ** 2).mean()

Why no requires_grad?

Because the Laplacian and time derivative are closed form, you do not need an autograd graph through the network to get u_xx. That is the whole speed and stability win for high-order PDEs.

Step 2 — initial and boundary conditions

# Initial condition u(x, 0) = sin(πx)
x0 = torch.rand(512, 1)
t0 = torch.zeros_like(x0)
ic = model(torch.cat([x0, t0], dim=1)) - torch.sin(torch.pi * x0)
loss_ic = (ic ** 2).mean()

# Dirichlet boundaries u(0, t) = u(1, t) = 0
tb = torch.rand(512, 1)
xb0 = torch.zeros_like(tb); xb1 = torch.ones_like(tb)
bc = model(torch.cat([xb0, tb], 1)) ** 2 + model(torch.cat([xb1, tb], 1)) ** 2
loss_bc = bc.mean()

Step 3 — train

opt = torch.optim.Adam(model.parameters(), lr=1e-3)

for step in range(2000):
    opt.zero_grad()
    loss = loss_pde + 10.0 * loss_ic + 10.0 * loss_bc
    loss.backward()
    opt.step()

(In a real run, resample collocation points each step and track the residual on a held-out grid.)

Step 4 — diagnostics

omnibias-pinn provides forecast-horizon, relative-L2-per-time, and spectral diagnostics so you can see where the residual is large, not just its mean.

from omnibias.pinn.torch import diagnostics
report = diagnostics.relative_l2_per_time(model, reference_solution)

Going further

• Higher-order PDEs. Swap Heat for Biharmonic, KuramotoSivashinsky, or CahnHilliard. The closed-form path keeps Δᵏ flat — see High-order PDEs.
• Hard conservation. Enforce invariants by construction with a cage instead of a penalty — see Structural cages.
• Certify the residual. Turn a trained surrogate into a certified statement — see Proof-carrying PDE.