Allen-Cahn 2D

Download .py All tutorials

This example solves a manufactured 2D Allen-Cahn problem and introduces a nonlinear cubic reaction term.

Problem Setup

The PDE has the Allen-Cahn structure u_t = epsilon^2 Delta u + u - u^3 + f, with a known exact solution used to build the forcing term.

Step 1: Build a Manufactured Nonlinear Problem

The exact solution is substituted into the PDE to derive a forcing term that makes validation straightforward.

eps = 0.1
T_end = 1.0

domain = jno.Shape.rect(0, 0, 1, 1, size=0.05).domain(time=(0, T_end, 4))
x, y, t = domain.variable("interior")

S = sin(π * x) * sin(π * y)
u_exact = exp(-t) * S

coeff  = 2 * eps**2 * π**2 - 2
source = exp(-t) * S * coeff + exp(-3 * t) * S**3

Step 2: Set Up the Space-Time Network

The model learns a field over space and time while respecting the chosen boundary handling. The optimizer is scaled by jno.fn.adaptive.dlrs, a loss-adaptive dynamic learning-rate scheduler that shrinks the step when the stiff Allen-Cahn interface stalls the loss and grows it when the loss can still descend.

net = jno.nn(
    foundax.deeponet(
        n_sensors=1, coord_dim=2, n_outputs=1,
        n_layers=3, basis_functions=64, hidden_dim=40,
        key=jax.random.PRNGKey(42),
    )
)
net.optimizer(optax.adam(1)).scale(jno.fn.adaptive.dlrs(lr0=1e-3, window=10))

xy = jno.np.concat([x, y])
u  = net(t, xy) * x * (1 - x) * y * (1 - y)

Step 3: Encode the Nonlinear Residual

The key change relative to the heat equation is the nonlinear reaction term u - u^3.

pde = u.d(t) - eps**2 * jno.np.laplacian(u, [x, y]) - u + u**3 - source

Step 4: Impose the Initial Condition

The script uses the same PDE infrastructure but anchors the solution at the initial time with an additional loss.

u_at_0 = net(0 * t, xy) * x * (1 - x) * y * (1 - y)
ini     = u_at_0 - sin(π * x) * sin(π * y)

crux    = jno.core([pde.mse, ini.mse])
history = crux.solve(5000)

Result

The network's own field is evaluated on a finer time grid and animated above; the single central bump decays like \(e^{-t}\), as the manufactured solution prescribes.

At the final time the prediction matches the manufactured solution to rel-\(L^2 \approx 1.5\times10^{-3}\) (signed-error panel, right, centered at 0).

What To Notice

• Nonlinear reaction terms are easy to express once the field is available symbolically.
• Manufactured solutions are especially valuable for nonlinear PDEs.
• This example is a good template for phase-field style problems.

Download full script All tutorials

Script Snippet

"""03 — 2-D Allen–Cahn equation (manufactured-solution verification)"""

import foundax
import jax
import optax

import jno

π = jno.np.pi
ε = 0.1
T_end = 1.0

# Time-dependent unit square: the Shape one-liner forwards ``time=`` to the domain.
domain = jno.Shape.rect(0, 0, 1, 1, size=0.05).domain(time=(0, T_end, 4))
x, y, t = domain.variable("interior")

S = jno.np.sin(π * x) * jno.np.sin(π * y)
u_exact = jno.np.exp(-t) * S
source = jno.np.exp(-t) * S * (2 * ε**2 * π**2 - 2) + jno.np.exp(-3 * t) * S**3

# Network with hard Dirichlet BCs in space; t is fed via the trunk input.
net = jno.nn(
    foundax.deeponet(
        n_sensors=1,
        coord_dim=2,
        n_outputs=1,
        n_layers=3,
        basis_functions=64,
        hidden_dim=40,
        key=jax.random.PRNGKey(42),
    )
)
# Loss-adaptive LR: the DLRS schedule raises/lowers the step size from the recent loss
# slope, so training slows when the stiff ε-thin Allen–Cahn interface makes the loss
# stagnate and speeds up when it can descend — no hand-tuned decay curve required.
net.optimizer(optax.adam(1)).scale(jno.fn.adaptive.dlrs(lr0=1e-3, window=10))

xy = jno.np.concat([x, y])
# Bind names so partials read like the math:  u.t, u.xx, u.yy, u.xy, ...
u = (net(t, xy) * x * (1 - x) * y * (1 - y)).scalar.bind(x=x, y=y, t=t)

pde = u.t - ε**2 * (u.xx + u.yy) - u + u**3 - source

# Initial condition  (t=0 via 0*t trick)
u_at_0 = net(0 * t, xy) * x * (1 - x) * y * (1 - y)
ini = u_at_0 - S

grad_norms = jno.trackers.gradient_norms(interval=500)
crux = jno.core([pde.mse, ini.mse])

print(f"Allen–Cahn 2-D  (ε = {ε})")
crux.solve(5000, callbacks=[grad_norms])

_u, _u_exact = crux.eval([u, u_exact])
rel_l2 = float(jax.numpy.linalg.norm(_u - _u_exact) / (jax.numpy.linalg.norm(_u_exact) + 1e-8))
print(f"Relative L2 error: {rel_l2:.4e}")
if grad_norms.value is not None:
    print(f"Final ∇L norms (pde, ini): {grad_norms.value['norms']}")

assert rel_l2 < 1e-1, f"relative L2 error too large: {rel_l2:.3e}"