xref: /petsc/src/ts/tests/ex7.c (revision fbf9dbe564678ed6eff1806adbc4c4f01b9743f4)
1 static char help[] = "Solves DAE with integrator only on non-algebraic terms \n";
2 
3 #include <petscts.h>
4 
5 /*
6         \dot{U} = f(U,V)
7         F(U,V)  = 0
8 
9     Same as ex6.c except the user provided functions take input values as a single vector instead of two vectors
10 */
11 
12 /*
13    f(U,V) = U + V
14 
15 */
16 PetscErrorCode f(PetscReal t, Vec UV, Vec F)
17 {
18   const PetscScalar *u, *v;
19   PetscScalar       *f;
20   PetscInt           n, i;
21 
22   PetscFunctionBeginUser;
23   PetscCall(VecGetLocalSize(UV, &n));
24   n = n / 2;
25   PetscCall(VecGetArrayRead(UV, &u));
26   v = u + n;
27   PetscCall(VecGetArrayWrite(F, &f));
28   for (i = 0; i < n; i++) f[i] = u[i] + v[i];
29   PetscCall(VecRestoreArrayRead(UV, &u));
30   PetscCall(VecRestoreArrayWrite(F, &f));
31   PetscFunctionReturn(PETSC_SUCCESS);
32 }
33 
34 /*
35    F(U,V) = U - V
36 */
37 PetscErrorCode F(PetscReal t, Vec UV, Vec F)
38 {
39   const PetscScalar *u, *v;
40   PetscScalar       *f;
41   PetscInt           n, i;
42 
43   PetscFunctionBeginUser;
44   PetscCall(VecGetLocalSize(UV, &n));
45   n = n / 2;
46   PetscCall(VecGetArrayRead(UV, &u));
47   v = u + n;
48   PetscCall(VecGetArrayWrite(F, &f));
49   for (i = 0; i < n; i++) f[i] = u[i] - v[i];
50   PetscCall(VecRestoreArrayRead(UV, &u));
51   PetscCall(VecRestoreArrayWrite(F, &f));
52   PetscFunctionReturn(PETSC_SUCCESS);
53 }
54 
55 typedef struct {
56   PetscReal  t;
57   SNES       snes;
58   Vec        UV, V;
59   VecScatter scatterU, scatterV;
60   PetscErrorCode (*f)(PetscReal, Vec, Vec);
61   PetscErrorCode (*F)(PetscReal, Vec, Vec);
62 } AppCtx;
63 
64 extern PetscErrorCode TSFunction(TS, PetscReal, Vec, Vec, void *);
65 extern PetscErrorCode SNESFunction(SNES, Vec, Vec, void *);
66 
67 int main(int argc, char **argv)
68 {
69   AppCtx      ctx;
70   TS          ts;
71   Vec         tsrhs, U;
72   IS          is;
73   PetscInt    i;
74   PetscMPIInt rank;
75 
76   PetscFunctionBeginUser;
77   PetscCall(PetscInitialize(&argc, &argv, (char *)0, help));
78   PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));
79   PetscCall(TSCreate(PETSC_COMM_WORLD, &ts));
80   PetscCall(TSSetProblemType(ts, TS_NONLINEAR));
81   PetscCall(TSSetType(ts, TSEULER));
82   PetscCall(TSSetFromOptions(ts));
83   PetscCall(VecCreateMPI(PETSC_COMM_WORLD, 1, PETSC_DETERMINE, &tsrhs));
84   PetscCall(VecDuplicate(tsrhs, &U));
85   PetscCall(TSSetRHSFunction(ts, tsrhs, TSFunction, &ctx));
86   PetscCall(TSSetMaxTime(ts, 1.0));
87   ctx.f = f;
88 
89   PetscCall(SNESCreate(PETSC_COMM_WORLD, &ctx.snes));
90   PetscCall(SNESSetFromOptions(ctx.snes));
91   PetscCall(SNESSetFunction(ctx.snes, NULL, SNESFunction, &ctx));
92   PetscCall(SNESSetJacobian(ctx.snes, NULL, NULL, SNESComputeJacobianDefault, &ctx));
93   ctx.F = F;
94   PetscCall(VecCreateMPI(PETSC_COMM_WORLD, 1, PETSC_DETERMINE, &ctx.V));
95 
96   /* Create scatters to move between separate U and V representation and UV representation of solution */
97   PetscCall(VecCreateMPI(PETSC_COMM_WORLD, 2, PETSC_DETERMINE, &ctx.UV));
98   i = 2 * rank;
99   PetscCall(ISCreateGeneral(PETSC_COMM_WORLD, 1, &i, PETSC_COPY_VALUES, &is));
100   PetscCall(VecScatterCreate(U, NULL, ctx.UV, is, &ctx.scatterU));
101   PetscCall(ISDestroy(&is));
102   i = 2 * rank + 1;
103   PetscCall(ISCreateGeneral(PETSC_COMM_WORLD, 1, &i, PETSC_COPY_VALUES, &is));
104   PetscCall(VecScatterCreate(ctx.V, NULL, ctx.UV, is, &ctx.scatterV));
105   PetscCall(ISDestroy(&is));
106 
107   PetscCall(VecSet(U, 1.0));
108   PetscCall(TSSolve(ts, U));
109 
110   PetscCall(VecDestroy(&ctx.V));
111   PetscCall(VecDestroy(&ctx.UV));
112   PetscCall(VecScatterDestroy(&ctx.scatterU));
113   PetscCall(VecScatterDestroy(&ctx.scatterV));
114   PetscCall(VecDestroy(&tsrhs));
115   PetscCall(VecDestroy(&U));
116   PetscCall(SNESDestroy(&ctx.snes));
117   PetscCall(TSDestroy(&ts));
118   PetscCall(PetscFinalize());
119   return 0;
120 }
121 
122 /*
123    Defines the RHS function that is passed to the time-integrator.
124 
125    Solves F(U,V) for V and then computes f(U,V)
126 */
127 PetscErrorCode TSFunction(TS ts, PetscReal t, Vec U, Vec F, void *actx)
128 {
129   AppCtx *ctx = (AppCtx *)actx;
130 
131   PetscFunctionBeginUser;
132   ctx->t = t;
133   PetscCall(VecScatterBegin(ctx->scatterU, U, ctx->UV, INSERT_VALUES, SCATTER_FORWARD));
134   PetscCall(VecScatterEnd(ctx->scatterU, U, ctx->UV, INSERT_VALUES, SCATTER_FORWARD));
135   PetscCall(SNESSolve(ctx->snes, NULL, ctx->V));
136   PetscCall(VecScatterBegin(ctx->scatterV, ctx->V, ctx->UV, INSERT_VALUES, SCATTER_FORWARD));
137   PetscCall(VecScatterEnd(ctx->scatterV, ctx->V, ctx->UV, INSERT_VALUES, SCATTER_FORWARD));
138   PetscCall((*ctx->f)(t, ctx->UV, F));
139   PetscFunctionReturn(PETSC_SUCCESS);
140 }
141 
142 /*
143    Defines the nonlinear function that is passed to the nonlinear solver
144 
145 */
146 PetscErrorCode SNESFunction(SNES snes, Vec V, Vec F, void *actx)
147 {
148   AppCtx *ctx = (AppCtx *)actx;
149 
150   PetscFunctionBeginUser;
151   PetscCall(VecScatterBegin(ctx->scatterV, V, ctx->UV, INSERT_VALUES, SCATTER_FORWARD));
152   PetscCall(VecScatterEnd(ctx->scatterV, V, ctx->UV, INSERT_VALUES, SCATTER_FORWARD));
153   PetscCall((*ctx->F)(ctx->t, ctx->UV, F));
154   PetscFunctionReturn(PETSC_SUCCESS);
155 }
156 
157 /*TEST
158 
159     test:
160       args: -ts_monitor -ts_view
161 
162 TEST*/
163