盒子
盒子
文章目录
  1. Initialization
  2. Options
  3. Communicator

PETSc[01] petsc4py的基础知识

2026.01.01 Zhuoyu Chen petsc  热度 ℃

Initialization

在petsc4py中，需要通过下面代码进行初始化：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
import sys, petsc4py
petsc4py.init(sys.argv)
from petsc4py import PETSc

opt = PETSc.Options()                # get PETSc options
comm = PETSc.COMM_WORLD              # get the global communicator
rank = comm.getRank()                # get the rank of the process
size = comm.getSize()                # get the size of the communicator
PETSc.Sys.setDefaultComm(comm)       # set default communicator

PETSc.Sys.Print(f"{opt.getAll()=}")  # see all PETSc options, only rank=0 prints
n = opt.getInt("n", 10)              # get integer option -n <value>, default 10
a = opt.getReal("a", 1.5)            # get real option -a <value>, default 1.5
s = opt.getString("s", "hello")      # get string option -s <value>, default "hello"

opt_pfx = PETSc.Options("test_")
test_x = opt_pfx.getInt("x", 42)     # get integer option -test_x <value>, default 42

PETSc.Sys.Print(f"{rank=}, {size=}")
PETSc.Sys.Print(f"{n=}, {a=}, {s=}")

# collect prints from all ranks and display them in order
PETSc.Sys.syncPrint(f"Hello from rank {rank}", comm=comm)
PETSc.Sys.syncFlush(comm=comm)

可以通过下面两种方式，来从命令行中传入参数（保存到了opt变量中），通过PETSc.Options().getAll()获取字典保存的结果，或者通过opt.get*方法，来读取（并设置默认值）命令行参数内容。更全面的函数，在PETSc的文档里可以查看到。

下面两种方法都可以进行含有参数的调用:

1
2
3
4
5
6
7
8
9
10
mpirun -np 4 python test.py -n 20 -test_x 30
PETSC_OPTIONS="-n 20 -test_x 30" mpirun -np 4 python test.py

# opt.getAll()={'n': '20', 'test_x': '30'}
# rank=0, size=4
# n=20, a=1.5, s='hello'
# Hello from rank 0
# Hello from rank 1
# Hello from rank 2
# Hello from rank 3

下面几个开关可以方便参数调试：

  • -help：查看PETSc所有的帮助选项
  • -options_view：启动时输出全部的 options
  • -options_left：输出没有被使用的参数列表
  • -options_file $FILENAME：从$FILENAME中读取参数列表，并加入到PETSc程序中，但是命令行参数的优先级更高

Options

PETSc里的一些常用对象（Vec, Mat, KSP ...），都可以通过.setFromOptions()方法，来通过命令行参数进行初始化。

如果需要特定的前缀进行区分，那么可以使用.setOptionsPrefix(...)，之后再进行.setFromOptions()。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
import sys, petsc4py
petsc4py.init(sys.argv)
from petsc4py import PETSc
comm = PETSc.COMM_WORLD
opt = PETSc.Options()
pprint = PETSc.Sys.Print

x = PETSc.Vec().create(comm=comm)
x.setSizes(1)
x.setFromOptions()        # -vec_type mpi
x.assemble()
pprint(f"{x.getType()=}")

y = PETSc.Vec().create(comm=PETSc.COMM_SELF)
y.setSizes(1)
y.setOptionsPrefix("y_")  # -y_vec_type seq
y.setFromOptions()
y.assemble()
pprint(f"{y.getType()=}")

运行结果：

1
2
3
4
5
6
7
python test.py -vec_type mpi -y_vec_type seqcuda -vec_view ascii:stdout
#  Vec Object: 1 MPI process
#   type: mpi
# Process [0]
# 0.
# x.getType()='mpi'
# y.getType()='seqcuda'

Communicator

在PETSc中，几乎所有核心对象(Vec, Mat, KSP, ...) 都绑定在某一个 MPI communicator 上。communicator 决定了哪些进程参与这个对象/数据如何分布，以及对象之间能否一起工作。

PETSc定义了常用的两个comm，分别为：

  • PETSc.COMM_WORLD：所有 MPI 进程共同参与
  • PETSc.COMM_SELF：每个进程各自独立（每个 rank 都有一个“单进程 communicator”）

下面代码，给出了在COMM_WORLD和在COMM_SELF上创建相同大小的Vec对象时，PETSc分配的向量类型和每个处理器上的向量大小：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
import sys, petsc4py
petsc4py.init(sys.argv)
from petsc4py import PETSc

rank = PETSc.COMM_WORLD.getRank()
size = PETSc.COMM_WORLD.getSize()

x = PETSc.Vec().create(comm=PETSc.COMM_WORLD)
x.setSizes(8)
x.setFromOptions()
x.assemble()
PETSc.Sys.syncPrint(f"[{rank}/{size}]: {x.getType()=}, {x.getLocalSize()=}", comm=PETSc.COMM_WORLD)
PETSc.Sys.syncFlush()

x = PETSc.Vec().create(comm=PETSc.COMM_SELF)
x.setSizes(8)
x.setFromOptions()
x.assemble()
PETSc.Sys.syncPrint(f"[{rank}/{size}]: {x.getType()=}, {x.getLocalSize()=}", comm=PETSc.COMM_WORLD)
PETSc.Sys.syncFlush()

运行结果为：

1
2
3
4
5
mpirun -np 2 python test.py
# [0/2]: x.getType()='mpi', x.getLocalSize()=4
# [1/2]: x.getType()='mpi', x.getLocalSize()=4
# [0/2]: x.getType()='seq', x.getLocalSize()=8
# [1/2]: x.getType()='seq', x.getLocalSize()=8

可以看到，在COMM_WORLD中创建向量对象，每个处理器会平均分配到向量的一部分内容（类型为mpi）；而在COMM_SELF上创建向量对象，每个处理器会分配一个完整的向量（同时，类型会被设置为序列对象seq）。

为了在处理器上，正确地进行数据间通信，PETSc要求进行相同操作的对象位于同一个communicator中，比如：

  • MatMult(A, x, y)：A/x/y 处在相同 comm 中
  • KSP.setOperators(A)：ksp 与 A 处在相同 comm 中
  • SNES/TS 与其内部用到的 Vec/Mat/DM：处在相同 comm 中

在创建PETSc对象时，可以通过.create(comm=PETSc.COMM_WORLD)来确定communicator环境。

利用communicator环境，可以方便的进行代码调试和输出，比如：

  • PETSc.Sys.Print("hello", comm=PETSc.COMM_WORLD)只会在comm中rank=0的处理器上进行输出

  • PETSc.Sys.syncPrint(f"hello {rank}", comm=comm)会按照rank的顺序，依次进行输出

  • PETSc.Sys.syncFlush(comm=comm)需要在执行完syncPrint之后，刷新输出缓存

在更复杂的应用里，可能希望把 COMM_WORLD 的进程拆成多个组，每个组运行一个独立的 PETSc 子问题，例如：

  • 同一脚本里同时求解多个互不耦合的线性系统
  • 多物理场：不同子模型用不同进程数
  • 一个组负责计算，另一个组负责 I/O 或后处理

通过不同的communicator分组，可以将没有耦合关系的求解过程进行解耦，减少额外的同步等待时间。

通过PETSc.Comm.tompi4py()可以将PETSc中的communicator,转换成为mpi4py中comm, 并进行进一步的切割，比如：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
import sys, petsc4py
petsc4py.init(sys.argv)
from petsc4py import PETSc

comm = PETSc.COMM_WORLD
rank = comm.getRank()
size = comm.getSize()
mpi_world = comm.tompi4py()

color = rank % 2
comm_sub = mpi_world.Split(color=color, key=rank)

rank_sub = comm_sub.Get_rank()
size_sub = comm_sub.Get_size()
PETSc.Sys.syncPrint(f"[{rank}/{size}]: local({rank_sub}/{size_sub}), {color=}", comm=comm)
PETSc.Sys.syncFlush(comm=comm)

运行结果如下：

1
2
3
4
5
mpirun -np 4 python test.py
# [0/4]: local(0/2), color=0
# [1/4]: local(0/2), color=1
# [2/4]: local(1/2), color=0
# [3/4]: local(1/2), color=1

其中，我们通过comm.Split(color=color, key=key)进行分割，保证：

  • color 决定属于哪一组，即相同 color 的进程会进入同一个子 communicator
  • key 决定子 communicator 内的 rank 排序，一般直接用原 rank 即可
  • 如果 color = MPI.UNDEFINED，它会得到 MPI_COMM_NULL，表示“不属于任何子 communicator”，可以通过这个设置来将某个处理器单独剥离出来