xref: /petsc/src/ts/tests/ex10.c (revision 58d68138c660dfb4e9f5b03334792cd4f2ffd7cc)
1 static char help[] = "Simple wrapper object to solve DAE of the form:\n\
2                              \\dot{U} = f(U,V)\n\
3                              F(U,V) = 0\n\n";
4 
5 #include <petscts.h>
6 
7 /* ----------------------------------------------------------------------------*/
8 
9 typedef struct _p_TSDAESimple *TSDAESimple;
10 struct _p_TSDAESimple {
11   MPI_Comm comm;
12   PetscErrorCode (*setfromoptions)(TSDAESimple, PetscOptionItems *);
13   PetscErrorCode (*solve)(TSDAESimple, Vec);
14   PetscErrorCode (*destroy)(TSDAESimple);
15   Vec U, V;
16   PetscErrorCode (*f)(PetscReal, Vec, Vec, Vec, void *);
17   PetscErrorCode (*F)(PetscReal, Vec, Vec, Vec, void *);
18   void *fctx, *Fctx;
19   void *data;
20 };
21 
22 PetscErrorCode TSDAESimpleCreate(MPI_Comm comm, TSDAESimple *tsdae) {
23   PetscFunctionBeginUser;
24   PetscCall(PetscNew(tsdae));
25   (*tsdae)->comm = comm;
26   PetscFunctionReturn(0);
27 }
28 
29 PetscErrorCode TSDAESimpleSetRHSFunction(TSDAESimple tsdae, Vec U, PetscErrorCode (*f)(PetscReal, Vec, Vec, Vec, void *), void *ctx) {
30   PetscFunctionBeginUser;
31   tsdae->f = f;
32   tsdae->U = U;
33   PetscCall(PetscObjectReference((PetscObject)U));
34   tsdae->fctx = ctx;
35   PetscFunctionReturn(0);
36 }
37 
38 PetscErrorCode TSDAESimpleSetIFunction(TSDAESimple tsdae, Vec V, PetscErrorCode (*F)(PetscReal, Vec, Vec, Vec, void *), void *ctx) {
39   PetscFunctionBeginUser;
40   tsdae->F = F;
41   tsdae->V = V;
42   PetscCall(PetscObjectReference((PetscObject)V));
43   tsdae->Fctx = ctx;
44   PetscFunctionReturn(0);
45 }
46 
47 PetscErrorCode TSDAESimpleDestroy(TSDAESimple *tsdae) {
48   PetscFunctionBeginUser;
49   PetscCall((*(*tsdae)->destroy)(*tsdae));
50   PetscCall(VecDestroy(&(*tsdae)->U));
51   PetscCall(VecDestroy(&(*tsdae)->V));
52   PetscCall(PetscFree(*tsdae));
53   PetscFunctionReturn(0);
54 }
55 
56 PetscErrorCode TSDAESimpleSolve(TSDAESimple tsdae, Vec Usolution) {
57   PetscFunctionBeginUser;
58   PetscCall((*tsdae->solve)(tsdae, Usolution));
59   PetscFunctionReturn(0);
60 }
61 
62 PetscErrorCode TSDAESimpleSetFromOptions(TSDAESimple tsdae) {
63   PetscFunctionBeginUser;
64   PetscCall((*tsdae->setfromoptions)(PetscOptionsObject, tsdae));
65   PetscFunctionReturn(0);
66 }
67 
68 /* ----------------------------------------------------------------------------*/
69 /*
70       Integrates the system by integrating the reduced ODE system and solving the
71    algebraic constraints at each stage with a separate SNES solve.
72 */
73 
74 typedef struct {
75   PetscReal t;
76   TS        ts;
77   SNES      snes;
78   Vec       U;
79 } TSDAESimple_Reduced;
80 
81 /*
82    Defines the RHS function that is passed to the time-integrator.
83 
84    Solves F(U,V) for V and then computes f(U,V)
85 
86 */
87 PetscErrorCode TSDAESimple_Reduced_TSFunction(TS ts, PetscReal t, Vec U, Vec F, void *actx) {
88   TSDAESimple          tsdae = (TSDAESimple)actx;
89   TSDAESimple_Reduced *red   = (TSDAESimple_Reduced *)tsdae->data;
90 
91   PetscFunctionBeginUser;
92   red->t = t;
93   red->U = U;
94   PetscCall(SNESSolve(red->snes, NULL, tsdae->V));
95   PetscCall((*tsdae->f)(t, U, tsdae->V, F, tsdae->fctx));
96   PetscFunctionReturn(0);
97 }
98 
99 /*
100    Defines the nonlinear function that is passed to the nonlinear solver
101 
102 */
103 PetscErrorCode TSDAESimple_Reduced_SNESFunction(SNES snes, Vec V, Vec F, void *actx) {
104   TSDAESimple          tsdae = (TSDAESimple)actx;
105   TSDAESimple_Reduced *red   = (TSDAESimple_Reduced *)tsdae->data;
106 
107   PetscFunctionBeginUser;
108   PetscCall((*tsdae->F)(red->t, red->U, V, F, tsdae->Fctx));
109   PetscFunctionReturn(0);
110 }
111 
112 PetscErrorCode TSDAESimpleSolve_Reduced(TSDAESimple tsdae, Vec U) {
113   TSDAESimple_Reduced *red = (TSDAESimple_Reduced *)tsdae->data;
114 
115   PetscFunctionBeginUser;
116   PetscCall(TSSolve(red->ts, U));
117   PetscFunctionReturn(0);
118 }
119 
120 PetscErrorCode TSDAESimpleSetFromOptions_Reduced(TSDAESimple tsdae, PetscOptionItems *PetscOptionsObject) {
121   TSDAESimple_Reduced *red = (TSDAESimple_Reduced *)tsdae->data;
122 
123   PetscFunctionBeginUser;
124   PetscCall(TSSetFromOptions(red->ts));
125   PetscCall(SNESSetFromOptions(red->snes));
126   PetscFunctionReturn(0);
127 }
128 
129 PetscErrorCode TSDAESimpleDestroy_Reduced(TSDAESimple tsdae) {
130   TSDAESimple_Reduced *red = (TSDAESimple_Reduced *)tsdae->data;
131 
132   PetscFunctionBeginUser;
133   PetscCall(TSDestroy(&red->ts));
134   PetscCall(SNESDestroy(&red->snes));
135   PetscCall(PetscFree(red));
136   PetscFunctionReturn(0);
137 }
138 
139 PetscErrorCode TSDAESimpleSetUp_Reduced(TSDAESimple tsdae) {
140   TSDAESimple_Reduced *red;
141   Vec                  tsrhs;
142 
143   PetscFunctionBeginUser;
144   PetscCall(PetscNew(&red));
145   tsdae->data = red;
146 
147   tsdae->setfromoptions = TSDAESimpleSetFromOptions_Reduced;
148   tsdae->solve          = TSDAESimpleSolve_Reduced;
149   tsdae->destroy        = TSDAESimpleDestroy_Reduced;
150 
151   PetscCall(TSCreate(tsdae->comm, &red->ts));
152   PetscCall(TSSetProblemType(red->ts, TS_NONLINEAR));
153   PetscCall(TSSetType(red->ts, TSEULER));
154   PetscCall(TSSetExactFinalTime(red->ts, TS_EXACTFINALTIME_STEPOVER));
155   PetscCall(VecDuplicate(tsdae->U, &tsrhs));
156   PetscCall(TSSetRHSFunction(red->ts, tsrhs, TSDAESimple_Reduced_TSFunction, tsdae));
157   PetscCall(VecDestroy(&tsrhs));
158 
159   PetscCall(SNESCreate(tsdae->comm, &red->snes));
160   PetscCall(SNESSetOptionsPrefix(red->snes, "tsdaesimple_"));
161   PetscCall(SNESSetFunction(red->snes, NULL, TSDAESimple_Reduced_SNESFunction, tsdae));
162   PetscCall(SNESSetJacobian(red->snes, NULL, NULL, SNESComputeJacobianDefault, tsdae));
163   PetscFunctionReturn(0);
164 }
165 
166 /* ----------------------------------------------------------------------------*/
167 
168 /*
169       Integrates the system by integrating directly the entire DAE system
170 */
171 
172 typedef struct {
173   TS         ts;
174   Vec        UV, UF, VF;
175   VecScatter scatterU, scatterV;
176 } TSDAESimple_Full;
177 
178 /*
179    Defines the RHS function that is passed to the time-integrator.
180 
181    f(U,V)
182    0
183 
184 */
185 PetscErrorCode TSDAESimple_Full_TSRHSFunction(TS ts, PetscReal t, Vec UV, Vec F, void *actx) {
186   TSDAESimple       tsdae = (TSDAESimple)actx;
187   TSDAESimple_Full *full  = (TSDAESimple_Full *)tsdae->data;
188 
189   PetscFunctionBeginUser;
190   PetscCall(VecSet(F, 0.0));
191   PetscCall(VecScatterBegin(full->scatterU, UV, tsdae->U, INSERT_VALUES, SCATTER_REVERSE));
192   PetscCall(VecScatterEnd(full->scatterU, UV, tsdae->U, INSERT_VALUES, SCATTER_REVERSE));
193   PetscCall(VecScatterBegin(full->scatterV, UV, tsdae->V, INSERT_VALUES, SCATTER_REVERSE));
194   PetscCall(VecScatterEnd(full->scatterV, UV, tsdae->V, INSERT_VALUES, SCATTER_REVERSE));
195   PetscCall((*tsdae->f)(t, tsdae->U, tsdae->V, full->UF, tsdae->fctx));
196   PetscCall(VecScatterBegin(full->scatterU, full->UF, F, INSERT_VALUES, SCATTER_FORWARD));
197   PetscCall(VecScatterEnd(full->scatterU, full->UF, F, INSERT_VALUES, SCATTER_FORWARD));
198   PetscFunctionReturn(0);
199 }
200 
201 /*
202    Defines the nonlinear function that is passed to the nonlinear solver
203 
204    \dot{U}
205    F(U,V)
206 
207 */
208 PetscErrorCode TSDAESimple_Full_TSIFunction(TS ts, PetscReal t, Vec UV, Vec UVdot, Vec F, void *actx) {
209   TSDAESimple       tsdae = (TSDAESimple)actx;
210   TSDAESimple_Full *full  = (TSDAESimple_Full *)tsdae->data;
211 
212   PetscFunctionBeginUser;
213   PetscCall(VecCopy(UVdot, F));
214   PetscCall(VecScatterBegin(full->scatterU, UV, tsdae->U, INSERT_VALUES, SCATTER_REVERSE));
215   PetscCall(VecScatterEnd(full->scatterU, UV, tsdae->U, INSERT_VALUES, SCATTER_REVERSE));
216   PetscCall(VecScatterBegin(full->scatterV, UV, tsdae->V, INSERT_VALUES, SCATTER_REVERSE));
217   PetscCall(VecScatterEnd(full->scatterV, UV, tsdae->V, INSERT_VALUES, SCATTER_REVERSE));
218   PetscCall((*tsdae->F)(t, tsdae->U, tsdae->V, full->VF, tsdae->Fctx));
219   PetscCall(VecScatterBegin(full->scatterV, full->VF, F, INSERT_VALUES, SCATTER_FORWARD));
220   PetscCall(VecScatterEnd(full->scatterV, full->VF, F, INSERT_VALUES, SCATTER_FORWARD));
221   PetscFunctionReturn(0);
222 }
223 
224 PetscErrorCode TSDAESimpleSolve_Full(TSDAESimple tsdae, Vec U) {
225   TSDAESimple_Full *full = (TSDAESimple_Full *)tsdae->data;
226 
227   PetscFunctionBeginUser;
228   PetscCall(VecSet(full->UV, 1.0));
229   PetscCall(VecScatterBegin(full->scatterU, U, full->UV, INSERT_VALUES, SCATTER_FORWARD));
230   PetscCall(VecScatterEnd(full->scatterU, U, full->UV, INSERT_VALUES, SCATTER_FORWARD));
231   PetscCall(TSSolve(full->ts, full->UV));
232   PetscCall(VecScatterBegin(full->scatterU, full->UV, U, INSERT_VALUES, SCATTER_REVERSE));
233   PetscCall(VecScatterEnd(full->scatterU, full->UV, U, INSERT_VALUES, SCATTER_REVERSE));
234   PetscFunctionReturn(0);
235 }
236 
237 PetscErrorCode TSDAESimpleSetFromOptions_Full(TSDAESimple tsdae, PetscOptionItems *PetscOptionsObject) {
238   TSDAESimple_Full *full = (TSDAESimple_Full *)tsdae->data;
239 
240   PetscFunctionBeginUser;
241   PetscCall(TSSetFromOptions(full->ts));
242   PetscFunctionReturn(0);
243 }
244 
245 PetscErrorCode TSDAESimpleDestroy_Full(TSDAESimple tsdae) {
246   TSDAESimple_Full *full = (TSDAESimple_Full *)tsdae->data;
247 
248   PetscFunctionBeginUser;
249   PetscCall(TSDestroy(&full->ts));
250   PetscCall(VecDestroy(&full->UV));
251   PetscCall(VecDestroy(&full->UF));
252   PetscCall(VecDestroy(&full->VF));
253   PetscCall(VecScatterDestroy(&full->scatterU));
254   PetscCall(VecScatterDestroy(&full->scatterV));
255   PetscCall(PetscFree(full));
256   PetscFunctionReturn(0);
257 }
258 
259 PetscErrorCode TSDAESimpleSetUp_Full(TSDAESimple tsdae) {
260   TSDAESimple_Full *full;
261   Vec               tsrhs;
262   PetscInt          nU, nV, UVstart;
263   IS                is;
264 
265   PetscFunctionBeginUser;
266   PetscCall(PetscNew(&full));
267   tsdae->data = full;
268 
269   tsdae->setfromoptions = TSDAESimpleSetFromOptions_Full;
270   tsdae->solve          = TSDAESimpleSolve_Full;
271   tsdae->destroy        = TSDAESimpleDestroy_Full;
272 
273   PetscCall(TSCreate(tsdae->comm, &full->ts));
274   PetscCall(TSSetProblemType(full->ts, TS_NONLINEAR));
275   PetscCall(TSSetType(full->ts, TSROSW));
276   PetscCall(TSSetExactFinalTime(full->ts, TS_EXACTFINALTIME_STEPOVER));
277   PetscCall(VecDuplicate(tsdae->U, &full->UF));
278   PetscCall(VecDuplicate(tsdae->V, &full->VF));
279 
280   PetscCall(VecGetLocalSize(tsdae->U, &nU));
281   PetscCall(VecGetLocalSize(tsdae->V, &nV));
282   PetscCall(VecCreateMPI(tsdae->comm, nU + nV, PETSC_DETERMINE, &tsrhs));
283   PetscCall(VecDuplicate(tsrhs, &full->UV));
284 
285   PetscCall(VecGetOwnershipRange(tsrhs, &UVstart, NULL));
286   PetscCall(ISCreateStride(tsdae->comm, nU, UVstart, 1, &is));
287   PetscCall(VecScatterCreate(tsdae->U, NULL, tsrhs, is, &full->scatterU));
288   PetscCall(ISDestroy(&is));
289   PetscCall(ISCreateStride(tsdae->comm, nV, UVstart + nU, 1, &is));
290   PetscCall(VecScatterCreate(tsdae->V, NULL, tsrhs, is, &full->scatterV));
291   PetscCall(ISDestroy(&is));
292 
293   PetscCall(TSSetRHSFunction(full->ts, tsrhs, TSDAESimple_Full_TSRHSFunction, tsdae));
294   PetscCall(TSSetIFunction(full->ts, NULL, TSDAESimple_Full_TSIFunction, tsdae));
295   PetscCall(VecDestroy(&tsrhs));
296   PetscFunctionReturn(0);
297 }
298 
299 /* ----------------------------------------------------------------------------*/
300 
301 /*
302    Simple example:   f(U,V) = U + V
303 
304 */
305 PetscErrorCode f(PetscReal t, Vec U, Vec V, Vec F, void *ctx) {
306   PetscFunctionBeginUser;
307   PetscCall(VecWAXPY(F, 1.0, U, V));
308   PetscFunctionReturn(0);
309 }
310 
311 /*
312    Simple example: F(U,V) = U - V
313 
314 */
315 PetscErrorCode F(PetscReal t, Vec U, Vec V, Vec F, void *ctx) {
316   PetscFunctionBeginUser;
317   PetscCall(VecWAXPY(F, -1.0, V, U));
318   PetscFunctionReturn(0);
319 }
320 
321 int main(int argc, char **argv) {
322   TSDAESimple tsdae;
323   Vec         U, V, Usolution;
324 
325   PetscFunctionBeginUser;
326   PetscCall(PetscInitialize(&argc, &argv, (char *)0, help));
327   PetscCall(TSDAESimpleCreate(PETSC_COMM_WORLD, &tsdae));
328 
329   PetscCall(VecCreateMPI(PETSC_COMM_WORLD, 1, PETSC_DETERMINE, &U));
330   PetscCall(VecCreateMPI(PETSC_COMM_WORLD, 1, PETSC_DETERMINE, &V));
331   PetscCall(TSDAESimpleSetRHSFunction(tsdae, U, f, NULL));
332   PetscCall(TSDAESimpleSetIFunction(tsdae, V, F, NULL));
333 
334   PetscCall(VecDuplicate(U, &Usolution));
335   PetscCall(VecSet(Usolution, 1.0));
336 
337   /*  PetscCall(TSDAESimpleSetUp_Full(tsdae)); */
338   PetscCall(TSDAESimpleSetUp_Reduced(tsdae));
339 
340   PetscCall(TSDAESimpleSetFromOptions(tsdae));
341   PetscCall(TSDAESimpleSolve(tsdae, Usolution));
342   PetscCall(TSDAESimpleDestroy(&tsdae));
343 
344   PetscCall(VecDestroy(&U));
345   PetscCall(VecDestroy(&Usolution));
346   PetscCall(VecDestroy(&V));
347   PetscCall(PetscFinalize());
348   return 0;
349 }
350