Deep Ritz: solve a PDE by minimising an energy functional
A VPINN (previous tutorial) minimises the weak residual test-projected onto the FE basis. The Deep Ritz method takes the other classical route: it minimises the energy functional directly — the network is the trial, there are no test functions and no assembled system. For Poisson \(-\Delta u = f\) with \(u=0\) on \(\partial\Omega\) the governing energy is
whose minimiser over \(H^1_0\) is the weak solution. Lower derivative order (only \(\nabla u\), no \(\Delta u\)) and no linear solve, at the price of a non-convex optimisation.
Method: Deep Ritz — E & Yu, Commun. Math. Stat. 6:1 (2018), §2 (arXiv:1710.00211); the solid-mechanics form is the Deep Energy Method (Samaniego et al., CMAME 362, 2020; Nguyen-Thanh, Zhuang & Rabczuk, Eur. J. Mech. A/Solids 80, 2020).
The energy as a loss
expr.integrate() integrates a scalar expression over the mesh (auto-detecting volume vs boundary
from the coordinate variable), so the whole method is: build the network trial, write the energy
density, integrate it, and hand the single scalar to jno.core as the loss.
xi, yi, _ = dom.variable("interior", split=True)
net = jnn.nn.wrap(foundax.mlp(2, hidden_dims=48, num_layers=4, activation=jax.nn.tanh, key=key))
ansatz = xi * (1 - xi) * yi * (1 - yi) # hard-BC ansatz: vanishes on ∂Ω → u ∈ H¹₀
u = net(xi, yi) * ansatz
ux, uy = jnn.grad(u, xi), jnn.grad(u, yi)
f = 2.0 * (xi * (1 - xi) + yi * (1 - yi))
energy = (0.5 * (ux**2 + uy**2) - f * u).integrate(quadrature="gauss") # J[u] = ∫ (½|∇u|² − f u) dx
net.optimizer(optax.adam(3e-3))
jno.core([energy], domain=dom).solve(4000) # minimise the (signed) energy directly
The loss is the signed energy — .integrate() reduces to a scalar and jno.core minimises it as
is (no square), so the reported loss converges to \(J_{\min}<0\). General geometries that cannot use a
hard-BC ansatz add a boundary penalty (β * (net(xb, yb) - g)).mse as a second constraint.
Quadrature consistency — why quadrature="gauss"
The default .integrate() uses the P1 nodal-volume (vertex) rule: the energy is sampled only at
the mesh nodes. A network expressive enough to develop structure between nodes can drive the
quadrature-estimated energy below the true minimum — a variational crime in which the discrete
energy stops being a faithful estimate, and the solution degrades even as the reported loss keeps
falling. .integrate(quadrature="gauss") samples the per-element Gauss points (weighting by
\(J\,W\)) instead — many points inside each element, far harder to alias — so a capable network can
minimise the energy safely. Pass quadrature=<int> to set the Gauss degree ("gauss" = 4).
With this same 48×4 network the two rules diverge sharply: Gauss holds the true minimum and reaches rel-\(L^2 \approx 8\times10^{-5}\), while the nodal rule lets the energy sink past the true minimum and the error blows up. (Gauss integration is currently volume-only; boundary integrals use the nodal rule.)
The result
The trained network matches the analytic \(x(1-x)y(1-y)\) to rel-\(L^2 \approx 8\times10^{-5}\). The energy plateaus at the true minimum within a few hundred epochs and stays there.
What to notice
- Energy, not residual. The loss is the functional \(J[u]\) itself, integrated over the mesh with
.integrate()— no test functions, no assembled operator. Contrast the VPINN, which needs the FE test space. - Signed objective.
jno.coreminimises the raw reduced value, so a signed energy that goes negative is exactly right — no.msewrapper. - Quadrature matters for energy losses. A fixed nodal rule can be aliased by an expressive
integrand;
quadrature="gauss"samples inside each element and resists it. If a loss dips below the analytic energy while the error grows, that is the quadrature being exploited.
Full script
"""Deep Ritz / Deep Energy Method: solve a PDE by **minimising an energy functional** with a network.
Where a VPINN (`vpinn_poisson_2d.py`) minimises the *weak residual* test-projected onto the FE basis,
the Deep Ritz method minimises the **energy functional** directly — no test functions, no assembly.
For Poisson ``-Δu = f`` with ``u = 0`` on ``∂Ω`` the governing energy is
J[u] = ∫_Ω ( ½ |∇u|² − f·u ) dx ,
whose minimiser over ``H¹₀`` is the weak solution. The neural network *is* the trial ``u``; we integrate
the energy density over the mesh with ``expr.integrate()`` and hand the single scalar ``J`` to
``jno.core`` as the loss. Lower derivative order and no linear solve, at the price of a
non-convex optimisation.
Method: **Deep Ritz** — E & Yu, *Commun. Math. Stat.* 6:1 (2018), §2 (arXiv:1710.00211); the solid-
mechanics form is the **Deep Energy Method**, Samaniego et al., *CMAME* 362 (2020) and Nguyen-Thanh,
Zhuang & Rabczuk, *Eur. J. Mech. A/Solids* 80 (2020).
We solve ``-Δu = f`` on the unit square with exact ``u = x(1-x)y(1-y)`` (so ``f = 2[x(1-x)+y(1-y)]``).
Quadrature consistency (why we use Gauss): the default ``.integrate()`` uses the P1 nodal-volume
(vertex) rule — the energy is sampled *only at mesh nodes*. A network expressive enough to develop
structure *between* nodes can then drive the discrete energy *below* the true minimum (a variational
crime: the quadrature is no longer a faithful estimate) and the solution degrades even as the
reported loss keeps dropping. ``.integrate(quadrature="gauss")`` samples the per-element Gauss points
instead, which is far harder to alias — it lets a *capable* network minimise the energy safely and
converge to ``rel-L2 ≈ 1e-4`` here. (Swap in the default nodal rule with this same network and the
energy sinks past the true minimum while the error blows up.)
"""
import os
os.environ["MPLBACKEND"] = "Agg"
from pathlib import Path # noqa: E402
import foundax # noqa: E402
import jax # noqa: E402 (jax.nn / jax.random for the network)
import matplotlib.pyplot as plt # noqa: E402
import matplotlib.tri as mtri # noqa: E402
import numpy as np # noqa: E402
import optax # noqa: E402
import jno # noqa: E402
import jno.jnp_ops as jnn # noqa: E402
jax.config.update("jax_enable_x64", True) # the mesh integral accumulates in float64
# ---- domain, network trial, energy functional ----------------------------------------------
dom = jno.shape.rect(0, 0, 1, 1, size=0.05).domain()
xi, yi, _ = dom.variable("interior", split=True) # .integrate() re-evaluates at the quadrature points
# a capable network — the Gauss quadrature below keeps its energy honest (see the consistency note)
net = jnn.nn.wrap(foundax.mlp(2, hidden_dims=48, num_layers=4, 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 ∈ H¹₀
u = net(xi, yi) * ansatz # the network trial
ux, uy = jnn.grad(u, xi), jnn.grad(u, yi)
f = 2.0 * (xi * (1 - xi) + yi * (1 - yi)) # -Δ[x(1-x)y(1-y)]
# energy functional J[u] = ∫ (½|∇u|² − f u) dx, integrated with element GAUSS quadrature
# (alias-resistant — a nodal rule would let this network cheat the integral; see the docstring)
energy = (0.5 * (ux**2 + uy**2) - f * u).integrate(quadrature="gauss")
# ---- train the network by minimising J[u] through jno.core ---------------------------------
net.optimizer(optax.adam(3e-3))
crux = jno.core([energy], domain=dom) # the *signed* energy is the loss (it converges to J_min < 0)
crux.solve(4000)
# ---- verify the trained network against the analytic solution (on a fresh grid) ------------
test_dom = jno.shape.rect(0, 0, 1, 1, size=0.035).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"\nDeep Ritz Poisson 2D: minimised energy functional; dofs={dom.mesh.points.shape[0]}")
print(f" trained Deep-Ritz vs analytic x(1-x)y(1-y): rel-L2 = {rel:.3e}")