portfolio / projects / pinn-playground

Watch a neural network learn physics.

This page trains a Physics-Informed Neural Network in your browser to solve the 1D viscous Burgers equation — the canonical benchmark from Raissi et al. (2019). No data labels, no ground truth: the network learns purely from the PDE residual, boundary conditions, and initial condition. Running entirely on your device via TensorFlow.js with WebGL acceleration.

PINN Burgers equation TensorFlow.js WebGL Automatic differentiation Zero data

[01] Controls

Viscosity ν 0.01/π

smaller ν → sharper shock

Hidden width 32

Hidden layers 4

Learning rate 1e-3

Collocation points 2000

PDE loss weight λ 1.0

[02] The PDE

∂u/∂t + u·∂u/∂x = ν·∂²u/∂x²

x ∈ [−1, 1] t ∈ [0, 1]
u(0, x) = −sin(πx)
u(t, ±1) = 0

Idle — press Train to start

Iteration

Total loss

—

PDE residual

—

Boundary + IC

—

Predicted solution u(t, x)

x → horizontal, t → vertical (0 top, 1 bottom)

u = −1

u = +1

Loss curve (log scale)

— total · — pde · — boundary

What's a Physics-Informed Neural Network?

A standard neural network learns from input/output pairs. A PINN learns from a differential equation. Instead of asking "what's the right answer for this input?" it asks "how badly does my prediction violate the laws of physics?" and uses that violation as the loss.

Concretely, we build a network u_θ(t, x) that takes a point in spacetime and predicts the solution value. We then use automatic differentiation to compute its partial derivatives — exactly, not numerically — and plug them into the PDE residual:

r(t, x) = ∂u/∂t + u · ∂u/∂x − ν · ∂²u/∂x²

A perfect solution satisfies r = 0 everywhere. So the training loss is just the mean squared residual at a sample of random collocation points in the domain, plus penalty terms for the initial and boundary conditions:

L = λ · ‖r‖² + ‖u(0, x) + sin(πx)‖² + ‖u(t, ±1)‖²

Why this is interesting

No training data is required. The "supervision" is the physics itself. This makes PINNs valuable for inverse problems (recovering hidden parameters from sparse measurements), high-dimensional PDEs where mesh-based solvers explode, and problems where governing equations are known but measurements are scarce — like cardiac biomechanics, my MSc research area.

Things to try

· Crank viscosity down (ν = 0.001/π) and watch a shock form near t ≈ 0.4.
· Set the PDE weight λ to zero and watch the network only learn the boundaries.
· Drop hidden width to 8 and see the solution become unstable.
· Increase collocation points to 5000 for cleaner residuals at higher cost per step.

Implementation notes

The network is a fully-connected MLP with tanh activations. First derivatives (u_t, u_x) are computed by automatic differentiation via tf.grad. The second derivative u_xx is approximated by a central difference of u_x evaluated at x ± ε — TensorFlow.js doesn't support differentiating through tf.grad like PyTorch's create_graph=True or JAX's jacfwd/jacrev, so we use a hybrid scheme. Gradients still flow correctly through the FD stencil because each u_x evaluation is itself a full autodiff result. Training uses Adam. Solution rendered to a 400×240 canvas with the RdBu_r diverging colormap on a 100×80 grid. Everything runs locally on your device's GPU via WebGL.