xref: /petsc/src/ts/utils/dmlocalts.c (revision 40badf4fbc550ac1f60bd080eaff6de6d55b946d)
1 #include <petsc/private/dmimpl.h>
2 #include <petsc/private/tsimpl.h>   /*I "petscts.h" I*/
3 
4 typedef struct {
5   PetscErrorCode (*boundarylocal)(DM,PetscReal,Vec,Vec,void*);
6   PetscErrorCode (*ifunctionlocal)(DM,PetscReal,Vec,Vec,Vec,void*);
7   PetscErrorCode (*ijacobianlocal)(DM,PetscReal,Vec,Vec,PetscReal,Mat,Mat,void*);
8   PetscErrorCode (*rhsfunctionlocal)(DM,PetscReal,Vec,Vec,void*);
9   void *boundarylocalctx;
10   void *ifunctionlocalctx;
11   void *ijacobianlocalctx;
12   void *rhsfunctionlocalctx;
13   Vec   lumpedmassinv;
14   Mat   mass;
15   KSP   kspmass;
16 } DMTS_Local;
17 
18 static PetscErrorCode DMTSDestroy_DMLocal(DMTS tdm)
19 {
20   PetscFunctionBegin;
21   CHKERRQ(PetscFree(tdm->data));
22   PetscFunctionReturn(0);
23 }
24 
25 static PetscErrorCode DMTSDuplicate_DMLocal(DMTS oldtdm, DMTS tdm)
26 {
27   PetscFunctionBegin;
28   CHKERRQ(PetscNewLog(tdm, (DMTS_Local **) &tdm->data));
29   if (oldtdm->data) CHKERRQ(PetscMemcpy(tdm->data, oldtdm->data, sizeof(DMTS_Local)));
30   PetscFunctionReturn(0);
31 }
32 
33 static PetscErrorCode DMLocalTSGetContext(DM dm, DMTS tdm, DMTS_Local **dmlocalts)
34 {
35   PetscFunctionBegin;
36   *dmlocalts = NULL;
37   if (!tdm->data) {
38     CHKERRQ(PetscNewLog(dm, (DMTS_Local **) &tdm->data));
39 
40     tdm->ops->destroy   = DMTSDestroy_DMLocal;
41     tdm->ops->duplicate = DMTSDuplicate_DMLocal;
42   }
43   *dmlocalts = (DMTS_Local *) tdm->data;
44   PetscFunctionReturn(0);
45 }
46 
47 static PetscErrorCode TSComputeIFunction_DMLocal(TS ts, PetscReal time, Vec X, Vec X_t, Vec F, void *ctx)
48 {
49   DM             dm;
50   Vec            locX, locX_t, locF;
51   DMTS_Local    *dmlocalts = (DMTS_Local *) ctx;
52 
53   PetscFunctionBegin;
54   PetscValidHeaderSpecific(ts,TS_CLASSID,1);
55   PetscValidHeaderSpecific(X,VEC_CLASSID,3);
56   PetscValidHeaderSpecific(X_t,VEC_CLASSID,4);
57   PetscValidHeaderSpecific(F,VEC_CLASSID,5);
58   CHKERRQ(TSGetDM(ts, &dm));
59   CHKERRQ(DMGetLocalVector(dm, &locX));
60   CHKERRQ(DMGetLocalVector(dm, &locX_t));
61   CHKERRQ(DMGetLocalVector(dm, &locF));
62   CHKERRQ(VecZeroEntries(locX));
63   CHKERRQ(VecZeroEntries(locX_t));
64   if (dmlocalts->boundarylocal) CHKERRQ((*dmlocalts->boundarylocal)(dm, time, locX, locX_t,dmlocalts->boundarylocalctx));
65   CHKERRQ(DMGlobalToLocalBegin(dm, X, INSERT_VALUES, locX));
66   CHKERRQ(DMGlobalToLocalEnd(dm, X, INSERT_VALUES, locX));
67   CHKERRQ(DMGlobalToLocalBegin(dm, X_t, INSERT_VALUES, locX_t));
68   CHKERRQ(DMGlobalToLocalEnd(dm, X_t, INSERT_VALUES, locX_t));
69   CHKERRQ(VecZeroEntries(locF));
70   CHKMEMQ;
71   CHKERRQ((*dmlocalts->ifunctionlocal)(dm, time, locX, locX_t, locF, dmlocalts->ifunctionlocalctx));
72   CHKMEMQ;
73   CHKERRQ(VecZeroEntries(F));
74   CHKERRQ(DMLocalToGlobalBegin(dm, locF, ADD_VALUES, F));
75   CHKERRQ(DMLocalToGlobalEnd(dm, locF, ADD_VALUES, F));
76   CHKERRQ(DMRestoreLocalVector(dm, &locX));
77   CHKERRQ(DMRestoreLocalVector(dm, &locX_t));
78   CHKERRQ(DMRestoreLocalVector(dm, &locF));
79   PetscFunctionReturn(0);
80 }
81 
82 static PetscErrorCode TSComputeRHSFunction_DMLocal(TS ts, PetscReal time, Vec X, Vec F, void *ctx)
83 {
84   DM             dm;
85   Vec            locX, locF;
86   DMTS_Local    *dmlocalts = (DMTS_Local *) ctx;
87 
88   PetscFunctionBegin;
89   PetscValidHeaderSpecific(ts,TS_CLASSID,1);
90   PetscValidHeaderSpecific(X,VEC_CLASSID,3);
91   PetscValidHeaderSpecific(F,VEC_CLASSID,4);
92   CHKERRQ(TSGetDM(ts, &dm));
93   CHKERRQ(DMGetLocalVector(dm, &locX));
94   CHKERRQ(DMGetLocalVector(dm, &locF));
95   CHKERRQ(VecZeroEntries(locX));
96   if (dmlocalts->boundarylocal) CHKERRQ((*dmlocalts->boundarylocal)(dm,time,locX,NULL,dmlocalts->boundarylocalctx));
97   CHKERRQ(DMGlobalToLocalBegin(dm, X, INSERT_VALUES, locX));
98   CHKERRQ(DMGlobalToLocalEnd(dm, X, INSERT_VALUES, locX));
99   CHKERRQ(VecZeroEntries(locF));
100   CHKMEMQ;
101   CHKERRQ((*dmlocalts->rhsfunctionlocal)(dm, time, locX, locF, dmlocalts->rhsfunctionlocalctx));
102   CHKMEMQ;
103   CHKERRQ(VecZeroEntries(F));
104   CHKERRQ(DMLocalToGlobalBegin(dm, locF, ADD_VALUES, F));
105   CHKERRQ(DMLocalToGlobalEnd(dm, locF, ADD_VALUES, F));
106   if (dmlocalts->lumpedmassinv) {
107     CHKERRQ(VecPointwiseMult(F, dmlocalts->lumpedmassinv, F));
108   } else if (dmlocalts->kspmass) {
109     Vec tmp;
110 
111     CHKERRQ(DMGetGlobalVector(dm, &tmp));
112     CHKERRQ(KSPSolve(dmlocalts->kspmass, F, tmp));
113     CHKERRQ(VecCopy(tmp, F));
114     CHKERRQ(DMRestoreGlobalVector(dm, &tmp));
115   }
116   CHKERRQ(DMRestoreLocalVector(dm, &locX));
117   CHKERRQ(DMRestoreLocalVector(dm, &locF));
118   PetscFunctionReturn(0);
119 }
120 
121 static PetscErrorCode TSComputeIJacobian_DMLocal(TS ts, PetscReal time, Vec X, Vec X_t, PetscReal a, Mat A, Mat B, void *ctx)
122 {
123   DM             dm;
124   Vec            locX, locX_t;
125   DMTS_Local    *dmlocalts = (DMTS_Local *) ctx;
126 
127   PetscFunctionBegin;
128   CHKERRQ(TSGetDM(ts, &dm));
129   if (dmlocalts->ijacobianlocal) {
130     CHKERRQ(DMGetLocalVector(dm, &locX));
131     CHKERRQ(DMGetLocalVector(dm, &locX_t));
132     CHKERRQ(VecZeroEntries(locX));
133     CHKERRQ(VecZeroEntries(locX_t));
134     if (dmlocalts->boundarylocal) CHKERRQ((*dmlocalts->boundarylocal)(dm,time,locX,locX_t,dmlocalts->boundarylocalctx));
135     CHKERRQ(DMGlobalToLocalBegin(dm, X, INSERT_VALUES, locX));
136     CHKERRQ(DMGlobalToLocalEnd(dm, X, INSERT_VALUES, locX));
137     CHKERRQ(DMGlobalToLocalBegin(dm, X_t, INSERT_VALUES, locX_t));
138     CHKERRQ(DMGlobalToLocalEnd(dm, X_t, INSERT_VALUES, locX_t));
139     CHKMEMQ;
140     CHKERRQ((*dmlocalts->ijacobianlocal)(dm, time, locX, locX_t, a, A, B, dmlocalts->ijacobianlocalctx));
141     CHKMEMQ;
142     CHKERRQ(DMRestoreLocalVector(dm, &locX));
143     CHKERRQ(DMRestoreLocalVector(dm, &locX_t));
144   } else {
145     MatFDColoring fdcoloring;
146     CHKERRQ(PetscObjectQuery((PetscObject) dm, "DMDASNES_FDCOLORING", (PetscObject *) &fdcoloring));
147     if (!fdcoloring) {
148       ISColoring coloring;
149 
150       CHKERRQ(DMCreateColoring(dm, dm->coloringtype, &coloring));
151       CHKERRQ(MatFDColoringCreate(B, coloring, &fdcoloring));
152       CHKERRQ(ISColoringDestroy(&coloring));
153       switch (dm->coloringtype) {
154       case IS_COLORING_GLOBAL:
155         CHKERRQ(MatFDColoringSetFunction(fdcoloring, (PetscErrorCode (*)(void)) TSComputeIFunction_DMLocal, dmlocalts));
156         break;
157       default: SETERRQ(PetscObjectComm((PetscObject) ts), PETSC_ERR_SUP, "No support for coloring type '%s'", ISColoringTypes[dm->coloringtype]);
158       }
159       CHKERRQ(PetscObjectSetOptionsPrefix((PetscObject) fdcoloring, ((PetscObject) dm)->prefix));
160       CHKERRQ(MatFDColoringSetFromOptions(fdcoloring));
161       CHKERRQ(MatFDColoringSetUp(B, coloring, fdcoloring));
162       CHKERRQ(PetscObjectCompose((PetscObject) dm, "DMDASNES_FDCOLORING", (PetscObject) fdcoloring));
163       CHKERRQ(PetscObjectDereference((PetscObject) fdcoloring));
164 
165       /* The following breaks an ugly reference counting loop that deserves a paragraph. MatFDColoringApply() will call
166        * VecDuplicate() with the state Vec and store inside the MatFDColoring. This Vec will duplicate the Vec, but the
167        * MatFDColoring is composed with the DM. We dereference the DM here so that the reference count will eventually
168        * drop to 0. Note the code in DMDestroy() that exits early for a negative reference count. That code path will be
169        * taken when the PetscObjectList for the Vec inside MatFDColoring is destroyed.
170        */
171       CHKERRQ(PetscObjectDereference((PetscObject) dm));
172     }
173     CHKERRQ(MatFDColoringApply(B, fdcoloring, X, ts));
174   }
175   /* This will be redundant if the user called both, but it's too common to forget. */
176   if (A != B) {
177     CHKERRQ(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
178     CHKERRQ(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
179   }
180   PetscFunctionReturn(0);
181 }
182 
183 /*@C
184   DMTSSetBoundaryLocal - set the function for essential boundary data for a local implicit function evaluation.
185     It should set the essential boundary data for the local portion of the solution X, as well its time derivative X_t (if it is not NULL).
186     Vectors are initialized to zero before this function, so it is only needed for non homogeneous data.
187 
188   Note that this function is somewhat optional: boundary data could potentially be inserted by a function passed to
189   DMTSSetIFunctionLocal().  The use case for this function is for discretizations with constraints (see
190   DMGetDefaultConstraints()): this function inserts boundary values before constraint interpolation.
191 
192   Logically Collective
193 
194   Input Parameters:
195 + dm   - DM to associate callback with
196 . func - local function evaluation
197 - ctx  - context for function evaluation
198 
199   Level: intermediate
200 
201 .seealso: DMTSSetIFunction(), DMTSSetIJacobianLocal()
202 @*/
203 PetscErrorCode DMTSSetBoundaryLocal(DM dm, PetscErrorCode (*func)(DM, PetscReal, Vec, Vec, void *), void *ctx)
204 {
205   DMTS           tdm;
206   DMTS_Local    *dmlocalts;
207 
208   PetscFunctionBegin;
209   PetscValidHeaderSpecific(dm,DM_CLASSID,1);
210   CHKERRQ(DMGetDMTSWrite(dm, &tdm));
211   CHKERRQ(DMLocalTSGetContext(dm, tdm, &dmlocalts));
212 
213   dmlocalts->boundarylocal    = func;
214   dmlocalts->boundarylocalctx = ctx;
215 
216   PetscFunctionReturn(0);
217 }
218 
219 /*@C
220   DMTSSetIFunctionLocal - set a local implicit function evaluation function. This function is called with local vector
221       containing the local vector information PLUS ghost point information. It should compute a result for all local
222       elements and DMTS will automatically accumulate the overlapping values.
223 
224   Logically Collective
225 
226   Input Parameters:
227 + dm   - DM to associate callback with
228 . func - local function evaluation
229 - ctx  - context for function evaluation
230 
231   Level: beginner
232 
233 .seealso: DMTSSetIFunction(), DMTSSetIJacobianLocal()
234 @*/
235 PetscErrorCode DMTSSetIFunctionLocal(DM dm, PetscErrorCode (*func)(DM, PetscReal, Vec, Vec, Vec, void *), void *ctx)
236 {
237   DMTS           tdm;
238   DMTS_Local    *dmlocalts;
239 
240   PetscFunctionBegin;
241   PetscValidHeaderSpecific(dm,DM_CLASSID,1);
242   CHKERRQ(DMGetDMTSWrite(dm, &tdm));
243   CHKERRQ(DMLocalTSGetContext(dm, tdm, &dmlocalts));
244 
245   dmlocalts->ifunctionlocal    = func;
246   dmlocalts->ifunctionlocalctx = ctx;
247 
248   CHKERRQ(DMTSSetIFunction(dm, TSComputeIFunction_DMLocal, dmlocalts));
249   if (!tdm->ops->ijacobian) {  /* Call us for the Jacobian too, can be overridden by the user. */
250     CHKERRQ(DMTSSetIJacobian(dm, TSComputeIJacobian_DMLocal, dmlocalts));
251   }
252   PetscFunctionReturn(0);
253 }
254 
255 /*@C
256   DMTSSetIJacobianLocal - set a local Jacobian evaluation function
257 
258   Logically Collective
259 
260   Input Parameters:
261 + dm - DM to associate callback with
262 . func - local Jacobian evaluation
263 - ctx - optional context for local Jacobian evaluation
264 
265   Level: beginner
266 
267 .seealso: DMTSSetIFunctionLocal(), DMTSSetIJacobian(), DMTSSetIFunction()
268 @*/
269 PetscErrorCode DMTSSetIJacobianLocal(DM dm, PetscErrorCode (*func)(DM, PetscReal, Vec, Vec, PetscReal, Mat, Mat, void *), void *ctx)
270 {
271   DMTS           tdm;
272   DMTS_Local    *dmlocalts;
273 
274   PetscFunctionBegin;
275   PetscValidHeaderSpecific(dm,DM_CLASSID,1);
276   CHKERRQ(DMGetDMTSWrite(dm, &tdm));
277   CHKERRQ(DMLocalTSGetContext(dm, tdm, &dmlocalts));
278 
279   dmlocalts->ijacobianlocal    = func;
280   dmlocalts->ijacobianlocalctx = ctx;
281 
282   CHKERRQ(DMTSSetIJacobian(dm, TSComputeIJacobian_DMLocal, dmlocalts));
283   PetscFunctionReturn(0);
284 }
285 
286 /*@C
287   DMTSSetRHSFunctionLocal - set a local rhs function evaluation function. This function is called with local vector
288       containing the local vector information PLUS ghost point information. It should compute a result for all local
289       elements and DMTS will automatically accumulate the overlapping values.
290 
291   Logically Collective
292 
293   Input Parameters:
294 + dm   - DM to associate callback with
295 . func - local function evaluation
296 - ctx  - context for function evaluation
297 
298   Level: beginner
299 
300 .seealso: DMTSSetRHSFunction(), DMTSSetIFunction(), DMTSSetIJacobianLocal()
301 @*/
302 PetscErrorCode DMTSSetRHSFunctionLocal(DM dm, PetscErrorCode (*func)(DM, PetscReal, Vec, Vec, void *), void *ctx)
303 {
304   DMTS           tdm;
305   DMTS_Local    *dmlocalts;
306 
307   PetscFunctionBegin;
308   PetscValidHeaderSpecific(dm,DM_CLASSID,1);
309   CHKERRQ(DMGetDMTSWrite(dm, &tdm));
310   CHKERRQ(DMLocalTSGetContext(dm, tdm, &dmlocalts));
311 
312   dmlocalts->rhsfunctionlocal    = func;
313   dmlocalts->rhsfunctionlocalctx = ctx;
314 
315   CHKERRQ(DMTSSetRHSFunction(dm, TSComputeRHSFunction_DMLocal, dmlocalts));
316   PetscFunctionReturn(0);
317 }
318 
319 /*@C
320   DMTSCreateRHSMassMatrix - This creates the mass matrix associated with the given DM, and a solver to invert it, and stores them in the DMTS context.
321 
322   Collective on dm
323 
324   Input Parameters:
325 . dm   - DM providing the mass matrix
326 
327   Note: The idea here is that an explicit system can be given a mass matrix, based on the DM, which is inverted on the RHS at each step.
328 
329   Level: developer
330 
331 .seealso: DMTSCreateRHSMassMatrixLumped(), DMTSDestroyRHSMassMatrix(), DMCreateMassMatrix(), DMTS
332 @*/
333 PetscErrorCode DMTSCreateRHSMassMatrix(DM dm)
334 {
335   DMTS           tdm;
336   DMTS_Local    *dmlocalts;
337   const char    *prefix;
338 
339   PetscFunctionBegin;
340   PetscValidHeaderSpecific(dm,DM_CLASSID,1);
341   CHKERRQ(DMGetDMTSWrite(dm, &tdm));
342   CHKERRQ(DMLocalTSGetContext(dm, tdm, &dmlocalts));
343   CHKERRQ(DMCreateMassMatrix(dm, dm, &dmlocalts->mass));
344   CHKERRQ(KSPCreate(PetscObjectComm((PetscObject) dm), &dmlocalts->kspmass));
345   CHKERRQ(PetscObjectGetOptionsPrefix((PetscObject) dm, &prefix));
346   CHKERRQ(KSPSetOptionsPrefix(dmlocalts->kspmass, prefix));
347   CHKERRQ(KSPAppendOptionsPrefix(dmlocalts->kspmass, "mass_"));
348   CHKERRQ(KSPSetFromOptions(dmlocalts->kspmass));
349   CHKERRQ(KSPSetOperators(dmlocalts->kspmass, dmlocalts->mass, dmlocalts->mass));
350   PetscFunctionReturn(0);
351 }
352 
353 /*@C
354   DMTSCreateRHSMassMatrixLumped - This creates the lumped mass matrix associated with the given DM, and a solver to invert it, and stores them in the DMTS context.
355 
356   Collective on dm
357 
358   Input Parameters:
359 . dm   - DM providing the mass matrix
360 
361   Note: The idea here is that an explicit system can be given a mass matrix, based on the DM, which is inverted on the RHS at each step.
362   Since the matrix is lumped, inversion is trivial.
363 
364   Level: developer
365 
366 .seealso: DMTSCreateRHSMassMatrix(), DMTSDestroyRHSMassMatrix(), DMCreateMassMatrix(), DMTS
367 @*/
368 PetscErrorCode DMTSCreateRHSMassMatrixLumped(DM dm)
369 {
370   DMTS           tdm;
371   DMTS_Local    *dmlocalts;
372 
373   PetscFunctionBegin;
374   PetscValidHeaderSpecific(dm,DM_CLASSID,1);
375   CHKERRQ(DMGetDMTSWrite(dm, &tdm));
376   CHKERRQ(DMLocalTSGetContext(dm, tdm, &dmlocalts));
377   CHKERRQ(DMCreateMassMatrixLumped(dm, &dmlocalts->lumpedmassinv));
378   CHKERRQ(VecReciprocal(dmlocalts->lumpedmassinv));
379   CHKERRQ(VecViewFromOptions(dmlocalts->lumpedmassinv, NULL, "-lumped_mass_inv_view"));
380   PetscFunctionReturn(0);
381 }
382 
383 /*@C
384   DMTSDestroyRHSMassMatrix - Destroys the mass matrix and solver stored in the DMTS context, if they exist.
385 
386   Logically Collective
387 
388   Input Parameters:
389 . dm   - DM providing the mass matrix
390 
391   Level: developer
392 
393 .seealso: DMTSCreateRHSMassMatrixLumped(), DMCreateMassMatrix(), DMCreateMassMatrix(), DMTS
394 @*/
395 PetscErrorCode DMTSDestroyRHSMassMatrix(DM dm)
396 {
397   DMTS           tdm;
398   DMTS_Local    *dmlocalts;
399 
400   PetscFunctionBegin;
401   PetscValidHeaderSpecific(dm,DM_CLASSID,1);
402   CHKERRQ(DMGetDMTSWrite(dm, &tdm));
403   CHKERRQ(DMLocalTSGetContext(dm, tdm, &dmlocalts));
404   CHKERRQ(VecDestroy(&dmlocalts->lumpedmassinv));
405   CHKERRQ(MatDestroy(&dmlocalts->mass));
406   CHKERRQ(KSPDestroy(&dmlocalts->kspmass));
407   PetscFunctionReturn(0);
408 }
409