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