Neural Networks for Differential Equations in Physics: Full Guide

Physics is written in the language of differential equations. Neural networks can now solve them — from the heat equation to the Schrödinger equation — with a technique called Physics-Informed Neural Networks (PINNs). This is the complete guide: theory, implementation, worked examples, and where the field is going next.

1. Why Neural Networks for DEs?
2. The Mathematical Foundation
3. Automatic Differentiation Explained
4. Building Your First PINN from Scratch
5. Solving the Heat Equation
6. Solving the Schrödinger Equation
7. Inverse Problems: Learning Physics from Data
8. DeepXDE: The PINN Framework
9. Neural ODEs
10. When PINNs Fail — and What To Do

import torch
import torch.nn as nn

# Any differentiable function — could be a neural network
def u_network(x, t):
    return torch.sin(x) * torch.exp(-t)   # toy example

# Create input tensors with requires_grad=True
# This tells PyTorch to track all operations for differentiation
x = torch.linspace(0, 1, 100, requires_grad=True).reshape(-1,1)
t = torch.linspace(0, 1, 100, requires_grad=True).reshape(-1,1)

u = u_network(x, t)

# First derivative ∂u/∂x
du_dx = torch.autograd.grad(
    u, x,
    grad_outputs=torch.ones_like(u),
    create_graph=True   # CRITICAL: lets us differentiate again
)[0]

# Second derivative ∂²u/∂x²  (differentiate du_dx again)
d2u_dx2 = torch.autograd.grad(
    du_dx, x,
    grad_outputs=torch.ones_like(du_dx),
    create_graph=True
)[0]

# First derivative ∂u/∂t
du_dt = torch.autograd.grad(
    u, t,
    grad_outputs=torch.ones_like(u),
    create_graph=True
)[0]

print(f"u shape:       {u.shape}")
print(f"du_dx shape:   {du_dx.shape}")
print(f"d2u_dx2 shape: {d2u_dx2.shape}")
print(f"du_dt shape:   {du_dt.shape}")

import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt

# ── 1. Generate synthetic "experimental" data ─────────────────
k_true = 2.5                          # true decay rate (unknown to the PINN)
t_data = torch.tensor([[0.0],[0.5],[1.0],[1.5],[2.0]])
u_data = torch.exp(-k_true * t_data) + 0.02*torch.randn_like(t_data)

# Collocation points — where we enforce the ODE
t_phys = torch.linspace(0, 2, 1000).reshape(-1,1)

# ── 2. Define the PINN architecture ──────────────────────────
class PINN(nn.Module):
    def __init__(self):
        super().__init__()
        # Network approximates u(t)
        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, 1)
        )
        # k is a LEARNABLE PARAMETER — this is the inverse problem part
        self.k = nn.Parameter(torch.tensor([1.0]))

    def forward(self, t):
        return self.net(t)

    def ode_residual(self, t):
        t = t.clone().requires_grad_(True)
        u = self.forward(t)
        du_dt = torch.autograd.grad(
            u, t, torch.ones_like(u), create_graph=True
        )[0]
        return du_dt + self.k * u   # du/dt + k*u = 0

# ── 3. Training loop ─────────────────────────────────────────
model  = PINN()
optim  = torch.optim.Adam(model.parameters(), lr=1e-3)
mse    = nn.MSELoss()
losses = []

for epoch in range(8000):
    optim.zero_grad()

    # Data loss — match measurements
    loss_data = mse(model(t_data), u_data)

    # Physics loss — enforce ODE at collocation points
    residual   = model.ode_residual(t_phys)
    loss_phys  = torch.mean(residual**2)

    # IC loss — enforce u(0) = 1
    t0         = torch.tensor([[0.0]])
    loss_ic    = (model(t0) - 1.0)**2

    loss = loss_data + 0.1*loss_phys + 10.0*loss_ic
    loss.backward()
    optim.step()
    losses.append(loss.item())

    if (epoch+1) % 2000 == 0:
        print(f"Epoch {epoch+1}: loss={loss.item():.6f}  k={model.k.item():.4f}")

print(f"\nFinal: k = {model.k.item():.4f}  (true k = {k_true})")

import torch, torch.nn as nn
import numpy as np

alpha = 0.4   # thermal diffusivity

class HeatPINN(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(2, 64), nn.Tanh(),   # 2 inputs: x and t
            nn.Linear(64, 128), nn.Tanh(),
            nn.Linear(128, 128), nn.Tanh(),
            nn.Linear(128, 64), nn.Tanh(),
            nn.Linear(64, 1)
        )

    def forward(self, x, t):
        xt = torch.cat([x, t], dim=1)
        return self.net(xt)

    def pde_residual(self, x, t):
        x = x.clone().requires_grad_(True)
        t = t.clone().requires_grad_(True)
        u        = self.forward(x, t)
        du_dt    = torch.autograd.grad(u, t, torch.ones_like(u), create_graph=True)[0]
        du_dx    = torch.autograd.grad(u, x, torch.ones_like(u), create_graph=True)[0]
        d2u_dx2  = torch.autograd.grad(du_dx, x, torch.ones_like(du_dx), create_graph=True)[0]
        return du_dt - alpha * d2u_dx2   # should = 0

# ── Collocation points (random in domain) ───────────────────
N_col  = 5000
x_col  = torch.rand(N_col, 1)
t_col  = torch.rand(N_col, 1)

# ── Boundary conditions: u=0 at x=0 and x=1 ────────────────
N_bc   = 200
t_bc   = torch.rand(N_bc, 1)
x_bc0  = torch.zeros(N_bc, 1)    # x=0
x_bc1  = torch.ones(N_bc, 1)     # x=1

# ── Initial condition: u(x,0) = sin(pi*x) ───────────────────
N_ic   = 200
x_ic   = torch.rand(N_ic, 1)
t_ic   = torch.zeros(N_ic, 1)
u_ic   = torch.sin(np.pi * x_ic)

# ── Train ────────────────────────────────────────────────────
model = HeatPINN()
optim = torch.optim.Adam(model.parameters(), lr=1e-3)
mse   = nn.MSELoss()

for epoch in range(15000):
    optim.zero_grad()
    loss_pde = torch.mean(model.pde_residual(x_col, t_col)**2)
    loss_bc  = mse(model(x_bc0, t_bc), torch.zeros_like(t_bc)) + \
               mse(model(x_bc1, t_bc), torch.zeros_like(t_bc))
    loss_ic  = mse(model(x_ic, t_ic), u_ic)
    loss = loss_pde + 10*loss_bc + 10*loss_ic
    loss.backward()
    optim.step()

# ── Verify against analytic solution ────────────────────────
x_test = torch.linspace(0,1,100).reshape(-1,1)
t_test = torch.full_like(x_test, 0.5)   # evaluate at t=0.5
u_pred = model(x_test, t_test).detach().numpy()
u_exact= np.sin(np.pi * x_test.numpy()) * np.exp(-alpha * 0.5 * np.pi**2)
err    = np.max(np.abs(u_pred - u_exact))
print(f"Max error at t=0.5: {err:.6f}")

# Schrödinger PINN — particle in a box (V=0 inside, infinite outside)
# Units: hbar=1, 2m=1. Domain: x in [0,1], t in [0,T]

class SchrodingerPINN(nn.Module):
    def __init__(self):
        super().__init__()
        # Two outputs: real (u) and imaginary (v) parts of psi
        self.net = nn.Sequential(
            nn.Linear(2, 100), nn.Tanh(),
            nn.Linear(100, 100), nn.Tanh(),
            nn.Linear(100, 100), nn.Tanh(),
            nn.Linear(100, 2)   # output: [u, v]
        )

    def forward(self, x, t):
        xt = torch.cat([x, t], dim=1)
        return self.net(xt)   # returns [u, v] per point

    def pde_residual(self, x, t):
        x = x.clone().requires_grad_(True)
        t = t.clone().requires_grad_(True)
        psi   = self.forward(x, t)
        u, v  = psi[:, 0:1], psi[:, 1:2]

        def D(f, var): return torch.autograd.grad(f, var, torch.ones_like(f), create_graph=True)[0]

        u_t   = D(u, t);  v_t  = D(v, t)
        u_xx  = D(D(u, x), x);  v_xx = D(D(v, x), x)

        # Coupled real PDEs from Schrodinger (V=0, hbar=1, 2m=1)
        res_u = u_t + 0.5 * v_xx   # ∂u/∂t = -(1/2)∂²v/∂x²
        res_v = v_t - 0.5 * u_xx   # ∂v/∂t = +(1/2)∂²u/∂x²
        return res_u, res_v

    def probability_conservation(self, x, t):
        # |psi|^2 should integrate to 1 at every time step
        psi  = self.forward(x, t)
        prob = psi[:, 0]**2 + psi[:, 1]**2
        return (prob.mean() - 1.0)**2   # normalization constraint

# pip install deepxde
import deepxde as dde
import numpy as np

# ── Solve 1D diffusion-reaction: du/dt = D*d2u/dx2 - k*u ────
D, k = 1.0, 1.0

def pde(x, y):
    # x[:,0]=position, x[:,1]=time.  y=u.  dde computes gradients automatically
    dy_t  = dde.grad.jacobian(y, x, i=0, j=1)    # du/dt
    dy_xx = dde.grad.hessian(y, x, i=0, j=0)     # d2u/dx2
    return dy_t - D*dy_xx + k*y

# ── Domain: x in [-1,1], t in [0,1] ─────────────────────────
geom    = dde.geometry.Interval(-1, 1)
timedomain = dde.geometry.TimeDomain(0, 1)
geomtime    = dde.geometry.GeometryXTime(geom, timedomain)

# ── Boundary and initial conditions ─────────────────────────
def boundary(x, on_boundary): return on_boundary
def ic_func(x): return np.sin(np.pi * x[:, 0:1])

bc = dde.DirichletBC(geomtime, lambda x: 0, boundary)
ic = dde.IC(geomtime, ic_func, lambda x, on_initial: on_initial)

# ── Assemble problem and train ───────────────────────────────
data  = dde.data.TimePDE(
    geomtime, pde, [bc, ic],
    num_domain=3000,      # collocation points inside domain
    num_boundary=300,    # boundary points
    num_initial=300      # initial condition points
)
net   = dde.nn.FNN([2, 64, 64, 64, 1], "tanh", "Glorot normal")
model = dde.Model(data, net)

model.compile("adam", lr=1e-3)
model.train(iterations=20000)
model.compile("L-BFGS")        # refine with L-BFGS after Adam
model.train()

print("Training complete. Predict with model.predict(X_test)")

# pip install torchdiffeq
from torchdiffeq import odeint
import torch, torch.nn as nn

# ── Neural ODE that learns the dynamics of a damped oscillator ─
class ODEFunc(nn.Module):
    """Neural network that approximates the RHS of dh/dt = f(h, t)"""
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(2, 64),  nn.Tanh(),
            nn.Linear(64, 64),  nn.Tanh(),
            nn.Linear(64, 2)    # output: [dx/dt, dv/dt]
        )

    def forward(self, t, h):
        return self.net(h)   # autonomous system (no explicit t)

# ── Generate training data from a known damped oscillator ──────
def true_dynamics(t, h):
    x, v = h[:, 0], h[:, 1]
    return torch.stack([v, -0.3*v - 4.0*x], dim=1)   # γ=0.3, ω₀²=4

t_span    = torch.linspace(0, 5, 100)
h0        = torch.tensor([[1.0, 0.0]])   # x(0)=1, v(0)=0
traj_true = odeint(true_dynamics, h0, t_span)  # [T, batch, 2]

# ── Train the Neural ODE to match the trajectory ───────────────
odefunc = ODEFunc()
optim   = torch.optim.Adam(odefunc.parameters(), lr=1e-3)

for epoch in range(3000):
    optim.zero_grad()
    traj_pred = odeint(odefunc, h0, t_span)   # predict trajectory
    loss      = nn.MSELoss()(traj_pred, traj_true)
    loss.backward()
    optim.step()
    if (epoch+1) % 1000 == 0:
        print(f"Epoch {epoch+1}: loss = {loss.item():.6f}")

# The learned odefunc.net now encodes the damped oscillator dynamics
# Extrapolate beyond training range — physics generalises!
t_extrap  = torch.linspace(0, 10, 200)  # 2x training range
traj_ext  = odeint(odefunc, h0, t_extrap)

class FourierFeatureNet(nn.Module):
    def __init__(self, sigma=10.0, n_freq=128):
        super().__init__()
        # Random frequency matrix — fixed, not trained
        self.register_buffer('B', sigma * torch.randn(2, n_freq))  # 2 inputs: x,t
        self.net = nn.Sequential(
            nn.Linear(2*n_freq, 128), nn.Tanh(),
            nn.Linear(128, 128),    nn.Tanh(),
            nn.Linear(128, 1)
        )

    def encode(self, x, t):
        xt  = torch.cat([x, t], dim=1)         # [N, 2]
        proj = 2 * np.pi * xt @ self.B          # [N, n_freq]
        return torch.cat([proj.sin(), proj.cos()], dim=1)  # [N, 2*n_freq]

    def forward(self, x, t):
        return self.net(self.encode(x, t))

• Raissi, Perdikaris & Karniadakis (2019) — Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear PDEs. Journal of Computational Physics. doi.org/10.1016/j.jcp.2018.10.045 — The original PINN paper. Start here.
• Chen et al. (2018) — Neural Ordinary Differential Equations. NeurIPS. arXiv:1806.07366 — The Neural ODE paper. One of the most cited ML papers of the decade.
• Lu et al. (2021) — DeepXDE: A deep learning library for solving differential equations. SIAM Review. doi.org/10.1137/19M1274067 — The DeepXDE framework paper.
• Wang, Teng & Perdikaris (2022) — Understanding and mitigating gradient pathologies in physics-informed neural networks. SIAM J. Sci. Comput. arXiv:2001.04536 — Essential reading on PINN failure modes and NTK-based fixes.
• Wang et al. (2022) — Respecting causality is all you need for training physics-informed neural networks. arXiv:2203.07404 — Causal training for time-dependent PDEs.
• Tancik et al. (2020) — Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains. NeurIPS. arXiv:2006.10739 — The Fourier feature embedding trick that fixes spectral bias.
• DeepXDE documentation — deepxde.readthedocs.io — Comprehensive tutorials for 40+ PDE types, inverse problems, and complex geometries.

• PINNs solve the forward problem (given a PDE, find u(x,t)) and the inverse problem (given data + a PDE structure, find unknown parameters) simultaneously.
• Three loss terms. Data loss + physics residual + boundary conditions. Weighting matters enormously.
• Automatic differentiation via PyTorch autograd computes exact PDE residuals — no finite differences needed.
• DeepXDE abstracts all the boilerplate. Use it for production. Write from scratch to understand internals.
• Neural ODEs parameterise the dynamics themselves as a neural network — powerful for learning equations of motion from data.
• Know the failure modes: loss imbalance, spectral bias, causality violation. All have known fixes.
• Don't use PINNs for simple problems. Use classical solvers when they're sufficient. PINNs earn their complexity on inverse problems, irregular domains, and high-dimensional PDEs.