Diffusion

Two transient problems in examples/diffusion/: a linear heat equation with implicit Euler time stepping (heat/heat.py) and a nonlinear phase-field problem solved by Newton’s method at each timestep (allen-cahn/). Together they exercise both the linear time-stepping pattern from Time Integration and the nonlinear / Newton pattern from Sparse Solvers. The Allen-Cahn problem ships in three flavours: the hand-written backward-Euler Newton loop (ac.py), the same step driven by torch_sla.nonlinear_solve (ac_torch_sla.py), and an unconditionally energy-stable convex–concave splitting (ac_convex_concave.py).

2D heat equation — heat/heat.py

The strong form is

\[u_t \;=\; D^2\, \Delta u \quad \text{in } \Omega, \qquad u = 0 \text{ on } \partial\Omega,\]

on the unit square with a multi-frequency Fourier-series initial condition from HeatMultiFrequency.

Backward (implicit) Euler turns each step into a linear solve:

\[(M + \Delta t\, D^2\, A)\, U^{n+1} \;=\; M\, U^{n},\]

where \(M\) is the mass matrix and \(A\) is the Laplace stiffness. The script assembles both once and reuses them across steps:

Listing 5 examples/diffusion/heat/heat.py
mesh = Mesh.gen_rectangle(chara_length=0.02, order=2,
                          element_type="tri")
dataset = HeatMultiFrequency(d=16)

class AAssembler(ElementAssembler):
    def forward(self, gradu, gradv): return gradu @ gradv
class MAssembler(ElementAssembler):
    def forward(self, u, v):         return u * v

M = MAssembler.from_mesh(mesh, quadrature_order=2)()
A = AAssembler.from_mesh(mesh, quadrature_order=2)()
condenser = Condenser(mesh.boundary_mask)

dt, D, n = 5e-5, 1.0, 100
K  = M + dt * D * D * A
K_ = condenser(K)[0]                  # condense once

U = dataset.initial_condition(mesh.points)
Us = [U]
for _ in range(n - 1):
    F  = M @ U                        # explicit RHS
    F_ = condenser.condense_rhs(F)
    U_ = K_.solve(F_)
    U  = condenser.recover(U_)
    Us.append(U)

A few choices reflect the user-guide patterns:

  • Quadratic triangles (order=2, element_type="tri") — heat diffusion is smooth on the interior, so order-2 elements give a much sharper match to the analytical solution at the same mesh size.

  • One assembly, many solves. K and K_ are built before the loop. Each step only does a matrix-vector M @ U, a condense-RHS, and a solve — so the bulk of the cost is the per-step linear solve.

  • Comparison with truth. solution() returns the analytical solution at each time, and mesh.plot({"prediction": Us, "ground truth": Us_gt}, save_path="heat.mp4", dt=dt) writes a side-by-side animation.

The same scheme expressed with TensorMesh’s ImplicitLinearEuler integrator ships alongside as heat/heat_ode.py — same problem, with Condenser wired through the three boundary-condition hooks instead of the manual loop. The two scripts produce identical snapshots to machine precision; Time Integration walks through the integrator version, and this script is the lower-level “by hand” one.

Output of heat.py: FEM prediction (left) vs analytical truth (right) over the time window. The two stay visually identical to the eye, with the maximum nodal error decaying smoothly as the backward-Euler step damps the highest Fourier mode.

Allen-Cahn phase field — allen-cahn/ac.py

The Allen-Cahn equation is the textbook nonlinear phase-field model:

\[u_t \;=\; \Delta u + \varepsilon^2\, u\, (1 - u^2),\]

with natural (no-flux) boundary conditions and a smooth initial condition that evolves toward the binary phases \(u \in \{-1, 1\}\). With \(\varepsilon = 220\) and \(\Delta t = 10^{-6}\), the nonlinearity is stiff enough that an explicit integrator would require infeasibly small steps; the script uses implicit Euler plus Newton iteration.

At each time step Newton solves the residual

\[R(c) \;=\; \frac{c - c_\text{old}}{\Delta t}\, v \;+\; D(c)\, \nabla v \cdot \nabla c \;-\; f(c)\, v \;\xrightarrow{!}\; 0,\]

with Jacobian \(K = \partial R / \partial c\) linearized around the current iterate. Both are written as TensorMesh assemblers so the pattern is the same as a linear solve:

Listing 6 examples/diffusion/allen-cahn/ac.py (essence)
class KAssembler(ElementAssembler):
    def forward(self, u, v, gradu, gradv, c, gradc, cold):
        return -1.0 * (
            (1.0/dt) * (u * v)
            + self.dD(c) * u * (gradv @ gradc)
            + self.D(c) * (gradu @ gradv)
            - self.df(c) * (u * v)
        )

class RAssembler(NodeAssembler):
    def forward(self, v, gradv, c, gradc, cold):
        cdot = (c - cold) / dt
        return cdot * v + self.D(c) * (gradv @ gradc) - self.f(c) * v

for step in range(steps):
    c = cold
    while True:
        pd = {"c": c, "cold": cold}
        K  = K_asm(mesh.points, point_data=pd)
        R  = R_asm(mesh.points, point_data=pd)
        K_, R_ = condenser(K, R)
        dC_ = K_.solve(R_)
        c   = c + condenser.recover(dC_)
        if torch.linalg.norm(R_) < 1e-10:
            break
    cold = c

Two things that make this easy to write:

  • Per-iterate parameters via point_data. The current iterate c and the previous timestep’s solution cold are passed by name into both forward methods — exactly the argument-dispatch contract documented in Forms. There is no per-quadrature-point Python loop.

  • Reassemble each iteration. Because K depends on c (through D(c), df(c)), it has to be rebuilt every Newton iteration. That is intentional — assembly is cheap on the GPU, and the alternative (computing analytical Jacobians by hand) would require pages of additional code.

If you would rather hand the whole step to a packaged driver, allen-cahn/ac_torch_sla.py solves the same problem with torch-sla’s nonlinear_solve (Newton / Picard / Anderson with Armijo line search; see Sparse Solvers). The per-step while loop above collapses to one call:

Listing 7 examples/diffusion/allen-cahn/ac_torch_sla.py (essence)
from torch_sla import nonlinear_solve

def residual_fn(c, cold):                       # F(c) = 0 is the BE step
    return R_asm(mesh.points, point_data={"c": c, "cold": cold})

def jacobian_fn(c, cold):                        # KAssembler gives K = -J,
    K = K_asm(mesh.points, point_data={"c": c, "cold": cold})
    return (-K.values, K.row, K.col, K.shape)    # so hand back J = -K

for step in range(steps):
    cold = nonlinear_solve(residual_fn, cold, cold, jacobian_fn=jacobian_fn,
                           method="newton", linear_solver="auto")

Two points worth noting. First, nonlinear_solve can build the Jacobian by autograd, but here we pass the FEM consistent tangent explicitly as jacobian_fn — it is already assembled, so this is both cheaper and exact. Second, the sign: KAssembler assembles the negative tangent \(K = -\partial R/\partial c\), so jacobian_fn returns -K.values to hand back the true Jacobian \(J\), which nonlinear_solve then steps with \(J\,\delta u = -R\). The two scripts agree to round-off, and nonlinear_solve additionally gives a correct adjoint backward through the converged solution.

Output of ac.py: the order parameter \(\varphi\) evolving on the unit square. Starting from random noise it rapidly coarsens into the two equilibrium phases ( \(\varphi=\pm1\) ), with diffuse interfaces of width set by \(\varepsilon\). Each frame is one resolved Newton solve.

Convex–concave splitting — allen-cahn/ac_convex_concave.py

ac.py and ac_torch_sla.py solve the same fully-implicit backward-Euler step and only differ in how the Newton solve is driven. ac_convex_concave.py solves a different time discretisation: Eyre’s convex splitting. It exploits the fact that Allen-Cahn is the \(L^2\) gradient flow of the double-well energy

\[E(c) \;=\; \int_\Omega \tfrac12\,|\nabla c|^2 \;+\; \varepsilon^2\, W(c)\, d\Omega, \qquad W(c) = \tfrac14\,(c^2 - 1)^2,\]

whose nonlinear term \(W'(c) = c^3 - c\) splits into a convex part \(c^3\) and a concave part \(-c\). Taking the convex part implicitly and the concave part explicitly gives

\[\frac{c - c_\text{old}}{\Delta t} \;=\; \Delta c \;-\; \varepsilon^2\, c^3 \;+\; \varepsilon^2\, c_\text{old}.\]

This scheme is unconditionally energy-stable — it decreases \(E\) for any time-step size — and the per-step Newton system is uniquely solvable independent of \(\Delta t\), so the residual and its tangent are

\[\begin{split}R(c)\,v &= \int_\Omega \tfrac{1}{\Delta t}\,(c - c_\text{old})\,v \;+\; \nabla c \cdot \nabla v \;+\; \varepsilon^2\, c^3\, v \;-\; \varepsilon^2\, c_\text{old}\, v \; d\Omega, \\ R'(c)[u]\,v &= \int_\Omega \tfrac{1}{\Delta t}\,u\,v \;+\; \nabla u \cdot \nabla v \;+\; 3\,\varepsilon^2\, c^2\, u\,v \; d\Omega.\end{split}\]

The tangent is symmetric positive-definite (the reaction term \(3\varepsilon^2 c^2 \ge 0\) only adds to the diagonal), which is the algebraic reflection of the unconditional stability — there is no step-size limit from the nonlinearity:

Listing 8 examples/diffusion/allen-cahn/ac_convex_concave.py (essence)
class RAssembler(NodeAssembler):           # residual, convex/concave split
    def forward(self, v, gradv, c, gradc, cold):
        c_t = (c - cold) / self.dt
        return c_t * v + (gradc @ gradv) + self.fi(c) * v - self.fe(cold) * v

class KAssembler(ElementAssembler):        # SPD tangent R'(c)
    def forward(self, c, u, v, gradu, gradv):
        return 1.0 / self.dt * (u * v) + (gradu @ gradv) - self.dfi(c) * (u * v)

for _ in range(steps):
    c = cold
    for newton_step in range(max_iter):
        pd = {"c": c, "cold": cold}
        K = K_asm(mesh.points, point_data=pd)
        R = R_asm(mesh.points, point_data=pd)
        c = c + K.solve(-R)
        if torch.linalg.norm(R) < 1e-10:
            break
    cold = c

Why this is worth a separate script:

  • Bigger time steps. The fully-implicit ac.py is stable but its Newton solve can stall when \(\Delta t\) is pushed too far; the convex split stays well-behaved, so --dt can be raised well beyond the fully-implicit limit while the energy still decays monotonically.

  • Same assembler contract. Both the residual NodeAssembler and the SPD-tangent ElementAssembler read the current iterate c and the previous step cold from point_data — the identical argument-dispatch pattern as ac.py; only the discretisation changed, not the API.

python ac_convex_concave.py                    # default dt
python ac_convex_concave.py --steps 200 --dt 1e-4   # larger step, still stable

The output animation is written to Allen-Cahn-convex-concave-dt<dt>.mp4 (the file name encodes the --dt used, so runs at different step sizes do not overwrite each other).

Running the examples

cd examples/diffusion/heat
python heat.py            # manual loop -> heat.mp4
python heat_ode.py        # integrator hooks -> heat_ode.mp4

cd ../allen-cahn
python ac.py                # hand-written Newton loop -> Allen-Cahn.mp4
python ac_torch_sla.py      # via torch_sla.nonlinear_solve -> Allen-Cahn-torch-sla.mp4
python ac_convex_concave.py # Eyre convex splitting -> Allen-Cahn-convex-concave-dt<dt>.mp4

What’s next