xref: /petsc/src/ts/interface/tsrhssplit.c (revision 89928cc5142e867cb7b0dd1ff74e0acffcd6b4b5)
1 #include <petsc/private/tsimpl.h> /*I "petscts.h"  I*/
2 #include <petscdm.h>
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(0);
16 }
17 
18 /*@C
19    TSRHSSplitSetIS - Set the index set for the specified split
20 
21    Logically Collective on TS
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: `TSRHSSplitGetIS()`
31 
32 @*/
33 PetscErrorCode TSRHSSplitSetIS(TS ts, const char splitname[], IS is)
34 {
35   TS_RHSSplitLink newsplit, next = ts->tsrhssplit;
36   char            prefix[128];
37 
38   PetscFunctionBegin;
39   PetscValidHeaderSpecific(ts, TS_CLASSID, 1);
40   PetscValidHeaderSpecific(is, IS_CLASSID, 3);
41 
42   PetscCall(PetscNew(&newsplit));
43   if (splitname) {
44     PetscCall(PetscStrallocpy(splitname, &newsplit->splitname));
45   } else {
46     PetscCall(PetscMalloc1(8, &newsplit->splitname));
47     PetscCall(PetscSNPrintf(newsplit->splitname, 7, "%" PetscInt_FMT, ts->num_rhs_splits));
48   }
49   PetscCall(PetscObjectReference((PetscObject)is));
50   newsplit->is = is;
51   PetscCall(TSCreate(PetscObjectComm((PetscObject)ts), &newsplit->ts));
52 
53   PetscCall(PetscObjectIncrementTabLevel((PetscObject)newsplit->ts, (PetscObject)ts, 1));
54   PetscCall(PetscSNPrintf(prefix, sizeof(prefix), "%srhsplit_%s_", ((PetscObject)ts)->prefix ? ((PetscObject)ts)->prefix : "", newsplit->splitname));
55   PetscCall(TSSetOptionsPrefix(newsplit->ts, prefix));
56   if (!next) ts->tsrhssplit = newsplit;
57   else {
58     while (next->next) next = next->next;
59     next->next = newsplit;
60   }
61   ts->num_rhs_splits++;
62   PetscFunctionReturn(0);
63 }
64 
65 /*@C
66    TSRHSSplitGetIS - Retrieves the elements for a split as an IS
67 
68    Logically Collective on TS
69 
70    Input Parameters:
71 +  ts        - the TS context obtained from TSCreate()
72 -  splitname - name of this split
73 
74    Output Parameters:
75 -  is        - the index set for part of the solution vector
76 
77    Level: intermediate
78 
79 .seealso: `TSRHSSplitSetIS()`
80 
81 @*/
82 PetscErrorCode TSRHSSplitGetIS(TS ts, const char splitname[], IS *is)
83 {
84   TS_RHSSplitLink isplit;
85 
86   PetscFunctionBegin;
87   PetscValidHeaderSpecific(ts, TS_CLASSID, 1);
88   *is = NULL;
89   /* look up the split */
90   PetscCall(TSRHSSplitGetRHSSplit(ts, splitname, &isplit));
91   if (isplit) *is = isplit->is;
92   PetscFunctionReturn(0);
93 }
94 
95 /*@C
96    TSRHSSplitSetRHSFunction - Set the split right-hand-side functions.
97 
98    Logically Collective on TS
99 
100    Input Parameters:
101 +  ts        - the TS context obtained from TSCreate()
102 .  splitname - name of this split
103 .  r         - vector to hold the residual (or NULL to have it created internally)
104 .  rhsfunc   - the RHS function evaluation routine
105 -  ctx       - user-defined context for private data for the split function evaluation routine (may be NULL)
106 
107  Calling sequence of fun:
108 $  rhsfunc(TS ts,PetscReal t,Vec u,Vec f,ctx);
109 
110 +  t    - time at step/stage being solved
111 .  u    - state vector
112 .  f    - function vector
113 -  ctx  - [optional] user-defined context for matrix evaluation routine (may be NULL)
114 
115  Level: beginner
116 
117 @*/
118 PetscErrorCode TSRHSSplitSetRHSFunction(TS ts, const char splitname[], Vec r, TSRHSFunction rhsfunc, void *ctx)
119 {
120   TS_RHSSplitLink isplit;
121   DM              dmc;
122   Vec             subvec, ralloc = NULL;
123 
124   PetscFunctionBegin;
125   PetscValidHeaderSpecific(ts, TS_CLASSID, 1);
126   if (r) PetscValidHeaderSpecific(r, VEC_CLASSID, 3);
127 
128   /* look up the split */
129   PetscCall(TSRHSSplitGetRHSSplit(ts, splitname, &isplit));
130   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);
131 
132   if (!r && ts->vec_sol) {
133     PetscCall(VecGetSubVector(ts->vec_sol, isplit->is, &subvec));
134     PetscCall(VecDuplicate(subvec, &ralloc));
135     r = ralloc;
136     PetscCall(VecRestoreSubVector(ts->vec_sol, isplit->is, &subvec));
137   }
138 
139   if (ts->dm) {
140     PetscInt dim;
141 
142     PetscCall(DMGetDimension(ts->dm, &dim));
143     if (dim != -1) {
144       PetscCall(DMClone(ts->dm, &dmc));
145       PetscCall(TSSetDM(isplit->ts, dmc));
146       PetscCall(DMDestroy(&dmc));
147     }
148   }
149 
150   PetscCall(TSSetRHSFunction(isplit->ts, r, rhsfunc, ctx));
151   PetscCall(VecDestroy(&ralloc));
152   PetscFunctionReturn(0);
153 }
154 
155 /*@C
156    TSRHSSplitGetSubTS - Get the sub-TS by split name.
157 
158    Logically Collective on TS
159 
160    Input Parameter:
161 .  ts - the TS context obtained from TSCreate()
162 
163    Output Parameters:
164 +  splitname - the number of the split
165 -  subts - the array of TS contexts
166 
167    Level: advanced
168 
169 .seealso: `TSGetRHSSplitFunction()`
170 @*/
171 PetscErrorCode TSRHSSplitGetSubTS(TS ts, const char splitname[], TS *subts)
172 {
173   TS_RHSSplitLink isplit;
174 
175   PetscFunctionBegin;
176   PetscValidHeaderSpecific(ts, TS_CLASSID, 1);
177   PetscValidPointer(subts, 3);
178   *subts = NULL;
179   /* look up the split */
180   PetscCall(TSRHSSplitGetRHSSplit(ts, splitname, &isplit));
181   if (isplit) *subts = isplit->ts;
182   PetscFunctionReturn(0);
183 }
184 
185 /*@C
186    TSRHSSplitGetSubTSs - Get an array of all sub-TS contexts.
187 
188    Logically Collective on TS
189 
190    Input Parameter:
191 .  ts - the TS context obtained from TSCreate()
192 
193    Output Parameters:
194 +  n - the number of splits
195 -  subksp - the array of TS contexts
196 
197    Note:
198    After TSRHSSplitGetSubTS() the array of TSs is to be freed by the user with PetscFree()
199    (not the TS just the array that contains them).
200 
201    Level: advanced
202 
203 .seealso: `TSGetRHSSplitFunction()`
204 @*/
205 PetscErrorCode TSRHSSplitGetSubTSs(TS ts, PetscInt *n, TS *subts[])
206 {
207   TS_RHSSplitLink ilink = ts->tsrhssplit;
208   PetscInt        i     = 0;
209 
210   PetscFunctionBegin;
211   PetscValidHeaderSpecific(ts, TS_CLASSID, 1);
212   if (subts) {
213     PetscCall(PetscMalloc1(ts->num_rhs_splits, subts));
214     while (ilink) {
215       (*subts)[i++] = ilink->ts;
216       ilink         = ilink->next;
217     }
218   }
219   if (n) *n = ts->num_rhs_splits;
220   PetscFunctionReturn(0);
221 }
222