Actual source code: dmdasnes.c

  1: #include <petscdmda.h>
  2: #include <petsc/private/dmimpl.h>
  3: #include <petsc/private/snesimpl.h>

  5: /* This structure holds the user-provided DMDA callbacks */
  6: typedef struct {
  7:   /* array versions for vector data */
  8:   DMDASNESFunctionFn  *residuallocal;
  9:   DMDASNESJacobianFn  *jacobianlocal;
 10:   DMDASNESObjectiveFn *objectivelocal;

 12:   /* Vec version for vector data */
 13:   DMDASNESFunctionVecFn  *residuallocalvec;
 14:   DMDASNESJacobianVecFn  *jacobianlocalvec;
 15:   DMDASNESObjectiveVecFn *objectivelocalvec;

 17:   /* user contexts */
 18:   void      *residuallocalctx;
 19:   void      *jacobianlocalctx;
 20:   void      *objectivelocalctx;
 21:   InsertMode residuallocalimode;

 23:   /*   For Picard iteration defined locally */
 24:   PetscErrorCode (*rhsplocal)(DMDALocalInfo *, void *, void *, PetscCtx);
 25:   PetscErrorCode (*jacobianplocal)(DMDALocalInfo *, void *, Mat, Mat, PetscCtx);
 26:   void *picardlocalctx;
 27: } DMSNES_DA;

 29: static PetscErrorCode DMSNESDestroy_DMDA(DMSNES sdm)
 30: {
 31:   PetscFunctionBegin;
 32:   PetscCall(PetscFree(sdm->data));
 33:   PetscFunctionReturn(PETSC_SUCCESS);
 34: }

 36: static PetscErrorCode DMSNESDuplicate_DMDA(DMSNES oldsdm, DMSNES sdm)
 37: {
 38:   PetscFunctionBegin;
 39:   PetscCall(PetscNew((DMSNES_DA **)&sdm->data));
 40:   if (oldsdm->data) PetscCall(PetscMemcpy(sdm->data, oldsdm->data, sizeof(DMSNES_DA)));
 41:   PetscFunctionReturn(PETSC_SUCCESS);
 42: }

 44: static PetscErrorCode DMDASNESGetContext(DM dm, DMSNES sdm, DMSNES_DA **dmdasnes)
 45: {
 46:   PetscFunctionBegin;
 47:   *dmdasnes = NULL;
 48:   if (!sdm->data) {
 49:     PetscCall(PetscNew((DMSNES_DA **)&sdm->data));
 50:     sdm->ops->destroy   = DMSNESDestroy_DMDA;
 51:     sdm->ops->duplicate = DMSNESDuplicate_DMDA;
 52:   }
 53:   *dmdasnes = (DMSNES_DA *)sdm->data;
 54:   PetscFunctionReturn(PETSC_SUCCESS);
 55: }

 57: static PetscErrorCode SNESComputeFunction_DMDA(SNES snes, Vec X, Vec F, PetscCtx ctx)
 58: {
 59:   DM            dm;
 60:   DMSNES_DA    *dmdasnes = (DMSNES_DA *)ctx;
 61:   DMDALocalInfo info;
 62:   Vec           Xloc;
 63:   void         *x, *f, *rctx;

 65:   PetscFunctionBegin;
 69:   PetscCheck(dmdasnes->residuallocal || dmdasnes->residuallocalvec, PetscObjectComm((PetscObject)snes), PETSC_ERR_PLIB, "Corrupt context");
 70:   PetscCall(SNESGetDM(snes, &dm));
 71:   PetscCall(DMGetLocalVector(dm, &Xloc));
 72:   PetscCall(DMGlobalToLocalBegin(dm, X, INSERT_VALUES, Xloc));
 73:   PetscCall(DMGlobalToLocalEnd(dm, X, INSERT_VALUES, Xloc));
 74:   PetscCall(DMDAGetLocalInfo(dm, &info));
 75:   rctx = dmdasnes->residuallocalctx ? dmdasnes->residuallocalctx : snes->ctx;
 76:   switch (dmdasnes->residuallocalimode) {
 77:   case INSERT_VALUES: {
 78:     PetscCall(PetscLogEventBegin(SNES_FunctionEval, snes, X, F, 0));
 79:     if (dmdasnes->residuallocalvec) PetscCallBack("SNES DMDA local callback function", (*dmdasnes->residuallocalvec)(&info, Xloc, F, rctx));
 80:     else {
 81:       PetscCall(DMDAVecGetArray(dm, Xloc, &x));
 82:       PetscCall(DMDAVecGetArray(dm, F, &f));
 83:       PetscCallBack("SNES DMDA local callback function", (*dmdasnes->residuallocal)(&info, x, f, rctx));
 84:       PetscCall(DMDAVecRestoreArray(dm, Xloc, &x));
 85:       PetscCall(DMDAVecRestoreArray(dm, F, &f));
 86:     }
 87:     PetscCall(PetscLogEventEnd(SNES_FunctionEval, snes, X, F, 0));
 88:   } break;
 89:   case ADD_VALUES: {
 90:     Vec Floc;
 91:     PetscCall(DMGetLocalVector(dm, &Floc));
 92:     PetscCall(VecZeroEntries(Floc));
 93:     PetscCall(PetscLogEventBegin(SNES_FunctionEval, snes, X, F, 0));
 94:     if (dmdasnes->residuallocalvec) PetscCallBack("SNES DMDA local callback function", (*dmdasnes->residuallocalvec)(&info, Xloc, Floc, rctx));
 95:     else {
 96:       PetscCall(DMDAVecGetArray(dm, Xloc, &x));
 97:       PetscCall(DMDAVecGetArray(dm, Floc, &f));
 98:       PetscCallBack("SNES DMDA local callback function", (*dmdasnes->residuallocal)(&info, x, f, rctx));
 99:       PetscCall(DMDAVecRestoreArray(dm, Xloc, &x));
100:       PetscCall(DMDAVecRestoreArray(dm, Floc, &f));
101:     }
102:     PetscCall(PetscLogEventEnd(SNES_FunctionEval, snes, X, F, 0));
103:     PetscCall(VecZeroEntries(F));
104:     PetscCall(DMLocalToGlobalBegin(dm, Floc, ADD_VALUES, F));
105:     PetscCall(DMLocalToGlobalEnd(dm, Floc, ADD_VALUES, F));
106:     PetscCall(DMRestoreLocalVector(dm, &Floc));
107:   } break;
108:   default:
109:     SETERRQ(PetscObjectComm((PetscObject)snes), PETSC_ERR_ARG_INCOMP, "Cannot use imode=%d", (int)dmdasnes->residuallocalimode);
110:   }
111:   PetscCall(DMRestoreLocalVector(dm, &Xloc));
112:   PetscCall(VecFlag(F, snes->functiondomainerror));
113:   PetscFunctionReturn(PETSC_SUCCESS);
114: }

116: static PetscErrorCode SNESComputeObjective_DMDA(SNES snes, Vec X, PetscReal *ob, PetscCtx ctx)
117: {
118:   DM            dm;
119:   DMSNES_DA    *dmdasnes = (DMSNES_DA *)ctx;
120:   DMDALocalInfo info;
121:   Vec           Xloc;
122:   void         *x, *octx;

124:   PetscFunctionBegin;
127:   PetscAssertPointer(ob, 3);
128:   PetscCheck(dmdasnes->objectivelocal || dmdasnes->objectivelocalvec, PetscObjectComm((PetscObject)snes), PETSC_ERR_PLIB, "Corrupt context");
129:   PetscCall(SNESGetDM(snes, &dm));
130:   PetscCall(DMGetLocalVector(dm, &Xloc));
131:   PetscCall(DMGlobalToLocalBegin(dm, X, INSERT_VALUES, Xloc));
132:   PetscCall(DMGlobalToLocalEnd(dm, X, INSERT_VALUES, Xloc));
133:   PetscCall(DMDAGetLocalInfo(dm, &info));
134:   octx = dmdasnes->objectivelocalctx ? dmdasnes->objectivelocalctx : snes->ctx;
135:   if (dmdasnes->objectivelocalvec) PetscCallBack("SNES DMDA local callback objective", (*dmdasnes->objectivelocalvec)(&info, Xloc, ob, octx));
136:   else {
137:     PetscCall(DMDAVecGetArray(dm, Xloc, &x));
138:     PetscCallBack("SNES DMDA local callback objective", (*dmdasnes->objectivelocal)(&info, x, ob, octx));
139:     PetscCall(DMDAVecRestoreArray(dm, Xloc, &x));
140:   }
141:   PetscCall(DMRestoreLocalVector(dm, &Xloc));
142:   PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, ob, 1, MPIU_REAL, MPIU_SUM, PetscObjectComm((PetscObject)snes)));
143:   PetscFunctionReturn(PETSC_SUCCESS);
144: }

146: /* Routine is called by example, hence must be labeled PETSC_EXTERN */
147: PETSC_EXTERN PetscErrorCode SNESComputeJacobian_DMDA(SNES snes, Vec X, Mat A, Mat B, PetscCtx ctx)
148: {
149:   DM            dm;
150:   DMSNES_DA    *dmdasnes = (DMSNES_DA *)ctx;
151:   DMDALocalInfo info;
152:   Vec           Xloc;
153:   void         *x, *jctx;

155:   PetscFunctionBegin;
156:   PetscCheck(dmdasnes->residuallocal || dmdasnes->residuallocalvec, PetscObjectComm((PetscObject)snes), PETSC_ERR_PLIB, "Corrupt context");
157:   PetscCall(SNESGetDM(snes, &dm));
158:   jctx = dmdasnes->jacobianlocalctx ? dmdasnes->jacobianlocalctx : snes->ctx;
159:   if (dmdasnes->jacobianlocal || dmdasnes->jacobianlocalvec) {
160:     PetscCall(DMGetLocalVector(dm, &Xloc));
161:     PetscCall(DMGlobalToLocalBegin(dm, X, INSERT_VALUES, Xloc));
162:     PetscCall(DMGlobalToLocalEnd(dm, X, INSERT_VALUES, Xloc));
163:     PetscCall(DMDAGetLocalInfo(dm, &info));
164:     if (dmdasnes->jacobianlocalvec) PetscCallBack("SNES DMDA local callback Jacobian", (*dmdasnes->jacobianlocalvec)(&info, Xloc, A, B, jctx));
165:     else {
166:       PetscCall(DMDAVecGetArray(dm, Xloc, &x));
167:       PetscCallBack("SNES DMDA local callback Jacobian", (*dmdasnes->jacobianlocal)(&info, x, A, B, jctx));
168:       PetscCall(DMDAVecRestoreArray(dm, Xloc, &x));
169:     }
170:     PetscCall(DMRestoreLocalVector(dm, &Xloc));
171:   } else {
172:     MatFDColoring fdcoloring;
173:     PetscCall(PetscObjectQuery((PetscObject)dm, "DMDASNES_FDCOLORING", (PetscObject *)&fdcoloring));
174:     if (!fdcoloring) {
175:       ISColoring coloring;

177:       PetscCall(DMCreateColoring(dm, dm->coloringtype, &coloring));
178:       PetscCall(MatFDColoringCreate(B, coloring, &fdcoloring));
179:       switch (dm->coloringtype) {
180:       case IS_COLORING_GLOBAL:
181:         PetscCall(MatFDColoringSetFunction(fdcoloring, (MatFDColoringFn *)SNESComputeFunction_DMDA, dmdasnes));
182:         break;
183:       default:
184:         SETERRQ(PetscObjectComm((PetscObject)snes), PETSC_ERR_SUP, "No support for coloring type '%s'", ISColoringTypes[dm->coloringtype]);
185:       }
186:       PetscCall(PetscObjectSetOptionsPrefix((PetscObject)fdcoloring, ((PetscObject)dm)->prefix));
187:       PetscCall(MatFDColoringSetFromOptions(fdcoloring));
188:       PetscCall(MatFDColoringSetUp(B, coloring, fdcoloring));
189:       PetscCall(ISColoringDestroy(&coloring));
190:       PetscCall(PetscObjectCompose((PetscObject)dm, "DMDASNES_FDCOLORING", (PetscObject)fdcoloring));
191:       PetscCall(PetscObjectDereference((PetscObject)fdcoloring));

193:       /* The following breaks an ugly reference counting loop that deserves a paragraph. MatFDColoringApply() will call
194:        * VecDuplicate() with the state Vec and store inside the MatFDColoring. This Vec will duplicate the Vec, but the
195:        * MatFDColoring is composed with the DM. We dereference the DM here so that the reference count will eventually
196:        * drop to 0. Note the code in DMDestroy() that exits early for a negative reference count. That code path will be
197:        * taken when the PetscObjectList for the Vec inside MatFDColoring is destroyed.
198:        */
199:       PetscCall(PetscObjectDereference((PetscObject)dm));
200:     }
201:     PetscCall(MatFDColoringApply(B, fdcoloring, X, snes));
202:   }
203:   /* This will be redundant if the user called both, but it's too common to forget. */
204:   if (A != B) {
205:     PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
206:     PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
207:   }
208:   PetscFunctionReturn(PETSC_SUCCESS);
209: }

211: /*@C
212:   DMDASNESSetFunctionLocal - set a local residual evaluation function for use with `DMDA`

214:   Logically Collective

216:   Input Parameters:
217: + dm    - `DM` to associate callback with
218: . imode - `INSERT_VALUES` if local function computes owned part, `ADD_VALUES` if it contributes to ghosted part
219: . func  - local residual evaluation
220: - ctx   - optional context for local residual evaluation

222:   Level: beginner

224: .seealso: [](ch_snes), `DMDA`, `DMDASNESFunctionFn`, `DMDASNESSetJacobianLocal()`, `DMSNESSetFunction()`, `DMDACreate1d()`, `DMDACreate2d()`, `DMDACreate3d()`
225: @*/
226: PetscErrorCode DMDASNESSetFunctionLocal(DM dm, InsertMode imode, DMDASNESFunctionFn *func, PetscCtx ctx)
227: {
228:   DMSNES     sdm;
229:   DMSNES_DA *dmdasnes;

231:   PetscFunctionBegin;
233:   PetscCall(DMGetDMSNESWrite(dm, &sdm));
234:   PetscCall(DMDASNESGetContext(dm, sdm, &dmdasnes));

236:   dmdasnes->residuallocalimode = imode;
237:   dmdasnes->residuallocal      = func;
238:   dmdasnes->residuallocalctx   = ctx;

240:   PetscCall(DMSNESSetFunction(dm, SNESComputeFunction_DMDA, dmdasnes));
241:   if (!sdm->ops->computejacobian) { /* Call us for the Jacobian too, can be overridden by the user. */
242:     PetscCall(DMSNESSetJacobian(dm, SNESComputeJacobian_DMDA, dmdasnes));
243:   }
244:   PetscFunctionReturn(PETSC_SUCCESS);
245: }

247: /*@C
248:   DMDASNESSetFunctionLocalVec - set a local residual evaluation function that operates on a local vector for `DMDA`

250:   Logically Collective

252:   Input Parameters:
253: + dm    - `DM` to associate callback with
254: . imode - `INSERT_VALUES` if local function computes owned part, `ADD_VALUES` if it contributes to ghosted part
255: . func  - local residual evaluation
256: - ctx   - optional context for local residual evaluation

258:   Level: beginner

260: .seealso: [](ch_snes), `DMDA`, `DMDASNESFunctionVecFn`, `DMDASNESSetFunctionLocal()`, `DMDASNESSetJacobianLocalVec()`, `DMSNESSetFunction()`, `DMDACreate1d()`, `DMDACreate2d()`, `DMDACreate3d()`
261: @*/
262: PetscErrorCode DMDASNESSetFunctionLocalVec(DM dm, InsertMode imode, DMDASNESFunctionVecFn *func, PetscCtx ctx)
263: {
264:   DMSNES     sdm;
265:   DMSNES_DA *dmdasnes;

267:   PetscFunctionBegin;
269:   PetscCall(DMGetDMSNESWrite(dm, &sdm));
270:   PetscCall(DMDASNESGetContext(dm, sdm, &dmdasnes));

272:   dmdasnes->residuallocalimode = imode;
273:   dmdasnes->residuallocalvec   = func;
274:   dmdasnes->residuallocalctx   = ctx;

276:   PetscCall(DMSNESSetFunction(dm, SNESComputeFunction_DMDA, dmdasnes));
277:   if (!sdm->ops->computejacobian) { /* Call us for the Jacobian too, can be overridden by the user. */
278:     PetscCall(DMSNESSetJacobian(dm, SNESComputeJacobian_DMDA, dmdasnes));
279:   }
280:   PetscFunctionReturn(PETSC_SUCCESS);
281: }

283: /*@C
284:   DMDASNESSetJacobianLocal - set a local Jacobian evaluation function for use with `DMDA`

286:   Logically Collective

288:   Input Parameters:
289: + dm   - `DM` to associate callback with
290: . func - local Jacobian evaluation function
291: - ctx  - optional context for local Jacobian evaluation

293:   Level: beginner

295:   Note:
296:   The `J` and `M` matrices are created internally by `DMCreateMatrix()`

298: .seealso: [](ch_snes), `DMDA`, `DMDASNESJacobianFn`, `DMDASNESSetFunctionLocal()`, `DMSNESSetJacobian()`, `DMDACreate1d()`, `DMDACreate2d()`, `DMDACreate3d()`
299: @*/
300: PetscErrorCode DMDASNESSetJacobianLocal(DM dm, DMDASNESJacobianFn *func, PetscCtx ctx)
301: {
302:   DMSNES     sdm;
303:   DMSNES_DA *dmdasnes;

305:   PetscFunctionBegin;
307:   PetscCall(DMGetDMSNESWrite(dm, &sdm));
308:   PetscCall(DMDASNESGetContext(dm, sdm, &dmdasnes));

310:   dmdasnes->jacobianlocal    = func;
311:   dmdasnes->jacobianlocalctx = ctx;

313:   PetscCall(DMSNESSetJacobian(dm, SNESComputeJacobian_DMDA, dmdasnes));
314:   PetscFunctionReturn(PETSC_SUCCESS);
315: }

317: /*@C
318:   DMDASNESSetJacobianLocalVec - set a local Jacobian evaluation function that operates on a local vector with `DMDA`

320:   Logically Collective

322:   Input Parameters:
323: + dm   - `DM` to associate callback with
324: . func - local Jacobian evaluation
325: - ctx  - optional context for local Jacobian evaluation

327:   Level: beginner

329: .seealso: [](ch_snes), `DMDA`, `DMDASNESJacobianVecFn`, `DMDASNESSetJacobianLocal()`, `DMDASNESSetFunctionLocalVec()`, `DMSNESSetJacobian()`, `DMDACreate1d()`, `DMDACreate2d()`, `DMDACreate3d()`
330: @*/
331: PetscErrorCode DMDASNESSetJacobianLocalVec(DM dm, DMDASNESJacobianVecFn *func, PetscCtx ctx)
332: {
333:   DMSNES     sdm;
334:   DMSNES_DA *dmdasnes;

336:   PetscFunctionBegin;
338:   PetscCall(DMGetDMSNESWrite(dm, &sdm));
339:   PetscCall(DMDASNESGetContext(dm, sdm, &dmdasnes));

341:   dmdasnes->jacobianlocalvec = func;
342:   dmdasnes->jacobianlocalctx = ctx;

344:   PetscCall(DMSNESSetJacobian(dm, SNESComputeJacobian_DMDA, dmdasnes));
345:   PetscFunctionReturn(PETSC_SUCCESS);
346: }

348: /*@C
349:   DMDASNESSetObjectiveLocal - set a local residual evaluation function to used with a `DMDA`

351:   Logically Collective

353:   Input Parameters:
354: + dm   - `DM` to associate callback with
355: . func - local objective evaluation, see `DMDASNESSetObjectiveLocal` for the calling sequence
356: - ctx  - optional context for local residual evaluation

358:   Level: beginner

360: .seealso: [](ch_snes), `DMDA`, `DMDASNESObjectiveFn`, `DMSNESSetFunction()`, `DMDASNESSetJacobianLocal()`, `DMDACreate1d()`, `DMDACreate2d()`, `DMDACreate3d()`, `DMDASNESFunctionFn`
361: @*/
362: PetscErrorCode DMDASNESSetObjectiveLocal(DM dm, DMDASNESObjectiveFn *func, PetscCtx ctx)
363: {
364:   DMSNES     sdm;
365:   DMSNES_DA *dmdasnes;

367:   PetscFunctionBegin;
369:   PetscCall(DMGetDMSNESWrite(dm, &sdm));
370:   PetscCall(DMDASNESGetContext(dm, sdm, &dmdasnes));

372:   dmdasnes->objectivelocal    = func;
373:   dmdasnes->objectivelocalctx = ctx;

375:   PetscCall(DMSNESSetObjective(dm, SNESComputeObjective_DMDA, dmdasnes));
376:   PetscFunctionReturn(PETSC_SUCCESS);
377: }

379: /*@C
380:   DMDASNESSetObjectiveLocalVec - set a local residual evaluation function that operates on a local vector with `DMDA`

382:   Logically Collective

384:   Input Parameters:
385: + dm   - `DM` to associate callback with
386: . func - local objective evaluation, see `DMDASNESSetObjectiveLocalVec` for the calling sequence
387: - ctx  - optional context for local residual evaluation

389:   Level: beginner

391: .seealso: [](ch_snes), `DMDA`, `DMDASNESObjectiveVecFn`, `DMDASNESSetObjectiveLocal()`, `DMSNESSetFunction()`, `DMDASNESSetJacobianLocalVec()`, `DMDACreate1d()`, `DMDACreate2d()`, `DMDACreate3d()`, `DMDASNESObjectiveFn`
392: @*/
393: PetscErrorCode DMDASNESSetObjectiveLocalVec(DM dm, DMDASNESObjectiveVecFn *func, PetscCtx ctx)
394: {
395:   DMSNES     sdm;
396:   DMSNES_DA *dmdasnes;

398:   PetscFunctionBegin;
400:   PetscCall(DMGetDMSNESWrite(dm, &sdm));
401:   PetscCall(DMDASNESGetContext(dm, sdm, &dmdasnes));

403:   dmdasnes->objectivelocalvec = func;
404:   dmdasnes->objectivelocalctx = ctx;

406:   PetscCall(DMSNESSetObjective(dm, SNESComputeObjective_DMDA, dmdasnes));
407:   PetscFunctionReturn(PETSC_SUCCESS);
408: }

410: static PetscErrorCode SNESComputePicard_DMDA(SNES snes, Vec X, Vec F, PetscCtx ctx)
411: {
412:   DM            dm;
413:   DMSNES_DA    *dmdasnes = (DMSNES_DA *)ctx;
414:   DMDALocalInfo info;
415:   Vec           Xloc;
416:   void         *x, *f;

418:   PetscFunctionBegin;
422:   PetscCheck(dmdasnes->rhsplocal, PetscObjectComm((PetscObject)snes), PETSC_ERR_PLIB, "Corrupt context");
423:   PetscCall(SNESGetDM(snes, &dm));
424:   PetscCall(DMGetLocalVector(dm, &Xloc));
425:   PetscCall(DMGlobalToLocalBegin(dm, X, INSERT_VALUES, Xloc));
426:   PetscCall(DMGlobalToLocalEnd(dm, X, INSERT_VALUES, Xloc));
427:   PetscCall(DMDAGetLocalInfo(dm, &info));
428:   PetscCall(DMDAVecGetArray(dm, Xloc, &x));
429:   switch (dmdasnes->residuallocalimode) {
430:   case INSERT_VALUES: {
431:     PetscCall(DMDAVecGetArray(dm, F, &f));
432:     PetscCallBack("SNES Picard DMDA local callback function", (*dmdasnes->rhsplocal)(&info, x, f, dmdasnes->picardlocalctx));
433:     PetscCall(DMDAVecRestoreArray(dm, F, &f));
434:   } break;
435:   case ADD_VALUES: {
436:     Vec Floc;
437:     PetscCall(DMGetLocalVector(dm, &Floc));
438:     PetscCall(VecZeroEntries(Floc));
439:     PetscCall(DMDAVecGetArray(dm, Floc, &f));
440:     PetscCallBack("SNES Picard DMDA local callback function", (*dmdasnes->rhsplocal)(&info, x, f, dmdasnes->picardlocalctx));
441:     PetscCall(DMDAVecRestoreArray(dm, Floc, &f));
442:     PetscCall(VecZeroEntries(F));
443:     PetscCall(DMLocalToGlobalBegin(dm, Floc, ADD_VALUES, F));
444:     PetscCall(DMLocalToGlobalEnd(dm, Floc, ADD_VALUES, F));
445:     PetscCall(DMRestoreLocalVector(dm, &Floc));
446:   } break;
447:   default:
448:     SETERRQ(PetscObjectComm((PetscObject)snes), PETSC_ERR_ARG_INCOMP, "Cannot use imode=%d", (int)dmdasnes->residuallocalimode);
449:   }
450:   PetscCall(DMDAVecRestoreArray(dm, Xloc, &x));
451:   PetscCall(DMRestoreLocalVector(dm, &Xloc));
452:   PetscFunctionReturn(PETSC_SUCCESS);
453: }

455: static PetscErrorCode SNESComputePicardJacobian_DMDA(SNES snes, Vec X, Mat A, Mat B, PetscCtx ctx)
456: {
457:   DM            dm;
458:   DMSNES_DA    *dmdasnes = (DMSNES_DA *)ctx;
459:   DMDALocalInfo info;
460:   Vec           Xloc;
461:   void         *x;

463:   PetscFunctionBegin;
464:   PetscCheck(dmdasnes->jacobianplocal, PetscObjectComm((PetscObject)snes), PETSC_ERR_PLIB, "Corrupt context");
465:   PetscCall(SNESGetDM(snes, &dm));

467:   PetscCall(DMGetLocalVector(dm, &Xloc));
468:   PetscCall(DMGlobalToLocalBegin(dm, X, INSERT_VALUES, Xloc));
469:   PetscCall(DMGlobalToLocalEnd(dm, X, INSERT_VALUES, Xloc));
470:   PetscCall(DMDAGetLocalInfo(dm, &info));
471:   PetscCall(DMDAVecGetArray(dm, Xloc, &x));
472:   PetscCallBack("SNES Picard DMDA local callback Jacobian", (*dmdasnes->jacobianplocal)(&info, x, A, B, dmdasnes->picardlocalctx));
473:   PetscCall(DMDAVecRestoreArray(dm, Xloc, &x));
474:   PetscCall(DMRestoreLocalVector(dm, &Xloc));
475:   if (A != B) {
476:     PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
477:     PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
478:   }
479:   PetscFunctionReturn(PETSC_SUCCESS);
480: }

482: /*@C
483:   DMDASNESSetPicardLocal - set a local right-hand side and matrix evaluation function for Picard iteration with `DMDA`

485:   Logically Collective

487:   Input Parameters:
488: + dm    - `DM` to associate callback with
489: . imode - `INSERT_VALUES` if local function computes owned part, `ADD_VALUES` if it contributes to ghosted part
490: . func  - local residual evaluation
491: . jac   - function to compute Jacobian
492: - ctx   - optional context for local residual evaluation

494:   Level: beginner

496:   Note:
497:   The user must use `SNESSetFunction`(`snes`,`NULL`,`SNESPicardComputeFunction`,&user));
498:   in their code before calling this routine.

500: .seealso: [](ch_snes), `SNES`, `DMDA`, `DMDASNESFunctionFn`, `DMDASNESJacobianFn`, `DMSNESSetFunction()`, `DMDASNESSetJacobian()`, `DMDACreate1d()`, `DMDACreate2d()`, `DMDACreate3d()`
501: @*/
502: PetscErrorCode DMDASNESSetPicardLocal(DM dm, InsertMode imode, DMDASNESFunctionFn *func, DMDASNESJacobianFn *jac, PetscCtx ctx)
503: {
504:   DMSNES     sdm;
505:   DMSNES_DA *dmdasnes;

507:   PetscFunctionBegin;
509:   PetscCall(DMGetDMSNESWrite(dm, &sdm));
510:   PetscCall(DMDASNESGetContext(dm, sdm, &dmdasnes));

512:   dmdasnes->residuallocalimode = imode;
513:   dmdasnes->rhsplocal          = func;
514:   dmdasnes->jacobianplocal     = jac;
515:   dmdasnes->picardlocalctx     = ctx;

517:   PetscCall(DMSNESSetPicard(dm, SNESComputePicard_DMDA, SNESComputePicardJacobian_DMDA, dmdasnes));
518:   PetscCall(DMSNESSetMFFunction(dm, SNESComputeFunction_DMDA, dmdasnes));
519:   PetscFunctionReturn(PETSC_SUCCESS);
520: }