扩散¶
examples/diffusion/ 中的两个瞬态问题:采用隐式 Euler 时间步进的线性热方程(heat/heat.py),以及每个时间步用牛顿法求解的非线性相场问题(allen-cahn/)。二者一起演练了来自 时间积分 的线性时间步进模式,以及来自 稀疏求解器 的非线性 / 牛顿模式。Allen-Cahn 问题提供三种风味:手写的后向 Euler 牛顿循环(ac.py)、把同一步交给 torch_sla.nonlinear_solve 驱动(ac_torch_sla.py),以及无条件能量稳定的凸-凹分裂(ac_convex_concave.py)。
二维热方程——heat/heat.py¶
强形式为
在单位正方形上，初始条件取自 HeatMultiFrequency 的多频傅里叶级数。
后向（隐式）Euler 将每一步都转化为一个线性求解：
其中 \(M\) 是质量矩阵，\(A\) 是拉普拉斯刚度。脚本将两者各装配一次，并在各步之间重复使用：
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 \in \{-1, 1\}\) 演化的光滑初始条件。在 \(\varepsilon = 220\) 和 \(\Delta t = 10^{-6}\) 下，非线性项足够刚性，以至于显式积分器需要小到不可行的步长；因此脚本采用隐式 Euler 加牛顿迭代。
每个时间步，牛顿法求解残差
其雅可比矩阵 \(K = \partial R / \partial c\) 围绕当前迭代值线性化。两者都写成 TensorMesh 装配器，因此其模式与线性求解相同：
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 循环会浓缩为一次调用：
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\) 决定。每一帧对应一次收敛的牛顿求解。
凸-凹分裂 —— allen-cahn/ac_convex_concave.py¶
ac.py 和 ac_torch_sla.py 求解 相同 的全隐式后向 Euler 步,只在如何驱动牛顿求解上不同。ac_convex_concave.py 求解一种 不同的时间离散:Eyre 凸分裂。它利用了 Allen-Cahn 是双井能量的 \(L^2\) 梯度流这一事实
其非线性项 \(W'(c) = c^3 - c\) 分裂为 凸 部分 \(c^3\) 和 凹 部分 \(-c\)。将凸部分取隐式、凹部分取显式,得到
该格式 无条件能量稳定 —— 它对 任意 时间步长都减少 \(E\) —— 且每步牛顿系统的唯一可解性与 \(\Delta t\) 无关,因此残差及其切线为
该切线是对称正定的(反应项 \(3\varepsilon^2 c^2 \ge 0\) 只往对角线上叠加),这是无条件稳定性在代数上的反映 —— 非线性不带来任何步长限制:
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
为什么值得单列一个脚本:
更大的时间步。 全隐式的
ac.py虽稳定,但当 \(\Delta t\) 推得太远时其牛顿求解可能停滞;凸分裂保持良好,因此--dt可以远超全隐式上限,而能量仍单调衰减。相同的装配器约定。 残差
NodeAssembler和 SPD 切线ElementAssembler都从point_data读取当前迭代值c和上一步cold—— 与ac.py完全相同的参数分派模式;改变的只是 离散格式,而非 API。
python ac_convex_concave.py # default dt
python ac_convex_concave.py --steps 200 --dt 1e-4 # larger step, still stable
输出动画写入 Allen-Cahn-convex-concave-dt<dt>.mp4``(文件名编码了所用的 ``--dt,因此不同步长的运行不会相互覆盖)。
运行示例¶
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
接下来¶
时间积分——同一个热传导问题，以
ImplicitLinearEuler重写。稀疏求解器——
SparseMatrix.nonlinear_solve，提供封装好的牛顿 / Picard / Anderson 迭代循环，并支持伴随反向传播。波动方程——下一个进阶的瞬态偏微分方程：双曲型，显式中心差分。
机器学习数据集——用于机器学习训练数据的批量热传导求解。