Diffusion

Two transient examples 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/ac.py). Together they exercise both the linear time-stepping pattern from Time Integration and the nonlinear / Newton pattern from Sparse Solvers.

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.

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

What’s next