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