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