Variational PINN: Poisson with a network trial

Download .py All tutorials

A variational PINN keeps the FEM test space but replaces the trial with a neural network: instead of solving a linear system for FE coefficients, you minimise the weak-form residual test-projected onto the FE basis. In jNO it is authored exactly like any other jno.fem problem — write the weak form with u = net(x,y) and hand it to jno.fem; no init_fem, no weak.assemble.

We solve Poisson \(-\Delta u = f\) on the unit square with exact \(u = x(1-x)y(1-y)\) (so \(f = 2[x(1-x)+y(1-y)]\)).

A network trial in the weak form

dom.fem_symbols() gives the trial symbol u and the FE test function phi. Build the trial from a network with a hard-BC ansatz that vanishes on the boundary, write the standard weak form, and jno.fem detects the network (a ModelCall) and returns a trainable test-projected residual:

u, phi = dom.fem_symbols()
xi, yi, _ = dom.variable("interior", split=True)
xb, yb, _ = dom.variable("boundary", split=True)

u_net = net(xi, yi) * (xi * (1 - xi) * yi * (1 - yi))   # network trial, vanishes on the boundary
vi    = phi.bind(x=xi, y=yi)
f     = 2.0 * (xi * (1 - xi) + yi * (1 - yi))

pde = jno.fem([
    jnn.grad(u_net, xi) * jnn.grad(vi, xi) + jnn.grad(u_net, yi) * jnn.grad(vi, yi) - f * vi,
    u(xb, yb) - 0.0,                                     # Dirichlet declaration (see below)
])
jno.core([pde.mse], domain=dom).solve(2500)             # minimise the test-projected residual

The Dirichlet condition is not optional

u(boundary) - 0 looks redundant — the hard-BC ansatz already vanishes on the boundary — but it is required. It tells jno.fem which FE test functions live on the boundary so their residual is masked. A test function that does not vanish on the boundary carries the exact solution's \(\partial u/\partial n\) flux — an irreducible term that makes the loss minimum not the PDE solution, so training would diverge from it. With the declaration the residual is tested on the interior only and the network converges.

The result

The trained network matches the analytic \(x(1-x)y(1-y)\) to rel-\(L^2 \approx 7\times10^{-4}\).

What to notice

• Same entry as FEM. A network trial vs an FE trial is the only difference; jno.fem routes by detecting the ModelCall, and the returned .mse is an ordinary jNO loss for jno.core.
• Declare the Dirichlet boundary (u(boundary) - g) so its test functions are masked.
• Single-field 2-D/3-D for now (1-D and coupled multi-field raise a clear error).

Full script

"""Variational PINN (VPINN): the trial is a **neural network**, the test functions are the **FE
basis**. ``jno.fem`` detects the network trial written into the weak form and test-projects it onto
the FE test space -- a trainable residual loss, trained through ``jno.core``. No ``init_fem``, no
``weak.assemble``: a VPINN is authored exactly like any other ``jno.fem`` problem.

Poisson  -Δu = f  on the unit square, exact ``u = x(1-x)y(1-y)``  (so ``f = 2[x(1-x)+y(1-y)]``).
The network trial uses a hard-BC ansatz  ``u = net(x,y) · x(1-x)y(1-y)``  that vanishes on the
boundary; the Dirichlet condition ``u(boundary) - 0`` tells ``jno.fem`` which test functions vanish
on the boundary, so their (irreducible ``∂u/∂n``-flux) residual is masked -- without it the loss
minimum is not the PDE solution and training would diverge from it.
"""

import foundax
import jax  # (jax.nn / jax.random for the network)
import numpy as np
import optax

import jno
import jno.jnp_ops as jnn

jax.config.update("jax_enable_x64", True)  # the assembler builds in float64

# ---- domain, network trial, weak form -------------------------------------------------------
dom = jno.Shape.rect(0, 0, 1, 1, size=0.07).domain()
u, phi = dom.fem_symbols()
xi, yi, _ = dom.variable("interior", split=True)
xb, yb, _ = dom.variable("boundary", split=True)

net = jnn.nn.wrap(foundax.mlp(2, hidden_dims=32, num_layers=3, activation=jax.nn.tanh, key=jax.random.PRNGKey(0)))
ansatz = xi * (1 - xi) * yi * (1 - yi)  # hard-BC ansatz: vanishes on the [0,1]^2 boundary
u_net = net(xi, yi) * ansatz  # the network trial
vi = phi.bind(x=xi, y=yi)  # FE test function
f = 2.0 * (xi * (1 - xi) + yi * (1 - yi))  # -Δ[x(1-x)y(1-y)]

# weak form with the NETWORK trial + the Dirichlet declaration (masks the boundary test functions)
pde = jno.fem([jnn.grad(u_net, xi) * jnn.grad(vi, xi) + jnn.grad(u_net, yi) * jnn.grad(vi, yi) - f * vi, u(xb, yb) - 0.0])
print(f"\nVPINN Poisson 2D: {type(pde).__name__} (test-projected residual); dofs={dom.mesh.points.shape[0]}")

# ---- train the network through jno.core (minimise the test-projected residual) --------------
net.optimizer(optax.adam(1e-2))
crux = jno.core([pde.mse], domain=dom)
crux.solve(2500)

# ---- verify the trained network against the analytic solution (on a fresh grid) ------------
test_dom = jno.Shape.rect(0, 0, 1, 1, size=0.04).domain()
xt, yt, _ = test_dom.variable("interior", split=True)
exact_expr = xt * (1 - xt) * yt * (1 - yt)
pred = np.asarray(crux.eval([net(xt, yt) * exact_expr], domain=test_dom)).reshape(-1)
exact = np.asarray(crux.eval([exact_expr], domain=test_dom)).reshape(-1)
rel = float(np.linalg.norm(pred - exact) / np.linalg.norm(exact))
print(f"  trained VPINN vs analytic x(1-x)y(1-y):  rel-L2 = {rel:.3e}")

assert rel < 1e-2, f"VPINN did not solve Poisson: rel-L2={rel:.3e}"