xref: /petsc/src/ts/tutorials/ex41.c (revision 4e278199b78715991f5c71ebbd945c1489263e6c)
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   PetscErrorCode    ierr;
20   const PetscScalar *u;
21 
22   PetscFunctionBegin;
23   /* Event for ball height */
24   ierr = VecGetArrayRead(U,&u);CHKERRQ(ierr);
25   fvalue[0] = u[0];
26   ierr = VecRestoreArrayRead(U,&u);CHKERRQ(ierr);
27   PetscFunctionReturn(0);
28 }
29 
30 PetscErrorCode PostEventFunction(TS ts,PetscInt nevents,PetscInt event_list[],PetscReal t,Vec U,PetscBool forwardsolve,void* ctx)
31 {
32   PetscErrorCode ierr;
33   PetscScalar    *u;
34   PetscMPIInt    rank;
35 
36   PetscFunctionBegin;
37   ierr = MPI_Comm_rank(PETSC_COMM_WORLD,&rank);CHKERRMPI(ierr);
38   if (nevents) {
39     ierr = PetscPrintf(PETSC_COMM_SELF,"Ball hit the ground at t = %5.2f seconds -> Processor[%d]\n",(double)t,rank);CHKERRQ(ierr);
40     /* Set new initial conditions with .9 attenuation */
41     ierr = VecGetArray(U,&u);CHKERRQ(ierr);
42     u[0] =  1.0*rank;
43     u[1] = -0.9*u[1];
44     ierr = VecRestoreArray(U,&u);CHKERRQ(ierr);
45   }
46   PetscFunctionReturn(0);
47 }
48 
49 /*
50      Defines the ODE passed to the ODE solver in explicit form: U_t = F(U)
51 */
52 static PetscErrorCode RHSFunction(TS ts,PetscReal t,Vec U,Vec F,void *ctx)
53 {
54   PetscErrorCode    ierr;
55   PetscScalar       *f;
56   const PetscScalar *u;
57 
58   PetscFunctionBegin;
59   /*  The next three lines allow us to access the entries of the vectors directly */
60   ierr = VecGetArrayRead(U,&u);CHKERRQ(ierr);
61   ierr = VecGetArray(F,&f);CHKERRQ(ierr);
62 
63   f[0] = u[1];
64   f[1] = - 9.8;
65 
66   ierr = VecRestoreArrayRead(U,&u);CHKERRQ(ierr);
67   ierr = VecRestoreArray(F,&f);CHKERRQ(ierr);
68   PetscFunctionReturn(0);
69 }
70 
71 /*
72      Defines the Jacobian of the ODE passed to the ODE solver. See TSSetRHSJacobian() for the meaning the Jacobian.
73 */
74 static PetscErrorCode RHSJacobian(TS ts,PetscReal t,Vec U,Mat A,Mat B,void *ctx)
75 {
76   PetscErrorCode    ierr;
77   PetscInt          rowcol[2],rstart;
78   PetscScalar       J[2][2];
79   const PetscScalar *u;
80 
81   PetscFunctionBegin;
82   ierr = VecGetArrayRead(U,&u);CHKERRQ(ierr);
83 
84   ierr = MatGetOwnershipRange(B,&rstart,NULL);CHKERRQ(ierr);
85   rowcol[0] = rstart; rowcol[1] = rstart+1;
86 
87   J[0][0] = 0.0;      J[0][1] = 1.0;
88   J[1][0] = 0.0;      J[1][1] = 0.0;
89   ierr = MatSetValues(B,2,rowcol,2,rowcol,&J[0][0],INSERT_VALUES);CHKERRQ(ierr);
90 
91   ierr = VecRestoreArrayRead(U,&u);CHKERRQ(ierr);
92   ierr = MatAssemblyBegin(B,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
93   ierr = MatAssemblyEnd(B,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
94   if (A != B) {
95     ierr = MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
96     ierr = MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
97   }
98   PetscFunctionReturn(0);
99 }
100 
101 /*
102      Defines the ODE passed to the ODE solver in implicit form: F(U_t,U) = 0
103 */
104 static PetscErrorCode IFunction(TS ts,PetscReal t,Vec U,Vec Udot,Vec F,void *ctx)
105 {
106   PetscErrorCode    ierr;
107   PetscScalar       *f;
108   const PetscScalar *u,*udot;
109 
110   PetscFunctionBegin;
111   /*  The next three lines allow us to access the entries of the vectors directly */
112   ierr = VecGetArrayRead(U,&u);CHKERRQ(ierr);
113   ierr = VecGetArrayRead(Udot,&udot);CHKERRQ(ierr);
114   ierr = VecGetArray(F,&f);CHKERRQ(ierr);
115 
116   f[0] = udot[0] - u[1];
117   f[1] = udot[1] + 9.8;
118 
119   ierr = VecRestoreArrayRead(U,&u);CHKERRQ(ierr);
120   ierr = VecRestoreArrayRead(Udot,&udot);CHKERRQ(ierr);
121   ierr = VecRestoreArray(F,&f);CHKERRQ(ierr);
122   PetscFunctionReturn(0);
123 }
124 
125 /*
126      Defines the Jacobian of the ODE passed to the ODE solver. See TSSetIJacobian() for the meaning of a and the Jacobian.
127 */
128 static PetscErrorCode IJacobian(TS ts,PetscReal t,Vec U,Vec Udot,PetscReal a,Mat A,Mat B,void *ctx)
129 {
130   PetscErrorCode    ierr;
131   PetscInt          rowcol[2],rstart;
132   PetscScalar       J[2][2];
133   const PetscScalar *u,*udot;
134 
135   PetscFunctionBegin;
136   ierr = VecGetArrayRead(U,&u);CHKERRQ(ierr);
137   ierr = VecGetArrayRead(Udot,&udot);CHKERRQ(ierr);
138 
139   ierr = MatGetOwnershipRange(B,&rstart,NULL);CHKERRQ(ierr);
140   rowcol[0] = rstart; rowcol[1] = rstart+1;
141 
142   J[0][0] = a;        J[0][1] = -1.0;
143   J[1][0] = 0.0;      J[1][1] = a;
144   ierr = MatSetValues(B,2,rowcol,2,rowcol,&J[0][0],INSERT_VALUES);CHKERRQ(ierr);
145 
146   ierr = VecRestoreArrayRead(U,&u);CHKERRQ(ierr);
147   ierr = VecRestoreArrayRead(Udot,&udot);CHKERRQ(ierr);
148 
149   ierr = MatAssemblyBegin(B,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
150   ierr = MatAssemblyEnd(B,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
151   if (A != B) {
152     ierr = MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
153     ierr = MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
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   PetscErrorCode ierr;
163   PetscMPIInt    rank;
164   PetscInt       n = 2;
165   PetscScalar    *u;
166   PetscInt       direction=-1;
167   PetscBool      terminate=PETSC_FALSE;
168   PetscBool      rhs_form=PETSC_FALSE,hist=PETSC_TRUE;
169   TSAdapt        adapt;
170 
171   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
172      Initialize program
173      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
174   ierr = PetscInitialize(&argc,&argv,(char*)0,help);if (ierr) return ierr;
175   ierr = MPI_Comm_rank(PETSC_COMM_WORLD,&rank);CHKERRMPI(ierr);
176 
177   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
178      Create timestepping solver context
179      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
180   ierr = TSCreate(PETSC_COMM_WORLD,&ts);CHKERRQ(ierr);
181   ierr = TSSetType(ts,TSROSW);CHKERRQ(ierr);
182 
183   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
184      Set ODE routines
185    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
186   ierr = TSSetProblemType(ts,TS_NONLINEAR);CHKERRQ(ierr);
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   ierr = PetscOptionsGetBool(NULL,NULL,"-rhs-form",&rhs_form,NULL);CHKERRQ(ierr);
193   if (rhs_form) {
194     ierr = TSSetRHSFunction(ts,NULL,RHSFunction,NULL);CHKERRQ(ierr);
195     ierr = TSSetRHSJacobian(ts,NULL,NULL,RHSJacobian,NULL);CHKERRQ(ierr);
196   } else {
197     ierr = TSSetIFunction(ts,NULL,IFunction,NULL);CHKERRQ(ierr);
198     ierr = TSSetIJacobian(ts,NULL,NULL,IJacobian,NULL);CHKERRQ(ierr);
199   }
200 
201   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
202      Set initial conditions
203    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
204   ierr = VecCreate(PETSC_COMM_WORLD,&U);CHKERRQ(ierr);
205   ierr = VecSetSizes(U,n,PETSC_DETERMINE);CHKERRQ(ierr);
206   ierr = VecSetUp(U);CHKERRQ(ierr);
207   ierr = VecGetArray(U,&u);CHKERRQ(ierr);
208   u[0] = 1.0*rank;
209   u[1] = 20.0;
210   ierr = VecRestoreArray(U,&u);CHKERRQ(ierr);
211   ierr = TSSetSolution(ts,U);CHKERRQ(ierr);
212 
213   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
214      Set solver options
215    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
216   ierr = TSSetSaveTrajectory(ts);CHKERRQ(ierr);
217   ierr = TSSetMaxTime(ts,30.0);CHKERRQ(ierr);
218   ierr = TSSetExactFinalTime(ts,TS_EXACTFINALTIME_STEPOVER);CHKERRQ(ierr);
219   ierr = TSSetTimeStep(ts,0.1);CHKERRQ(ierr);
220   /* The adapative time step controller could take very large timesteps resulting in
221      the same event occuring multiple times in the same interval. A maximum step size
222      limit is enforced here to avoid this issue. */
223   ierr = TSGetAdapt(ts,&adapt);CHKERRQ(ierr);
224   ierr = TSAdaptSetType(adapt,TSADAPTBASIC);CHKERRQ(ierr);
225   ierr = TSAdaptSetStepLimits(adapt,0.0,0.5);CHKERRQ(ierr);
226 
227   /* Set direction and terminate flag for the event */
228   ierr = TSSetEventHandler(ts,1,&direction,&terminate,EventFunction,PostEventFunction,NULL);CHKERRQ(ierr);
229 
230   ierr = TSSetFromOptions(ts);CHKERRQ(ierr);
231 
232   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
233      Run timestepping solver
234      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
235   ierr = TSSolve(ts,U);CHKERRQ(ierr);
236 
237   if (hist) { /* replay following history */
238     TSTrajectory tj;
239     PetscReal    tf,t0,dt;
240 
241     ierr = TSGetTime(ts,&tf);CHKERRQ(ierr);
242     ierr = TSSetMaxTime(ts,tf);CHKERRQ(ierr);
243     ierr = TSSetStepNumber(ts,0);CHKERRQ(ierr);
244     ierr = TSRestartStep(ts);CHKERRQ(ierr);
245     ierr = TSSetExactFinalTime(ts,TS_EXACTFINALTIME_MATCHSTEP);CHKERRQ(ierr);
246     ierr = TSSetFromOptions(ts);CHKERRQ(ierr);
247     ierr = TSSetEventHandler(ts,1,&direction,&terminate,EventFunction,PostEventFunction,NULL);CHKERRQ(ierr);
248     ierr = TSGetAdapt(ts,&adapt);CHKERRQ(ierr);
249     ierr = TSAdaptSetType(adapt,TSADAPTHISTORY);CHKERRQ(ierr);
250     ierr = TSGetTrajectory(ts,&tj);CHKERRQ(ierr);
251     ierr = TSAdaptHistorySetTrajectory(adapt,tj,PETSC_FALSE);CHKERRQ(ierr);
252     ierr = TSAdaptHistoryGetStep(adapt,0,&t0,&dt);CHKERRQ(ierr);
253     /* this example fails with single (or smaller) precision */
254 #if defined(PETSC_USE_REAL_SINGLE) || defined(PETSC_USE_REAL__FP16)
255     ierr = TSAdaptSetType(adapt,TSADAPTBASIC);CHKERRQ(ierr);
256     ierr = TSAdaptSetStepLimits(adapt,0.0,0.5);CHKERRQ(ierr);
257     ierr = TSSetFromOptions(ts);CHKERRQ(ierr);
258 #endif
259     ierr = TSSetTime(ts,t0);CHKERRQ(ierr);
260     ierr = TSSetTimeStep(ts,dt);CHKERRQ(ierr);
261     ierr = TSResetTrajectory(ts);CHKERRQ(ierr);
262     ierr = VecGetArray(U,&u);CHKERRQ(ierr);
263     u[0] = 1.0*rank;
264     u[1] = 20.0;
265     ierr = VecRestoreArray(U,&u);CHKERRQ(ierr);
266     ierr = TSSolve(ts,U);CHKERRQ(ierr);
267   }
268   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
269      Free work space.  All PETSc objects should be destroyed when they are no longer needed.
270      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
271   ierr = VecDestroy(&U);CHKERRQ(ierr);
272   ierr = TSDestroy(&ts);CHKERRQ(ierr);
273 
274   ierr = PetscFinalize();
275   return ierr;
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