xref: /petsc/src/ts/utils/dmlocalts.c (revision 2c9ec6dfe874b911fa49ef7e759d29a8430d6aff)
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   CHKMEMQ;
114   ierr = (*dmlocalts->rhsfunctionlocal)(dm, time, locX, F, dmlocalts->rhsfunctionlocalctx);CHKERRQ(ierr);
115   CHKMEMQ;
116   ierr = DMRestoreLocalVector(dm, &locX);CHKERRQ(ierr);
117   PetscFunctionReturn(0);
118 }
119 
120 #undef __FUNCT__
121 #define __FUNCT__ "TSComputeIJacobian_DMLocal"
122 static PetscErrorCode TSComputeIJacobian_DMLocal(TS ts, PetscReal time, Vec X, Vec X_t, PetscReal a, Mat A, Mat B, void *ctx)
123 {
124   DM             dm;
125   Vec            locX, locX_t;
126   DMTS_Local    *dmlocalts = (DMTS_Local *) ctx;
127   PetscErrorCode ierr;
128 
129   PetscFunctionBegin;
130   ierr = TSGetDM(ts, &dm);CHKERRQ(ierr);
131   if (dmlocalts->ijacobianlocal) {
132     ierr = DMGetLocalVector(dm, &locX);CHKERRQ(ierr);
133     ierr = DMGetLocalVector(dm, &locX_t);CHKERRQ(ierr);
134     ierr = VecZeroEntries(locX);CHKERRQ(ierr);
135     ierr = VecZeroEntries(locX_t);CHKERRQ(ierr);
136     ierr = DMGlobalToLocalBegin(dm, X, INSERT_VALUES, locX);CHKERRQ(ierr);
137     ierr = DMGlobalToLocalEnd(dm, X, INSERT_VALUES, locX);CHKERRQ(ierr);
138     ierr = DMGlobalToLocalBegin(dm, X_t, INSERT_VALUES, locX_t);CHKERRQ(ierr);
139     ierr = DMGlobalToLocalEnd(dm, X_t, INSERT_VALUES, locX_t);CHKERRQ(ierr);
140     CHKMEMQ;
141     ierr = (*dmlocalts->ijacobianlocal)(dm, time, locX, locX_t, a, A, B, dmlocalts->ijacobianlocalctx);CHKERRQ(ierr);
142     CHKMEMQ;
143     ierr = DMRestoreLocalVector(dm, &locX);CHKERRQ(ierr);
144     ierr = DMRestoreLocalVector(dm, &locX_t);CHKERRQ(ierr);
145   } else {
146     MatFDColoring fdcoloring;
147     ierr = PetscObjectQuery((PetscObject) dm, "DMDASNES_FDCOLORING", (PetscObject *) &fdcoloring);CHKERRQ(ierr);
148     if (!fdcoloring) {
149       ISColoring coloring;
150 
151       ierr = DMCreateColoring(dm, dm->coloringtype, &coloring);CHKERRQ(ierr);
152       ierr = MatFDColoringCreate(B, coloring, &fdcoloring);CHKERRQ(ierr);
153       ierr = ISColoringDestroy(&coloring);CHKERRQ(ierr);
154       switch (dm->coloringtype) {
155       case IS_COLORING_GLOBAL:
156         ierr = MatFDColoringSetFunction(fdcoloring, (PetscErrorCode (*)(void)) TSComputeIFunction_DMLocal, dmlocalts);CHKERRQ(ierr);
157         break;
158       default: SETERRQ1(PetscObjectComm((PetscObject) ts), PETSC_ERR_SUP, "No support for coloring type '%s'", ISColoringTypes[dm->coloringtype]);
159       }
160       ierr = PetscObjectSetOptionsPrefix((PetscObject) fdcoloring, ((PetscObject) dm)->prefix);CHKERRQ(ierr);
161       ierr = MatFDColoringSetFromOptions(fdcoloring);CHKERRQ(ierr);
162       ierr = MatFDColoringSetUp(B, coloring, fdcoloring);CHKERRQ(ierr);
163       ierr = PetscObjectCompose((PetscObject) dm, "DMDASNES_FDCOLORING", (PetscObject) fdcoloring);CHKERRQ(ierr);
164       ierr = PetscObjectDereference((PetscObject) fdcoloring);CHKERRQ(ierr);
165 
166       /* The following breaks an ugly reference counting loop that deserves a paragraph. MatFDColoringApply() will call
167        * VecDuplicate() with the state Vec and store inside the MatFDColoring. This Vec will duplicate the Vec, but the
168        * MatFDColoring is composed with the DM. We dereference the DM here so that the reference count will eventually
169        * drop to 0. Note the code in DMDestroy() that exits early for a negative reference count. That code path will be
170        * taken when the PetscObjectList for the Vec inside MatFDColoring is destroyed.
171        */
172       ierr = PetscObjectDereference((PetscObject) dm);CHKERRQ(ierr);
173     }
174     ierr = MatFDColoringApply(B, fdcoloring, X, ts);CHKERRQ(ierr);
175   }
176   /* This will be redundant if the user called both, but it's too common to forget. */
177   if (A != B) {
178     ierr = MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
179     ierr = MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
180   }
181   PetscFunctionReturn(0);
182 }
183 
184 #undef __FUNCT__
185 #define __FUNCT__ "DMTSSetIFunctionLocal"
186 /*@C
187   DMTSSetIFunctionLocal - set a local implicit function evaluation function. This function is called with local vector
188       containing the local vector information PLUS ghost point information. It should compute a result for all local
189       elements and DMTS will automatically accumulate the overlapping values.
190 
191   Logically Collective
192 
193   Input Arguments:
194 + dm   - DM to associate callback with
195 . func - local function evaluation
196 - ctx  - context for function evaluation
197 
198   Level: beginner
199 
200 .seealso: DMTSSetIFunction(), DMTSSetIJacobianLocal()
201 @*/
202 PetscErrorCode DMTSSetIFunctionLocal(DM dm, PetscErrorCode (*func)(DM, PetscReal, Vec, Vec, Vec, void *), void *ctx)
203 {
204   DMTS           tdm;
205   DMTS_Local    *dmlocalts;
206   PetscErrorCode ierr;
207 
208   PetscFunctionBegin;
209   PetscValidHeaderSpecific(dm,DM_CLASSID,1);
210   ierr = DMGetDMTSWrite(dm, &tdm);CHKERRQ(ierr);
211   ierr = DMLocalTSGetContext(dm, tdm, &dmlocalts);CHKERRQ(ierr);
212 
213   dmlocalts->ifunctionlocal    = func;
214   dmlocalts->ifunctionlocalctx = ctx;
215 
216   ierr = DMTSSetIFunction(dm, TSComputeIFunction_DMLocal, dmlocalts);CHKERRQ(ierr);
217   if (!tdm->ops->ijacobian) {  /* Call us for the Jacobian too, can be overridden by the user. */
218     ierr = DMTSSetIJacobian(dm, TSComputeIJacobian_DMLocal, dmlocalts);CHKERRQ(ierr);
219   }
220   PetscFunctionReturn(0);
221 }
222 
223 #undef __FUNCT__
224 #define __FUNCT__ "DMTSSetIJacobianLocal"
225 /*@C
226   DMTSSetIJacobianLocal - set a local Jacobian evaluation function
227 
228   Logically Collective
229 
230   Input Arguments:
231 + dm - DM to associate callback with
232 . func - local Jacobian evaluation
233 - ctx - optional context for local Jacobian evaluation
234 
235   Level: beginner
236 
237 .seealso: DMTSSetIFunctionLocal(), DMTSSetIJacobian(), DMTSSetIFunction()
238 @*/
239 PetscErrorCode DMTSSetIJacobianLocal(DM dm, PetscErrorCode (*func)(DM, PetscReal, Vec, Vec, PetscReal, Mat, Mat, void *), void *ctx)
240 {
241   DMTS           tdm;
242   DMTS_Local    *dmlocalts;
243   PetscErrorCode ierr;
244 
245   PetscFunctionBegin;
246   PetscValidHeaderSpecific(dm,DM_CLASSID,1);
247   ierr = DMGetDMTSWrite(dm, &tdm);CHKERRQ(ierr);
248   ierr = DMLocalTSGetContext(dm, tdm, &dmlocalts);CHKERRQ(ierr);
249 
250   dmlocalts->ijacobianlocal    = func;
251   dmlocalts->ijacobianlocalctx = ctx;
252 
253   ierr = DMTSSetIJacobian(dm, TSComputeIJacobian_DMLocal, dmlocalts);CHKERRQ(ierr);
254   PetscFunctionReturn(0);
255 }
256 
257 #undef __FUNCT__
258 #define __FUNCT__ "DMTSSetRHSFunctionLocal"
259 /*@C
260   DMTSSetRHSFunctionLocal - set a local rhs function evaluation function. This function is called with local vector
261       containing the local vector information PLUS ghost point information. It should compute a result for all local
262       elements and DMTS will automatically accumulate the overlapping values.
263 
264   Logically Collective
265 
266   Input Arguments:
267 + dm   - DM to associate callback with
268 . func - local function evaluation
269 - ctx  - context for function evaluation
270 
271   Level: beginner
272 
273 .seealso: DMTSSetRHSFunction(), DMTSSetIFunction(), DMTSSetIJacobianLocal()
274 @*/
275 PetscErrorCode DMTSSetRHSFunctionLocal(DM dm, PetscErrorCode (*func)(DM, PetscReal, Vec, Vec, 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->rhsfunctionlocal    = func;
287   dmlocalts->rhsfunctionlocalctx = ctx;
288 
289   ierr = DMTSSetRHSFunction(dm, TSComputeRHSFunction_DMLocal, dmlocalts);CHKERRQ(ierr);
290   PetscFunctionReturn(0);
291 }
292