Training a physics-informed neural field

Solve u''(x) = -2 on [-1, 1], with u(-1) = u(1) = 0. The exact solution is 1 - x², but training only sees the equation and boundary conditions, not samples of that solution. A tanh MLP represents the field. The objective is half the mean squared equation residual plus half the mean squared boundary residual. derivativeVjp differentiates the second coordinate derivative with respect to all model parameters; ordinary SGD then updates the same typed state.

This is a small CPU example of the residual method of Raissi, Perdikaris, and Karniadakis, "Physics-informed neural networks", Journal of Computational Physics 378 (2019), 686–707. The reported residuals are sampled diagnostics, not a uniform PDE certificate.

def NN.Examples.Models.Operators.Pinn.model : TorchLean.nn.Sequential [1] [1]

A smooth neural field; tanh permits the second derivatives in the equation residual.

def NN.Examples.Models.Operators.Pinn.collocation : Array Float

The equation is imposed at five interior collocation points.

def NN.Examples.Models.Operators.Pinn.lossGradient (state : TorchLean.nn.State Float (Runtime.Autograd.Model.Layers.Seq.stateShapes model)) : IO (Float × TorchLean.nn.State Float (Runtime.Autograd.Model.Layers.Seq.stateShapes model))

Evaluate the residual objective and accumulate its parameter gradient.

def NN.Examples.Models.Operators.Pinn.run (steps : ℕ := 1500) : IO Unit

Fit the equation and boundary conditions, then compare with the held-out exact solution.

def NN.Examples.Models.Operators.Pinn.main (args : List String) : IO Unit

Run with the default step count or one explicit natural-number argument.