PETSc 科学计算库实战指南

#include <petsc.h>

PetscErrorCode CreateLaplacian2D(MPI_Comm comm, PetscInt N, Mat *A) {
    Mat            mat;
    PetscInt       n = N * N;          // 总自由度
    PetscInt       rank, size;
    PetscInt       istart, iend;

    MPI_Comm_rank(comm, &rank);
    MPI_Comm_size(comm, &size);

    // 行范围划分
    istart = rank * (n / size);
    iend   = (rank == size - 1) ? n : (rank + 1) * (n / size);

    MatCreate(comm, &mat);
    MatSetSizes(mat, iend - istart, iend - istart, n, n);
    MatSetType(mat, MATAIJ);           // 自动选择 SEQ or MPI
    MatMPIAIJSetPreallocation(mat, 5, NULL, 2, NULL);  // 5 非零/行(本地)，2 非零/行(远端)
    MatSetOption(mat, MAT_NEW_NONZERO_LOCATIONS, PETSC_TRUE);

    // 遍历本地行，填充 5 点 stencil
    for (PetscInt idx = istart; idx < iend; idx++) {
        PetscInt i = idx % N;
        PetscInt j = idx / N;
        PetscScalar v = -1.0;
        PetscScalar center = 4.0;

        // 中心点
        MatSetValue(mat, idx, idx, center, INSERT_VALUES);

        // 邻居：上下左右
        if (i > 0)    MatSetValue(mat, idx, idx - 1, v, INSERT_VALUES);
        if (i < N-1)  MatSetValue(mat, idx, idx + 1, v, INSERT_VALUES);
        if (j > 0)    MatSetValue(mat, idx, idx - N, v, INSERT_VALUES);
        if (j < N-1)  MatSetValue(mat, idx, idx + N, v, INSERT_VALUES);
    }

    MatAssemblyBegin(mat, MAT_FINAL_ASSEMBLY);
    MatAssemblyEnd(mat, MAT_FINAL_ASSEMBLY);
    *A = mat;
    return 0;
}

KSP ksp;
Vec x, b;

KSPCreate(PETSC_COMM_WORLD, &ksp);
KSPSetOperators(ksp, A, A);         // 系统矩阵与预处理器矩阵

// 设置求解器类型（也可通过命令行 -ksp_type cg）
KSPSetType(ksp, KSPCG);              // 对称正定问题用 CG
// KSPSetType(ksp, KSPGMRES);        // 非对称用 GMRES
// KSPSetType(ksp, KSPBCGS);         // 非对称替代：BiCGSTAB

// 设置预处理
KSPGetPC(ksp, &pc);
PCSetType(pc, PCJACOBI);             // 简单并行可用 Jacobi
// PCSetType(pc, PCILU);             // 串行/小并行有效
// PCSetType(pc, PCHYPRE);           // BoomerAMG，大规模并行推荐

KSPSetFromOptions(ksp);              // 覆盖：优先采用命令行参数
KSPSolve(ksp, b, x);

DM da;
DMDACreate2d(PETSC_COMM_WORLD,
             DM_BOUNDARY_NONE, DM_BOUNDARY_NONE,  // 边界条件
             DMDA_STENCIL_STAR,                   // 5 点 stencil
             N, N, PETSC_DECIDE, PETSC_DECIDE,    // 全局网格
             1, 2, NULL, NULL, &da);              // 自由度 1, 重叠层数 2
DMSetFromOptions(da);
DMSetUp(da);

// 从 DM 创建矩阵，PETSc 自动处理并行分布
DMCreateMatrix(da, &A);
// 通过 DMDA 遍历本地网格点，填充矩阵

SNES snes;
Mat  J;      // Jacobian 矩阵
Vec  r;      // 残差向量

SNESCreate(PETSC_COMM_WORLD, &snes);
SNESSetFunction(snes, r, FormFunction, &user_ctx);   // 用户提供 F(x)
SNESSetJacobian(snes, J, J, FormJacobian, &user_ctx); // 用户提供 J = dF/dx

SNESSetType(snes, SNESNEWTONLS);      // 牛顿法 + 线搜索
SNESSetFromOptions(snes);
SNESSolve(snes, NULL, x);

// 可在运行时切换：-snes_type ksponly （不求 Jacobian，纯线性迭代）

# 运行并输出详细计时日志
mpiexec -n 64 ./poisson -log_view :poisson.log

# 通过 PETSc 内置 event 标记自定义代码段
PetscLogEvent  event;
PetscLogEventRegister("MyStencilCompute", 0, &event);
PetscLogEventBegin(event, 0, 0, 0, 0);
// ... 自定义计算 ...
PetscLogEventEnd(event, 0, 0, 0, 0);