1: static const char help[] = "Tests PCMG setup with DMKSP{Create|Compute}Operators.\n\n";

  3: #include <petscksp.h>
  4: #include <petscdmda.h>

  6: typedef struct {
  7:   PetscBool same_operator;
  8:   PetscInt  ncreate;
  9:   PetscInt  ncompute;
 10: } AppCtx;

 12: static PetscErrorCode AssembleDiagonal(Mat A, PetscScalar diag)
 13: {
 14:   PetscInt rstart, rend;

 16:   PetscFunctionBeginUser;
 17:   PetscCall(MatZeroEntries(A));
 18:   PetscCall(MatGetOwnershipRange(A, &rstart, &rend));
 19:   for (PetscInt row = rstart; row < rend; row++) PetscCall(MatSetValue(A, row, row, diag, INSERT_VALUES));
 20:   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
 21:   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
 22:   PetscFunctionReturn(PETSC_SUCCESS);
 23: }

 25: static PetscErrorCode CreateOperators(KSP ksp, Mat *A, Mat *P, void *ctx)
 26: {
 27:   AppCtx *user = (AppCtx *)ctx;
 28:   DM      dm;

 30:   PetscFunctionBeginUser;
 31:   PetscCall(KSPGetDM(ksp, &dm));
 32:   PetscCall(DMCreateMatrix(dm, A));
 33:   if (!user->same_operator) PetscCall(DMCreateMatrix(dm, P));
 34:   else *P = *A; /* No need to increment ref count of A */
 35:   user->ncreate++;
 36:   PetscFunctionReturn(PETSC_SUCCESS);
 37: }

 39: static PetscErrorCode ComputeOperators(KSP ksp, Mat A, Mat P, void *ctx)
 40: {
 41:   AppCtx *user = (AppCtx *)ctx;

 43:   PetscFunctionBeginUser;
 44:   PetscCall(AssembleDiagonal(P, 3.0));
 45:   if (A != P) PetscCall(AssembleDiagonal(A, 2.0));
 46:   user->ncompute++;
 47:   PetscFunctionReturn(PETSC_SUCCESS);
 48: }

 50: int main(int argc, char **argv)
 51: {
 52:   AppCtx    user;
 53:   DM        dm;
 54:   KSP       ksp, cksp;
 55:   PC        pc;
 56:   Mat       A, P;
 57:   PetscBool user_finest = PETSC_TRUE, expected_same_operator, use_amat;
 58:   PetscInt  nl, expected_ncreate, expected_ncompute;

 60:   PetscFunctionBeginUser;
 61:   user.same_operator = PETSC_FALSE;
 62:   user.ncreate       = 0;
 63:   user.ncompute      = 0;

 65:   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
 66:   PetscCall(PetscOptionsGetBool(NULL, NULL, "-user_finest", &user_finest, NULL));
 67:   PetscCall(PetscOptionsGetBool(NULL, NULL, "-same_operator", &user.same_operator, NULL));

 69:   PetscCall(DMDACreate1d(PETSC_COMM_WORLD, DM_BOUNDARY_NONE, 3, 1, 1, NULL, &dm));
 70:   PetscCall(DMSetFromOptions(dm));
 71:   PetscCall(DMSetUp(dm));

 73:   PetscCall(DMKSPSetCreateOperators(dm, CreateOperators, &user));
 74:   PetscCall(DMKSPSetComputeOperators(dm, ComputeOperators, &user));

 76:   PetscCall(KSPCreate(PETSC_COMM_WORLD, &ksp));
 77:   PetscCall(KSPSetDM(ksp, dm));
 78:   PetscCall(KSPSetDMActive(ksp, KSP_DMACTIVE_OPERATOR, (PetscBool)!user_finest));
 79:   PetscCall(KSPSetType(ksp, KSPPREONLY));
 80:   PetscCall(KSPGetPC(ksp, &pc));
 81:   PetscCall(PCSetType(pc, PCMG));
 82:   if (user_finest) {
 83:     PetscCall(DMCreateMatrix(dm, &A));
 84:     PetscCall(AssembleDiagonal(A, 2.0));
 85:     P = A;
 86:     if (!user.same_operator) {
 87:       PetscCall(DMCreateMatrix(dm, &P));
 88:       PetscCall(AssembleDiagonal(P, 3.0));
 89:     }
 90:     PetscCall(KSPSetOperators(ksp, A, P));
 91:     PetscCall(MatDestroy(&A));
 92:     if (!user.same_operator) PetscCall(MatDestroy(&P));
 93:   }
 94:   PetscCall(KSPSetFromOptions(ksp));
 95:   PetscCall(KSPSetUp(ksp));
 96:   PetscCall(KSPViewFromOptions(ksp, NULL, "-ksp_view"));

 98:   PetscCall(PCMGGetLevels(pc, &nl));
 99:   PetscCall(PCGetUseAmat(pc, &use_amat));
100:   expected_ncreate       = user_finest ? nl - 1 : nl;
101:   expected_ncompute      = user_finest ? nl - 1 : nl;
102:   expected_same_operator = (PetscBool)(!use_amat || user.same_operator);
103:   PetscCheck(user.ncreate == expected_ncreate, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "Expected PCMG to create %" PetscInt_FMT " operators: found %" PetscInt_FMT, expected_ncreate, user.ncreate);
104:   PetscCheck(user.ncompute == expected_ncompute, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "Expected PCMG to compute %" PetscInt_FMT " operators: found %" PetscInt_FMT, expected_ncompute, user.ncompute);

106:   for (PetscInt l = nl - 1; l > 0; l--) {
107:     KSP sksp;

109:     PetscCall(PCMGGetSmootherDown(pc, l, &sksp));
110:     PetscCall(KSPGetOperators(sksp, &A, &P));
111:     if (!expected_same_operator) PetscCheck(A != P, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "Smooth KSP operators were not distinct");
112:     else PetscCheck(A == P, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "Smooth KSP operators were not Pmat-only");
113:   }
114:   PetscCall(PCMGGetCoarseSolve(pc, &cksp));
115:   PetscCall(KSPGetOperators(cksp, &A, &P));
116:   if (!expected_same_operator) PetscCheck(A != P, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "Coarse KSP operators were not distinct");
117:   else PetscCheck(A == P, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "Coarse KSP operators were not Pmat-only");
118:   PetscCall(KSPDestroy(&ksp));
119:   PetscCall(DMDestroy(&dm));
120:   PetscCall(PetscFinalize());
121:   return 0;
122: }

124: /*TEST

126:    test:
127:       args: -pc_mg_galerkin none -user_finest {{0 1}} -same_operator {{0 1}} -da_refine {{0 1 2}} -pc_use_amat {{0 1}}
128:       output_file: output/empty.out

130: TEST*/