Vec

在PETSc中，创建Vec对象往往需要以下几个步骤：

x = PETSc.Vec().create(comm=comm)
x.setSizes(100)
x.setFromOptions()
x.assemble()

其中：

create()负责在对应的communicator上创建Vec对象，或者通过createSeq, createMPI来指定创建方式
setSizes() 定义的是逻辑大小，并行情况下 PETSc 会根据 communicator 自动决定每个 rank 的本地大小
setFromOptions() 让 Vec 的相关配置(如VecType)可以由命令行控制
assemble() 是必要的结束步骤，在Vec.pyx文件中，是通过进行assemblyBegin, assemblyEnd实现的，目的是在所有处理器上完成Vec对象的组装；注意，如果在命令行指定了-vec_view，那么需要调用assemble()后，才会执行view()方法。

一般来说，还需要通过x.setUp()来完成创建，但是如果使用了setFromOptions()，或者进行了set()，又或者使用了assemble()，那么PETSc将会自动进行setUp()部分。

对于-vec_type的设置，在PETSc的VecType文档中，可以查询到不同的类型和含义：

VecType	值	含义
`VECSEQ`	seq	基本的顺序（单 MPI 进程）向量（CPU）
`VECMPI`	mpi	基本的并行（多 MPI 进程分布式）向量（CPU）
`VECSTANDARD`	standard	seq on one process and mpi on multiple
`VECSHARED`	shared	共享内存并行向量（历史上只在特定平台真正是“shared”，很多情况下等同 `VECMPI`）。目前 `VecCreateShared()` 仅在 SGI 上可用；否则，此例程与 VecCreateMPI () 相同。
`VECSEQVIENNACL`	seqviennacl	顺序向量，使用 ViennaCL（OpenCL）后端
`VECMPIVIENNACL`	mpiviennacl	并行向量，使用 ViennaCL（OpenCL）后端
`VECVIENNACL`	viennacl	seqviennacl on one process and mpiviennacl on multiple
`VECSEQCUDA`	seqcuda	顺序向量，使用 CUDA 后端
`VECMPICUDA`	mpicuda	并行向量，使用 CUDA 后端
`VECCUDA`	cuda	seqcuda on one process and mpicuda on multiple
`VECSEQHIP`	seqhip	顺序向量，使用 HIP 后端
`VECMPIHIP`	mpihip	并行向量，使用 HIP 后端
`VECHIP`	hip	seqhip on one process and mpihip on multiple
`VECNEST`	nest	由多个子向量“拼起来”的嵌套向量，每个子向量独立存储
`VECSEQKOKKOS`	seqkokkos	顺序向量，使用 Kokkos 后端
`VECMPIKOKKOS`	mpikokkos	并行向量，使用 Kokkos 后端
`VECKOKKOS`	kokkos	seqkokkos on one process and mpikokkos on multiple

其中，standard, viennacl, cuda, hip, kokkos会根据处理器的数量，来自动选择串行版本或并行版本。

在另一个文档中，可以看到具体的外部包实现：

Format	Vector Types	External Packages	Details
Dense array	`VECSTANDARD`	BLAS
	`VECCUDA`	NVIDIA’s cuBLAS	NVIDIA GPUs
	`VECHIP`	AMD’s RocBLAS	AMD GPUs
	`VECKOKKOS`	Kokkos	GPUs, CPUs, OpenMP
Nested	`VECNEST`		Provides efficient access to inner vectors

这些计算后端的区别在于：

CUDA
- NVIDIA 的专有 GPU 编程体系
- PETSc 的 CUDA Vec 和 Mat 直接基于 CUDA runtime
HIP
- AMD 的 GPU 编程体系
- PETSc 的 HIP Vec 和 Mat 对应 AMD GPU。
ViennaCL
- 基于 OpenCL 的线性代数库
- 理论上可以跑在很多设备上，但在实践中，可能生态和性能都不如 CUDA/HIP/Kokkos
Kokkos
- 这是一个抽象层，不是某一家 GPU 的底层 API
- Kokkos 可以在下面接 CUDA、HIP、OpenMP、Threads 等后端。

在多处理器上，并行 Vec 中的每个 rank 只拥有一段连续的全局索引区间。可以通过：

1	lo, hi = x.getOwnershipRange()

来查询当前 rank 负责的全局索引范围 [lo, hi)。

还可以通过下面的函数，来获得不同程度的分配大小信息：

getLocalSize() -> (nl)：返回当前处理器上的长度
getSize() -> (ng)：返回总长度
getSizes() -> (nl, ng)：返回当前处理器上的长度、以及总长度

在Vec.pyx文件中，实现了__enter__, __exit__这两个成员函数，因此可以通过with vec as vec_arr:来将PETSc中的向量对象转化为可以直接访问和修改的内存表示。更具体的实现函数在petscvec.pxi文件中实现。

如果要在全局环境下，设置Vec的具体数值，可以使用下面的函数：

x.set(value)：将Vec的所有位置设置为同一个值value，这个不需要assemble()，直接调用就可以完成初始值设置
x.setValue(index:int, value:Scalar, addv:InsertModeSpec = None)
x.setValues(indices: Sequence[int], values: Sequence[Scalar], addv:InsertModeSpec = None)

以及对应的getValue, getValues方法。注意，使用setValue, setValues之后，需要立刻进行assemble()组装Vec，否则操作可能会被缓存，导致没有正确修改数组内容。在官方源代码中如此指出：

Notes
-----
The values may be cached so `assemblyBegin` and `assemblyEnd`
must be called after all calls of this method are completed.

Multiple calls to `setValue` cannot be made with different values
for ``addv`` without intermediate calls to `assemblyBegin` and
`assemblyEnd`.

另一方面，如果只是进行位置上的赋值操作，那么如下直接进行就可以：

1 2	x.setValue(9, 100.0, addv=PETSc.InsertMode.INSERT_VALUES) x.assemble()

但是，如果要进行累加操作，那么只能在单个处理器上进行，否则会被执行多遍：

1
2
3

if rank == 0:
    x.setValue(8, 100.0, addv=PETSc.InsertMode.ADD_VALUES)
x.assemble()

因此，最好避免全局性的修改，而是进行with ... as ...:的局部修改。

如果需要创建含有ghost boundary的Vec，可以考虑createGhost, ghostUpdate, localForm等一系列函数，不过后续我们会发现，可以使用DA系列更好的管理ghost，因此我们不必在这里深入研究这些函数。

对于VECNEST，可以当成一个Vec的包装器，用法如下：

# assume u, p are Vec objects
field = PETSc.Vec().createNest([u, p], comm=comm)
u_sub, p_sub = field.getNestSubVecs()
with u_sub as u_loc: ...
with p_sub as p_loc: ...

如果不想依赖DA等网格管理器，比如自己实现分块矩阵求解：
$$
\begin{pmatrix}
A & B^T \\
B & -C
\end{pmatrix}
\begin{pmatrix}
u \\ p
\end{pmatrix}
=\begin{pmatrix}
f \\ g
\end{pmatrix}
$$
那么可以用Vec().createNest来创建u, p和f, g，然后通过Mat().createNest来创建左侧的完整矩阵：

X = PETSc.Vec().createNest([u, p], comm=comm); X.assemble()
F = PETSc.Vec().createNest([f, g], comm=comm); F.assemble()
M = PETSc.Mat().createNest([[A, Bt], [B, minusC]], comm=comm); M.assemble()
# solve M*X=F with KSP

对于createNest部分的测试代码如下：

import os
os.environ["PETSC_OPTIONS"] = " ".join([
    "-ksp_type gmres",
    "-ksp_rtol 1e-10",
    "-ksp_monitor",
    "-pc_type fieldsplit",
    "-pc_fieldsplit_type schur",
    "-pc_fieldsplit_schur_fact_type full",
    "-fieldsplit_u_ksp_type preonly",
    "-fieldsplit_u_pc_type jacobi",
    "-fieldsplit_p_ksp_type preonly",
    "-fieldsplit_p_pc_type jacobi",
])

import sys, petsc4py
petsc4py.init(sys.argv)
from petsc4py import PETSc
comm = PETSc.COMM_WORLD
rank = comm.getRank()
opt = PETSc.Options()

Nu, Np = 8, 4

# create u, p
u = PETSc.Vec().create(comm=comm)
u.setSizes(Nu)
u.setFromOptions()
lu = u.getLocalSize()

p = PETSc.Vec().create(comm=comm)
p.setSizes(Np)
p.setFromOptions()
lp = p.getLocalSize()

# create f, g
f = u.duplicate()
g = p.duplicate()
f.set(1.0)
g.set(1.0)

# create nested vectors
X = PETSc.Vec().createNest([u, p], comm=comm)
F = PETSc.Vec().createNest([f, g], comm=comm)

# create block matrix
A = PETSc.Mat().createAIJ(size=((lu, Nu), (lu, Nu)), comm=comm)
A.setUp()
diag_u = u.duplicate()
diag_u.set(1.0)
A.setDiagonal(diag_u)
A.assemble()

C = PETSc.Mat().createAIJ(size=((lp, Np), (lp, Np)), comm=comm)
C.setUp()
diag_p = p.duplicate()
diag_p.set(1.0)
C.setDiagonal(diag_p)
C.assemble()
minusC = C.duplicate(copy=True)
minusC.scale(-1.0)

B = PETSc.Mat().createAIJ(size=((lp, Np), (lu, Nu)), comm=comm)
B.assemble() # 0.0
Bt = B.transpose()

# create nested matrix
M_nest = PETSc.Mat().createNest([[A, Bt], [B, minusC]], comm=comm)
M_nest.assemble()
M = M_nest.convert("aij")

# create KSP solver
ksp = PETSc.KSP().create(comm=comm)
ksp.setOperators(M)

# initialize PC as fieldsplit
pc = ksp.getPC()
pc.setType("fieldsplit")

# create index sets for fieldsplit
u_is = PETSc.IS().createStride(Nu, first=0, step=1, comm=comm)
p_is = PETSc.IS().createStride(Np, first=Nu, step=1, comm=comm)
pc.setFieldSplitIS(("u", u_is), ("p", p_is))

# use Schur complement for pressure block
pc.setFieldSplitType(PETSc.PC.CompositeType.SCHUR)
pc.setFieldSplitSchurFactType(PETSc.PC.SchurFactType.FULL)

# solve the system
ksp.setFromOptions()
ksp.solve(F, X)

u_sol, p_sol = X.getNestSubVecs()
PETSc.Sys.Print(f"Converged reason: {ksp.getConvergedReason()}")
PETSc.Sys.Print(f"Iterations: {ksp.getIterationNumber()}")
PETSc.Sys.syncPrint(f"[rank {rank}] |u|={u_sol.norm():.2e}, |p|={p_sol.norm():.2e}", comm=comm)
PETSc.Sys.syncFlush(comm=comm)

输出结果如下：

mpirun -np 4 python test.py
#   0 KSP Residual norm 3.464101615138e+00
#   1 KSP Residual norm 1.914854215513e+00
#   2 KSP Residual norm 1.058630546810e-15
# Converged reason: 2
# Iterations: 2
# [rank 0] |u|=2.83e+00, |p|=2.00e+00
# [rank 1] |u|=2.83e+00, |p|=2.00e+00
# [rank 2] |u|=2.83e+00, |p|=2.00e+00
# [rank 3] |u|=2.83e+00, |p|=2.00e+00

下面代码展示了关于Vec的大部分常用的使用方法：

import sys, petsc4py
petsc4py.init(sys.argv)
from petsc4py import PETSc
import numpy as np

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

x = PETSc.Vec().create(comm=comm)
x.setSizes(10)
x.setFromOptions()

with x as local_x:
    lo, hi = x.getOwnershipRange()
    local_x[:] = np.arange(lo, hi)
x.assemble() # if -vec_view is given, assemble() is required.

if rank == 0:
    x.setValue(8, 100.0, addv=PETSc.InsertMode.ADD_VALUES)
x.assemble()

x.setValue(9, 100.0, addv=PETSc.InsertMode.INSERT_VALUES)
x.assemble()

PETSc.Sys.Print(f"Full ownership ranges: {x.getOwnershipRanges()}")

PETSc.Sys.Print(f"Array ranges:")
PETSc.Sys.syncPrint(f"  [{rank}/{size}]: ownership ranges: {x.getOwnershipRange()}", comm=comm)
PETSc.Sys.syncFlush(comm=comm)

PETSc.Sys.Print(f"Array values:")
PETSc.Sys.syncPrint(f"  [{rank}/{size}]: {x.getArray(readonly=True)}", comm=comm)
PETSc.Sys.syncFlush(comm=comm)

# x0 is SEQ on rank=0, empty on others
scatter, x0 = PETSc.Scatter.toZero(x)
scatter.scatter(x, x0,
    addv=PETSc.InsertMode.INSERT,
    mode=PETSc.ScatterMode.FORWARD)
PETSc.Sys.Print(f"Gatherd array: {x0.getArray(readonly=True)}", comm=comm)

# print sizes: (local_size, global_size)
PETSc.Sys.Print(f"Array sizes(local_size, global_size):")
PETSc.Sys.syncPrint(f"  [{rank}/{size}]: x0:{x0.getSizes()}; x:{x.getSizes()}", comm=comm)
PETSc.Sys.syncFlush(comm=comm)

# remember to destroy objects if not needed anymore
scatter.destroy()
x0.destroy()
x.destroy()

最终得到输出结果：

mpirun -np 4 python test.py
# Full ownership ranges: [ 0  3  6  8 10]
# Array ranges:
#   [0/4]: ownership ranges: (0, 3)
#   [1/4]: ownership ranges: (3, 6)
#   [2/4]: ownership ranges: (6, 8)
#   [3/4]: ownership ranges: (8, 10)
# Array values:
#   [0/4]: [0. 1. 2.]
#   [1/4]: [3. 4. 5.]
#   [2/4]: [6. 7.]
#   [3/4]: [108. 100.]
# Gatherd array: [  0.   1.   2.   3.   4.   5.   6.   7. 108. 100.]
# Array sizes(local_size, global_size):
#   [0/4]: x0:(10, 10); x:(3, 10)
#   [1/4]: x0:(0, 0); x:(3, 10)
#   [2/4]: x0:(0, 0); x:(2, 10)
#   [3/4]: x0:(0, 0); x:(2, 10)