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