盒子
盒子
文章目录
  1. SNES
    1. Jacobian-free Newton-Krylov (JFNK)
    2. 图着色有限差分Newton Krylov
    3. Poisson equation

PETSc[05] petsc4py中的SNES

2026.01.01 Zhuoyu Chen petsc  热度 ℃

SNES

SNES（Scalable Nonlinear Equations Solvers）负责求解非线性方程组：
$$
F(x) = 0
$$
SNES并不直接求解线性系统，而是在每一步非线性迭代中，构造并求解一个线性化问题，这一线性问题由 KSP 完成，即：

  • SNES：决定非线性迭代策略（Newton、线搜索、收敛判据）
  • KSP：负责解当前一步的线性系统
  • PC：影响每一步线性系统的求解效率
  • DM：负责数据布局与装配回调

在 petsc4py 中，SNES 的 Python 接口主要定义在 SNES.pyx 中，而其与底层 C 接口的绑定细节集中在 petscsnes.pxi。

一个最小的 SNES 使用流程为：

1
2
3
4
5
6
7
8
snes = PETSc.SNES().create(comm)
snes.setFunction(form_function, F)
snes.setJacobian(form_jacobian, J, P) # P is the preconditioner of J
snes.setFromOptions()
snes.solve(None, x)
# `snes.solve(b, x)`:
#     b: The affine right-hand side or `None` to use zero.
#     x: The starting vector or `None` to use the vector stored internally.

其中：

  • x：是未知量（Vec）
  • F：是残差向量（Vec），表示 $F(x)$
  • J：是Jacobian 矩阵（Mat），表示 $\partial F / \partial x$

在 petsc4py 中，setFunction() 接受一个 Python callable 作为残差回调，其标准签名是：

1
2
3
def form_function(snes, x, F):
    # compute F(x)
    ...

语义上，这个回调必须完全覆盖写入 F，而不是返回一个新对象。

SNES 的每一步 Newton 型方法都需要构造一个线性化问题：
$$
J(x_k),\delta x = -F(x_k)
$$
在 petsc4py 中，Jacobian 回调的标准形式是：

1
2
3
4
5
def form_jacobian(snes, x, J, P):
    ...
    J.assemble()
    ...
    P.assemble()

直接进行牛顿迭代，计算效率往往不会很高，因此在 SNES 中，Newton 方法几乎总是与 线搜索（LineSearch） 绑定使用，即：
$$
x_{k+1} = x_k + \alpha_k ,\delta x_k
$$
其中 $\delta x_k$ 来自 KSP 解线性系统，而 $\alpha_k$ 由 LineSearch 决定。在 petsc4py 中，可以通过选项控制线搜索类型：

1
2
3
-snes_linesearch_type basic
-snes_linesearch_type bt
-snes_linesearch_type l2

SNES 的收敛判据与 KSP 是完全独立的，常用的诊断选项包括：

1
2
3
-snes_monitor
-snes_converged_reason
-snes_view

也可以通过snes.setErrorIfNotConverged(True)将非收敛视为错误。

Jacobian-free Newton-Krylov (JFNK)

如果显式构建 Jacobian 矩阵过于复杂，可以采用 matrix-free 方法，即 Jacobian-free Newton–Krylov（JFNK）。在这种方法中，SNES 不再显式装配 Jacobian，而是通过有限差分近似来计算 Jacobian–向量乘法，从而驱动内层的 Krylov 迭代。

在 Newton 型方法中，每一步非线性迭代都需要求解线性化问题
$$
J(x_k),\delta x = -F(x_k),
$$
其中 $J(x_k)$ 是当前状态下的 Jacobian。KSP 在求解该线性系统时，并不需要完整的 Jacobian 矩阵，而只需要能够计算算子–向量乘法 $Jv$。JFNK 正是利用这一点，通过有限差分近似
$$
J(x),v \approx \frac{F(x + \varepsilon v) - F(x)}{\varepsilon},
$$
从而避免显式构建 Jacobian。

在 PETSc 中，SNES 对 JFNK 的支持是内建的。在 petsc4py 里，只需打开 matrix-free 开关即可：

1
2
snes.setUseMF(True)
# or use `-snes_mf`

启用该选项后，SNES 会在内部创建一个 matrix-free 的线性算子，并将其传递给 KSP，此时：

  • FormFunction 是必需的，因为有限差分需要多次计算残差
  • FormJacobian 不再用于构造线性算子本身
  • KSP 所看到的算子是一个基于残差计算的近似解 $Jv$

需要注意的是，JFNK 只替代 Jacobian 的算子作用，并不自动替代预处理矩阵。在实际计算中，通常仍然需要提供一个显式装配的矩阵作为预处理子 $P$，例如一个近似 Jacobian 或简化算子。否则，许多常用的预处理方法（如 ILU、GAMG、fieldsplit）将无法使用，KSP 的收敛性也会明显变差。

因此，JFNK 在工程中的典型配置是：

  • Jacobian：matrix-free（仅用于 Krylov 迭代中的算子乘法）
  • 预处理矩阵 $P$：显式装配，用于加速线性求解

这种组合在 Jacobian 推导困难、但仍可给出合理近似结构的非线性问题中非常常见。

图着色有限差分Newton Krylov

如果无法给出 Jacobian 的解析表达式，但已知其非零结构，则可以使用图着色方法，将 Jacobian 的列按稀疏结构进行分组，并通过有限差分在一次残差计算中同时更新多列，从而高效地构建 Jacobian 的数值近似。

与 JFNK 不同，图着色方法的目标仍然是构建一个显式 Jacobian 矩阵，只是采用数值方式完成装配，并尽可能减少残差评估的次数。最朴素的有限差分 Jacobian 装配方法需要对每一个自由度分别扰动一次，即需要 $O(n)$ 次 FormFunction 调用，这在大规模问题中代价很高。图着色方法利用 Jacobian 的稀疏结构，将列按照在同一行上不同时非零的原则分组，同一组中的列可以同时扰动，从而在一次残差计算中更新多列。

在 PETSc 中，这一机制由 FD coloring 实现；在 petsc4py 中，可以通过以下方式启用：

1
2
snes.setUseFD(True)
# or use `-snes_fd_color`

在这种模式下，用户通常只需要提供 FormFunction，并为 Jacobian 准备一个预分配好非零结构的稀疏矩阵容器：

1
snes.setJacobian(None, J=J, P=J)

SNES 会根据矩阵的非零结构自动生成着色方案，并使用有限差分填充 Jacobian。需要注意的是，图着色方法对 Jacobian 的稀疏结构假设是强依赖的。如果矩阵的预分配结构与真实 Jacobian 不匹配，可能导致装配效率低下，甚至得到错误的数值结果。

Poisson equation

下面考虑在一个2D正方形$[0,1]\times [0,1]$上求解泊松方程：
$$
-\nabla^2 u=f
$$
边界条件设置为$u|_{\partial \Omega}=0$，分别用下面三种方法求解：

  • 解析 Jacobian：人工装配 $J=A$（稀疏五点格式）
  • JFNK（matrix-free）：SNES 用有限差分计算 $Jv$，但仍然装配一个稀疏矩阵 $P$ 给 PC 做预处理
  • 图着色 + 有限差分 Jacobian：不给解析 Jacobian，让 PETSc 用 FD + coloring 去装配显式 $J$

其中，构造了人造解 $u(x,y)=\sin(\pi x)\cdot \sin(\pi y)$，满足：
$$
f(x,y)=2\pi^2 \cdot \sin(\pi x) \cdot \sin(\pi y)
$$
代码如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
import sys, petsc4py
petsc4py.init(sys.argv)
from petsc4py import PETSc
import numpy as np
SNES_REASONS = {
    v: k for k, v in vars(PETSc.SNES.ConvergedReason).items()
    if isinstance(v, int)
}
comm = PETSc.COMM_WORLD
opt = PETSc.Options()

jac_type = opt.getString("jac_type", "analytic")
nx = opt.getInt("nx", 50)
ny = opt.getInt("ny", 50)
N = nx * ny
hx = 1.0 / (nx - 1)
hy = 1.0 / (ny - 1)

def get_coords(row, nx): return row % nx, row // nx
def exact_u(x, y): return np.sin(np.pi * x) * np.sin(np.pi * y)
def rhs_f(x, y): return 2.0 * (np.pi**2) * np.sin(np.pi * x) * np.sin(np.pi * y)

def form_function(snes, X, F):
    # compute F = A*X - R
    ctx = snes.getApplicationContext()
    A, R = ctx["A"], ctx["R"]
    A.mult(X, F)      # F = A * X
    F.axpy(-1.0, R)   # F = F - R

def form_jacobian(snes, X, J, P):
    # compute J and P (P=J)
    ctx = snes.getApplicationContext()
    A = ctx["A"]

    # copy A to P
    if P.handle != A.handle:
        A.copy(P, structure=PETSc.Mat.Structure.SAME_NONZERO_PATTERN)
        P.assemble()
    
    # copy A to J
    if J.handle != P.handle and J.handle != A.handle:
        if J.getType() != "mffd": # JFNK
            A.copy(J, structure=PETSc.Mat.Structure.SAME_NONZERO_PATTERN)
        J.assemble()

# create matrix A
A = PETSc.Mat().createAIJ(size=(N, N), comm=comm)
A.setFromOptions()
A.setPreallocationNNZ(5) # 5-point stencil
A.setUp()

rstart, rend = A.getOwnershipRange()
cstart, cend = A.getOwnershipRangeColumn()

for row in range(rstart, rend):
    i, j = get_coords(row, nx)
    if i == 0 or j == 0 or i == nx - 1 or j == ny - 1:
        A.setValue(row, row, 1.0) # Dirichlet boundary, u=0
    else:
        # -Δu discretization
        inv_hx2 = 1.0 / (hx**2)
        inv_hy2 = 1.0 / (hy**2)
        diag_val = 2.0 * (inv_hx2 + inv_hy2)
        A.setValues(row, [row,      row-1,    row+1,    row-nx,   row+nx], 
                         [diag_val, -inv_hx2, -inv_hx2, -inv_hy2, -inv_hy2])
A.assemble()

# create rhs `R`
R = A.createVecLeft() # F(X) = A*X - R
with R as r_arr:
    for row in range(rstart, rend):
        i, j = get_coords(row, nx)
        if i == 0 or j == 0 or i == nx - 1 or j == ny - 1:
            r_arr[row - rstart] = 0.0 # u = 0
        else:
            r_arr[row - rstart] = rhs_f(i * hx, j * hy)
R.assemble()

# create SNES solver
snes = PETSc.SNES().create(comm=comm)
snes.setApplicationContext({"A": A, "R": R})

F = A.createVecLeft() # F(X) = A*X - R
J = A.copy()
P = A.copy()

# setup SNES: form_function, form_jacobian, options
snes.setFunction(form_function, F)
if jac_type == "analytic":
    snes.setJacobian(form_jacobian, J, P)
elif jac_type == "jfnk":
    snes.setUseMF(True) # -snes_mf
    snes.setJacobian(form_jacobian, J, P)
elif jac_type == "color":
    snes.setUseFD(True)  # -snes_fd_color
    J.zeroEntries()
    snes.setJacobian(None, J) # let P=J
else: raise ValueError(f"Unknown {jac_type=}")
snes.setFromOptions()

# solve F(X)=0
X = A.createVecRight() # F(X) = A*X - R
X.set(0.0)
snes.solve(None, X) # `None` stands for `0` in `F(X)=0`

if jac_type == "color":
    # output the coloring info
    func_calls = snes.getFunctionEvaluations()
    PETSc.Sys.Print(f"SNES form_function calls = {func_calls}")
    # or use `-mat_coloring_view` to see the coloring info

# compute exact solution
U_ex = A.createVecRight()
with U_ex as ue_arr:
    for col in range(cstart, cend):
        i, j = get_coords(col, nx)
        ue_arr[col - cstart] = exact_u(i * hx, j * hy)
U_ex.assemble()

X.viewFromOptions("-vec_view")          # -vec_view binary:result_analytic.bin
U_ex.viewFromOptions("-vec_view_exact") # -vec_view_exact binary:result_exact.bin

# error analysis
E = X.copy()
E.axpy(-1.0, U_ex) # E = X - U_ex

err_2 = E.norm(PETSc.NormType.NORM_2)
err_inf = E.norm(PETSc.NormType.NORM_INFINITY)
reason = SNES_REASONS[snes.getConvergedReason()]
iters = snes.getIterationNumber()
norm = snes.getFunctionNorm()

PETSc.Sys.Print(f"SNES done: \t{reason}\n"
                f"  iters   =\t{iters}\n"
                f"  |F|     =\t{norm:.2e}\n"
                f"  |e|_2   =\t{err_2:.2e}\n"
                f"  |e|_inf =\t{err_inf:.2e}")

对于显式Jacobian, 使用下面方式运行：

1
2
3
4
mpirun -n 4 python test.py \
  -nx 100 -ny 100 -jac_type analytic \
  -snes_type newtonls -snes_converged_reason -snes_monitor \
  -ksp_type cg -pc_type jacobi

对于JFNK, 使用下面方式运行：

1
2
3
4
mpirun -n 4 python test.py \
  -nx 100 -ny 100 -jac_type jfnk \
  -snes_type newtonls -snes_converged_reason -snes_monitor \
  -ksp_type cg -pc_type jacobi -ksp_monitor

对于图着色（FD coloring 装配 Jacobian），使用下面方式运行：

1
2
3
4
mpirun -n 4 python test.py \
  -nx 100 -ny 100 -jac_type color \
  -snes_type newtonls -snes_converged_reason -snes_monitor \
  -ksp_type cg -pc_type jacobi -ksp_monitor