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