xref: /petsc/src/ts/utils/dmlocalts.c (revision 7eb36faedfe3d52783893769aae549936b218a00)
1 #include <petsc/private/dmimpl.h>
2 #include <petsc/private/tsimpl.h>   /*I "petscts.h" I*/
3 
4 typedef struct {
5   PetscErrorCode (*ifunctionlocal)(DM,PetscReal,Vec,Vec,Vec,void*);
6   PetscErrorCode (*ijacobianlocal)(DM,PetscReal,Vec,Vec,PetscReal,Mat,Mat,void*);
7   PetscErrorCode (*rhsfunctionlocal)(DM,PetscReal,Vec,Vec,void*);
8   PetscErrorCode (*boundarylocal)(DM,Vec,void*);
9   void *ifunctionlocalctx;
10   void *ijacobianlocalctx;
11   void *rhsfunctionlocalctx;
12   void *boundarylocalctx;
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   if (dmlocalts->boundarylocal) {ierr = (*dmlocalts->boundarylocal)(dm,locX,dmlocalts->boundarylocalctx);CHKERRQ(ierr);}
76   ierr = VecZeroEntries(locX_t);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,locX,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,locX,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 a local boundary value function. This function is called with local vector
190     containing the local vector information PLUS ghost point information. It should insert values into the
191     local vector that do not come from the global vector, such as essential boundary condition data.
192 
193   Logically Collective
194 
195   Input Arguments:
196 + dm   - DM to associate callback with
197 . func - local boundary value evaluation
198 - ctx  - optional context for local boundary value evaluation
199 
200   Level: intermediate
201 
202 .seealso: DMTSSetIFunctionLocal(), DMTSetIJacobianLocal(), DMSNESSetBoundaryLocal(), DMSNESSetFunctionLocal()
203 @*/
204 PetscErrorCode DMTSSetBoundaryLocal(DM dm, PetscErrorCode (*func)(DM, Vec, void *), void *ctx)
205 {
206   DMTS           tdm;
207   DMTS_Local    *dmlocalts;
208   PetscErrorCode ierr;
209 
210   PetscFunctionBegin;
211   PetscValidHeaderSpecific(dm, DM_CLASSID, 1);
212   ierr = DMGetDMTSWrite(dm, &tdm);CHKERRQ(ierr);
213   ierr = DMLocalTSGetContext(dm, tdm, &dmlocalts);CHKERRQ(ierr);
214 
215   dmlocalts->boundarylocal    = func;
216   dmlocalts->boundarylocalctx = ctx;
217   PetscFunctionReturn(0);
218 }
219 
220 #undef __FUNCT__
221 #define __FUNCT__ "DMTSSetIFunctionLocal"
222 /*@C
223   DMTSSetIFunctionLocal - set a local implicit function evaluation function. This function is called with local vector
224       containing the local vector information PLUS ghost point information. It should compute a result for all local
225       elements and DMTS will automatically accumulate the overlapping values.
226 
227   Logically Collective
228 
229   Input Arguments:
230 + dm   - DM to associate callback with
231 . func - local function evaluation
232 - ctx  - context for function evaluation
233 
234   Level: beginner
235 
236 .seealso: DMTSSetIFunction(), DMTSSetIJacobianLocal()
237 @*/
238 PetscErrorCode DMTSSetIFunctionLocal(DM dm, PetscErrorCode (*func)(DM, PetscReal, Vec, Vec, Vec, void *), void *ctx)
239 {
240   DMTS           tdm;
241   DMTS_Local    *dmlocalts;
242   PetscErrorCode ierr;
243 
244   PetscFunctionBegin;
245   PetscValidHeaderSpecific(dm,DM_CLASSID,1);
246   ierr = DMGetDMTSWrite(dm, &tdm);CHKERRQ(ierr);
247   ierr = DMLocalTSGetContext(dm, tdm, &dmlocalts);CHKERRQ(ierr);
248 
249   dmlocalts->ifunctionlocal    = func;
250   dmlocalts->ifunctionlocalctx = ctx;
251 
252   ierr = DMTSSetIFunction(dm, TSComputeIFunction_DMLocal, dmlocalts);CHKERRQ(ierr);
253   if (!tdm->ops->ijacobian) {  /* Call us for the Jacobian too, can be overridden by the user. */
254     ierr = DMTSSetIJacobian(dm, TSComputeIJacobian_DMLocal, dmlocalts);CHKERRQ(ierr);
255   }
256   PetscFunctionReturn(0);
257 }
258 
259 #undef __FUNCT__
260 #define __FUNCT__ "DMTSSetIJacobianLocal"
261 /*@C
262   DMTSSetIJacobianLocal - set a local Jacobian evaluation function
263 
264   Logically Collective
265 
266   Input Arguments:
267 + dm - DM to associate callback with
268 . func - local Jacobian evaluation
269 - ctx - optional context for local Jacobian evaluation
270 
271   Level: beginner
272 
273 .seealso: DMTSSetIFunctionLocal(), DMTSSetIJacobian(), DMTSSetIFunction()
274 @*/
275 PetscErrorCode DMTSSetIJacobianLocal(DM dm, PetscErrorCode (*func)(DM, PetscReal, Vec, Vec, PetscReal, Mat, Mat, void *), void *ctx)
276 {
277   DMTS           tdm;
278   DMTS_Local    *dmlocalts;
279   PetscErrorCode ierr;
280 
281   PetscFunctionBegin;
282   PetscValidHeaderSpecific(dm,DM_CLASSID,1);
283   ierr = DMGetDMTSWrite(dm, &tdm);CHKERRQ(ierr);
284   ierr = DMLocalTSGetContext(dm, tdm, &dmlocalts);CHKERRQ(ierr);
285 
286   dmlocalts->ijacobianlocal    = func;
287   dmlocalts->ijacobianlocalctx = ctx;
288 
289   ierr = DMTSSetIJacobian(dm, TSComputeIJacobian_DMLocal, dmlocalts);CHKERRQ(ierr);
290   PetscFunctionReturn(0);
291 }
292 
293 #undef __FUNCT__
294 #define __FUNCT__ "DMTSSetRHSFunctionLocal"
295 /*@C
296   DMTSSetRHSFunctionLocal - set a local rhs function evaluation function. This function is called with local vector
297       containing the local vector information PLUS ghost point information. It should compute a result for all local
298       elements and DMTS will automatically accumulate the overlapping values.
299 
300   Logically Collective
301 
302   Input Arguments:
303 + dm   - DM to associate callback with
304 . func - local function evaluation
305 - ctx  - context for function evaluation
306 
307   Level: beginner
308 
309 .seealso: DMTSSetRHSFunction(), DMTSSetIFunction(), DMTSSetIJacobianLocal()
310 @*/
311 PetscErrorCode DMTSSetRHSFunctionLocal(DM dm, PetscErrorCode (*func)(DM, PetscReal, Vec, Vec, void *), void *ctx)
312 {
313   DMTS           tdm;
314   DMTS_Local    *dmlocalts;
315   PetscErrorCode ierr;
316 
317   PetscFunctionBegin;
318   PetscValidHeaderSpecific(dm,DM_CLASSID,1);
319   ierr = DMGetDMTSWrite(dm, &tdm);CHKERRQ(ierr);
320   ierr = DMLocalTSGetContext(dm, tdm, &dmlocalts);CHKERRQ(ierr);
321 
322   dmlocalts->rhsfunctionlocal    = func;
323   dmlocalts->rhsfunctionlocalctx = ctx;
324 
325   ierr = DMTSSetRHSFunction(dm, TSComputeRHSFunction_DMLocal, dmlocalts);CHKERRQ(ierr);
326   PetscFunctionReturn(0);
327 }
328