xref: /petsc/src/ts/interface/tsrhssplit.c (revision 545aaa6f248f4d02efcfa901a11a18c680409921)
1 #include <petsc/private/tsimpl.h>        /*I "petscts.h"  I*/
2 
3 /*@C
4    TSRHSSplitSetIS - Set the index set for the specified split
5 
6    Logically Collective on TS
7 
8    Input Parameters:
9 +  ts        - the TS context obtained from TSCreate()
10 .  splitname - name of this split, if NULL the number of the split is used
11 -  is        - the index set for part of the solution vector
12 
13    Level: intermediate
14 
15 .seealso: TSRHSSplitGetIS
16 
17 .keywords: TS, TSRHSSplit
18 @*/
19 PetscErrorCode TSRHSSplitSetIS(TS ts,const char splitname[],IS is)
20 {
21   TS_RHSSplit    newsplit;
22   char           prefix[128];
23   PetscErrorCode ierr;
24 
25   PetscFunctionBegin;
26   PetscValidHeaderSpecific(ts,TS_CLASSID,1);
27   PetscValidHeaderSpecific(is,IS_CLASSID,3);
28   if (ts->num_rhs_splits == MAXRHSSPLITS) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_USER,"MAXIMUM number of splits reached");
29   ierr = PetscNew(&newsplit);CHKERRQ(ierr);
30   if (splitname) {
31     ierr = PetscStrallocpy(splitname,&newsplit->splitname);CHKERRQ(ierr);
32   } else {
33     ierr = PetscMalloc1(8,&newsplit->splitname);CHKERRQ(ierr);
34     ierr = PetscSNPrintf(newsplit->splitname,7,"%D",ts->num_rhs_splits);CHKERRQ(ierr);
35   }
36   ierr = PetscObjectReference((PetscObject)is);CHKERRQ(ierr);
37   newsplit->is = is;
38   ierr = TSCreate(PetscObjectComm((PetscObject)ts),&newsplit->ts);CHKERRQ(ierr);
39   ierr = PetscObjectIncrementTabLevel((PetscObject)newsplit->ts,(PetscObject)ts,1);CHKERRQ(ierr);
40   ierr = PetscLogObjectParent((PetscObject)ts,(PetscObject)newsplit->ts);CHKERRQ(ierr);
41   ierr = PetscSNPrintf(prefix,sizeof(prefix),"%srhsplit_%s_",((PetscObject)ts)->prefix ? ((PetscObject)ts)->prefix : "",newsplit->splitname);
42   ierr = TSSetOptionsPrefix(newsplit->ts,prefix);CHKERRQ(ierr);
43   ts->tsrhssplit[ts->num_rhs_splits++] = newsplit;
44   PetscFunctionReturn(0);
45 }
46 
47 /*@C
48    TSRHSSplitGetIS - Retrieves the elements for a split as an IS
49 
50    Logically Collective on TS
51 
52    Input Parameters:
53 +  ts        - the TS context obtained from TSCreate()
54 -  splitname - name of this split
55 
56    Output Parameters:
57 -  is        - the index set for part of the solution vector
58 
59    Level: intermediate
60 
61 .seealso: TSRHSSplitSetIS
62 
63 .keywords: TS, TSRHSSplit
64 @*/
65 PetscErrorCode TSRHSSplitGetIS(TS ts,const char splitname[],IS *is)
66 {
67   PetscInt       i = 0;
68   PetscBool      found = PETSC_FALSE;
69   PetscErrorCode ierr;
70 
71   PetscFunctionBegin;
72   PetscValidHeaderSpecific(ts,TS_CLASSID,1);
73   /* look up the split */
74   while (i<ts->num_rhs_splits) {
75     ierr = PetscStrcmp(ts->tsrhssplit[i]->splitname,splitname,&found);CHKERRQ(ierr);
76     if (found) {
77       *is = ts->tsrhssplit[i]->is;
78       break;
79     }
80     i++;
81   }
82   PetscFunctionReturn(0);
83 }
84 
85 /*@C
86    TSRHSSplitSetRHSFunction - Set the split right-hand-side functions.
87 
88    Logically Collective on TS
89 
90    Input Parameters:
91 +  ts        - the TS context obtained from TSCreate()
92 .  splitname - name of this split
93 .  r         - vector to hold the residual (or NULL to have it created internally)
94 .  rhsfunc   - the RHS function evaluation routine
95 -  ctx       - user-defined context for private data for the split function evaluation routine (may be NULL)
96 
97  Calling sequence of fun:
98 $  rhsfunc(TS ts,PetscReal t,Vec u,Vec f,ctx);
99 
100 +  t    - time at step/stage being solved
101 .  u    - state vector
102 .  f    - function vector
103 -  ctx  - [optional] user-defined context for matrix evaluation routine (may be NULL)
104 
105  Level: beginner
106 
107 .keywords: TS, timestep, set, ODE, Hamiltonian, Function
108 
109 .seealso: TSGetRHSSplitFunction()
110 @*/
111 PetscErrorCode TSRHSSplitSetRHSFunction(TS ts,const char splitname[],Vec r,TSRHSFunction rhsfunc,void *ctx)
112 {
113   DM             dm;
114   SNES           snes;
115   Vec            subvec,ralloc = NULL;
116   PetscBool      found = PETSC_FALSE;
117   PetscInt       i = 0;
118   TS             subts;
119   PetscErrorCode ierr;
120 
121   PetscFunctionBegin;
122   PetscValidHeaderSpecific(ts,TS_CLASSID,1);
123   if (r) PetscValidHeaderSpecific(r,VEC_CLASSID,2);
124 
125   /* look up the split */
126   while (i<ts->num_rhs_splits) {
127      ierr = PetscStrcmp(ts->tsrhssplit[i]->splitname,splitname,&found);CHKERRQ(ierr);
128      if (found) break;
129      i++;
130   }
131   if (!found) SETERRQ1(PETSC_COMM_SELF,PETSC_ERR_USER,"The split %s is not created, check the split name or call TSRHSSplitSetIS() to create one",splitname);
132 
133   subts = ts->tsrhssplit[i]->ts;
134 
135   ierr = TSGetDM(subts,&dm);CHKERRQ(ierr);
136   ierr = DMTSSetRHSFunction(dm,rhsfunc,ctx);CHKERRQ(ierr);
137 
138   ierr = TSGetSNES(subts,&snes);CHKERRQ(ierr);
139   if (!r && !subts->dm && ts->vec_sol) {
140     ierr = VecGetSubVector(ts->vec_sol,ts->tsrhssplit[i]->is,&subvec);CHKERRQ(ierr);
141     ierr = VecDuplicate(subvec,&ralloc);CHKERRQ(ierr);
142     r = ralloc;
143     ierr = VecRestoreSubVector(ts->vec_sol,ts->tsrhssplit[i]->is,&subvec);CHKERRQ(ierr);
144   }
145   ierr = SNESSetFunction(snes,r,SNESTSFormFunction,subts);CHKERRQ(ierr);
146   ierr = VecDestroy(&ralloc);CHKERRQ(ierr);
147   PetscFunctionReturn(0);
148 }
149 
150 /*@C
151    TSRHSSplitGetSubTS - Get the split right-hand-side functions.
152 
153    Logically Collective on TS
154 
155    Output Parameters:
156 +  n - the number of splits
157 -  subksp - the array of TS contexts
158 
159    Note:
160    After TSRHSSplitGetSubTS() the array of TSs is to be freed by the user with PetscFree()
161    (not the TS just the array that contains them).
162 
163    Level: advanced
164 
165 .seealso: TSGetRHSSplitFunction()
166 @*/
167 PetscErrorCode TSRHSSplitGetSubTS(TS ts,PetscInt *n,TS **subts)
168 {
169   PetscInt       i = 0;
170   PetscErrorCode ierr;
171 
172   PetscFunctionBegin;
173   PetscValidHeaderSpecific(ts,TS_CLASSID,1);
174   if (subts) {
175     ierr = PetscMalloc1(ts->num_rhs_splits,subts);CHKERRQ(ierr);
176     while (i<ts->num_rhs_splits) {
177       (*subts)[i] = ts->tsrhssplit[i]->ts;
178       i++;
179     }
180   }
181   if (n) *n = ts->num_rhs_splits;
182   PetscFunctionReturn(0);
183 }
184 
185 
186