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