What it actually does
A neural network is a differentiable function, so you can compute its derivatives exactly with the same machinery that computes its gradients. That means you can write down a differential equation, evaluate how badly the network violates it at a scatter of points, and make that violation the loss. The result is a mesh-free solver whose solution is a closed-form function you can differentiate anywhere — and, more usefully, a framework where unknown coefficients in the equation are just more parameters to fit alongside the network.
- The governing equation is known and correct — the load-bearing assumption, and where it fails every residual is wrong in the same systematic direction
- The solution is smooth enough for a neural network to represent, which is false at a genuine discontinuity
- The domain of interest is fixed and known before training, since the network has no validity outside it
- The loss terms can be balanced so every constraint is meaningfully enforced, which is a tuning problem rather than a given
- Collocation points cover the domain densely enough that satisfying the equation there implies satisfying it everywhere
Loss function
Three losses summed, and the summation is the problem. These terms have different units, different scales and different curvature, so the weights are not a regularization choice but a determination of which constraint the optimizer actually respects. Get them wrong and the network satisfies the boundary conditions beautifully while ignoring the equation, or satisfies the equation with a solution that drifts from the boundary — both of which look like convergence. The second thing to notice is what the residual term needs: derivatives of the network with respect to its inputs, not its parameters. A second-order PDE requires second derivatives of the output with respect to the input, and then gradients of those with respect to the weights, which is why a PINN step costs several times what a same-sized regression step costs.
- the differential operator — the PDE, written as something that should equal zero
- the network, standing in for the solution field as a closed-form differentiable function
- collocation points: locations where the equation is enforced and nothing was ever measured
- the loss weights, and the single most consequential set of numbers in the method
- unknown physical coefficients, optimized jointly with the network — the inverse problem
How it is minimized
Adam followed by L-BFGS, with adaptive loss weighting and collocation resampling
- gradient of the residual term; its magnitude relative to the others is what the weights must equalize
- the weight-update rate, kept small so the balance moves slowly relative to the weights themselves
- derivatives of the network output with respect to its INPUTS, obtained by automatic differentiation
- a physical coefficient; known in the forward problem, and a fitted parameter in the inverse one
Two things make the optimization unusual and both are worth separating. The first is that gradients of the different loss terms routinely differ by orders of magnitude, so a fixed weighting silently selects which constraint gets satisfied; adaptive schemes rebalance the weights from observed gradient magnitudes, which converts a hyperparameter search into a running measurement. The second is the two-stage schedule, which is not folklore: Adam is robust enough to get out of the terrible initial region, and L-BFGS then exploits curvature to reach the sharp minima that a first-order method approaches only slowly. The gap between stopping after Adam and running L-BFGS to convergence is frequently an order of magnitude in residual. Underneath both sits spectral bias — networks fit low frequencies first — which is exactly the wrong inductive bias for a solution with sharp gradients, and it is why a PINN can converge smoothly to a plausible solution that is simply wrong in the shock region.
loss weights
Which constraint is actually enforced. The most consequential choice in the method, and adaptive schemes exist because fixed ones fail (1 to 1000, or adaptive)
collocation points
Where the equation is enforced. More improves accuracy and cost linearly; placement matters more than count (1e3 to 1e5)
network depth and width
Capacity for the solution field. Usually modest — 4 to 8 layers of 20 to 100 units (4-8 layers, 20-100 wide)
activation
Must be smooth to the order of the PDE. ReLU has zero second derivative and cannot express a second-order residual at all (tanh / sin / GELU)
Adam steps before L-BFGS
The handover point. Too early and L-BFGS stalls in a bad region; too late and Adam wastes time (1e4 to 5e4)
Fourier feature scale
Input encoding that counteracts spectral bias; frequently the difference between failure and success on sharp solutions (1 to 10)
No guarantee of anything, and a distinctive set of failures that all look like success. The most important is that the loss is a poor proxy for solution accuracy: a residual of 1e-5 is compatible with a solution that is qualitatively wrong, because the residual is measured at collocation points and a network can be nearly right there while being badly wrong between them. The second is failure on sharp features — spectral bias means high frequencies arrive last if at all, so shocks and boundary layers are smoothed into something plausible and incorrect, and the loss curve gives no hint. The third is loss-term imbalance, where one term dominates and the others are effectively unconstrained; the diagnostic is to track the terms separately, which is trivial and routinely skipped. Underlying all of it is the honest comparison: on a well-posed forward problem with a standard geometry, a classical finite-element or spectral solver is faster by orders of magnitude and comes with error bounds, and a PINN is the wrong tool.
O(N_c · d · P) per step, where the derivative order multiplies the cost of a plain forward pass: a second-order PDE needs second derivatives of the output with respect to the input, then gradients of those with respect to the weights, which is typically three to five times the cost of an equivalent regression step. Fully parallel across collocation points, and with no sequential structure anywhere.
How, where, and why it is used
Time-Series Forecasting
AdaptedTreat time as one of the input coordinates and solve on a space-time domain, so the network gives the field at any future instant inside that domain. Sparse measurements enter as a data term while the equation constrains everything between and beyond them, which is what lets the method work with far fewer observations than a purely data-driven forecaster would need.
- Forecasting a physical process with a known governing equation and very sparse sensors
- Digital twins where the model must respect conservation laws rather than merely fit history
- Filling gaps in space-time observation records under a physical constraint
- Hybrid settings where a known equation is augmented by a learned closure term
The argument for it is real but narrow: where a governing equation is genuinely known, it constrains the solution in regions with no data at all, which is something no statistical forecaster can do — and with a handful of sensors that is the difference between a solvable and an unsolvable problem. Against it, two things are decisive. Extrapolation beyond the trained space-time domain is not supported: the network was fitted on that domain and behaves arbitrarily outside it, so "forecasting" means solving on a domain that already includes the future, which must be decided before training. And where no governing equation is known, which is nearly all forecasting, the method has nothing to offer. This is a tool for physics, not for demand.
- Normalize the space and time coordinates to comparable ranges, or the derivative terms differ by orders of magnitude before training starts
- Include the forecast window in the training domain, since the network cannot extrapolate past the collocation region
- Weight the sparse data term against the residual explicitly and monitor both, since one silently dominates otherwise
- Use a smooth activation to the order of the equation — a ReLU network has zero second derivative and cannot represent a second-order residual at all
Error against held-out measurements inside the domain, and separately the residual on a dense grid of points never used for training — a low training residual with a high held-out residual means the network satisfied the equation only where it was asked to. Compare against a classical solver wherever one exists, because that comparison is usually unflattering and is the honest test of whether this approach was warranted.
- Treating the trained domain boundary as extrapolatable, which it is not in any sense
- A low residual masking a qualitatively wrong solution, which is the characteristic failure of the method
- Loss-term imbalance leaving the sparse data term effectively unconstrained
- Applying it where no governing equation is known, where it reduces to an ordinary and worse regressor
Anomaly Detection
AdaptedOnce a network is fitted to normal operation under a governing equation, the PDE residual evaluated at new measurements is a physically meaningful anomaly score: a reading that violates conservation of mass or energy is anomalous for a reason you can state, rather than because it is statistically unusual. That interpretability is the distinctive benefit, and it is rare.
- Sensor fault detection in instrumented physical systems, where a violated conservation law localizes the faulty instrument
- Leak and blockage detection in flow networks, where the governing equation is well established
- Structural health monitoring against a known mechanical model
- Process monitoring in chemical and thermal plants where the balance equations are documented
The score means something, which no statistical detector can claim: "this reading is inconsistent with conservation of mass" is an explanation an engineer can act on, where "this reading is in the 0.1st percentile of a learned density" is not. It also distinguishes a faulty sensor from a genuine process change, because the two violate the physics differently. Against it: it requires a governing equation you actually trust, and if the model is wrong every residual is wrong in the same direction; the network must be refitted when operating conditions move outside the trained domain; and on any system without a documented physical model this has nothing to offer that a much cheaper detector does not.
- Validate the governing equation against known-good operation first, since a wrong model makes every residual meaningless in the same direction
- Normalize the residual per equation term, so a large-magnitude term does not dominate the score by units alone
- Keep the per-term residual rather than the total, since which term is violated is what identifies the fault
- Refit when operating conditions leave the trained domain, because the network is not valid outside it
Precision and recall on confirmed faults, plus an explicit check that the implicated equation term matches the physical root cause — that second check is the whole value proposition, and a detector that flags the right events for the wrong reason has not delivered it.
- A wrong or incomplete governing equation making every residual systematically misleading
- Operating conditions drifting outside the trained domain, where the network is simply undefined
- Reporting a scalar residual and losing the per-term attribution that made the method worth using
- Using it where no trusted physical model exists, which removes the only advantage it has
Optimization
Primary fitA PINN is a PDE solver implemented as an optimization problem, so the subject is optimization almost end to end. The forward problem replaces a linear solve with gradient descent on a residual; the inverse problem, which is where the method genuinely earns its place, makes unknown physical coefficients parameters in that same optimization, so parameter estimation and field reconstruction happen in one fit rather than in a nested loop. Around this sit the multi-objective weighting problem, the two-stage Adam-then-L-BFGS schedule, and spectral bias as a structural limitation of the hypothesis class.
- Multi-objective loss balancing where the terms have incommensurable units — the defining difficulty here
- Inverse problems solved as joint optimization rather than as nested forward solves, which is the method’s strongest use
- First-order to second-order handover, where Adam escapes the bad region and L-BFGS exploits curvature
- Spectral bias as an inductive-bias limitation, and Fourier features as a targeted fix rather than a general one
Worth studying because it is the clearest example in this reference of a method whose headline framing and whose actual value diverge. As a forward solver it loses badly to classical numerical methods — slower by orders of magnitude, with no error bounds, on problems those methods solve reliably. As an inverse solver it is genuinely strong, because adding unknown coefficients to an existing optimization is nearly free while the classical equivalent requires an outer loop around a full solve. Knowing which problem you have is the whole decision, and the multi-objective weighting difficulty is the transferable lesson: when a loss sums terms with different units, the weights are choosing the answer.
- Track each loss term separately at all times; a total loss hides that one term is unconstrained, and this is the cheapest diagnostic available
- Use adaptive weighting from measured gradient magnitudes rather than searching fixed weights, which converts tuning into measurement
- Run L-BFGS after Adam — the residual improvement is frequently an order of magnitude and is routinely left on the table
- Add Fourier features when the solution has sharp gradients, since spectral bias otherwise smooths them into something plausible and wrong
Compare against a classical solver on any problem where one applies — that is the honest benchmark and the one most PINN evaluations omit. For inverse problems, report parameter recovery error rather than residual, since that is what the method is actually for. Always evaluate the residual on points never used in training.
- Using it as a forward solver where finite elements are faster, more accurate and come with error bounds
- Treating the loss as a proxy for solution accuracy, which it is not — a low residual is compatible with a qualitatively wrong solution
- Fixed loss weights silently determining which constraint is satisfied
- Stopping after Adam and leaving an order of magnitude of residual improvement unclaimed
Control & Operations
ViableUsed as a fast surrogate for a physical simulation inside a control or optimization loop: train once on the governing equations across a parameter range, then evaluate in microseconds where the classical solver takes seconds. The network is differentiable with respect to its inputs, so gradients of the objective with respect to control variables come for free.
- Real-time surrogates for thermal, flow or structural simulation inside a controller
- Design optimization where the physical solve sits inside the objective and must be evaluated thousands of times
- Digital twins requiring physically consistent state estimates between sparse measurements
- Model-predictive control over a system with a known continuous physical model
The economics are the argument: training is expensive and evaluation is nearly free, which is exactly the right shape when the solve sits inside a loop evaluated thousands of times. Differentiability with respect to inputs is a genuine second benefit, since it gives design sensitivities without finite differencing. Against it: the surrogate is valid only inside the parameter range it was trained on and degrades silently outside it, there is no error bound anywhere, and a controller acting on a wrong surrogate has no way to detect that it is wrong. In safety-critical settings that absence of a bound is usually disqualifying.
- Train across the full parameter range the controller will explore, since the surrogate has no validity outside it and no way to say so
- Validate against the classical solver on a dense grid, which is the only error estimate available
- Monitor whether inputs at serving time fall inside the trained range, and abstain rather than extrapolate
- Prefer a hybrid form where a known part of the physics is computed exactly and only the expensive part is surrogated
Error against the classical solver across the parameter range, reported at the worst case rather than the mean — a controller meets the worst case, not the average. Then closed-loop performance, which is the number that matters and which diverges from surrogate accuracy more than expected.
- Silent degradation outside the trained parameter range, with no in-model way to detect it
- No error bound, which is frequently disqualifying in safety-critical control
- Optimizing mean surrogate accuracy when the controller is exposed to the worst case
What it costs to run
Minutes to hours on a single GPU for a modest domain, but with a wide spread: derivative order multiplies the per-step cost, and difficult problems need very many steps. The L-BFGS stage is expensive per iteration and usually worth it.
One forward pass per query point and trivially batched — microseconds, and mesh-free, so the solution can be evaluated anywhere in the domain at any resolution without re-solving. This asymmetry between expensive training and nearly free evaluation is the whole argument for the surrogate use case.
Per problem rather than on a schedule. A new geometry, a new boundary condition or a new parameter range is a new training run — there is no notion of updating a PINN with new data the way a statistical model is updated.
- Track each loss term separately and permanently; a total loss hides that one constraint has become unconstrained, and this is the cheapest diagnostic there is
- Evaluate the residual on points never used for training, since a network can satisfy the equation only where it was asked to
- Check at serving time that query points lie inside the trained domain, and abstain rather than extrapolate
- Compare against a classical solver wherever one exists, as the only available error estimate
- The trained domain is a hard boundary. Outside it the network is not approximate, it is undefined — and it will still return a confident number
- The activation must be smooth to the order of the PDE. A ReLU network has zero second derivative everywhere, so a second-order residual is identically zero and the model trains to a meaningless optimum
- A low loss does not mean an accurate solution. This is the single most important thing to internalize about the method, and it catches experienced people
- Loss weights are part of the model. Different weights give a different solution to the same equation, and the weights must be versioned with the checkpoint
- Input normalization interacts with the derivative terms, since scaling a coordinate scales its derivatives by the same factor to the order of the equation
Pros and cons, in context
Inverse problems come nearly free
Unknown physical coefficients are just more parameters in the same optimization, so parameter estimation and field reconstruction happen in one fit. The classical equivalent needs an outer loop around a full solve, and this is where the method is genuinely strong rather than merely interesting.
Mesh-free, with a closed-form differentiable solution
No mesh generation, and the solution can be evaluated and differentiated anywhere in the domain at any resolution. Meshing is a real cost in complex geometries, and this sidesteps it entirely.
Physics constrains regions with no data
With a handful of sensors the equation determines the field between them, which no statistical model can do. This is the difference between a solvable and an unsolvable problem in sparsely instrumented systems.
Anomaly scores that mean something physical
"This violates conservation of mass" is an explanation an engineer can act on, and the violated term identifies the fault. Rare among detectors and genuinely valuable where a trusted model exists.
Expensive to train, nearly free to evaluate
Exactly the right shape for a surrogate inside an optimization or control loop evaluated thousands of times. The asymmetry is the whole argument for that use case.
A low loss does not mean an accurate solution
The residual is measured at collocation points and a network can be nearly right there while being badly wrong between them. The most important and most under-appreciated property of the method, and it catches experienced practitioners.
Loses badly to classical solvers on forward problems
On a well-posed problem with standard geometry, finite elements are orders of magnitude faster and come with error bounds. Choosing a PINN there is a mistake, and the literature does not always say so plainly.
Multi-objective weighting determines the answer
Terms with different units summed together mean the weights choose which constraint is satisfied. Not a regularization knob — a modelling decision that must be measured rather than guessed.
Spectral bias fails on sharp features
Networks fit low frequencies first, so shocks and boundary layers are smoothed into something plausible and wrong, with no signal in the loss curve. Fourier features help and do not eliminate it.
No validity outside the trained domain and no error bound anywhere
The network returns a confident number for any input, including ones it was never fitted on. In engineering contexts the absence of an error bound is frequently disqualifying on its own.
Intuitive, idiomatic, optimized
The same algorithm in three languages, each progressing through Kent Beck's three stages. Every optimization names the lever it pulls and the trade-off it accepts — cost profiles are illustrative, not measured benchmarks.
The math, transcribed literally. No dependencies.
"""A PINN for Burgers' equation, with the derivatives written out by hand.
The equation:
u_t + u * u_x - nu * u_xx = 0 on x in [-1, 1], t in [0, 1]
u(x, 0) = -sin(pi x) initial condition
u(-1, t) = u(1, t) = 0 boundary conditions
The loss is the residual of that equation at COLLOCATION POINTS - locations
where nothing was ever measured - plus the boundary and initial conditions.
In the forward problem there is no data anywhere.
A real implementation gets u_x and u_xx from automatic differentiation. Here
they are differentiated by hand, which is only possible because the network is
one tanh layer - and which makes visible exactly what the residual needs:
derivatives with respect to the INPUTS, not the parameters.
"""
import math
import random
SEED = 5
NU = 0.01 / math.pi
def forward(x, t, params):
"""u(x, t) through one tanh layer. Returns the value and the hidden state.
The hidden activations come back because every derivative below reuses
them - recomputing the forward pass for each derivative would triple the
cost of the residual.
"""
hidden = []
for i in range(len(params['w_x'])):
pre = params['w_x'][i] * x + params['w_t'][i] * t + params['b'][i]
hidden.append(math.tanh(pre))
value = params['b_out']
for i, h in enumerate(hidden):
value += params['w_out'][i] * h
return value, hidden
def derivatives(x, t, params, hidden):
"""u_t, u_x and u_xx, by the chain rule, using the cached activations.
For h = tanh(a): dh/da = 1 - h^2, d2h/da2 = -2h(1 - h^2)
Both are functions of h alone, which is why the hidden state is enough and
the pre-activations need not be stored. The ACTIVATION MUST BE SMOOTH to
the order of the equation: a ReLU network has zero second derivative
everywhere, so u_xx is identically zero and a second-order residual is
trivially satisfied by a meaningless solution.
"""
u_x = 0.0
u_t = 0.0
u_xx = 0.0
for i, h in enumerate(hidden):
first = 1.0 - h * h
second = -2.0 * h * first
u_x += params['w_out'][i] * first * params['w_x'][i]
u_t += params['w_out'][i] * first * params['w_t'][i]
u_xx += params['w_out'][i] * second * params['w_x'][i] * params['w_x'][i]
return u_t, u_x, u_xx
def residual(x, t, params):
"""How badly the network violates the equation at one point.
This IS the training signal. No label is involved anywhere.
"""
value, hidden = forward(x, t, params)
u_t, u_x, u_xx = derivatives(x, t, params, hidden)
return u_t + value * u_x - NU * u_xx
def initial_condition(x):
return -math.sin(math.pi * x)
def sample_collocation(count, rng):
"""Points where the equation is enforced. Nothing is measured here."""
return [(rng.uniform(-1.0, 1.0), rng.uniform(0.0, 1.0)) for _ in range(count)]
def total_loss(params, collocation, boundary, initial, weights):
"""Three terms with DIFFERENT UNITS, summed.
That summation is the central difficulty of the whole method. The residual
term carries units of du/dt; the boundary term carries units of u. Their
gradients routinely differ by orders of magnitude, so the weights are not
a regularization choice - they determine which constraint the optimizer
actually respects.
"""
residual_loss = 0.0
for x, t in collocation:
r = residual(x, t, params)
residual_loss += r * r
residual_loss /= len(collocation)
boundary_loss = 0.0
for x, t in boundary:
value, _ = forward(x, t, params)
boundary_loss += value * value # u = 0 on both boundaries
boundary_loss /= len(boundary)
initial_loss = 0.0
for x in initial:
value, _ = forward(x, 0.0, params)
difference = value - initial_condition(x)
initial_loss += difference * difference
initial_loss /= len(initial)
total = (
weights['residual'] * residual_loss
+ weights['boundary'] * boundary_loss
+ weights['initial'] * initial_loss
)
# The per-term values come back too. A total loss hides that one term has
# become unconstrained, and that is the cheapest diagnostic available.
return total, residual_loss, boundary_loss, initial_loss
def numerical_gradient(params, key, index, collocation, boundary, initial, weights,
epsilon=1e-6):
"""Finite differences over the parameters.
Enormously wasteful - two full loss evaluations per parameter - and here
only to keep the mechanism visible. Note the nesting this exposes: the
loss already contains derivatives with respect to the INPUTS, so a real
implementation differentiates through those to get parameter gradients,
which is why a PINN step costs several times a regression step.
"""
original = params[key][index] if isinstance(params[key], list) else params[key]
def set_value(value):
if isinstance(params[key], list):
params[key][index] = value
else:
params[key] = value
set_value(original + epsilon)
high, _, _, _ = total_loss(params, collocation, boundary, initial, weights)
set_value(original - epsilon)
low, _, _, _ = total_loss(params, collocation, boundary, initial, weights)
set_value(original)
return (high - low) / (2.0 * epsilon)
def train(width=20, steps=2000, learning_rate=0.01, num_collocation=200):
rng = random.Random(SEED)
params = {
'w_x': [rng.uniform(-1.0, 1.0) for _ in range(width)],
'w_t': [rng.uniform(-1.0, 1.0) for _ in range(width)],
'b': [0.0] * width,
'w_out': [rng.uniform(-1.0, 1.0) for _ in range(width)],
'b_out': 0.0,
}
collocation = sample_collocation(num_collocation, rng)
boundary = [(-1.0, rng.uniform(0.0, 1.0)) for _ in range(50)]
boundary += [(1.0, rng.uniform(0.0, 1.0)) for _ in range(50)]
initial = [rng.uniform(-1.0, 1.0) for _ in range(50)]
weights = {'residual': 1.0, 'boundary': 1.0, 'initial': 1.0}
for step in range(steps):
for key in ('w_x', 'w_t', 'b', 'w_out'):
for index in range(width):
gradient = numerical_gradient(
params, key, index, collocation, boundary, initial, weights
)
params[key][index] -= learning_rate * gradient
return params
def residual_on_held_out(params, rng, count=1000):
"""The check that matters, and the one people skip.
A network can satisfy the equation at its collocation points while being
badly wrong between them. A low TRAINING residual with a high held-out
residual is exactly that failure, and nothing else reveals it.
"""
points = sample_collocation(count, rng)
total = 0.0
for x, t in points:
r = residual(x, t, params)
total += r * r
return total / len(points)
O(N_c * W) per loss evaluation, multiplied by two evaluations per parameter for the finite-difference gradient — so O(N_c * W^2) per step. Illustrative, not a measured benchmark: the quadratic factor is entirely an artifact of numerical differentiation, and it is what makes the literal version unusable past a handful of hidden units.
How this language actually wants it written.
"""The same PINN, with the loss terms as a type and the weighting measured.
Two things change. The loss terms become a structure that keeps them
separate, because summing quantities with different units is the central
difficulty of the method and a scalar total hides which constraint is
actually being enforced. And forward-mode automatic differentiation replaces
hand-written derivatives, which is what makes the approach work for any
equation rather than only for the one it was derived for.
"""
from __future__ import annotations
import math
import random
from dataclasses import dataclass
from typing import Callable, NamedTuple, Sequence
class NonSmoothActivation(ValueError):
"""Raised when the activation cannot support the equation's order.
Its own type because the failure is silent and total: a ReLU network has
zero second derivative everywhere, so a second-order residual is
identically zero and the model trains happily to a meaningless optimum.
"""
class DomainViolation(ValueError):
"""Raised when a query point lies outside the trained domain.
Outside its collocation region the network is not approximate - it is
undefined. It will still return a confident number, which is why this is
checked rather than trusted.
"""
class ShapeMismatch(ValueError):
"""Raised on a dimension violation instead of computing nonsense."""
class Dual(NamedTuple):
"""A second-order dual number: value, first derivative, second derivative.
Forward-mode automatic differentiation, carried explicitly. This is what
replaces hand-differentiating the network: the arithmetic rules below
propagate derivatives through any composition, so the residual can be
written for ANY equation rather than only the one the derivatives were
derived for.
"""
value: float
first: float
second: float
def __add__(self, other: Dual | float) -> Dual:
if isinstance(other, Dual):
return Dual(self.value + other.value, self.first + other.first,
self.second + other.second)
return Dual(self.value + other, self.first, self.second)
def __mul__(self, other: Dual | float) -> Dual:
if not isinstance(other, Dual):
return Dual(self.value * other, self.first * other, self.second * other)
# Product rule, and the second derivative of a product needs the
# cross term - the place a hand-rolled implementation goes wrong.
return Dual(
self.value * other.value,
self.first * other.value + self.value * other.first,
self.second * other.value + 2.0 * self.first * other.first
+ self.value * other.second,
)
def tanh(self) -> Dual:
activation = math.tanh(self.value)
first = 1.0 - activation * activation
second = -2.0 * activation * first
# Chain rule to second order: (f o g)'' = f''(g) g'^2 + f'(g) g''
return Dual(
activation,
first * self.first,
second * self.first * self.first + first * self.second,
)
@dataclass(frozen=True)
class Domain:
"""The region the network is valid on. Frozen, because it is a hard edge."""
x_min: float
x_max: float
t_min: float
t_max: float
def __post_init__(self) -> None:
if self.x_min >= self.x_max or self.t_min >= self.t_max:
raise ShapeMismatch('domain bounds must be strictly increasing')
def contains(self, x: float, t: float) -> bool:
return self.x_min <= x <= self.x_max and self.t_min <= t <= self.t_max
def assert_contains(self, x: float, t: float) -> None:
"""Guard clause: outside the domain the network is undefined, not
approximate — and it will still return a confident number."""
if not self.contains(x, t):
raise DomainViolation(
f'({x}, {t}) lies outside the trained domain; the network has no '
'validity there and extrapolation is not supported'
)
def sample(self, count: int, rng: random.Random) -> list[tuple[float, float]]:
return [
(rng.uniform(self.x_min, self.x_max), rng.uniform(self.t_min, self.t_max))
for _ in range(count)
]
class LossTerms(NamedTuple):
"""Each term kept separate, with its weight.
Returning a scalar total is the mistake that makes PINNs hard to debug: a
falling total is compatible with one term being entirely unconstrained,
and tracking the terms separately is the cheapest diagnostic available.
"""
residual: float
boundary: float
initial: float
data: float = 0.0
def weighted_total(self, weights: LossWeights) -> float:
return (
weights.residual * self.residual
+ weights.boundary * self.boundary
+ weights.initial * self.initial
+ weights.data * self.data
)
def imbalance_ratio(self) -> float:
"""How far apart the terms are. Above a couple of orders of magnitude
the smaller ones are effectively unconstrained."""
values = [v for v in (self.residual, self.boundary, self.initial, self.data) if v > 0.0]
return max(values) / min(values) if len(values) > 1 else 1.0
@dataclass
class LossWeights:
"""Not a regularization knob — a determination of which constraint holds.
The terms carry different units: the residual is in du/dt, the boundary
term is in u. Their gradients routinely differ by orders of magnitude, so
the weights choose the answer rather than tuning it.
"""
residual: float = 1.0
boundary: float = 1.0
initial: float = 1.0
data: float = 1.0
def rebalance(self, gradient_norms: LossTerms, rate: float = 0.1) -> None:
"""Adaptive weighting from measured gradient magnitudes.
Converts a hyperparameter search into a running measurement: scale
each weight so every term contributes comparable gradient magnitude.
The rate is kept small so the balance moves slowly relative to the
weights it is balancing.
"""
reference = gradient_norms.residual
if reference <= 0.0:
return
for name in ('boundary', 'initial', 'data'):
observed = getattr(gradient_norms, name)
if observed <= 0.0:
continue
target = reference / observed
current = getattr(self, name)
setattr(self, name, (1.0 - rate) * current + rate * target)
class PinnNetwork:
"""One tanh layer, evaluated through dual numbers for exact derivatives."""
SMOOTH_ACTIVATIONS = frozenset({'tanh', 'sin', 'gelu'})
def __init__(
self,
width: int,
domain: Domain,
activation: str = 'tanh',
seed: int = 5,
) -> None:
# Guard clause for the failure with no other symptom.
if activation not in self.SMOOTH_ACTIVATIONS:
raise NonSmoothActivation(
f'{activation} is not smooth enough for a second-order equation; '
'a ReLU network has zero second derivative and the residual '
'becomes identically zero'
)
rng = random.Random(seed)
self._domain = domain
self._w_x = [rng.uniform(-1.0, 1.0) for _ in range(width)]
self._w_t = [rng.uniform(-1.0, 1.0) for _ in range(width)]
self._bias = [0.0] * width
self._w_out = [rng.uniform(-1.0, 1.0) for _ in range(width)]
self._b_out = 0.0
def evaluate(self, x: float, t: float, differentiate_in_x: bool = True) -> Dual:
"""u(x, t) as a dual number, carrying derivatives in x OR t.
The seed determines which variable is differentiated: first derivative
one, second derivative zero for the variable of interest, both zero
for the other. That is the whole of forward-mode AD.
"""
x_dual = Dual(x, 1.0 if differentiate_in_x else 0.0, 0.0)
t_dual = Dual(t, 0.0 if differentiate_in_x else 1.0, 0.0)
total = Dual(self._b_out, 0.0, 0.0)
for index in range(len(self._w_x)):
pre = (x_dual * self._w_x[index]) + (t_dual * self._w_t[index]) + self._bias[index]
total = total + pre.tanh() * self._w_out[index]
return total
def residual(self, x: float, t: float, viscosity: float) -> float:
"""u_t + u u_x - nu u_xx, from two dual evaluations.
Two passes rather than one: forward-mode AD carries derivatives with
respect to one input at a time, so a two-variable equation costs two
forward passes. That is the structural cost of the method, and it is
why reverse mode is preferred for many inputs and forward mode here.
"""
self._domain.assert_contains(x, t)
in_x = self.evaluate(x, t, differentiate_in_x=True)
in_t = self.evaluate(x, t, differentiate_in_x=False)
return in_t.first + in_x.value * in_x.first - viscosity * in_x.second
def compute_loss(
network: PinnNetwork,
collocation: Sequence[tuple[float, float]],
boundary: Sequence[tuple[float, float]],
initial: Sequence[float],
initial_condition: Callable[[float], float],
viscosity: float,
) -> LossTerms:
"""Each term computed and returned separately, never pre-summed."""
if not collocation:
raise ShapeMismatch('a PINN needs at least one collocation point')
residual_loss = math.fsum(
network.residual(x, t, viscosity) ** 2 for x, t in collocation
) / len(collocation)
boundary_loss = math.fsum(
network.evaluate(x, t).value ** 2 for x, t in boundary
) / max(len(boundary), 1)
initial_loss = math.fsum(
(network.evaluate(x, 0.0).value - initial_condition(x)) ** 2 for x in initial
) / max(len(initial), 1)
return LossTerms(residual=residual_loss, boundary=boundary_loss, initial=initial_loss)
def held_out_residual(
network: PinnNetwork,
domain: Domain,
viscosity: float,
rng: random.Random,
count: int = 1000,
) -> float:
"""The check that matters, and the one people skip.
A network can satisfy the equation at its collocation points while being
badly wrong between them. A low training residual with a high held-out
residual is exactly that failure, and nothing else reveals it — the loss
curve looks identical either way.
"""
points = domain.sample(count, rng)
return math.fsum(network.residual(x, t, viscosity) ** 2 for x, t in points) / count
What changed: Two changes carry the weight. Forward-mode automatic differentiation through second-order dual numbers replaces the hand-written derivatives, which is what makes the method general rather than tied to one equation — and writing the arithmetic rules out shows where a hand-rolled version goes wrong, since the second derivative of a product needs a cross term that is easy to omit. The dual seed makes explicit that forward mode carries derivatives with respect to one input at a time, so a two-variable equation costs two forward passes: that is the structural cost of the approach rather than an implementation detail. The second change is that the loss terms are kept separate in a structure with an imbalance ratio, because summing quantities with different units is the central difficulty of the method and a scalar total hides that one constraint has become unconstrained; the adaptive rebalancing then converts a hyperparameter search into a running measurement of gradient magnitudes. Around those, the domain becomes a frozen type that refuses queries outside it — the network is undefined there rather than approximate, and it will still return a confident number — and the activation is checked against the equation order, because a ReLU network makes a second-order residual identically zero and trains happily to a meaningless optimum.
Forward-mode AD costs a constant factor over a plain forward pass — roughly three times for second order — against the two full loss evaluations per parameter that numerical differentiation needed. Illustrative, not a measured benchmark: the change is asymptotic in parameter count, not constant.
Faster, and the justification for each change.
"""Batched collocation and vectorized derivatives. The point loop disappears.
The structural observation: collocation points are entirely independent, so
the whole residual evaluation is one batched forward pass. And because the
network is a composition of smooth functions, the derivative rules vectorize
exactly like the forward pass does - a second derivative is another array
expression, not another loop.
Three changes:
1. Every point at once. The hidden pre-activations become a matrix and the
derivative chain rules apply elementwise across it.
2. Derivative reuse: tanh, its first derivative and its second derivative
are all functions of the activation alone, so one tanh evaluation serves
all three and the transcendental is computed once.
3. Gradient magnitudes per loss term are measured, not guessed, because
that is what adaptive weighting needs and it is nearly free here.
"""
from __future__ import annotations
import numpy as np
from numpy.typing import NDArray
FLOAT = np.float64 # PINNs are genuinely precision-sensitive; see the trade-off
class BatchedPinn:
"""One tanh layer, evaluated for every collocation point simultaneously."""
def __init__(self, width: int, seed: int = 5) -> None:
rng = np.random.default_rng(seed)
self.w_x = rng.uniform(-1.0, 1.0, width).astype(FLOAT)
self.w_t = rng.uniform(-1.0, 1.0, width).astype(FLOAT)
self.bias = np.zeros(width, dtype=FLOAT)
self.w_out = rng.uniform(-1.0, 1.0, width).astype(FLOAT)
self.b_out = FLOAT(0.0)
def forward_with_derivatives(
self, x: NDArray[np.float64], t: NDArray[np.float64]
) -> tuple[NDArray[np.float64], ...]:
"""u, u_t, u_x and u_xx for every point, in ONE pass.
The key reuse: for h = tanh(a), both dh/da = 1 - h^2 and
d2h/da2 = -2h(1 - h^2) are functions of h ALONE. So one tanh
evaluation - the only transcendental here - serves the value and both
derivative orders, and the chain rules become elementwise products.
"""
# (points, width) pre-activations as one outer-product-style GEMM.
pre = np.outer(x, self.w_x)
pre += np.outer(t, self.w_t)
pre += self.bias
activation = np.tanh(pre)
# In place, reusing the activation buffer's sibling rather than
# allocating: first = 1 - h^2, second = -2 h first.
first = np.square(activation)
np.subtract(1.0, first, out=first)
second = activation * first
second *= -2.0
value = activation @ self.w_out
value += self.b_out
# Each derivative is one contraction: the chain-rule factors are
# elementwise and the sum over hidden units is a matrix-vector product.
u_x = (first * self.w_x) @ self.w_out
u_t = (first * self.w_t) @ self.w_out
u_xx = (second * self.w_x * self.w_x) @ self.w_out
return value, u_t, u_x, u_xx
def residual(
self, x: NDArray[np.float64], t: NDArray[np.float64], viscosity: float
) -> NDArray[np.float64]:
"""u_t + u u_x - nu u_xx, for every collocation point at once.
This IS the training signal. No label is involved anywhere.
"""
value, u_t, u_x, u_xx = self.forward_with_derivatives(x, t)
out = u_t
out += value * u_x
out -= viscosity * u_xx
return out
def loss_terms(
network: BatchedPinn,
collocation: tuple[NDArray[np.float64], NDArray[np.float64]],
boundary: tuple[NDArray[np.float64], NDArray[np.float64]],
initial_x: NDArray[np.float64],
initial_values: NDArray[np.float64],
viscosity: float,
) -> tuple[float, float, float]:
"""Each term computed separately and never pre-summed.
Returning a scalar total is the mistake that makes PINNs hard to debug: a
falling total is compatible with one term being entirely unconstrained.
"""
residual = network.residual(*collocation, viscosity)
# einsum contracts the squared norm in one pass rather than squaring into
# a temporary and summing it.
residual_loss = np.einsum('i,i->', residual, residual) / residual.size
boundary_value, *_ = network.forward_with_derivatives(*boundary)
boundary_loss = np.einsum('i,i->', boundary_value, boundary_value) / boundary_value.size
initial_value, *_ = network.forward_with_derivatives(
initial_x, np.zeros_like(initial_x)
)
initial_value -= initial_values
initial_loss = np.einsum('i,i->', initial_value, initial_value) / initial_value.size
return float(residual_loss), float(boundary_loss), float(initial_loss)
def gradient_magnitudes(
network: BatchedPinn,
collocation: tuple[NDArray[np.float64], NDArray[np.float64]],
boundary: tuple[NDArray[np.float64], NDArray[np.float64]],
viscosity: float,
) -> tuple[float, float]:
"""Typical gradient magnitude per loss term, for adaptive weighting.
This is what turns loss weighting from a hyperparameter search into a
measurement. The terms carry different units - the residual is in du/dt,
the boundary term is in u - so their gradients routinely differ by orders
of magnitude and a fixed weighting silently chooses which one is enforced.
"""
residual = network.residual(*collocation, viscosity)
boundary_value, *_ = network.forward_with_derivatives(*boundary)
# Proxy: the RMS of each residual scaled by the output-layer magnitude.
# A full implementation takes the real parameter gradients; the point is
# that the ratio is measured rather than assumed.
scale = float(np.abs(network.w_out).max())
return (
float(np.sqrt(np.mean(np.square(residual)))) * scale,
float(np.sqrt(np.mean(np.square(boundary_value)))) * scale,
)
def rebalance_weights(
weights: dict[str, float], magnitudes: tuple[float, float], rate: float = 0.1
) -> dict[str, float]:
"""Move each weight toward equalizing gradient contribution.
The rate is kept small so the balance moves slowly relative to the
weights it is balancing — a fast rebalance oscillates and is worse than
no rebalancing at all.
"""
residual_magnitude, boundary_magnitude = magnitudes
if boundary_magnitude <= 0.0:
return weights
target = residual_magnitude / boundary_magnitude
weights['boundary'] = (1.0 - rate) * weights['boundary'] + rate * target
return weights
def latin_hypercube(count: int, bounds: NDArray[np.float64], rng: np.random.Generator) -> NDArray[np.float64]:
"""Stratified collocation sampling rather than uniform.
Uniform sampling leaves clumps and gaps; a stratified design covers the
domain far more evenly for the same point count, which matters because
the residual is only enforced where points are - and the network can be
badly wrong in a gap while reporting an excellent loss.
"""
dimensions = bounds.shape[0]
# One permuted stratum per dimension: the defining property of the design.
samples = np.empty((count, dimensions), dtype=FLOAT)
for dimension in range(dimensions):
strata = (rng.permutation(count) + rng.random(count)) / count
low, high = bounds[dimension]
samples[:, dimension] = low + strata * (high - low)
return samples
def held_out_residual(
network: BatchedPinn,
bounds: NDArray[np.float64],
viscosity: float,
rng: np.random.Generator,
count: int = 10_000,
) -> float:
"""The check that matters, and the one people skip.
A network can satisfy the equation at its collocation points while being
badly wrong between them. A low training residual with a high held-out
residual is exactly that failure, and the loss curve looks identical
either way — this is the only thing that reveals it.
"""
points = latin_hypercube(count, bounds, rng)
residual = network.residual(points[:, 0], points[:, 1], viscosity)
return float(np.einsum('i,i->', residual, residual) / count)
What changed: The structural observation is that collocation points are entirely independent, so the whole residual evaluation is one batched pass — and because the network is a composition of smooth functions, the derivative chain rules vectorize exactly like the forward pass does, so a second derivative is another array expression rather than another loop. The reuse that matters is that for a tanh unit both the first and second derivatives are functions of the activation alone, so one tanh evaluation — the only transcendental in the whole residual — serves the value and both derivative orders, and each derivative reduces to one elementwise product followed by a matrix-vector contraction. Squared norms use einsum to contract in a single pass rather than squaring into a temporary. Two non-arithmetic changes matter as much: gradient magnitudes per loss term are measured rather than guessed, which is what turns loss weighting from a hyperparameter search into a running measurement, and collocation sampling becomes stratified rather than uniform, because uniform sampling leaves gaps and the network can be badly wrong in a gap while reporting an excellent loss. The dtype stays float64 deliberately, which is the one place this model differs from almost everything else in this reference.
Vectorize to NumPy/BLAS (@, np.dot, einsum)
Pre-activations for every collocation point are one outer-product construction, each derivative is one contraction over the hidden axis, and the squared norms contract in a single einsum pass.
Trade-off: The pre-activation matrix is points times width in float64, so a hundred thousand collocation points against a wide network is the dominant memory cost and forces chunking — which then reintroduces a Python loop over chunks.
Fuse adjacent operations so intermediates are never materialized
One tanh evaluation serves the value and both derivative orders, since each derivative is a function of the activation alone, and the residual accumulates in place into the u_t buffer.
Trade-off: The u_t array is destroyed to become the residual, so a diagnostic that wants to see which term of the equation dominates the violation — the natural question when a PINN converges to something wrong — needs an unfused pass.
Pre-allocate output arrays and use in-place operations
The derivative factors are computed through existing buffers with out parameters, so a residual evaluation allocates only the pre-activation matrix.
Trade-off: The activation buffer is consumed by the derivative computation, so the raw activations are unavailable afterwards — and a saturated tanh, which is why a PINN stops learning, is only visible in those.
Ensure contiguous memory layout and a single dtype
A single dtype throughout keeps every contraction on the BLAS fast path, and points-by-width row-major matches how the derivative reductions read.
Trade-off: The dtype is float64 rather than the float32 used elsewhere in this reference, and deliberately: the residual is a difference of derivative terms that nearly cancel, so float32 loses most of its significant digits exactly where the signal is — the memory and bandwidth cost is doubled and it is not optional.
A residual evaluation becomes one outer product, one tanh and three contractions for every collocation point at once. Illustrative, not a measured benchmark: against numerical differentiation over parameters this is not a constant-factor improvement but an asymptotic one, since that version cost two full loss evaluations per parameter.