Actual source code: snesj2.c

  1: #include <petsc/private/snesimpl.h>
  2: #include <petscdm.h>

  4: /*
  5:    MatFDColoringSetFunction() takes a function with four arguments, we want to use SNESComputeFunction()
  6:    since it logs function computation information.
  7: */
  8: static PetscErrorCode SNESComputeFunctionCtx(SNES snes, Vec x, Vec f, void *ctx)
  9: {
 10:   return SNESComputeFunction(snes, x, f);
 11: }
 12: static PetscErrorCode SNESComputeMFFunctionCtx(SNES snes, Vec x, Vec f, void *ctx)
 13: {
 14:   return SNESComputeMFFunction(snes, x, f);
 15: }

 17: /*@C
 18:   SNESComputeJacobianDefaultColor - Computes the Jacobian using
 19:   finite differences and coloring to exploit matrix sparsity.

 21:   Collective

 23:   Input Parameters:
 24: + snes - nonlinear solver object
 25: . x1   - location at which to evaluate Jacobian
 26: - ctx  - `MatFDColoring` context or `NULL`

 28:   Output Parameters:
 29: + J - Jacobian matrix (not altered in this routine)
 30: - B - newly computed Jacobian matrix to use with preconditioner (generally the same as `J`)

 32:   Level: intermediate

 34:   Options Database Keys:
 35: + -snes_fd_color_use_mat       - use a matrix coloring from the explicit matrix nonzero pattern instead of from the `DM` providing the matrix
 36: . -snes_fd_color               - Activates `SNESComputeJacobianDefaultColor()` in `SNESSetFromOptions()`
 37: . -mat_fd_coloring_err <err>   - Sets <err> (square root of relative error in the function)
 38: . -mat_fd_coloring_umin <umin> - Sets umin, the minimum allowable u-value magnitude
 39: . -mat_fd_type                 - Either wp or ds (see `MATMFFD_WP` or `MATMFFD_DS`)
 40: . -snes_mf_operator            - Use matrix-free application of Jacobian
 41: - -snes_mf                     - Use matrix-free Jacobian with no explicit Jacobian representation

 43:   Notes:
 44:   If the coloring is not provided through the context, this will first try to get the
 45:   coloring from the `DM`.  If the `DM` has no coloring routine, then it will try to
 46:   get the coloring from the matrix.  This requires that the matrix have its nonzero locations already provided.

 48:   `SNES` supports three approaches for computing (approximate) Jacobians: user provided via `SNESSetJacobian()`, matrix-free via `SNESSetUseMatrixFree()`,
 49:   and computing explicitly with finite differences and coloring using `MatFDColoring`. It is also possible to use automatic differentiation
 50:   and the `MatFDColoring` object, see src/ts/tutorials/autodiff/ex16adj_tl.cxx

 52:   This function can be provided to `SNESSetJacobian()` along with an appropriate sparse matrix to hold the Jacobian

 54: .seealso: `SNES`, `SNESSetJacobian()`, `SNESTestJacobian()`, `SNESComputeJacobianDefault()`, `SNESSetUseMatrixFree()`,
 55:           `MatFDColoringCreate()`, `MatFDColoringSetFunction()`
 56: @*/
 57: PetscErrorCode SNESComputeJacobianDefaultColor(SNES snes, Vec x1, Mat J, Mat B, void *ctx)
 58: {
 59:   MatFDColoring color = (MatFDColoring)ctx;
 60:   DM            dm;
 61:   MatColoring   mc;
 62:   ISColoring    iscoloring;
 63:   PetscBool     hascolor;
 64:   PetscBool     solvec, matcolor = PETSC_FALSE;
 65:   DMSNES        dms;

 67:   PetscFunctionBegin;
 69:   if (!color) PetscCall(PetscObjectQuery((PetscObject)B, "SNESMatFDColoring", (PetscObject *)&color));

 71:   if (!color) {
 72:     PetscCall(SNESGetDM(snes, &dm));
 73:     PetscCall(DMHasColoring(dm, &hascolor));
 74:     matcolor = PETSC_FALSE;
 75:     PetscCall(PetscOptionsGetBool(((PetscObject)snes)->options, ((PetscObject)snes)->prefix, "-snes_fd_color_use_mat", &matcolor, NULL));
 76:     if (hascolor && !matcolor) {
 77:       PetscCall(DMCreateColoring(dm, IS_COLORING_GLOBAL, &iscoloring));
 78:     } else {
 79:       PetscCall(MatColoringCreate(B, &mc));
 80:       PetscCall(MatColoringSetDistance(mc, 2));
 81:       PetscCall(MatColoringSetType(mc, MATCOLORINGSL));
 82:       PetscCall(MatColoringSetFromOptions(mc));
 83:       PetscCall(MatColoringApply(mc, &iscoloring));
 84:       PetscCall(MatColoringDestroy(&mc));
 85:     }
 86:     PetscCall(MatFDColoringCreate(B, iscoloring, &color));
 87:     PetscCall(DMGetDMSNES(dm, &dms));
 88:     if (dms->ops->computemffunction) {
 89:       PetscCall(MatFDColoringSetFunction(color, (PetscErrorCode(*)(void))SNESComputeMFFunctionCtx, NULL));
 90:     } else {
 91:       PetscCall(MatFDColoringSetFunction(color, (PetscErrorCode(*)(void))SNESComputeFunctionCtx, NULL));
 92:     }
 93:     PetscCall(MatFDColoringSetFromOptions(color));
 94:     PetscCall(MatFDColoringSetUp(B, iscoloring, color));
 95:     PetscCall(ISColoringDestroy(&iscoloring));
 96:     PetscCall(PetscObjectCompose((PetscObject)B, "SNESMatFDColoring", (PetscObject)color));
 97:     PetscCall(PetscObjectDereference((PetscObject)color));
 98:   }

100:   /* F is only usable if there is no RHS on the SNES and the full solution corresponds to x1 */
101:   PetscCall(VecEqual(x1, snes->vec_sol, &solvec));
102:   if (!snes->vec_rhs && solvec) {
103:     Vec F;
104:     PetscCall(SNESGetFunction(snes, &F, NULL, NULL));
105:     PetscCall(MatFDColoringSetF(color, F));
106:   }
107:   PetscCall(MatFDColoringApply(B, color, x1, snes));
108:   if (J != B) {
109:     PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY));
110:     PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY));
111:   }
112:   PetscFunctionReturn(PETSC_SUCCESS);
113: }

115: /*@C
116:   SNESPruneJacobianColor - Remove nondiagonal zeros in the Jacobian matrix and update the `MatMFFD` coloring information.

118:   Collective

120:   Input Parameters:
121: + snes - the `SNES` context
122: . J    - Jacobian matrix (not altered in this routine)
123: - B    - newly computed Jacobian matrix to use with preconditioner (generally the same as `J`)

125:   Level: intermediate

127:   Notes:
128:   This function improves the `MatMFFD` coloring performance when the Jacobian matrix is overallocated or contains
129:   many constant zeros entries, which is typically the case when the matrix is generated by a `DM`
130:   and multiple fields are involved.

132:   Users need to make sure that the Jacobian matrix is properly filled to reflect the sparsity
133:   structure. For `MatMFFD` coloring, the values of nonzero entries are not important. So one can
134:   usually call `SNESComputeJacobian()` with randomized input vectors to generate a dummy Jacobian.
135:   `SNESComputeJacobian()` should be called before `SNESSolve()` but after `SNESSetUp()`.

137: .seealso: `SNESComputeJacobianDefaultColor()`, `MatEliminateZeros()`, `MatFDColoringCreate()`, `MatFDColoringSetFunction()`
138: @*/
139: PetscErrorCode SNESPruneJacobianColor(SNES snes, Mat J, Mat B)
140: {
141:   DM            dm;
142:   DMSNES        dms;
143:   MatColoring   mc;
144:   ISColoring    iscoloring;
145:   MatFDColoring matfdcoloring;

147:   PetscFunctionBegin;
148:   /* Generate new coloring after eliminating zeros in the matrix */
149:   PetscCall(MatEliminateZeros(B, PETSC_TRUE));
150:   PetscCall(MatColoringCreate(B, &mc));
151:   PetscCall(MatColoringSetDistance(mc, 2));
152:   PetscCall(MatColoringSetType(mc, MATCOLORINGSL));
153:   PetscCall(MatColoringSetFromOptions(mc));
154:   PetscCall(MatColoringApply(mc, &iscoloring));
155:   PetscCall(MatColoringDestroy(&mc));
156:   /* Replace the old coloring with the new one */
157:   PetscCall(MatFDColoringCreate(B, iscoloring, &matfdcoloring));
158:   PetscCall(SNESGetDM(snes, &dm));
159:   PetscCall(DMGetDMSNES(dm, &dms));
160:   // Comment out the following branch to bypass the coverage test. You can uncomment it when needed.
161:   //if (dms->ops->computemffunction) {
162:   //  PetscCall(MatFDColoringSetFunction(matfdcoloring, (PetscErrorCode(*)(void))SNESComputeMFFunctionCtx, NULL));
163:   //} else {
164:   PetscCall(MatFDColoringSetFunction(matfdcoloring, (PetscErrorCode(*)(void))SNESComputeFunctionCtx, NULL));
165:   //}
166:   PetscCall(MatFDColoringSetFromOptions(matfdcoloring));
167:   PetscCall(MatFDColoringSetUp(B, iscoloring, matfdcoloring));
168:   PetscCall(PetscObjectCompose((PetscObject)B, "SNESMatFDColoring", (PetscObject)matfdcoloring));
169:   PetscCall(PetscObjectDereference((PetscObject)matfdcoloring));
170:   PetscCall(ISColoringDestroy(&iscoloring));
171:   PetscFunctionReturn(PETSC_SUCCESS);
172: }