← Back to Scimba Jax Module · scimba_jax

PINNs & projector

A Physics-Informed Neural Network replaces the local, mesh-tied bases used elsewhere in the library with a single network, and turns solving a PDE into minimizing the residual it leaves behind at a batch of sampled points. In scimba_jax that training loop is a Projector — the same object a collocation basis or any other approximation space is trained through.

1. A network, a model, a sampler — and a Projector

ApproximationSpace wraps the network (an MLP here) with a post_processing map that bakes the boundary condition into every output — a strong BC, exact for any weights at all. LaplacianDirichletND is the physical model, and Projector ties model, space and sampler into one object whose .project(...) runs the training loop and returns the improved space. 

from scimba_jax.domains.meshless_domains.domains_2d import Square2D
from scimba_jax.nonlinear_approximation.approximation_spaces.approximation_spaces import (  # noqa: E501
    ApproximationSpace,
)
from scimba_jax.nonlinear_approximation.integration.monte_carlo import (
    DomainSampler,
    TensorizedSampler,
)
from scimba_jax.nonlinear_approximation.networks.mlp import MLP
from scimba_jax.nonlinear_approximation.numerical_solvers.projectors import Projector
from scimba_jax.physical_models.elliptic_pde.laplacians import LaplacianDirichletND

domain = Square2D([(-1.0, 1.0), (-1.0, 1.0)], is_main_domain=True)
sampler = TensorizedSampler([DomainSampler(domain)], bc=False)

nn = MLP(in_size=2, out_size=1, hidden_sizes=[16, 16], key=key)
# post_processing bakes the BC into the network's output -- a "strong" BC:
# u_h is exactly 0 on the boundary for any weights at all.
space = ApproximationSpace(
    {"x": 2}, [(nn, "scalar", None)], model_type="x",
    post_processing=lambda approx, xy: approx * (1 - xy[0] ** 2) * (1 - xy[1] ** 2),
)
model = LaplacianDirichletND(domain, lambda x: source(x), bc="strong", model_type="x")

# A PINN is a Projector: the same space -> better space contract every
# approximation space in scimba_jax trains against.
pinn = Projector(model, space, sampler)  # default optimizer: see below
key, pinn = pinn.project(key, space, n_epochs=200, n_colloc=1000)
# loss 98.7 -> falling; see the Optimizer section for what changes it fastest.

2. The math, briefly

Training minimizes a sum of squared residuals, one term per piece of the problem — the interior PDE, and one more per boundary label:

\(\mathcal{L}(\theta) = \sum_{\text{label}} \frac{1}{N_{\text{label}}} \sum_{i=1}^{N_{\text{label}}} \left\lVert R_{\text{label}}(u_\theta, x_i) \right\rVert^2\)

exactly as many residual terms as sub-domains and boundary pieces the model declares — a Dirichlet edge and a Neumann edge each get their own label and their own term, summed into one scalar loss. The network itself is scalar-valued for a single unknown (temperature, pressure) or vector-valued when the problem couples several fields at once — ApproximationSpace's "scalar"/"vec" tag says which; a 2D Navier–Stokes velocity–pressure system is a length-3 vector output [u_x, u_y, p] from the very same network class.

3. Optimizer: first order vs. second order

Projector's training rule is a single keyword, swappable without touching the model, the space or the sampler. A first-order optimizer only ever uses the gradient:

\(\theta_{k+1} = \theta_k - \eta\, \nabla \mathcal{L}(\theta_k)\)

— Adam, AdamW and Muon fall here (Muon just reshapes the step by the weight matrix's own geometry instead of a plain learning rate). A second-order optimizer also uses some model of curvature C to rescale that gradient:

\(\theta_{k+1} = \theta_k - \eta\, C^{-1} \nabla \mathcal{L}(\theta_k)\)

L-BFGS, SS-BFGS and SS-Broyden build C from successive gradients (quasi-Newton); ENG, its matrix-free twin MFENG, ANaGRAM, RATATAM and RATATAM-LM build it instead from the PDE's own natural-gradient/Gauss–Newton metric — curvature that comes from the physics, not from a weight history. The default is ENG (Energy Natural Gradient): it converges on problems where plain Adam stalls, which is why it is the default rather than a special case. 

# `Projector`'s `optimizer=` argument, unchanged everywhere else:
#
#   "ENG"          Energy Natural Gradient -- the DEFAULT. Follows the
#                   gradient in the metric the PDE's own energy induces,
#                   not the Euclidean one; converges where Adam stalls on
#                   the elliptic and parabolic problems in this library.
#   "MFENG"        Matrix-free ENG -- same idea, without assembling the
#                   (n_params x n_params) Gram matrix explicitly.
#   "Adam"         The default choice everywhere else in deep learning; a
#                   safe fallback when ENG's Gram matrix is too costly.
#   "AdamW"        Adam with DECOUPLED weight decay (Loshchilov & Hutter).
#   "Muon"         Rescales the update's SPECTRUM instead of each
#                   coordinate: every singular value becomes 1.
#   "L-BFGS"       Quasi-Newton, second order without a Hessian.
#   "SS-BFGS"      A self-scaled BFGS variant, tuned for PINN loss surfaces.
#   "SS-Broyden"   Same family, Broyden's update instead of BFGS's.
#   "ANaGRAM"      A natural-gradient method sized for large networks.
#   "RATATAM"      A Levenberg-Marquardt-flavoured PINN optimizer.
#
# Swapping is a one-line change -- nothing about `model`, `space` or
# `sampler` moves:
pinn = Projector(model, space, sampler, optimizer="Adam")
key, pinn = pinn.project(key, space, n_epochs=1000, n_colloc=1000)
print(pinn.best_loss["total"])  # a dict, one entry per residual plus "total"
# 20 Adam epochs on the same problem: loss 98.7 -> 86.1 -- compare against
# ENG's 98.7 -> 52.2 in 5 epochs above; which one wins depends on the
# problem, which is exactly why the argument exists.

A PINN and a mesh-based solver answer the same question with opposite trade-offs: no mesh to build, at the cost of a nonconvex training problem instead of a linear system. When a mesh already fits the geometry well, see Mesh-based solvers instead.