xref: /petsc/src/snes/utils/dmlocalsnes.c (revision ccfb0f9f40a0131988d7995ed9679700dae2a75a)
1 #include <petsc/private/dmimpl.h>
2 #include <petsc/private/snesimpl.h> /*I "petscsnes.h" I*/
3 
4 typedef struct {
5   PetscErrorCode (*objectivelocal)(DM, Vec, PetscReal *, void *);
6   PetscErrorCode (*residuallocal)(DM, Vec, Vec, void *);
7   PetscErrorCode (*jacobianlocal)(DM, Vec, Mat, Mat, void *);
8   PetscErrorCode (*boundarylocal)(DM, Vec, void *);
9   void *objectivelocalctx;
10   void *residuallocalctx;
11   void *jacobianlocalctx;
12   void *boundarylocalctx;
13 } DMSNES_Local;
14 
15 static PetscErrorCode DMSNESDestroy_DMLocal(DMSNES sdm)
16 {
17   PetscFunctionBegin;
18   PetscCall(PetscFree(sdm->data));
19   sdm->data = NULL;
20   PetscFunctionReturn(PETSC_SUCCESS);
21 }
22 
23 static PetscErrorCode DMSNESDuplicate_DMLocal(DMSNES oldsdm, DMSNES sdm)
24 {
25   PetscFunctionBegin;
26   if (sdm->data != oldsdm->data) {
27     PetscCall(PetscFree(sdm->data));
28     PetscCall(PetscNew((DMSNES_Local **)&sdm->data));
29     if (oldsdm->data) PetscCall(PetscMemcpy(sdm->data, oldsdm->data, sizeof(DMSNES_Local)));
30   }
31   PetscFunctionReturn(PETSC_SUCCESS);
32 }
33 
34 static PetscErrorCode DMLocalSNESGetContext(DM dm, DMSNES sdm, DMSNES_Local **dmlocalsnes)
35 {
36   PetscFunctionBegin;
37   *dmlocalsnes = NULL;
38   if (!sdm->data) {
39     PetscCall(PetscNew((DMSNES_Local **)&sdm->data));
40 
41     sdm->ops->destroy   = DMSNESDestroy_DMLocal;
42     sdm->ops->duplicate = DMSNESDuplicate_DMLocal;
43   }
44   *dmlocalsnes = (DMSNES_Local *)sdm->data;
45   PetscFunctionReturn(PETSC_SUCCESS);
46 }
47 
48 static PetscErrorCode SNESComputeObjective_DMLocal(SNES snes, Vec X, PetscReal *obj, void *ctx)
49 {
50   DMSNES_Local *dmlocalsnes = (DMSNES_Local *)ctx;
51   DM            dm;
52   Vec           Xloc;
53   PetscBool     transform;
54 
55   PetscFunctionBegin;
56   PetscValidHeaderSpecific(snes, SNES_CLASSID, 1);
57   PetscValidHeaderSpecific(X, VEC_CLASSID, 2);
58   PetscCall(SNESGetDM(snes, &dm));
59   PetscCall(DMGetLocalVector(dm, &Xloc));
60   PetscCall(VecZeroEntries(Xloc));
61   /* Non-conforming routines needs boundary values before G2L */
62   if (dmlocalsnes->boundarylocal) PetscCall((*dmlocalsnes->boundarylocal)(dm, Xloc, dmlocalsnes->boundarylocalctx));
63   PetscCall(DMGlobalToLocalBegin(dm, X, INSERT_VALUES, Xloc));
64   PetscCall(DMGlobalToLocalEnd(dm, X, INSERT_VALUES, Xloc));
65   /* Need to reset boundary values if we transformed */
66   PetscCall(DMHasBasisTransform(dm, &transform));
67   if (transform && dmlocalsnes->boundarylocal) PetscCall((*dmlocalsnes->boundarylocal)(dm, Xloc, dmlocalsnes->boundarylocalctx));
68   CHKMEMQ;
69   PetscCall((*dmlocalsnes->objectivelocal)(dm, Xloc, obj, dmlocalsnes->objectivelocalctx));
70   CHKMEMQ;
71   PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, obj, 1, MPIU_REAL, MPIU_SUM, PetscObjectComm((PetscObject)snes)));
72   PetscCall(DMRestoreLocalVector(dm, &Xloc));
73   PetscFunctionReturn(PETSC_SUCCESS);
74 }
75 
76 static PetscErrorCode SNESComputeFunction_DMLocal(SNES snes, Vec X, Vec F, void *ctx)
77 {
78   DMSNES_Local *dmlocalsnes = (DMSNES_Local *)ctx;
79   DM            dm;
80   Vec           Xloc, Floc;
81   PetscBool     transform;
82 
83   PetscFunctionBegin;
84   PetscValidHeaderSpecific(snes, SNES_CLASSID, 1);
85   PetscValidHeaderSpecific(X, VEC_CLASSID, 2);
86   PetscValidHeaderSpecific(F, VEC_CLASSID, 3);
87   PetscCall(SNESGetDM(snes, &dm));
88   PetscCall(DMGetLocalVector(dm, &Xloc));
89   PetscCall(DMGetLocalVector(dm, &Floc));
90   PetscCall(VecZeroEntries(Xloc));
91   PetscCall(VecZeroEntries(Floc));
92   /* Non-conforming routines needs boundary values before G2L */
93   if (dmlocalsnes->boundarylocal) PetscCall((*dmlocalsnes->boundarylocal)(dm, Xloc, dmlocalsnes->boundarylocalctx));
94   PetscCall(DMGlobalToLocalBegin(dm, X, INSERT_VALUES, Xloc));
95   PetscCall(DMGlobalToLocalEnd(dm, X, INSERT_VALUES, Xloc));
96   /* Need to reset boundary values if we transformed */
97   PetscCall(DMHasBasisTransform(dm, &transform));
98   if (transform && dmlocalsnes->boundarylocal) PetscCall((*dmlocalsnes->boundarylocal)(dm, Xloc, dmlocalsnes->boundarylocalctx));
99   CHKMEMQ;
100   PetscCall((*dmlocalsnes->residuallocal)(dm, Xloc, Floc, dmlocalsnes->residuallocalctx));
101   CHKMEMQ;
102   PetscCall(VecZeroEntries(F));
103   PetscCall(DMLocalToGlobalBegin(dm, Floc, ADD_VALUES, F));
104   PetscCall(DMLocalToGlobalEnd(dm, Floc, ADD_VALUES, F));
105   PetscCall(DMRestoreLocalVector(dm, &Floc));
106   PetscCall(DMRestoreLocalVector(dm, &Xloc));
107   {
108     char        name[PETSC_MAX_PATH_LEN];
109     char        oldname[PETSC_MAX_PATH_LEN];
110     const char *tmp;
111     PetscInt    it;
112 
113     PetscCall(SNESGetIterationNumber(snes, &it));
114     PetscCall(PetscSNPrintf(name, PETSC_MAX_PATH_LEN, "Solution, Iterate %" PetscInt_FMT, it));
115     PetscCall(PetscObjectGetName((PetscObject)X, &tmp));
116     PetscCall(PetscStrncpy(oldname, tmp, PETSC_MAX_PATH_LEN - 1));
117     PetscCall(PetscObjectSetName((PetscObject)X, name));
118     PetscCall(VecViewFromOptions(X, (PetscObject)snes, "-dmsnes_solution_vec_view"));
119     PetscCall(PetscObjectSetName((PetscObject)X, oldname));
120     PetscCall(PetscSNPrintf(name, PETSC_MAX_PATH_LEN, "Residual, Iterate %" PetscInt_FMT, it));
121     PetscCall(PetscObjectSetName((PetscObject)F, name));
122     PetscCall(VecViewFromOptions(F, (PetscObject)snes, "-dmsnes_residual_vec_view"));
123   }
124   PetscFunctionReturn(PETSC_SUCCESS);
125 }
126 
127 static PetscErrorCode SNESComputeJacobian_DMLocal(SNES snes, Vec X, Mat A, Mat B, void *ctx)
128 {
129   DMSNES_Local *dmlocalsnes = (DMSNES_Local *)ctx;
130   DM            dm;
131   Vec           Xloc;
132   PetscBool     transform;
133 
134   PetscFunctionBegin;
135   PetscCall(SNESGetDM(snes, &dm));
136   if (dmlocalsnes->jacobianlocal) {
137     PetscCall(DMGetLocalVector(dm, &Xloc));
138     PetscCall(VecZeroEntries(Xloc));
139     /* Non-conforming routines needs boundary values before G2L */
140     if (dmlocalsnes->boundarylocal) PetscCall((*dmlocalsnes->boundarylocal)(dm, Xloc, dmlocalsnes->boundarylocalctx));
141     PetscCall(DMGlobalToLocalBegin(dm, X, INSERT_VALUES, Xloc));
142     PetscCall(DMGlobalToLocalEnd(dm, X, INSERT_VALUES, Xloc));
143     /* Need to reset boundary values if we transformed */
144     PetscCall(DMHasBasisTransform(dm, &transform));
145     if (transform && dmlocalsnes->boundarylocal) PetscCall((*dmlocalsnes->boundarylocal)(dm, Xloc, dmlocalsnes->boundarylocalctx));
146     CHKMEMQ;
147     PetscCall((*dmlocalsnes->jacobianlocal)(dm, Xloc, A, B, dmlocalsnes->jacobianlocalctx));
148     CHKMEMQ;
149     PetscCall(DMRestoreLocalVector(dm, &Xloc));
150   } else {
151     MatFDColoring fdcoloring;
152     PetscCall(PetscObjectQuery((PetscObject)dm, "DMDASNES_FDCOLORING", (PetscObject *)&fdcoloring));
153     if (!fdcoloring) {
154       ISColoring coloring;
155 
156       PetscCall(DMCreateColoring(dm, dm->coloringtype, &coloring));
157       PetscCall(MatFDColoringCreate(B, coloring, &fdcoloring));
158       PetscCall(ISColoringDestroy(&coloring));
159       switch (dm->coloringtype) {
160       case IS_COLORING_GLOBAL:
161         PetscCall(MatFDColoringSetFunction(fdcoloring, (MatFDColoringFn *)SNESComputeFunction_DMLocal, dmlocalsnes));
162         break;
163       default:
164         SETERRQ(PetscObjectComm((PetscObject)snes), PETSC_ERR_SUP, "No support for coloring type '%s'", ISColoringTypes[dm->coloringtype]);
165       }
166       PetscCall(PetscObjectSetOptionsPrefix((PetscObject)fdcoloring, ((PetscObject)dm)->prefix));
167       PetscCall(MatFDColoringSetFromOptions(fdcoloring));
168       PetscCall(MatFDColoringSetUp(B, coloring, fdcoloring));
169       PetscCall(PetscObjectCompose((PetscObject)dm, "DMDASNES_FDCOLORING", (PetscObject)fdcoloring));
170       PetscCall(PetscObjectDereference((PetscObject)fdcoloring));
171 
172       /* The following breaks an ugly reference counting loop that deserves a paragraph. MatFDColoringApply() will call
173        * VecDuplicate() with the state Vec and store inside the MatFDColoring. This Vec will duplicate the Vec, but the
174        * MatFDColoring is composed with the DM. We dereference the DM here so that the reference count will eventually
175        * drop to 0. Note the code in DMDestroy() that exits early for a negative reference count. That code path will be
176        * taken when the PetscObjectList for the Vec inside MatFDColoring is destroyed.
177        */
178       PetscCall(PetscObjectDereference((PetscObject)dm));
179     }
180     PetscCall(MatFDColoringApply(B, fdcoloring, X, snes));
181   }
182   /* This will be redundant if the user called both, but it's too common to forget. */
183   if (A != B) {
184     PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
185     PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
186   }
187   PetscFunctionReturn(PETSC_SUCCESS);
188 }
189 
190 /*@C
191   DMSNESSetObjectiveLocal - set a local objective evaluation function. This function is called with local vector
192   containing the local vector information PLUS ghost point information. It should compute a result for all local
193   elements and `DMSNES` will automatically accumulate the overlapping values.
194 
195   Logically Collective
196 
197   Input Parameters:
198 + dm   - `DM` to associate callback with
199 . func - local objective evaluation
200 - ctx  - optional context for local residual evaluation
201 
202   Level: advanced
203 
204 .seealso: `DMSNESSetFunctionLocal()`, `DMSNESSetJacobianLocal()`
205 @*/
206 PetscErrorCode DMSNESSetObjectiveLocal(DM dm, PetscErrorCode (*func)(DM, Vec, PetscReal *, void *), void *ctx)
207 {
208   DMSNES        sdm;
209   DMSNES_Local *dmlocalsnes;
210 
211   PetscFunctionBegin;
212   PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
213   PetscCall(DMGetDMSNESWrite(dm, &sdm));
214   PetscCall(DMLocalSNESGetContext(dm, sdm, &dmlocalsnes));
215 
216   dmlocalsnes->objectivelocal    = func;
217   dmlocalsnes->objectivelocalctx = ctx;
218 
219   PetscCall(DMSNESSetObjective(dm, SNESComputeObjective_DMLocal, dmlocalsnes));
220   PetscFunctionReturn(PETSC_SUCCESS);
221 }
222 
223 /*@C
224   DMSNESSetFunctionLocal - set a local residual evaluation function. This function is called with local vector
225   containing the local vector information PLUS ghost point information. It should compute a result for all local
226   elements and `DMSNES` will automatically accumulate the overlapping values.
227 
228   Logically Collective
229 
230   Input Parameters:
231 + dm   - `DM` to associate callback with
232 . func - local residual evaluation
233 - ctx  - optional context for local residual evaluation
234 
235   Calling sequence of `func`:
236 + dm  - `DM` for the function
237 . x   - vector to state at which to evaluate residual
238 . f   - vector to hold the function evaluation
239 - ctx - optional context passed above
240 
241   Level: advanced
242 
243 .seealso: [](ch_snes), `DMSNESSetFunction()`, `DMSNESSetJacobianLocal()`
244 @*/
245 PetscErrorCode DMSNESSetFunctionLocal(DM dm, PetscErrorCode (*func)(DM dm, Vec x, Vec f, void *ctx), void *ctx)
246 {
247   DMSNES        sdm;
248   DMSNES_Local *dmlocalsnes;
249 
250   PetscFunctionBegin;
251   PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
252   PetscCall(DMGetDMSNESWrite(dm, &sdm));
253   PetscCall(DMLocalSNESGetContext(dm, sdm, &dmlocalsnes));
254 
255   dmlocalsnes->residuallocal    = func;
256   dmlocalsnes->residuallocalctx = ctx;
257 
258   PetscCall(DMSNESSetFunction(dm, SNESComputeFunction_DMLocal, dmlocalsnes));
259   if (!sdm->ops->computejacobian) { /* Call us for the Jacobian too, can be overridden by the user. */
260     PetscCall(DMSNESSetJacobian(dm, SNESComputeJacobian_DMLocal, dmlocalsnes));
261   }
262   PetscFunctionReturn(PETSC_SUCCESS);
263 }
264 
265 /*@C
266   DMSNESSetBoundaryLocal - set a function to insert, for example, essential boundary conditions into a ghosted solution vector
267 
268   Logically Collective
269 
270   Input Parameters:
271 + dm   - `DM` to associate callback with
272 . func - local boundary value evaluation
273 - ctx  - optional context for local boundary value evaluation
274 
275   Calling sequence of `func`:
276 + dm  - the `DM` context
277 . X   - ghosted solution vector, appropriate locations (such as essential boundary condition nodes) should be filled
278 - ctx - option context passed in `DMSNESSetBoundaryLocal()`
279 
280   Level: advanced
281 
282 .seealso: [](ch_snes), `DMSNESSetObjectiveLocal()`, `DMSNESSetFunctionLocal()`, `DMSNESSetJacobianLocal()`
283 @*/
284 PetscErrorCode DMSNESSetBoundaryLocal(DM dm, PetscErrorCode (*func)(DM dm, Vec X, void *ctx), void *ctx)
285 {
286   DMSNES        sdm;
287   DMSNES_Local *dmlocalsnes;
288 
289   PetscFunctionBegin;
290   PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
291   PetscCall(DMGetDMSNESWrite(dm, &sdm));
292   PetscCall(DMLocalSNESGetContext(dm, sdm, &dmlocalsnes));
293 
294   dmlocalsnes->boundarylocal    = func;
295   dmlocalsnes->boundarylocalctx = ctx;
296   PetscFunctionReturn(PETSC_SUCCESS);
297 }
298 
299 /*@C
300   DMSNESSetJacobianLocal - set a local Jacobian evaluation function
301 
302   Logically Collective
303 
304   Input Parameters:
305 + dm   - `DM` to associate callback with
306 . func - local Jacobian evaluation
307 - ctx  - optional context for local Jacobian evaluation
308 
309   Calling sequence of `func`:
310 + dm  - the `DM` context
311 . X   - current solution vector (ghosted or not?)
312 . J   - the Jacobian
313 . Jp  - approximate Jacobian used to compute the preconditioner, often `J`
314 - ctx - a user provided context
315 
316   Level: advanced
317 
318 .seealso: [](ch_snes), `DMSNESSetObjectiveLocal()`, `DMSNESSetFunctionLocal()`, `DMSNESSetBoundaryLocal()`
319 @*/
320 PetscErrorCode DMSNESSetJacobianLocal(DM dm, PetscErrorCode (*func)(DM dm, Vec X, Mat J, Mat Jp, void *ctx), void *ctx)
321 {
322   DMSNES        sdm;
323   DMSNES_Local *dmlocalsnes;
324 
325   PetscFunctionBegin;
326   PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
327   PetscCall(DMGetDMSNESWrite(dm, &sdm));
328   PetscCall(DMLocalSNESGetContext(dm, sdm, &dmlocalsnes));
329 
330   dmlocalsnes->jacobianlocal    = func;
331   dmlocalsnes->jacobianlocalctx = ctx;
332 
333   PetscCall(DMSNESSetJacobian(dm, SNESComputeJacobian_DMLocal, dmlocalsnes));
334   PetscFunctionReturn(PETSC_SUCCESS);
335 }
336 
337 /*@C
338   DMSNESGetObjectiveLocal - get the local objective evaluation function information set with `DMSNESSetObjectiveLocal()`.
339 
340   Not Collective
341 
342   Input Parameter:
343 . dm - `DM` with the associated callback
344 
345   Output Parameters:
346 + func - local objective evaluation
347 - ctx  - context for local residual evaluation
348 
349   Level: beginner
350 
351 .seealso: `DMSNESSetObjective()`, `DMSNESSetObjectiveLocal()`, `DMSNESSetFunctionLocal()`
352 @*/
353 PetscErrorCode DMSNESGetObjectiveLocal(DM dm, PetscErrorCode (**func)(DM, Vec, PetscReal *, void *), void **ctx)
354 {
355   DMSNES        sdm;
356   DMSNES_Local *dmlocalsnes;
357 
358   PetscFunctionBegin;
359   PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
360   PetscCall(DMGetDMSNES(dm, &sdm));
361   PetscCall(DMLocalSNESGetContext(dm, sdm, &dmlocalsnes));
362   if (func) *func = dmlocalsnes->objectivelocal;
363   if (ctx) *ctx = dmlocalsnes->objectivelocalctx;
364   PetscFunctionReturn(PETSC_SUCCESS);
365 }
366 
367 /*@C
368   DMSNESGetFunctionLocal - get the local residual evaluation function information set with `DMSNESSetFunctionLocal()`.
369 
370   Not Collective
371 
372   Input Parameter:
373 . dm - `DM` with the associated callback
374 
375   Output Parameters:
376 + func - local residual evaluation
377 - ctx  - context for local residual evaluation
378 
379   Level: beginner
380 
381 .seealso: [](ch_snes), `DMSNESSetFunction()`, `DMSNESSetFunctionLocal()`, `DMSNESSetJacobianLocal()`
382 @*/
383 PetscErrorCode DMSNESGetFunctionLocal(DM dm, PetscErrorCode (**func)(DM, Vec, Vec, void *), void **ctx)
384 {
385   DMSNES        sdm;
386   DMSNES_Local *dmlocalsnes;
387 
388   PetscFunctionBegin;
389   PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
390   PetscCall(DMGetDMSNES(dm, &sdm));
391   PetscCall(DMLocalSNESGetContext(dm, sdm, &dmlocalsnes));
392   if (func) *func = dmlocalsnes->residuallocal;
393   if (ctx) *ctx = dmlocalsnes->residuallocalctx;
394   PetscFunctionReturn(PETSC_SUCCESS);
395 }
396 
397 /*@C
398   DMSNESGetBoundaryLocal - get the local boundary value function set with `DMSNESSetBoundaryLocal()`.
399 
400   Not Collective
401 
402   Input Parameter:
403 . dm - `DM` with the associated callback
404 
405   Output Parameters:
406 + func - local boundary value evaluation
407 - ctx  - context for local boundary value evaluation
408 
409   Level: intermediate
410 
411 .seealso: [](ch_snes), `DMSNESSetFunctionLocal()`, `DMSNESSetBoundaryLocal()`, `DMSNESSetJacobianLocal()`
412 @*/
413 PetscErrorCode DMSNESGetBoundaryLocal(DM dm, PetscErrorCode (**func)(DM, Vec, void *), void **ctx)
414 {
415   DMSNES        sdm;
416   DMSNES_Local *dmlocalsnes;
417 
418   PetscFunctionBegin;
419   PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
420   PetscCall(DMGetDMSNES(dm, &sdm));
421   PetscCall(DMLocalSNESGetContext(dm, sdm, &dmlocalsnes));
422   if (func) *func = dmlocalsnes->boundarylocal;
423   if (ctx) *ctx = dmlocalsnes->boundarylocalctx;
424   PetscFunctionReturn(PETSC_SUCCESS);
425 }
426 
427 /*@C
428   DMSNESGetJacobianLocal - the local Jacobian evaluation function set with `DMSNESSetJacobianLocal()`.
429 
430   Logically Collective
431 
432   Input Parameter:
433 . dm - `DM` with the associated callback
434 
435   Output Parameters:
436 + func - local Jacobian evaluation
437 - ctx  - context for local Jacobian evaluation
438 
439   Level: beginner
440 
441 .seealso: [](ch_snes), `DMSNESSetJacobianLocal()`, `DMSNESSetJacobian()`
442 @*/
443 PetscErrorCode DMSNESGetJacobianLocal(DM dm, PetscErrorCode (**func)(DM, Vec, Mat, Mat, void *), void **ctx)
444 {
445   DMSNES        sdm;
446   DMSNES_Local *dmlocalsnes;
447 
448   PetscFunctionBegin;
449   PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
450   PetscCall(DMGetDMSNES(dm, &sdm));
451   PetscCall(DMLocalSNESGetContext(dm, sdm, &dmlocalsnes));
452   if (func) *func = dmlocalsnes->jacobianlocal;
453   if (ctx) *ctx = dmlocalsnes->jacobianlocalctx;
454   PetscFunctionReturn(PETSC_SUCCESS);
455 }
456