扩散

examples/diffusion/ 中有两个瞬态示例：采用隐式 Euler 时间步进的线性热方程（heat/heat.py），以及每个时间步用牛顿法求解的非线性相场问题（allen-cahn/ac.py）。两者合起来既演练了 时间积分 中的线性时间步进模式，也演练了 稀疏求解器 中的非线性 / 牛顿迭代模式。

二维热方程——heat/heat.py

强形式为

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

在单位正方形上，初始条件取自 HeatMultiFrequency 的多频傅里叶级数。

后向（隐式）Euler 将每一步都转化为一个线性求解：

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

其中 \(M\) 是质量矩阵，\(A\) 是拉普拉斯刚度。脚本将两者各装配一次，并在各步之间重复使用：

列表 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)

其中有几处选择体现了用户指南中的模式：

• 二次三角形（order=2, element_type="tri"）——热扩散在内部是光滑的，因此在相同网格尺寸下，2 阶单元能与解析解高度吻合得多。

• 装配一次，多次求解。 K 和 K_ 在循环之前构建。每一步只做一次矩阵-向量乘 M @ U、一次右端项凝聚和一次求解——因此大部分开销都在每步的线性求解上。

• 与真值对比。 solution() 返回每个时刻的解析解，而 mesh.plot({"prediction": Us, "ground truth": Us_gt}, save_path="heat.mp4", dt=dt) 会写出一段并排对比的动画。

用 TensorMesh 的 ImplicitLinearEuler 积分器表达的同一格式，作为 heat/heat_ode.py 一并提供——同一问题，只是将 Condenser 接入三个边界条件钩子，而非手动循环。两个脚本生成的快照在机器精度内完全一致；时间积分 详解了积分器版本，而本脚本则是更底层的“手动”版本。

``heat.py`` 的输出：在时间窗口内有限元预测（左）与解析真值（右）的对比。两者在肉眼看来始终完全相同，随着后向 Euler 步阻尼掉最高的傅里叶模态，最大节点误差平滑衰减。

Allen-Cahn 相场——allen-cahn/ac.py

Allen-Cahn 方程是教科书中的非线性相场模型：

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

采用自然（无通量）边界条件，以及一个向二元相 \(u \in \{-1, 1\}\) 演化的光滑初始条件。在 \(\varepsilon = 220\) 和 \(\Delta t = 10^{-6}\) 下，非线性项足够刚性，以至于显式积分器需要小到不可行的步长；因此脚本采用隐式 Euler 加牛顿迭代。

每个时间步，牛顿法求解残差

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

其雅可比矩阵 \(K = \partial R / \partial c\) 围绕当前迭代值线性化。两者都写成 TensorMesh 装配器，因此其模式与线性求解相同：

列表 6 examples/diffusion/allen-cahn/ac.py（核心）

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

有两点让这段代码易于编写：

• 通过 point_data 传入逐迭代参数。 当前迭代值 c 和上一时间步的解 cold 按名称传入两个 forward 方法——这正是 形式 中所记载的参数分派约定。这里没有逐求积点的 Python 循环。

• 每次迭代重新装配。 由于 K 依赖于 c（通过 D(c)、df(c)），它必须在每次牛顿迭代时重建。这是有意为之——在 GPU 上装配开销很低，而另一种做法（手动计算解析雅可比矩阵）则需要数页额外代码。

如果你更愿意把整个步骤交给一个封装好的驱动器，allen-cahn/ac_torch_sla.py 用 torch-sla 的 nonlinear_solve（牛顿 / Picard / Anderson，并带 Armijo 线搜索；参见 稀疏求解器）求解同一问题。上面每步的 while 循环会浓缩为一次调用：

列表 7 examples/diffusion/allen-cahn/ac_torch_sla.py（核心）

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")

有两点值得注意。第一，nonlinear_solve 可以用 autograd 构建雅可比矩阵，但这里我们将有限元一致切线显式地作为 jacobian_fn 传入——它已经装配好，因此这样既更省又精确。第二，符号问题：KAssembler 装配的是负切线 \(K = -\partial R/\partial c\)，因此 jacobian_fn 返回 -K.values 以交回真正的雅可比矩阵 \(J\)，随后 nonlinear_solve 据此以 \(J\,\delta u = -R\) 进行迭代。两个脚本在舍入误差内一致，并且 nonlinear_solve 还能给出穿过收敛解的正确伴随反向传播。

``ac.py`` 的输出：序参量 \(\varphi\) 在单位正方形上的演化。从随机噪声出发，它迅速粗化为两个平衡相（ \(\varphi=\pm1\) ），界面呈弥散状，其宽度由 \(\varepsilon\) 决定。每一帧对应一次收敛的牛顿求解。

运行示例

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

接下来

• 时间积分——同一个热传导问题，以 ImplicitLinearEuler 重写。

• 稀疏求解器——SparseMatrix.nonlinear_solve，提供封装好的牛顿 / Picard / Anderson 迭代循环，并支持伴随反向传播。

• 波动方程——下一个进阶的瞬态偏微分方程：双曲型，显式中心差分。

• 机器学习数据集——用于机器学习训练数据的批量热传导求解。