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