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