xref: /petsc/src/ts/interface/tsrhssplit.c (revision 4e8208cbcbc709572b8abe32f33c78b69c819375)
1 #include <petsc/private/tsimpl.h> /*I "petscts.h"  I*/
2 #include <petscdm.h>
TSRHSSplitGetRHSSplit(TS ts,const char splitname[],TS_RHSSplitLink * isplit)3 static PetscErrorCode TSRHSSplitGetRHSSplit(TS ts, const char splitname[], TS_RHSSplitLink *isplit)
4 {
5   PetscBool found = PETSC_FALSE;
6 
7   PetscFunctionBegin;
8   *isplit = ts->tsrhssplit;
9   /* look up the split */
10   while (*isplit) {
11     PetscCall(PetscStrcmp((*isplit)->splitname, splitname, &found));
12     if (found) break;
13     *isplit = (*isplit)->next;
14   }
15   PetscFunctionReturn(PETSC_SUCCESS);
16 }
17 
18 /*@
19   TSRHSSplitSetIS - Set the index set for the specified split
20 
21   Logically Collective
22 
23   Input Parameters:
24 + ts        - the `TS` context obtained from `TSCreate()`
25 . splitname - name of this split, if `NULL` the number of the split is used
26 - is        - the index set for part of the solution vector
27 
28   Level: intermediate
29 
30 .seealso: [](ch_ts), `TS`, `IS`, `TSRHSSplitGetIS()`, `TSARKIMEXSetFastSlowSplit()`
31 @*/
TSRHSSplitSetIS(TS ts,const char splitname[],IS is)32 PetscErrorCode TSRHSSplitSetIS(TS ts, const char splitname[], IS is)
33 {
34   TS_RHSSplitLink newsplit, next = ts->tsrhssplit;
35   char            prefix[128];
36 
37   PetscFunctionBegin;
38   PetscValidHeaderSpecific(ts, TS_CLASSID, 1);
39   PetscValidHeaderSpecific(is, IS_CLASSID, 3);
40 
41   PetscCall(PetscNew(&newsplit));
42   if (splitname) {
43     PetscCall(PetscStrallocpy(splitname, &newsplit->splitname));
44   } else {
45     PetscCall(PetscMalloc1(8, &newsplit->splitname));
46     PetscCall(PetscSNPrintf(newsplit->splitname, 7, "%" PetscInt_FMT, ts->num_rhs_splits));
47   }
48   PetscCall(PetscObjectReference((PetscObject)is));
49   newsplit->is = is;
50   PetscCall(TSCreate(PetscObjectComm((PetscObject)ts), &newsplit->ts));
51 
52   PetscCall(PetscObjectIncrementTabLevel((PetscObject)newsplit->ts, (PetscObject)ts, 1));
53   PetscCall(PetscSNPrintf(prefix, sizeof(prefix), "%srhsplit_%s_", ((PetscObject)ts)->prefix ? ((PetscObject)ts)->prefix : "", newsplit->splitname));
54   PetscCall(TSSetOptionsPrefix(newsplit->ts, prefix));
55   if (!next) ts->tsrhssplit = newsplit;
56   else {
57     while (next->next) next = next->next;
58     next->next = newsplit;
59   }
60   ts->num_rhs_splits++;
61   PetscFunctionReturn(PETSC_SUCCESS);
62 }
63 
64 /*@
65   TSRHSSplitGetIS - Retrieves the elements for a split as an `IS`
66 
67   Logically Collective
68 
69   Input Parameters:
70 + ts        - the `TS` context obtained from `TSCreate()`
71 - splitname - name of this split
72 
73   Output Parameter:
74 . is - the index set for part of the solution vector
75 
76   Level: intermediate
77 
78 .seealso: [](ch_ts), `TS`, `IS`, `TSRHSSplitSetIS()`
79 @*/
TSRHSSplitGetIS(TS ts,const char splitname[],IS * is)80 PetscErrorCode TSRHSSplitGetIS(TS ts, const char splitname[], IS *is)
81 {
82   TS_RHSSplitLink isplit;
83 
84   PetscFunctionBegin;
85   PetscValidHeaderSpecific(ts, TS_CLASSID, 1);
86   *is = NULL;
87   /* look up the split */
88   PetscCall(TSRHSSplitGetRHSSplit(ts, splitname, &isplit));
89   if (isplit) *is = isplit->is;
90   PetscFunctionReturn(PETSC_SUCCESS);
91 }
92 
93 /*@C
94   TSRHSSplitSetRHSFunction - Set the split right-hand-side functions.
95 
96   Logically Collective
97 
98   Input Parameters:
99 + ts        - the `TS` context obtained from `TSCreate()`
100 . splitname - name of this split
101 . r         - vector to hold the residual (or `NULL` to have it created internally)
102 . rhsfunc   - the RHS function evaluation routine
103 - ctx       - user-defined context for private data for the split function evaluation routine (may be `NULL`)
104 
105   Level: intermediate
106 
107 .seealso: [](ch_ts), `TS`, `TSRHSFunctionFn`, `IS`, `TSRHSSplitSetIS()`
108 @*/
TSRHSSplitSetRHSFunction(TS ts,const char splitname[],Vec r,TSRHSFunctionFn * rhsfunc,PetscCtx ctx)109 PetscErrorCode TSRHSSplitSetRHSFunction(TS ts, const char splitname[], Vec r, TSRHSFunctionFn *rhsfunc, PetscCtx ctx)
110 {
111   TS_RHSSplitLink isplit;
112   DM              dmc;
113   Vec             subvec, ralloc = NULL;
114 
115   PetscFunctionBegin;
116   PetscValidHeaderSpecific(ts, TS_CLASSID, 1);
117   if (r) PetscValidHeaderSpecific(r, VEC_CLASSID, 3);
118 
119   /* look up the split */
120   PetscCall(TSRHSSplitGetRHSSplit(ts, splitname, &isplit));
121   PetscCheck(isplit, PETSC_COMM_SELF, PETSC_ERR_USER, "The split %s is not created, check the split name or call TSRHSSplitSetIS() to create one", splitname);
122 
123   if (!r && ts->vec_sol) {
124     PetscCall(VecGetSubVector(ts->vec_sol, isplit->is, &subvec));
125     PetscCall(VecDuplicate(subvec, &ralloc));
126     r = ralloc;
127     PetscCall(VecRestoreSubVector(ts->vec_sol, isplit->is, &subvec));
128   }
129 
130   if (ts->dm) {
131     PetscInt dim;
132 
133     PetscCall(DMGetDimension(ts->dm, &dim));
134     if (dim != -1) {
135       PetscCall(DMClone(ts->dm, &dmc));
136       PetscCall(TSSetDM(isplit->ts, dmc));
137       PetscCall(DMDestroy(&dmc));
138     }
139   }
140 
141   PetscCall(TSSetRHSFunction(isplit->ts, r, rhsfunc, ctx));
142   PetscCall(VecDestroy(&ralloc));
143   PetscFunctionReturn(PETSC_SUCCESS);
144 }
145 
146 /*@C
147   TSRHSSplitSetIFunction - Set the split implicit function for `TSARKIMEX`
148 
149   Logically Collective
150 
151   Input Parameters:
152 + ts        - the `TS` context obtained from `TSCreate()`
153 . splitname - name of this split
154 . r         - vector to hold the residual (or `NULL` to have it created internally)
155 . ifunc     - the implicit function evaluation routine
156 - ctx       - user-defined context for private data for the split function evaluation routine (may be `NULL`)
157 
158   Level: intermediate
159 
160 .seealso: [](ch_ts), `TS`, `TSIFunctionFn`, `IS`, `TSRHSSplitSetIS()`, `TSARKIMEX`, `TSARKIMEXSetFastSlowSplit()`
161 @*/
TSRHSSplitSetIFunction(TS ts,const char splitname[],Vec r,TSIFunctionFn * ifunc,PetscCtx ctx)162 PetscErrorCode TSRHSSplitSetIFunction(TS ts, const char splitname[], Vec r, TSIFunctionFn *ifunc, PetscCtx ctx)
163 {
164   TS_RHSSplitLink isplit;
165   DM              dmc;
166   Vec             subvec, ralloc = NULL;
167 
168   PetscFunctionBegin;
169   PetscValidHeaderSpecific(ts, TS_CLASSID, 1);
170   if (r) PetscValidHeaderSpecific(r, VEC_CLASSID, 3);
171 
172   /* look up the split */
173   PetscCall(TSRHSSplitGetRHSSplit(ts, splitname, &isplit));
174   PetscCheck(isplit, PETSC_COMM_SELF, PETSC_ERR_USER, "The split %s is not created, check the split name or call TSRHSSplitSetIS() to create one", splitname);
175 
176   if (!r && ts->vec_sol) {
177     PetscCall(VecGetSubVector(ts->vec_sol, isplit->is, &subvec));
178     PetscCall(VecDuplicate(subvec, &ralloc));
179     r = ralloc;
180     PetscCall(VecRestoreSubVector(ts->vec_sol, isplit->is, &subvec));
181   }
182 
183   if (ts->dm) {
184     PetscInt dim;
185 
186     PetscCall(DMGetDimension(ts->dm, &dim));
187     if (dim != -1) {
188       PetscCall(DMClone(ts->dm, &dmc));
189       PetscCall(TSSetDM(isplit->ts, dmc));
190       PetscCall(DMDestroy(&dmc));
191     }
192   }
193 
194   PetscCall(TSSetIFunction(isplit->ts, r, ifunc, ctx));
195   PetscCall(VecDestroy(&ralloc));
196   PetscFunctionReturn(PETSC_SUCCESS);
197 }
198 
199 /*@C
200   TSRHSSplitSetIJacobian - Set the Jacobian for the split implicit function with `TSARKIMEX`
201 
202   Logically Collective
203 
204   Input Parameters:
205 + ts        - the `TS` context obtained from `TSCreate()`
206 . splitname - name of this split
207 . Amat      - (approximate) matrix to store Jacobian entries computed by `f`
208 . Pmat      - matrix used to compute preconditioner (usually the same as `Amat`)
209 . ijac      - the Jacobian evaluation routine
210 - ctx       - user-defined context for private data for the split function evaluation routine (may be `NULL`)
211 
212   Level: intermediate
213 
214 .seealso: [](ch_ts), `TS`, `TSRHSSplitSetIFunction`, `TSIJacobianFn`, `IS`, `TSRHSSplitSetIS()`, `TSARKIMEXSetFastSlowSplit()`
215 @*/
TSRHSSplitSetIJacobian(TS ts,const char splitname[],Mat Amat,Mat Pmat,TSIJacobianFn * ijac,PetscCtx ctx)216 PetscErrorCode TSRHSSplitSetIJacobian(TS ts, const char splitname[], Mat Amat, Mat Pmat, TSIJacobianFn *ijac, PetscCtx ctx)
217 {
218   TS_RHSSplitLink isplit;
219   DM              dmc;
220 
221   PetscFunctionBegin;
222   PetscValidHeaderSpecific(ts, TS_CLASSID, 1);
223   if (Amat) PetscValidHeaderSpecific(Amat, MAT_CLASSID, 3);
224   if (Pmat) PetscValidHeaderSpecific(Pmat, MAT_CLASSID, 4);
225   if (Amat) PetscCheckSameComm(ts, 1, Amat, 3);
226   if (Pmat) PetscCheckSameComm(ts, 1, Pmat, 4);
227 
228   /* look up the split */
229   PetscCall(TSRHSSplitGetRHSSplit(ts, splitname, &isplit));
230   PetscCheck(isplit, PETSC_COMM_SELF, PETSC_ERR_USER, "The split %s is not created, check the split name or call TSRHSSplitSetIS() to create one", splitname);
231 
232   if (ts->dm) {
233     PetscInt dim;
234 
235     PetscCall(DMGetDimension(ts->dm, &dim));
236     if (dim != -1) {
237       PetscCall(DMClone(ts->dm, &dmc));
238       PetscCall(TSSetDM(isplit->ts, dmc));
239       PetscCall(DMDestroy(&dmc));
240     }
241   }
242 
243   PetscCall(TSSetIJacobian(isplit->ts, Amat, Pmat, ijac, ctx));
244   PetscFunctionReturn(PETSC_SUCCESS);
245 }
246 
247 /*@C
248   TSRHSSplitGetSubTS - Get the sub-`TS` by split name.
249 
250   Logically Collective
251 
252   Input Parameter:
253 . ts - the `TS` context obtained from `TSCreate()`
254 
255   Output Parameters:
256 + splitname - the number of the split
257 - subts     - the sub-`TS`
258 
259   Level: advanced
260 
261 .seealso: [](ch_ts), `TS`, `IS`, `TSGetRHSSplitFunction()`
262 @*/
TSRHSSplitGetSubTS(TS ts,const char splitname[],TS * subts)263 PetscErrorCode TSRHSSplitGetSubTS(TS ts, const char splitname[], TS *subts)
264 {
265   TS_RHSSplitLink isplit;
266 
267   PetscFunctionBegin;
268   PetscValidHeaderSpecific(ts, TS_CLASSID, 1);
269   PetscAssertPointer(subts, 3);
270   *subts = NULL;
271   /* look up the split */
272   PetscCall(TSRHSSplitGetRHSSplit(ts, splitname, &isplit));
273   if (isplit) *subts = isplit->ts;
274   PetscFunctionReturn(PETSC_SUCCESS);
275 }
276 
277 /*@C
278   TSRHSSplitGetSubTSs - Get an array of all sub-`TS` contexts.
279 
280   Logically Collective
281 
282   Input Parameter:
283 . ts - the `TS` context obtained from `TSCreate()`
284 
285   Output Parameters:
286 + n     - the number of splits
287 - subts - the array of `TS` contexts
288 
289   Level: advanced
290 
291   Note:
292   After `TSRHSSplitGetSubTS()` the array of `TS`s is to be freed by the user with `PetscFree()`
293   (not the `TS` in the array just the array that contains them).
294 
295 .seealso: [](ch_ts), `TS`, `IS`, `TSGetRHSSplitFunction()`
296 @*/
TSRHSSplitGetSubTSs(TS ts,PetscInt * n,TS * subts[])297 PetscErrorCode TSRHSSplitGetSubTSs(TS ts, PetscInt *n, TS *subts[])
298 {
299   TS_RHSSplitLink ilink = ts->tsrhssplit;
300   PetscInt        i     = 0;
301 
302   PetscFunctionBegin;
303   PetscValidHeaderSpecific(ts, TS_CLASSID, 1);
304   if (subts) {
305     PetscCall(PetscMalloc1(ts->num_rhs_splits, subts));
306     while (ilink) {
307       (*subts)[i++] = ilink->ts;
308       ilink         = ilink->next;
309     }
310   }
311   if (n) *n = ts->num_rhs_splits;
312   PetscFunctionReturn(PETSC_SUCCESS);
313 }
314 
315 /*@
316   TSRHSSplitGetSNES - Returns the `SNES` (nonlinear solver) associated with
317   a `TS` (timestepper) context when RHS splits are used.
318 
319   Not Collective, but `snes` is parallel if `ts` is parallel
320 
321   Input Parameter:
322 . ts - the `TS` context obtained from `TSCreate()`
323 
324   Output Parameter:
325 . snes - the nonlinear solver context
326 
327   Level: intermediate
328 
329   Note:
330   The returned `SNES` may have a different `DM` with the `TS` `DM`.
331 
332 .seealso: [](ch_ts), `TS`, `SNES`, `TSCreate()`, `TSRHSSplitSetSNES()`
333 @*/
TSRHSSplitGetSNES(TS ts,SNES * snes)334 PetscErrorCode TSRHSSplitGetSNES(TS ts, SNES *snes)
335 {
336   PetscFunctionBegin;
337   PetscValidHeaderSpecific(ts, TS_CLASSID, 1);
338   PetscAssertPointer(snes, 2);
339   if (!ts->snesrhssplit) {
340     PetscCall(SNESCreate(PetscObjectComm((PetscObject)ts), &ts->snesrhssplit));
341     PetscCall(PetscObjectSetOptions((PetscObject)ts->snesrhssplit, ((PetscObject)ts)->options));
342     PetscCall(SNESSetFunction(ts->snesrhssplit, NULL, SNESTSFormFunction, ts));
343     PetscCall(PetscObjectIncrementTabLevel((PetscObject)ts->snesrhssplit, (PetscObject)ts, 1));
344     if (ts->problem_type == TS_LINEAR) PetscCall(SNESSetType(ts->snesrhssplit, SNESKSPONLY));
345   }
346   *snes = ts->snesrhssplit;
347   PetscFunctionReturn(PETSC_SUCCESS);
348 }
349 
350 /*@
351   TSRHSSplitSetSNES - Set the `SNES` (nonlinear solver) to be used by the
352   timestepping context when RHS splits are used.
353 
354   Collective
355 
356   Input Parameters:
357 + ts   - the `TS` context obtained from `TSCreate()`
358 - snes - the nonlinear solver context
359 
360   Level: intermediate
361 
362   Note:
363   Most users should have the `TS` created by calling `TSRHSSplitGetSNES()`
364 
365 .seealso: [](ch_ts), `TS`, `SNES`, `TSCreate()`, `TSRHSSplitGetSNES()`
366 @*/
TSRHSSplitSetSNES(TS ts,SNES snes)367 PetscErrorCode TSRHSSplitSetSNES(TS ts, SNES snes)
368 {
369   PetscErrorCode (*func)(SNES, Vec, Mat, Mat, void *);
370 
371   PetscFunctionBegin;
372   PetscValidHeaderSpecific(ts, TS_CLASSID, 1);
373   PetscValidHeaderSpecific(snes, SNES_CLASSID, 2);
374   PetscCall(PetscObjectReference((PetscObject)snes));
375   PetscCall(SNESDestroy(&ts->snesrhssplit));
376 
377   ts->snesrhssplit = snes;
378 
379   PetscCall(SNESSetFunction(ts->snesrhssplit, NULL, SNESTSFormFunction, ts));
380   PetscCall(SNESGetJacobian(ts->snesrhssplit, NULL, NULL, &func, NULL));
381   if (func == SNESTSFormJacobian) PetscCall(SNESSetJacobian(ts->snesrhssplit, NULL, NULL, SNESTSFormJacobian, ts));
382   PetscFunctionReturn(PETSC_SUCCESS);
383 }
384