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