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