50.2 Physics-Informed Learning

1. The Inverse Problem Framework

When a jet engine blade cracks mid-flight, regulators need the thermal diffusivity of the failed alloy, but no sensor survives inside a running turbine to measure it. Pure curve-fitting on the sparse temperature readings will happily produce a smooth interpolation that violates the heat equation and returns a physically meaningless diffusivity. This is the problem that physics-informed inverse methods were built to solve.

A forward problem takes known parameters and initial/boundary conditions and computes the solution of a PDE. An inverse problem takes (possibly sparse and noisy) observations of the solution and infers unknown parameters, boundary conditions, or even the form of the governing equation.

An inverse problem deduces causes from observed effects: given measurements of a system's behavior (temperatures, displacements, concentrations), recover the hidden parameters or laws that produced that behavior. Scientists and engineers care most about quantities like material properties, reaction rates, and source locations, yet these are almost never measured directly. Indirect observations carry the signal. The core mechanism embeds known physics (a PDE or conservation law) as a constraint inside an optimization loop. The optimizer then searches only the subspace of solutions consistent with the governing equations. Prefer an inverse problem formulation over pure curve fitting whenever a known or partially known physical model exists. The physics constraint reduces the data needed and prevents nonphysical solutions. Resort to purely data-driven methods only when no reliable physical model is available.

Consider the heat equation in one spatial dimension:

In the forward problem, \(\alpha\) (the thermal diffusivity) is known, and we solve for \(u(x, t)\). In the inverse problem, \(\alpha\) is unknown, and we have temperature measurements \(\{(x_i, t_i, u_i)\}_{i=1}^N\) at scattered locations. The goal is to recover \(\alpha\) from the data.

The PINN approach treats both the solution \(u(x, t)\) and the unknown parameter \(\alpha\) as optimization targets. The network \(\hat{u}_\theta(x, t)\) approximates the temperature field, while \(\alpha\) is an additional trainable parameter. The total loss function combines three terms. Figure 50.2.1 illustrates PINN inverse problem architecture.

where:

The PDE residual is evaluated at collocation points \((x_j, t_j)\) sampled throughout the domain. These points do not require data; the PDE itself provides the supervision signal. Automatic differentiation (covered in Chapter 42) computes the required partial derivatives \(\partial \hat{u}/\partial t\) and \(\partial^2 \hat{u}/\partial x^2\) exactly through the network. In short: the PDE does not just constrain the answer; it teaches the network everywhere the data cannot reach. The diagram in Figure 50.2a below illustrates this architecture.

2. Inverse Heat Equation with DeepXDE

DeepXDE provides a high-level API for both forward and inverse PINN problems. The following code recovers the thermal diffusivity \(\alpha\) from sparse temperature measurements using the inverse heat equation.

import deepxde as dde
import numpy as np

# True diffusivity (what we want to recover)
alpha_true = 0.4

# Define the computational domain: x in [0, 1], t in [0, 1]
geom = dde.geometry.Interval(0, 1)
timedomain = dde.geometry.TimeDomain(0, 1)
geomtime = dde.geometry.GeometryXTime(geom, timedomain)

# The unknown parameter: initialized far from the true value
alpha = dde.Variable(1.0)  # trainable scalar

# Define the PDE residual
def heat_equation(x, u):
    """PDE residual: du/dt - alpha * d2u/dx2 = 0"""
    du_t = dde.grad.jacobian(u, x, i=0, j=1)   # du/dt
    du_xx = dde.grad.hessian(u, x, i=0, j=0)    # d2u/dx2
    return du_t - alpha * du_xx

# Exact solution for generating observation data
def exact_solution(x):
    """u(x,t) = sin(pi*x) * exp(-alpha_true * pi^2 * t)"""
    return np.sin(np.pi * x[:, 0:1]) * np.exp(
        -alpha_true * np.pi**2 * x[:, 1:2]
    )

# Boundary conditions: u(0,t) = u(1,t) = 0
bc = dde.icbc.DirichletBC(
    geomtime, lambda x: 0, lambda _, on_boundary: on_boundary
)

# Initial condition: u(x,0) = sin(pi*x)
ic = dde.icbc.IC(
    geomtime,
    lambda x: np.sin(np.pi * x[:, 0:1]),
    lambda _, on_initial: on_initial,
)

# Sparse observation data (simulating sensor measurements)
n_observations = 50
rng = np.random.default_rng(42)
observe_x = rng.uniform(0.1, 0.9, (n_observations, 1))
observe_t = rng.uniform(0.1, 0.9, (n_observations, 1))
observe_pts = np.hstack([observe_x, observe_t])
observe_u = exact_solution(observe_pts)

# Add 2% Gaussian noise to observations
observe_u += 0.02 * rng.standard_normal(observe_u.shape)

# PointSetBC enforces data matching at observation points
observe_bc = dde.icbc.PointSetBC(observe_pts, observe_u, component=0)

# Assemble the PDE problem
data = dde.data.TimePDE(
    geomtime,
    heat_equation,
    [bc, ic, observe_bc],
    num_domain=2000,     # collocation points for PDE residual
    num_boundary=100,
    num_initial=100,
    num_test=500,
)

# Network architecture: 2 inputs (x, t) -> 1 output (u)
net = dde.nn.FNN(
    layer_sizes=[2, 64, 64, 64, 1],
    activation="tanh",
    kernel_initializer="Glorot normal",
)

# Build and train the model
model = dde.Model(data, net)

# Phase 1: Adam optimizer for broad convergence
model.compile(
    "adam",
    lr=1e-3,
    external_trainable_variables=[alpha],
)
losshistory, train_state = model.train(
    iterations=10000,
    display_every=2000,
)

# Phase 2: L-BFGS (a quasi-Newton optimizer that approximates
# curvature for faster convergence) for fine-tuning
model.compile(
    "L-BFGS",
    external_trainable_variables=[alpha],
)
losshistory, train_state = model.train()

# Check recovered parameter
alpha_recovered = alpha.numpy()
print(f"True alpha:      {alpha_true:.4f}")
print(f"Recovered alpha: {alpha_recovered:.4f}")
print(f"Relative error:  {abs(alpha_recovered - alpha_true) / alpha_true:.2%}")

On this problem, the two-phase training strategy (Adam for 10,000 iterations followed by L-BFGS for fine-tuning) typically recovers \(\alpha\) with relative error below 1%, even with 2% observation noise and only 50 measurement points. The key is that 2,000 collocation points enforce the heat equation throughout the domain, providing orders of magnitude more supervision than the 50 data points alone. (Fifty noisy measurements plus one known PDE outperform thousands of label-only data points, because every collocation point is a free physics lesson.)

3. Inverse PINNs in JAX

DeepXDE provides convenience, but understanding the mechanics requires seeing the raw computation. Here is the same inverse problem implemented directly in JAX, exposing every gradient computation.

import jax
import jax.numpy as jnp
from jax import random, grad, jit, vmap
import optax

def init_network(key, layer_sizes):
    """Initialize a simple feedforward network."""
    params = []
    for in_dim, out_dim in zip(layer_sizes[:-1], layer_sizes[1:]):
        key, subkey = random.split(key)
        w = random.normal(subkey, (in_dim, out_dim)) * jnp.sqrt(2.0 / in_dim)
        key, subkey = random.split(key)
        b = jnp.zeros(out_dim)
        params.append((w, b))
    return params

def forward(params, x, t):
    """Network forward pass: (x, t) -> u."""
    h = jnp.stack([x, t], axis=-1)
    for w, b in params[:-1]:
        h = jnp.tanh(h @ w + b)
    w, b = params[-1]
    return (h @ w + b).squeeze(-1)

def pde_residual(params, alpha, x, t):
    """Compute du/dt - alpha * d2u/dx2 at a single point."""
    u_fn = lambda x_, t_: forward(params, x_, t_)

    # First derivatives
    du_dt = grad(u_fn, argnums=1)(x, t)

    # Second derivative w.r.t. x
    du_dx = grad(u_fn, argnums=0)
    d2u_dx2 = grad(du_dx, argnums=0)(x, t)

    return du_dt - alpha * d2u_dx2

# Vectorize over batches of points
pde_residual_batch = vmap(pde_residual, in_axes=(None, None, 0, 0))

@jit
def loss_fn(params, log_alpha, x_data, t_data, u_data,
            x_colloc, t_colloc, x_bc, t_bc):
    """Total PINN loss: data + PDE residual + boundary."""
    alpha = jnp.exp(log_alpha)  # ensure positivity

    # Data loss
    u_pred = vmap(forward, in_axes=(None, 0, 0))(params, x_data, t_data)
    loss_data = jnp.mean((u_pred - u_data)**2)

    # PDE residual loss
    residuals = pde_residual_batch(params, alpha, x_colloc, t_colloc)
    loss_pde = jnp.mean(residuals**2)

    # Boundary loss (u = 0 at x=0 and x=1)
    u_bc = vmap(forward, in_axes=(None, 0, 0))(params, x_bc, t_bc)
    loss_bc = jnp.mean(u_bc**2)

    return loss_data + 10.0 * loss_pde + 10.0 * loss_bc

# Training loop
key = random.PRNGKey(0)
params = init_network(key, [2, 64, 64, 64, 1])
log_alpha = jnp.log(jnp.array(1.0))  # initial guess

optimizer = optax.adam(1e-3)
opt_state = optimizer.init((params, log_alpha))

grad_fn = jit(grad(loss_fn, argnums=(0, 1)))

for step in range(10000):
    grads = grad_fn(
        params, log_alpha,
        x_data, t_data, u_data,
        x_colloc, t_colloc, x_bc, t_bc,
    )
    updates, opt_state = optimizer.update(grads, opt_state)
    params, log_alpha = optax.apply_updates((params, log_alpha), updates)

    if step % 2000 == 0:
        alpha_est = float(jnp.exp(log_alpha))
        print(f"Step {step:5d} | alpha = {alpha_est:.4f}")

print(f"\nRecovered alpha: {float(jnp.exp(log_alpha)):.4f}")
print(f"True alpha:      {alpha_true:.4f}")

The JAX implementation makes several design choices visible. We parameterize \(\alpha\) as \(\exp(\log \alpha)\) to enforce positivity without constrained optimization. The vmap transform vectorizes the per-point PDE residual over the collocation batch. And grad computes the exact derivatives \(\partial u / \partial t\) and \(\partial^2 u / \partial x^2\) through the network, with no finite differences and no discretization error.

4. Noether's Theorem as a Structural Prior

Recovering unknown coefficients is only half the story; we can also embed the fundamental symmetries of physics directly into the network architecture, so the model cannot produce solutions that violate conservation laws in the first place.

Noether's theorem is one of the deepest results in theoretical physics: every continuous symmetry of a physical system corresponds to a conserved quantity. Time-translation symmetry gives conservation of energy. Spatial-translation symmetry gives conservation of momentum. Rotational symmetry gives conservation of angular momentum. For discovery AI, Noether's theorem provides structural priors that constrain the space of possible physical laws.

The connection to machine learning is through Lagrangian and Hamiltonian neural networks. Instead of learning the dynamics \(\dot{x} = f(x)\) directly, we learn the Lagrangian \(L(q, \dot{q})\) or Hamiltonian \(H(q, p)\), and derive the dynamics through the Euler-Lagrange or Hamilton's equations (variational and symplectic formulations that express a system's time evolution purely in terms of energy functions rather than forces). This architectural choice automatically ensures energy conservation, because the Hamiltonian is a constant of motion by construction.

import jax
import jax.numpy as jnp
from jax import grad, jit, vmap

def hamiltonian_nn(params, q, p):
    """
    Neural network that outputs a scalar Hamiltonian H(q, p).
    The dynamics are derived from Hamilton's equations:
        dq/dt =  dH/dp
        dp/dt = -dH/dq
    """
    x = jnp.concatenate([q, p])
    for w, b in params[:-1]:
        x = jnp.tanh(x @ w + b)
    w, b = params[-1]
    H = (x @ w + b).squeeze()
    return H

def hamiltonian_dynamics(params, q, p):
    """
    Derive equations of motion from the Hamiltonian.
    Energy conservation is guaranteed by construction.
    """
    dH_dq = grad(hamiltonian_nn, argnums=1)(params, q, p)
    dH_dp = grad(hamiltonian_nn, argnums=2)(params, q, p)

    dq_dt = dH_dp    # Hamilton's first equation
    dp_dt = -dH_dq   # Hamilton's second equation

    return dq_dt, dp_dt

def rollout(params, q0, p0, dt, n_steps):
    """Symplectic Euler integration (a structure-preserving integrator
    that updates momentum and position in sequence to conserve energy
    over long time horizons) preserving Hamiltonian structure."""
    q, p = q0, p0
    trajectory = [(q, p)]

    for _ in range(n_steps):
        dq, dp = hamiltonian_dynamics(params, q, p)
        p = p + dp * dt      # update momentum first
        dq, _ = hamiltonian_dynamics(params, q, p)
        q = q + dq * dt      # then update position
        trajectory.append((q, p))

    return trajectory

The Hamiltonian Neural Network (HNN) of Greydanus et al. (2019) preserves energy by construction (up to integration error), producing closed orbits that match the true physics over thousands of time steps. A standard neural ODE (where a neural network directly parameterizes the time derivative \(\dot{x} = f_\theta(x)\) without any structural physics constraint), by contrast, drifts in energy and spirals into unphysical trajectories.

5. Simulation-Based Inference

Not every inverse problem in physics has a tractable likelihood function. Many settings provide a forward simulator (particle physics Monte Carlo, cosmological N-body simulation, climate model) that maps parameters \(\theta\) to observables \(x\). Yet evaluating \(p(x | \theta)\) analytically is impossible. Simulation-based inference (SBI) addresses this by learning the posterior \(p(\theta | x)\) directly from simulated parameter-observation pairs.

The Three Main SBI Approaches

Neural Posterior Estimation (NPE) trains a conditional density estimator (typically a normalizing flow) to approximate \(p(\theta | x)\) from simulated pairs \(\{(\theta_i, x_i)\}\). Given a new observation \(x_{\text{obs}}\), the trained network immediately produces the posterior without additional simulations.

Neural Likelihood Estimation (NLE) learns \(p(x | \theta)\) instead, then combines with a prior via Bayes' rule. This is useful when the prior changes between analyses but the likelihood does not.

Neural Ratio Estimation (NRE) learns the likelihood-to-evidence ratio \(r(x, \theta) = p(x | \theta) / p(x)\) using a binary classifier, avoiding explicit density estimation entirely.

import torch
import numpy as np

# Example: inferring parameters of a damped oscillator
# from noisy trajectory observations

def simulator(theta):
    """
    Forward simulator: damped oscillator.
    theta = [amplitude, damping, frequency]
    Returns a noisy trajectory (100 time points).
    """
    A, gamma, omega = theta
    t = np.linspace(0, 10, 100)
    x = A * np.exp(-gamma * t) * np.cos(omega * t)
    # Add observation noise
    x += 0.05 * np.random.randn(len(t))
    return x

def generate_training_data(n_simulations=10000):
    """Generate parameter-observation pairs for SBI training."""
    prior_low = np.array([0.5, 0.1, 1.0])
    prior_high = np.array([3.0, 1.0, 8.0])

    thetas = np.random.uniform(prior_low, prior_high,
                                size=(n_simulations, 3))
    observations = np.array([simulator(theta) for theta in thetas])

    return (
        torch.tensor(thetas, dtype=torch.float32),
        torch.tensor(observations, dtype=torch.float32),
    )

# Using the sbi library for Neural Posterior Estimation
# (sbi v0.22+; earlier versions used a slightly different API)
from sbi.inference import SNPE
from sbi.utils import BoxUniform

# Define prior
prior = BoxUniform(
    low=torch.tensor([0.5, 0.1, 1.0]),
    high=torch.tensor([3.0, 1.0, 8.0]),
)

# Generate training simulations
thetas, observations = generate_training_data(10000)

# Train the neural posterior estimator
inference = SNPE(prior=prior)
inference.append_simulations(thetas, observations)
density_estimator = inference.train()

# Build the posterior
posterior = inference.build_posterior(density_estimator)

# Given a real observation, sample from the posterior
x_observed = torch.tensor(simulator([2.0, 0.3, 4.0]),
                           dtype=torch.float32)
samples = posterior.sample((5000,), x=x_observed)

print(f"Posterior mean: {samples.mean(dim=0).numpy()}")
print(f"Posterior std:  {samples.std(dim=0).numpy()}")
print(f"True values:    [2.0, 0.3, 4.0]")

The SBI approach is particularly powerful for physics problems where the forward model is expensive (each simulation takes minutes to hours) but we need to perform inference for many different observations. Once the density estimator is trained, posterior evaluation is instantaneous, amortizing the cost of the training simulations across all future inferences.

6. Failure Modes and Mitigations

The power of PINNs and SBI comes with practical pitfalls that can silently undermine results if left unaddressed.

PINNs are not a universal solution. Several well-documented failure modes require awareness and mitigation strategies.

Spectral bias. Neural networks with smooth activation functions learn low-frequency components of the solution first, struggling with sharp gradients, shocks, and high-frequency oscillations. For the Burgers equation with viscosity approaching zero, standard PINNs tend to fail severely. Mitigation: use Fourier feature embeddings (Tancik et al., 2020) or multi-scale architectures that explicitly represent high-frequency content.

Imbalanced loss terms. The data loss, PDE residual, and boundary losses operate on different scales. If the PDE residual dominates, the network satisfies the equation approximately but ignores the data. If the data loss dominates, the network interpolates the data but violates the physics between measurement points. Mitigation: use adaptive loss weighting (Wang et al., 2021) or the neural tangent kernel perspective (a theoretical framework that analyzes neural network training dynamics through the lens of kernel methods, revealing how different loss components compete for gradient updates) to balance gradient magnitudes across loss terms.

Optimization landscape. The composite loss function can have many local minima, saddle points, and flat regions. The two-phase training strategy (Adam followed by L-BFGS) helps but does not guarantee convergence to the global minimum. Mitigation: use curriculum training (start with easier sub-problems), causal training (enforce temporal causality), or ensemble methods to explore the loss landscape more thoroughly.

# Adaptive loss weighting (Learning Rate Annealing)
# following Wang et al. (2021)

def adaptive_weights(loss_data, loss_pde, loss_bc,
                     prev_weights, eta=0.1):
    """
    Update loss weights based on gradient statistics.
    Balances gradient magnitudes across loss components.
    """
    # Compute gradient norms for each loss term
    grad_data = jax.grad(lambda p: loss_data(p))(params)
    grad_pde = jax.grad(lambda p: loss_pde(p))(params)
    grad_bc = jax.grad(lambda p: loss_bc(p))(params)

    norm_data = jnp.sqrt(sum(
        jnp.sum(g**2) for g in jax.tree_util.tree_leaves(grad_data)
    ))
    norm_pde = jnp.sqrt(sum(
        jnp.sum(g**2) for g in jax.tree_util.tree_leaves(grad_pde)
    ))
    norm_bc = jnp.sqrt(sum(
        jnp.sum(g**2) for g in jax.tree_util.tree_leaves(grad_bc)
    ))

    # Target: equal gradient contributions
    max_norm = jnp.maximum(norm_data, jnp.maximum(norm_pde, norm_bc))
    w_data = max_norm / (norm_data + 1e-8)
    w_pde = max_norm / (norm_pde + 1e-8)
    w_bc = max_norm / (norm_bc + 1e-8)

    # Exponential moving average for stability
    w_data = (1 - eta) * prev_weights[0] + eta * w_data
    w_pde = (1 - eta) * prev_weights[1] + eta * w_pde
    w_bc = (1 - eta) * prev_weights[2] + eta * w_bc

    return jnp.array([w_data, w_pde, w_bc])

7. When to Use Each Approach

The choice between PINNs, SBI, and classical inverse methods depends on the problem structure:

For physics discovery, PINNs and SBI are complementary rather than competing. PINNs excel when the governing PDE is known and the unknown is a small number of parameters. SBI excels when the forward model is a complex simulator and you need full posterior uncertainty quantification. In Section 50.4, we combine both: PySR discovers the equation form, a PINN recovers unknown coefficients, and Bayesian optimization explores the design space.