Actual source code: ex90.c

  1: static char help[] = "Tests KSPMatSolve() with KSPRICHARDSON and preconditioners that provide PCApplyRichardson().\n\
  2: Use -set_initial_guess to fill both blocks of solutions with the same nonzero initial guess and -compare false to skip the comparison against KSPSolve().\n\
  3: Use -nullspace to solve instead with a singular operator that has a constant null space, and -nullspace_attach false to leave that null space off the operator.\n\
  4: Use -shell to precondition with a PCSHELL that provides PCApplyRichardson() and -shell_block to make it provide PCMatApplyRichardson() as well.\n\
  5: Use -resize to solve a second system of twice the size with the same KSP, and -resize_uneven to grow it by one row instead, which changes the local size on some processes only.\n\
  6: Use -transpose to solve a second time with a nonsymmetric operator and KSPMatSolveTranspose().\n\
  7: Use -selfscale to solve a second time after turning KSPRichardsonSetSelfScale() on.\n\
  8: Use -n to set the size of the system and -nrhs to set the number of right-hand sides.\n\n";

 10: #include <petscksp.h>

 12: /*
 13:    Single Jacobi sweep with a zero initial guess, x <- D^{-1} b, the context is the reciprocal of the diagonal of the operator
 14: */
 15: static PetscErrorCode Apply_User(PC pc, Vec b, Vec x)
 16: {
 17:   Vec dinv;

 19:   PetscFunctionBeginUser;
 20:   PetscCall(PCShellGetContext(pc, &dinv));
 21:   PetscCall(VecPointwiseMult(x, b, dinv));
 22:   PetscFunctionReturn(PETSC_SUCCESS);
 23: }

 25: /*
 26:    Jacobi sweeps, x <- x + D^{-1} (b - A x), the context is the reciprocal of the diagonal of the operator
 27: */
 28: static PetscErrorCode ApplyRichardson_User(PC pc, Vec b, Vec x, Vec r, PetscReal rtol, PetscReal abstol, PetscReal dtol, PetscInt maxits, PetscBool guesszero, PetscInt *its, PCRichardsonConvergedReason *reason)
 29: {
 30:   Mat      A;
 31:   Vec      dinv;
 32:   PetscInt i;

 34:   PetscFunctionBeginUser;
 35:   PetscCall(PCShellGetContext(pc, &dinv));
 36:   PetscCall(PCGetOperators(pc, &A, NULL));
 37:   for (i = 0; i < maxits; i++) {
 38:     if (i == 0 && guesszero) PetscCall(VecCopy(b, r));
 39:     else {
 40:       PetscCall(MatMult(A, x, r));
 41:       PetscCall(VecAYPX(r, -1.0, b));
 42:     }
 43:     PetscCall(VecPointwiseMult(r, r, dinv));
 44:     if (i == 0 && guesszero) PetscCall(VecCopy(r, x));
 45:     else PetscCall(VecAXPY(x, 1.0, r));
 46:   }
 47:   *its    = maxits;
 48:   *reason = PCRICHARDSON_CONVERGED_ITS;
 49:   PetscFunctionReturn(PETSC_SUCCESS);
 50: }

 52: /*
 53:    Block analog of ApplyRichardson_User(), KSPRICHARDSON passes W = NULL so the work block of vectors is allocated here
 54: */
 55: static PetscErrorCode MatApplyRichardson_User(PC pc, Mat B, Mat X, Mat W, PetscReal rtol, PetscReal abstol, PetscReal dtol, PetscInt maxits, PetscBool guesszero, PetscInt *its, PCRichardsonConvergedReason *reason)
 56: {
 57:   Mat      A, R;
 58:   Vec      dinv;
 59:   PetscInt i;

 61:   PetscFunctionBeginUser;
 62:   PetscCall(PCShellGetContext(pc, &dinv));
 63:   PetscCall(PCGetOperators(pc, &A, NULL));
 64:   PetscCall(MatDuplicate(B, MAT_DO_NOT_COPY_VALUES, &R));
 65:   /* set the product A X up once so that only its numeric phase is run in the loop below */
 66:   PetscCall(MatProductCreateWithMat(A, X, NULL, R));
 67:   PetscCall(MatProductSetType(R, MATPRODUCT_AB));
 68:   PetscCall(MatProductSetFromOptions(R));
 69:   PetscCall(MatProductSymbolic(R));
 70:   for (i = 0; i < maxits; i++) {
 71:     if (i == 0 && guesszero) PetscCall(MatCopy(B, R, SAME_NONZERO_PATTERN));
 72:     else {
 73:       PetscCall(MatProductNumeric(R));
 74:       PetscCall(MatAYPX(R, -1.0, B, SAME_NONZERO_PATTERN));
 75:     }
 76:     PetscCall(MatDiagonalScale(R, dinv, NULL));
 77:     if (i == 0 && guesszero) PetscCall(MatCopy(R, X, SAME_NONZERO_PATTERN));
 78:     else PetscCall(MatAXPY(X, 1.0, R, SAME_NONZERO_PATTERN));
 79:   }
 80:   PetscCall(MatProductClear(R));
 81:   PetscCall(MatDestroy(&R));
 82:   *its    = maxits;
 83:   *reason = PCRICHARDSON_CONVERGED_ITS;
 84:   PetscFunctionReturn(PETSC_SUCCESS);
 85: }

 87: /*
 88:    Tridiagonal operator, singular with a constant null space when nullspace is PETSC_TRUE (1D Laplacian with Neumann ends), nonsymmetric when nonsymmetric is PETSC_TRUE so that a transpose solve differs from a forward one
 89: */
 90: static PetscErrorCode CreateOperator(PetscInt n, PetscBool nullspace, PetscBool nonsymmetric, Mat *A, MatNullSpace *nsp)
 91: {
 92:   PetscInt i, Istart, Iend;

 94:   PetscFunctionBeginUser;
 95:   *nsp = NULL;
 96:   PetscCall(MatCreate(PETSC_COMM_WORLD, A));
 97:   PetscCall(MatSetSizes(*A, PETSC_DECIDE, PETSC_DECIDE, n, n));
 98:   PetscCall(MatSetFromOptions(*A));
 99:   PetscCall(MatSetUp(*A));
100:   PetscCall(MatGetOwnershipRange(*A, &Istart, &Iend));
101:   for (i = Istart; i < Iend; i++) {
102:     if (nullspace) PetscCall(MatSetValue(*A, i, i, i == 0 || i == n - 1 ? 1.0 : 2.0, INSERT_VALUES));
103:     else PetscCall(MatSetValue(*A, i, i, 4.0, INSERT_VALUES));
104:     if (i > 0) PetscCall(MatSetValue(*A, i, i - 1, nonsymmetric ? -2.0 : -1.0, INSERT_VALUES));
105:     if (i < n - 1) PetscCall(MatSetValue(*A, i, i + 1, -1.0, INSERT_VALUES));
106:   }
107:   PetscCall(MatAssemblyBegin(*A, MAT_FINAL_ASSEMBLY));
108:   PetscCall(MatAssemblyEnd(*A, MAT_FINAL_ASSEMBLY));
109:   if (nullspace) {
110:     PetscCall(MatNullSpaceCreate(PETSC_COMM_WORLD, PETSC_TRUE, 0, NULL, nsp));
111:     PetscCall(MatSetNullSpace(*A, *nsp));
112:   }
113:   PetscFunctionReturn(PETSC_SUCCESS);
114: }

116: /*
117:    Solve with the whole block of right-hand sides at once, then one right-hand side at a time, and compare the two blocks of solutions
118: */
119: static PetscErrorCode SolveAndCompare(KSP ksp, Mat A, MatNullSpace nsp, PetscInt nrhs, PetscBool set_initial_guess, PetscBool compare, PetscBool transpose)
120: {
121:   Mat       B, X, Y;
122:   Vec       cb, cx;
123:   PetscInt  i, n, m;
124:   PetscReal nrm;
125:   PetscBool guess_nonzero = PETSC_FALSE;

127:   PetscFunctionBeginUser;
128:   /* block of right-hand sides and the two blocks of solutions to compare */
129:   PetscCall(MatGetSize(A, &n, NULL));
130:   PetscCall(MatGetLocalSize(A, &m, NULL));
131:   PetscCall(MatCreateDense(PETSC_COMM_WORLD, m, PETSC_DECIDE, n, nrhs, NULL, &B));
132:   PetscCall(MatSetRandom(B, NULL));
133:   if (nsp) { /* project the null space out of each right-hand side so that the singular systems are consistent */
134:     for (i = 0; i < nrhs; i++) {
135:       PetscCall(MatDenseGetColumnVec(B, i, &cb));
136:       PetscCall(MatNullSpaceRemove(nsp, cb));
137:       PetscCall(MatDenseRestoreColumnVec(B, i, &cb));
138:     }
139:   }
140:   PetscCall(MatDuplicate(B, MAT_DO_NOT_COPY_VALUES, &X));
141:   PetscCall(MatDuplicate(B, MAT_DO_NOT_COPY_VALUES, &Y));
142:   PetscCall(MatZeroEntries(X));
143:   PetscCall(MatZeroEntries(Y));
144:   if (set_initial_guess) { /* the same nonzero initial guess in both blocks of solutions */
145:     PetscCall(MatCopy(B, X, SAME_NONZERO_PATTERN));
146:     PetscCall(MatScale(X, 100.0));
147:     PetscCall(MatCopy(X, Y, SAME_NONZERO_PATTERN));
148:   }
149:   PetscCall(KSPGetInitialGuessNonzero(ksp, &guess_nonzero));

151:   if (transpose) PetscCall(KSPMatSolveTranspose(ksp, B, X));
152:   else PetscCall(KSPMatSolve(ksp, B, X));
153:   for (i = 0; i < nrhs; i++) {
154:     PetscCall(MatDenseGetColumnVecRead(B, i, &cb));
155:     /* KSPSolve() reads the column when it is used as the initial guess */
156:     if (guess_nonzero) PetscCall(MatDenseGetColumnVec(Y, i, &cx));
157:     else PetscCall(MatDenseGetColumnVecWrite(Y, i, &cx));
158:     if (transpose) PetscCall(KSPSolveTranspose(ksp, cb, cx));
159:     else PetscCall(KSPSolve(ksp, cb, cx));
160:     if (guess_nonzero) PetscCall(MatDenseRestoreColumnVec(Y, i, &cx));
161:     else PetscCall(MatDenseRestoreColumnVecWrite(Y, i, &cx));
162:     PetscCall(MatDenseRestoreColumnVecRead(B, i, &cb));
163:   }
164:   /* with tolerance-based stopping, block-aggregate and per-column convergence legitimately differ at the level of the relative tolerance, so the comparison is only meaningful for fixed-iteration runs */
165:   if (compare) {
166:     PetscCall(MatAXPY(Y, -1.0, X, SAME_NONZERO_PATTERN));
167:     PetscCall(MatNorm(Y, NORM_FROBENIUS, &nrm));
168:     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "KSPMatSolve() and KSPSolve() %s\n", nrm < PETSC_SQRT_MACHINE_EPSILON ? "agree" : "disagree"));
169:   }
170:   PetscCall(MatDestroy(&Y));
171:   PetscCall(MatDestroy(&X));
172:   PetscCall(MatDestroy(&B));
173:   PetscFunctionReturn(PETSC_SUCCESS);
174: }

176: int main(int argc, char **args)
177: {
178:   Mat          A;
179:   MatNullSpace nsp  = NULL;
180:   Vec          dinv = NULL;
181:   KSP          ksp;
182:   PC           pc;
183:   PetscInt     n = 20, nrhs = 5;
184:   PetscBool    shell = PETSC_FALSE, block = PETSC_FALSE, set_initial_guess = PETSC_FALSE, compare = PETSC_TRUE, nullspace = PETSC_FALSE, resize = PETSC_FALSE, resize_uneven = PETSC_FALSE, selfscale = PETSC_FALSE, transpose = PETSC_FALSE;
185:   PetscBool    nullspace_attach = PETSC_TRUE;

187:   PetscFunctionBeginUser;
188:   PetscCall(PetscInitialize(&argc, &args, NULL, help));
189:   PetscCall(PetscOptionsGetInt(NULL, NULL, "-n", &n, NULL));
190:   PetscCall(PetscOptionsGetInt(NULL, NULL, "-nrhs", &nrhs, NULL));
191:   PetscCall(PetscOptionsGetBool(NULL, NULL, "-shell", &shell, NULL));
192:   PetscCall(PetscOptionsGetBool(NULL, NULL, "-shell_block", &block, NULL));
193:   PetscCall(PetscOptionsGetBool(NULL, NULL, "-set_initial_guess", &set_initial_guess, NULL));
194:   PetscCall(PetscOptionsGetBool(NULL, NULL, "-compare", &compare, NULL));
195:   PetscCall(PetscOptionsGetBool(NULL, NULL, "-nullspace", &nullspace, NULL));
196:   PetscCall(PetscOptionsGetBool(NULL, NULL, "-nullspace_attach", &nullspace_attach, NULL));
197:   PetscCall(PetscOptionsGetBool(NULL, NULL, "-resize", &resize, NULL));
198:   PetscCall(PetscOptionsGetBool(NULL, NULL, "-resize_uneven", &resize_uneven, NULL));
199:   PetscCall(PetscOptionsGetBool(NULL, NULL, "-selfscale", &selfscale, NULL));
200:   PetscCall(PetscOptionsGetBool(NULL, NULL, "-transpose", &transpose, NULL));

202:   PetscCall(CreateOperator(n, nullspace, transpose, &A, &nsp));
203:   /* a failed preconditioner flags the block of solutions with infinities, and removing a null space from such a block turns them into NaN in complex arithmetic, so leaving the null space off the
204:      operator lets the failure path be exercised on a singular operator, the local MatNullSpace is still used to make the right-hand sides consistent */
205:   if (nsp && !nullspace_attach) PetscCall(MatSetNullSpace(A, NULL));

207:   PetscCall(KSPCreate(PETSC_COMM_WORLD, &ksp));
208:   PetscCall(KSPSetOperators(ksp, A, A));
209:   PetscCall(KSPSetType(ksp, KSPRICHARDSON));
210:   if (shell) {
211:     PetscCall(MatCreateVecs(A, NULL, &dinv));
212:     PetscCall(MatGetDiagonal(A, dinv));
213:     PetscCall(VecReciprocal(dinv));
214:     PetscCall(KSPGetPC(ksp, &pc));
215:     PetscCall(PCSetType(pc, PCSHELL));
216:     PetscCall(PCShellSetName(pc, "Jacobi sweeps"));
217:     PetscCall(PCShellSetContext(pc, dinv));
218:     PetscCall(PCShellSetApply(pc, Apply_User));
219:     PetscCall(PCShellSetApplyRichardson(pc, ApplyRichardson_User));
220:     if (block) PetscCall(PCShellSetMatApplyRichardson(pc, MatApplyRichardson_User));
221:   }
222:   PetscCall(KSPSetFromOptions(ksp));

224:   PetscCall(SolveAndCompare(ksp, A, nsp, nrhs, set_initial_guess, compare, PETSC_FALSE));

226:   /* the transpose of the operator is not the operator itself, so a transpose block solve on the same KSP must not reuse the product set up by the forward one */
227:   if (transpose) PetscCall(SolveAndCompare(ksp, A, nsp, nrhs, set_initial_guess, compare, PETSC_TRUE));

229:   /* the self-scaled variant needs more work vectors than the plain one, so turning it on after a first solve must run KSPSetUp() again */
230:   if (selfscale) {
231:     PetscCall(KSPRichardsonSetSelfScale(ksp, PETSC_TRUE));
232:     PetscCall(SolveAndCompare(ksp, A, nsp, nrhs, set_initial_guess, compare, PETSC_FALSE));
233:   }

235:   /* KSPReset() allows a second operator of a different size to be set on the KSP and ensures that its implementation and PC rebuild their work space,
236:      -resize_uneven grows the operator by a single row so that the local size changes on the first process only */
237:   if (resize || resize_uneven) {
238:     PetscCall(MatNullSpaceDestroy(&nsp));
239:     PetscCall(MatDestroy(&A));
240:     PetscCall(CreateOperator(resize_uneven ? n + 1 : 2 * n, nullspace, transpose, &A, &nsp));
241:     if (nsp && !nullspace_attach) PetscCall(MatSetNullSpace(A, NULL)); /* the second operator is built the same way, so it follows the same choice */
242:     PetscCall(KSPReset(ksp));
243:     PetscCall(KSPSetOperators(ksp, A, A));
244:     if (shell) { /* the context of the PCSHELL is the reciprocal of the diagonal of the operator, so it must be rebuilt as well */
245:       PetscCall(VecDestroy(&dinv));
246:       PetscCall(MatCreateVecs(A, NULL, &dinv));
247:       PetscCall(MatGetDiagonal(A, dinv));
248:       PetscCall(VecReciprocal(dinv));
249:       PetscCall(PCShellSetContext(pc, dinv));
250:     }
251:     PetscCall(SolveAndCompare(ksp, A, nsp, nrhs, set_initial_guess, compare, PETSC_FALSE));
252:   }

254:   PetscCall(MatNullSpaceDestroy(&nsp));
255:   PetscCall(VecDestroy(&dinv));
256:   PetscCall(KSPDestroy(&ksp));
257:   PetscCall(MatDestroy(&A));
258:   PetscCall(PetscFinalize());
259:   return 0;
260: }

262: /*TEST

264:    testset:
265:       output_file: output/ex90.out
266:       nsize: {{1 2}}
267:       args: -ksp_type richardson -ksp_max_it 5 -ksp_norm_type none

269:       test:
270:          suffix: shell_matapplyrichardson
271:          args: -shell -shell_block

273:       test:
274:          suffix: shell_applyrichardson
275:          args: -shell

277:       test:
278:          suffix: shell_matapplyrichardson_guess
279:          args: -shell -shell_block -ksp_initial_guess_nonzero

281:       test:
282:          suffix: shell_applyrichardson_guess
283:          args: -shell -ksp_initial_guess_nonzero

285:       test:
286:          suffix: sor
287:          args: -pc_type sor

289:       test:
290:          suffix: jacobi
291:          args: -pc_type jacobi

293:       test:
294:          suffix: nullspace
295:          args: -nullspace -pc_type jacobi

297:    # a second system of a different size solved after KSPReset(), which must rebuild the KSP and PC work space for the new operator
298:    testset:
299:       output_file: output/ex90_resize.out
300:       nsize: {{1 2}}
301:       args: -ksp_type richardson -ksp_max_it 5 -ksp_norm_type none -resize

303:       test:
304:          suffix: resize_sor
305:          args: -pc_type sor

307:       test:
308:          suffix: resize_shell_applyrichardson
309:          args: -shell

311:       test:
312:          suffix: resize_jacobi
313:          args: -pc_type jacobi

315:    # KSPConvergedDefault() removes the null space of the operator from the left-preconditioned block of right-hand sides, as in the Vec path
316:    test:
317:       suffix: nullspace_guess
318:       nsize: {{1 2}}
319:       output_file: output/ex90.out
320:       args: -ksp_type richardson -ksp_max_it 5 -ksp_rtol 1e-50 -ksp_norm_type preconditioned -nullspace -pc_type jacobi -set_initial_guess -ksp_initial_guess_nonzero

322:    # the singular operator makes the Cholesky factorization fail, so the block iteration flags the whole block of solutions with MatFlag() and stops with KSP_DIVERGED_PC_FAILED,
323:    # its null space is left off the operator since removing a null space from a block of infinities would turn them into NaN in complex arithmetic, and -fp_trap 0
324:    # overrides the option some harness configurations pass since the zero-pivot factorization and the flagged infinities raise floating point exceptions by design
325:    test:
326:       suffix: pc_failed
327:       nsize: 1
328:       args: -ksp_type richardson -nullspace -nullspace_attach false -pc_type cholesky -pc_factor_shift_type none -compare false -ksp_converged_reason -fp_trap 0

330:    # the same failure with batching, so that MatFlag() is called on a MatDenseGetSubMatrix() view
331:    test:
332:       suffix: pc_failed_batch
333:       nsize: 1
334:       args: -ksp_type richardson -nullspace -nullspace_attach false -pc_type cholesky -pc_factor_shift_type none -compare false -ksp_converged_reason -ksp_matsolve_batch_size 2 -fp_trap 0

336:    # two batches of the same global width but of different local column layouts, the cached work blocks of KSPMatSolve_Richardson() must be rebuilt for the second one
337:    test:
338:       suffix: batch_layout
339:       nsize: 2
340:       args: -ksp_type richardson -ksp_max_it 5 -ksp_norm_type none -pc_type jacobi -ksp_matsolve_batch_size 2

342:    # the local column layouts of the two batches match on the second process only, so the decision to rebuild the cached work blocks must be reduced over the processes
343:    test:
344:       suffix: batch_layout_uneven
345:       nsize: 3
346:       args: -ksp_type richardson -ksp_max_it 5 -ksp_norm_type none -pc_type jacobi -nrhs 6 -ksp_matsolve_batch_size 3

348:    # the local row count of the second operator changes on the first process only, so KSPReset() must allow setup with the new layout
349:    test:
350:       suffix: resize_uneven_sor
351:       nsize: 2
352:       args: -ksp_type richardson -ksp_max_it 5 -ksp_norm_type none -pc_type sor -resize_uneven

354:    # a forward block solve followed by a transpose one on the same KSP, the cached product of KSPMatSolve_Richardson() must be set up again for the other direction
355:    test:
356:       suffix: transpose
357:       nsize: {{1 2}}
358:       output_file: output/ex90_transpose.out
359:       args: -ksp_type richardson -ksp_max_it 5 -ksp_norm_type none -pc_type jacobi -transpose

361:    # KSPRichardsonSetSelfScale() after a first solve must run KSPSetUp() again, the self-scaled variant has no block analog so KSPMatSolve() falls back to solving one right-hand side at a time
362:    test:
363:       suffix: selfscale
364:       nsize: 1
365:       args: -ksp_type richardson -ksp_max_it 5 -ksp_norm_type none -pc_type jacobi -selfscale

367:    test:
368:       suffix: jacobi_guess
369:       args: -ksp_type richardson -pc_type jacobi -set_initial_guess -compare false -ksp_initial_guess_nonzero -ksp_norm_type {{unpreconditioned preconditioned}} -ksp_rtol 1e-4 -ksp_max_it 100 -ksp_converged_reason

371:    test:
372:       suffix: jacobi_guess_batch
373:       args: -ksp_type richardson -pc_type jacobi -set_initial_guess -compare false -ksp_initial_guess_nonzero -ksp_norm_type unpreconditioned -ksp_rtol 1e-4 -ksp_max_it 100 -ksp_converged_reason -ksp_matsolve_batch_size 2

375: TEST*/