Example: a PINN for the heat equation

A compact training loop for u_t = κ·u_xx on (x, t) ∈ [0,1]×[0,1]. For the step-by-step explanation, see the tutorial; this page is the condensed, copy-pasteable version.

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

KAPPA = 0.1
model = FieldMLP(in_dim=2, hidden=64, depth=4, base="tanh")
opt = torch.optim.Adam(model.parameters(), lr=1e-3)
heat = equations.Heat(kappa=KAPPA)          # residual: u_t - κ·u_xx

def sample_interior(n):
    xt = torch.rand(n, 2, requires_grad=True)
    return xt

def initial_condition(x):                   # u(x, 0) = sin(pi x)
    return torch.sin(torch.pi * x)

for step in range(2000):
    opt.zero_grad()

    xt = sample_interior(2048)
    r = heat(model, xt)                     # closed-form u_xx, autograd u_t
    loss_pde = (r ** 2).mean()

    x0 = torch.rand(256, 1)
    ic = torch.cat([x0, torch.zeros_like(x0)], dim=1)
    loss_ic = (model(ic) - initial_condition(x0)).pow(2).mean()

    tb = torch.rand(256, 1)
    left  = torch.cat([torch.zeros_like(tb), tb], dim=1)
    right = torch.cat([torch.ones_like(tb),  tb], dim=1)
    loss_bc = model(left).pow(2).mean() + model(right).pow(2).mean()

    loss = loss_pde + 10.0 * (loss_ic + loss_bc)
    loss.backward()
    opt.step()

    if step % 500 == 0:
        print(f"step {step:4d}  loss {loss.item():.3e}")

Why the spatial term is closed form

u_xx comes from the closed-form σ'' tower, not from a second autograd pass — this is what keeps the residual cheap and stable as you raise the spatial order.

PINN heat tutorial — the full explanation.
High-order PDEs guide.
Proof-carrying PDE — certify the trained residual.

Next​

Next