xref: /petsc/src/ts/tutorials/ex41.c (revision d71ae5a4db6382e7f06317b8d368875286fe9008)
1 static char help[] = "Parallel bouncing ball example to test TS event feature.\n";
2 
3 /*
4   The dynamics of the bouncing ball is described by the ODE
5                   u1_t = u2
6                   u2_t = -9.8
7 
8   Each processor is assigned one ball.
9 
10   The event function routine checks for the ball hitting the
11   ground (u1 = 0). Every time the ball hits the ground, its velocity u2 is attenuated by
12   a factor of 0.9 and its height set to 1.0*rank.
13 */
14 
15 #include <petscts.h>
16 
17 PetscErrorCode EventFunction(TS ts, PetscReal t, Vec U, PetscScalar *fvalue, void *ctx)
18 {
19   const PetscScalar *u;
20 
21   PetscFunctionBeginUser;
22   /* Event for ball height */
23   PetscCall(VecGetArrayRead(U, &u));
24   fvalue[0] = u[0];
25   PetscCall(VecRestoreArrayRead(U, &u));
26   PetscFunctionReturn(0);
27 }
28 
29 PetscErrorCode PostEventFunction(TS ts, PetscInt nevents, PetscInt event_list[], PetscReal t, Vec U, PetscBool forwardsolve, void *ctx)
30 {
31   PetscScalar *u;
32   PetscMPIInt  rank;
33 
34   PetscFunctionBeginUser;
35   PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));
36   if (nevents) {
37     PetscCall(PetscPrintf(PETSC_COMM_SELF, "Ball hit the ground at t = %5.2f seconds -> Processor[%d]\n", (double)t, rank));
38     /* Set new initial conditions with .9 attenuation */
39     PetscCall(VecGetArray(U, &u));
40     u[0] = 1.0 * rank;
41     u[1] = -0.9 * u[1];
42     PetscCall(VecRestoreArray(U, &u));
43   }
44   PetscFunctionReturn(0);
45 }
46 
47 /*
48      Defines the ODE passed to the ODE solver in explicit form: U_t = F(U)
49 */
50 static PetscErrorCode RHSFunction(TS ts, PetscReal t, Vec U, Vec F, void *ctx)
51 {
52   PetscScalar       *f;
53   const PetscScalar *u;
54 
55   PetscFunctionBeginUser;
56   /*  The next three lines allow us to access the entries of the vectors directly */
57   PetscCall(VecGetArrayRead(U, &u));
58   PetscCall(VecGetArray(F, &f));
59 
60   f[0] = u[1];
61   f[1] = -9.8;
62 
63   PetscCall(VecRestoreArrayRead(U, &u));
64   PetscCall(VecRestoreArray(F, &f));
65   PetscFunctionReturn(0);
66 }
67 
68 /*
69      Defines the Jacobian of the ODE passed to the ODE solver. See TSSetRHSJacobian() for the meaning the Jacobian.
70 */
71 static PetscErrorCode RHSJacobian(TS ts, PetscReal t, Vec U, Mat A, Mat B, void *ctx)
72 {
73   PetscInt           rowcol[2], rstart;
74   PetscScalar        J[2][2];
75   const PetscScalar *u;
76 
77   PetscFunctionBeginUser;
78   PetscCall(VecGetArrayRead(U, &u));
79 
80   PetscCall(MatGetOwnershipRange(B, &rstart, NULL));
81   rowcol[0] = rstart;
82   rowcol[1] = rstart + 1;
83 
84   J[0][0] = 0.0;
85   J[0][1] = 1.0;
86   J[1][0] = 0.0;
87   J[1][1] = 0.0;
88   PetscCall(MatSetValues(B, 2, rowcol, 2, rowcol, &J[0][0], INSERT_VALUES));
89 
90   PetscCall(VecRestoreArrayRead(U, &u));
91   PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
92   PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
93   if (A != B) {
94     PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
95     PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
96   }
97   PetscFunctionReturn(0);
98 }
99 
100 /*
101      Defines the ODE passed to the ODE solver in implicit form: F(U_t,U) = 0
102 */
103 static PetscErrorCode IFunction(TS ts, PetscReal t, Vec U, Vec Udot, Vec F, void *ctx)
104 {
105   PetscScalar       *f;
106   const PetscScalar *u, *udot;
107 
108   PetscFunctionBeginUser;
109   /*  The next three lines allow us to access the entries of the vectors directly */
110   PetscCall(VecGetArrayRead(U, &u));
111   PetscCall(VecGetArrayRead(Udot, &udot));
112   PetscCall(VecGetArray(F, &f));
113 
114   f[0] = udot[0] - u[1];
115   f[1] = udot[1] + 9.8;
116 
117   PetscCall(VecRestoreArrayRead(U, &u));
118   PetscCall(VecRestoreArrayRead(Udot, &udot));
119   PetscCall(VecRestoreArray(F, &f));
120   PetscFunctionReturn(0);
121 }
122 
123 /*
124      Defines the Jacobian of the ODE passed to the ODE solver. See TSSetIJacobian() for the meaning of a and the Jacobian.
125 */
126 static PetscErrorCode IJacobian(TS ts, PetscReal t, Vec U, Vec Udot, PetscReal a, Mat A, Mat B, void *ctx)
127 {
128   PetscInt           rowcol[2], rstart;
129   PetscScalar        J[2][2];
130   const PetscScalar *u, *udot;
131 
132   PetscFunctionBeginUser;
133   PetscCall(VecGetArrayRead(U, &u));
134   PetscCall(VecGetArrayRead(Udot, &udot));
135 
136   PetscCall(MatGetOwnershipRange(B, &rstart, NULL));
137   rowcol[0] = rstart;
138   rowcol[1] = rstart + 1;
139 
140   J[0][0] = a;
141   J[0][1] = -1.0;
142   J[1][0] = 0.0;
143   J[1][1] = a;
144   PetscCall(MatSetValues(B, 2, rowcol, 2, rowcol, &J[0][0], INSERT_VALUES));
145 
146   PetscCall(VecRestoreArrayRead(U, &u));
147   PetscCall(VecRestoreArrayRead(Udot, &udot));
148 
149   PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
150   PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
151   if (A != B) {
152     PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
153     PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
154   }
155   PetscFunctionReturn(0);
156 }
157 
158 int main(int argc, char **argv)
159 {
160   TS           ts; /* ODE integrator */
161   Vec          U;  /* solution will be stored here */
162   PetscMPIInt  rank;
163   PetscInt     n = 2;
164   PetscScalar *u;
165   PetscInt     direction = -1;
166   PetscBool    terminate = PETSC_FALSE;
167   PetscBool    rhs_form = PETSC_FALSE, hist = PETSC_TRUE;
168   TSAdapt      adapt;
169 
170   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
171      Initialize program
172      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
173   PetscFunctionBeginUser;
174   PetscCall(PetscInitialize(&argc, &argv, (char *)0, help));
175   PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));
176 
177   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
178      Create timestepping solver context
179      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
180   PetscCall(TSCreate(PETSC_COMM_WORLD, &ts));
181   PetscCall(TSSetType(ts, TSROSW));
182 
183   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
184      Set ODE routines
185    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
186   PetscCall(TSSetProblemType(ts, TS_NONLINEAR));
187   /* Users are advised against the following branching and code duplication.
188      For problems without a mass matrix like the one at hand, the RHSFunction
189      (and companion RHSJacobian) interface is enough to support both explicit
190      and implicit timesteppers. This tutorial example also deals with the
191      IFunction/IJacobian interface for demonstration and testing purposes. */
192   PetscCall(PetscOptionsGetBool(NULL, NULL, "-rhs-form", &rhs_form, NULL));
193   if (rhs_form) {
194     PetscCall(TSSetRHSFunction(ts, NULL, RHSFunction, NULL));
195     PetscCall(TSSetRHSJacobian(ts, NULL, NULL, RHSJacobian, NULL));
196   } else {
197     PetscCall(TSSetIFunction(ts, NULL, IFunction, NULL));
198     PetscCall(TSSetIJacobian(ts, NULL, NULL, IJacobian, NULL));
199   }
200 
201   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
202      Set initial conditions
203    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
204   PetscCall(VecCreate(PETSC_COMM_WORLD, &U));
205   PetscCall(VecSetSizes(U, n, PETSC_DETERMINE));
206   PetscCall(VecSetUp(U));
207   PetscCall(VecGetArray(U, &u));
208   u[0] = 1.0 * rank;
209   u[1] = 20.0;
210   PetscCall(VecRestoreArray(U, &u));
211   PetscCall(TSSetSolution(ts, U));
212 
213   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
214      Set solver options
215    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
216   PetscCall(TSSetSaveTrajectory(ts));
217   PetscCall(TSSetMaxTime(ts, 30.0));
218   PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_STEPOVER));
219   PetscCall(TSSetTimeStep(ts, 0.1));
220   /* The adaptive time step controller could take very large timesteps resulting in
221      the same event occurring multiple times in the same interval. A maximum step size
222      limit is enforced here to avoid this issue. */
223   PetscCall(TSGetAdapt(ts, &adapt));
224   PetscCall(TSAdaptSetType(adapt, TSADAPTBASIC));
225   PetscCall(TSAdaptSetStepLimits(adapt, 0.0, 0.5));
226 
227   /* Set direction and terminate flag for the event */
228   PetscCall(TSSetEventHandler(ts, 1, &direction, &terminate, EventFunction, PostEventFunction, NULL));
229 
230   PetscCall(TSSetFromOptions(ts));
231 
232   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
233      Run timestepping solver
234      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
235   PetscCall(TSSolve(ts, U));
236 
237   if (hist) { /* replay following history */
238     TSTrajectory tj;
239     PetscReal    tf, t0, dt;
240 
241     PetscCall(TSGetTime(ts, &tf));
242     PetscCall(TSSetMaxTime(ts, tf));
243     PetscCall(TSSetStepNumber(ts, 0));
244     PetscCall(TSRestartStep(ts));
245     PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP));
246     PetscCall(TSSetFromOptions(ts));
247     PetscCall(TSSetEventHandler(ts, 1, &direction, &terminate, EventFunction, PostEventFunction, NULL));
248     PetscCall(TSGetAdapt(ts, &adapt));
249     PetscCall(TSAdaptSetType(adapt, TSADAPTHISTORY));
250     PetscCall(TSGetTrajectory(ts, &tj));
251     PetscCall(TSAdaptHistorySetTrajectory(adapt, tj, PETSC_FALSE));
252     PetscCall(TSAdaptHistoryGetStep(adapt, 0, &t0, &dt));
253     /* this example fails with single (or smaller) precision */
254 #if defined(PETSC_USE_REAL_SINGLE) || defined(PETSC_USE_REAL___FP16)
255     PetscCall(TSAdaptSetType(adapt, TSADAPTBASIC));
256     PetscCall(TSAdaptSetStepLimits(adapt, 0.0, 0.5));
257     PetscCall(TSSetFromOptions(ts));
258 #endif
259     PetscCall(TSSetTime(ts, t0));
260     PetscCall(TSSetTimeStep(ts, dt));
261     PetscCall(TSResetTrajectory(ts));
262     PetscCall(VecGetArray(U, &u));
263     u[0] = 1.0 * rank;
264     u[1] = 20.0;
265     PetscCall(VecRestoreArray(U, &u));
266     PetscCall(TSSolve(ts, U));
267   }
268   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
269      Free work space.  All PETSc objects should be destroyed when they are no longer needed.
270      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
271   PetscCall(VecDestroy(&U));
272   PetscCall(TSDestroy(&ts));
273 
274   PetscCall(PetscFinalize());
275   return 0;
276 }
277 
278 /*TEST
279 
280    test:
281       suffix: a
282       nsize: 2
283       args: -ts_trajectory_type memory -snes_stol 1e-4
284       filter: sort -b
285 
286    test:
287       suffix: b
288       nsize: 2
289       args: -ts_trajectory_type memory -ts_type arkimex -snes_stol 1e-4
290       filter: sort -b
291 
292    test:
293       suffix: c
294       nsize: 2
295       args: -ts_trajectory_type memory -ts_type theta -ts_adapt_type basic -ts_atol 1e-1 -snes_stol 1e-4
296       filter: sort -b
297 
298    test:
299       suffix: d
300       nsize: 2
301       args: -ts_trajectory_type memory -ts_type alpha -ts_adapt_type basic -ts_atol 1e-1 -snes_stol 1e-4
302       filter: sort -b
303 
304    test:
305       suffix: e
306       nsize: 2
307       args: -ts_trajectory_type memory -ts_type bdf -ts_adapt_dt_max 0.015 -ts_max_steps 3000
308       filter: sort -b
309 
310    test:
311       suffix: f
312       nsize: 2
313       args: -ts_trajectory_type memory -rhs-form -ts_type rk -ts_rk_type 3bs
314       filter: sort -b
315 
316    test:
317       suffix: g
318       nsize: 2
319       args: -ts_trajectory_type memory -rhs-form -ts_type rk -ts_rk_type 5bs
320       filter: sort -b
321 
322    test:
323       suffix: h
324       nsize: 2
325       args: -ts_trajectory_type memory -rhs-form -ts_type rk -ts_rk_type 6vr
326       filter: sort -b
327       output_file: output/ex41_g.out
328 
329    test:
330       suffix: i
331       nsize: 2
332       args: -ts_trajectory_type memory -rhs-form -ts_type rk -ts_rk_type 7vr
333       filter: sort -b
334       output_file: output/ex41_g.out
335 
336    test:
337       suffix: j
338       nsize: 2
339       args: -ts_trajectory_type memory -rhs-form -ts_type rk -ts_rk_type 8vr
340       filter: sort -b
341       output_file: output/ex41_g.out
342 
343 TEST*/
344