xref: /petsc/src/ts/tutorials/ex8.c (revision 2ff79c18c26c94ed8cb599682f680f231dca6444)
1 static char help[] = "Nonlinear DAE benchmark problems.\n";
2 
3 /*
4    Include "petscts.h" so that we can use TS solvers.  Note that this
5    file automatically includes:
6      petscsys.h       - base PETSc routines   petscvec.h - vectors
7      petscmat.h - matrices
8      petscis.h     - index sets            petscksp.h - Krylov subspace methods
9      petscviewer.h - viewers               petscpc.h  - preconditioners
10      petscksp.h   - linear solvers
11 */
12 #include <petscts.h>
13 
14 typedef struct _Problem *Problem;
15 struct _Problem {
16   PetscErrorCode (*destroy)(Problem);
17   TSIFunctionFn *function;
18   TSIJacobianFn *jacobian;
19   PetscErrorCode (*solution)(PetscReal, Vec, void *);
20   MPI_Comm  comm;
21   PetscReal final_time;
22   PetscInt  n;
23   PetscBool hasexact;
24   void     *data;
25 };
26 
27 /*
28       Stiff 3-variable system from chemical reactions, due to Robertson (1966), problem ROBER in Hairer&Wanner, ODE 2, 1996
29 */
30 static PetscErrorCode RoberFunction(TS ts, PetscReal t, Vec X, Vec Xdot, Vec F, void *ctx)
31 {
32   PetscScalar       *f;
33   const PetscScalar *x, *xdot;
34 
35   PetscFunctionBeginUser;
36   PetscCall(VecGetArrayRead(X, &x));
37   PetscCall(VecGetArrayRead(Xdot, &xdot));
38   PetscCall(VecGetArray(F, &f));
39   f[0] = xdot[0] + 0.04 * x[0] - 1e4 * x[1] * x[2];
40   f[1] = xdot[1] - 0.04 * x[0] + 1e4 * x[1] * x[2] + 3e7 * PetscSqr(x[1]);
41   f[2] = xdot[2] - 3e7 * PetscSqr(x[1]);
42   PetscCall(VecRestoreArrayRead(X, &x));
43   PetscCall(VecRestoreArrayRead(Xdot, &xdot));
44   PetscCall(VecRestoreArray(F, &f));
45   PetscFunctionReturn(PETSC_SUCCESS);
46 }
47 
48 static PetscErrorCode RoberJacobian(TS ts, PetscReal t, Vec X, Vec Xdot, PetscReal a, Mat A, Mat B, void *ctx)
49 {
50   PetscInt           rowcol[] = {0, 1, 2};
51   PetscScalar        J[3][3];
52   const PetscScalar *x, *xdot;
53 
54   PetscFunctionBeginUser;
55   PetscCall(VecGetArrayRead(X, &x));
56   PetscCall(VecGetArrayRead(Xdot, &xdot));
57   J[0][0] = a + 0.04;
58   J[0][1] = -1e4 * x[2];
59   J[0][2] = -1e4 * x[1];
60   J[1][0] = -0.04;
61   J[1][1] = a + 1e4 * x[2] + 3e7 * 2 * x[1];
62   J[1][2] = 1e4 * x[1];
63   J[2][0] = 0;
64   J[2][1] = -3e7 * 2 * x[1];
65   J[2][2] = a;
66   PetscCall(MatSetValues(B, 3, rowcol, 3, rowcol, &J[0][0], INSERT_VALUES));
67   PetscCall(VecRestoreArrayRead(X, &x));
68   PetscCall(VecRestoreArrayRead(Xdot, &xdot));
69 
70   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
71   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
72   if (A != B) {
73     PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
74     PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
75   }
76   PetscFunctionReturn(PETSC_SUCCESS);
77 }
78 
79 static PetscErrorCode RoberSolution(PetscReal t, Vec X, void *ctx)
80 {
81   PetscScalar *x;
82 
83   PetscFunctionBeginUser;
84   PetscCheck(t == 0, PETSC_COMM_WORLD, PETSC_ERR_SUP, "not implemented");
85   PetscCall(VecGetArray(X, &x));
86   x[0] = 1;
87   x[1] = 0;
88   x[2] = 0;
89   PetscCall(VecRestoreArray(X, &x));
90   PetscFunctionReturn(PETSC_SUCCESS);
91 }
92 
93 static PetscErrorCode RoberCreate(Problem p)
94 {
95   PetscFunctionBeginUser;
96   p->destroy    = 0;
97   p->function   = &RoberFunction;
98   p->jacobian   = &RoberJacobian;
99   p->solution   = &RoberSolution;
100   p->final_time = 1e11;
101   p->n          = 3;
102   PetscFunctionReturn(PETSC_SUCCESS);
103 }
104 
105 /*
106      Stiff scalar valued problem
107 */
108 
109 typedef struct {
110   PetscReal lambda;
111 } CECtx;
112 
113 static PetscErrorCode CEDestroy(Problem p)
114 {
115   PetscFunctionBeginUser;
116   PetscCall(PetscFree(p->data));
117   PetscFunctionReturn(PETSC_SUCCESS);
118 }
119 
120 static PetscErrorCode CEFunction(TS ts, PetscReal t, Vec X, Vec Xdot, Vec F, void *ctx)
121 {
122   PetscReal          l = ((CECtx *)ctx)->lambda;
123   PetscScalar       *f;
124   const PetscScalar *x, *xdot;
125 
126   PetscFunctionBeginUser;
127   PetscCall(VecGetArrayRead(X, &x));
128   PetscCall(VecGetArrayRead(Xdot, &xdot));
129   PetscCall(VecGetArray(F, &f));
130   f[0] = xdot[0] + l * (x[0] - PetscCosReal(t));
131 #if 0
132   PetscCall(PetscPrintf(PETSC_COMM_WORLD," f(t=%g,x=%g,xdot=%g) = %g\n",(double)t,(double)x[0],(double)xdot[0],(double)f[0]));
133 #endif
134   PetscCall(VecRestoreArrayRead(X, &x));
135   PetscCall(VecRestoreArrayRead(Xdot, &xdot));
136   PetscCall(VecRestoreArray(F, &f));
137   PetscFunctionReturn(PETSC_SUCCESS);
138 }
139 
140 static PetscErrorCode CEJacobian(TS ts, PetscReal t, Vec X, Vec Xdot, PetscReal a, Mat A, Mat B, void *ctx)
141 {
142   PetscReal          l        = ((CECtx *)ctx)->lambda;
143   PetscInt           rowcol[] = {0};
144   PetscScalar        J[1][1];
145   const PetscScalar *x, *xdot;
146 
147   PetscFunctionBeginUser;
148   PetscCall(VecGetArrayRead(X, &x));
149   PetscCall(VecGetArrayRead(Xdot, &xdot));
150   J[0][0] = a + l;
151   PetscCall(MatSetValues(B, 1, rowcol, 1, rowcol, &J[0][0], INSERT_VALUES));
152   PetscCall(VecRestoreArrayRead(X, &x));
153   PetscCall(VecRestoreArrayRead(Xdot, &xdot));
154 
155   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
156   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
157   if (A != B) {
158     PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
159     PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
160   }
161   PetscFunctionReturn(PETSC_SUCCESS);
162 }
163 
164 static PetscErrorCode CESolution(PetscReal t, Vec X, void *ctx)
165 {
166   PetscReal    l = ((CECtx *)ctx)->lambda;
167   PetscScalar *x;
168 
169   PetscFunctionBeginUser;
170   PetscCall(VecGetArray(X, &x));
171   x[0] = l / (l * l + 1) * (l * PetscCosReal(t) + PetscSinReal(t)) - l * l / (l * l + 1) * PetscExpReal(-l * t);
172   PetscCall(VecRestoreArray(X, &x));
173   PetscFunctionReturn(PETSC_SUCCESS);
174 }
175 
176 static PetscErrorCode CECreate(Problem p)
177 {
178   CECtx *ce;
179 
180   PetscFunctionBeginUser;
181   PetscCall(PetscMalloc(sizeof(CECtx), &ce));
182   p->data = (void *)ce;
183 
184   p->destroy    = &CEDestroy;
185   p->function   = &CEFunction;
186   p->jacobian   = &CEJacobian;
187   p->solution   = &CESolution;
188   p->final_time = 10;
189   p->n          = 1;
190   p->hasexact   = PETSC_TRUE;
191 
192   ce->lambda = 10;
193   PetscOptionsBegin(p->comm, NULL, "CE options", "");
194   {
195     PetscCall(PetscOptionsReal("-problem_ce_lambda", "Parameter controlling stiffness: xdot + lambda*(x - cos(t))", "", ce->lambda, &ce->lambda, NULL));
196   }
197   PetscOptionsEnd();
198   PetscFunctionReturn(PETSC_SUCCESS);
199 }
200 
201 /*
202    Stiff 3-variable oscillatory system from chemical reactions. problem OREGO in Hairer&Wanner
203 */
204 static PetscErrorCode OregoFunction(TS ts, PetscReal t, Vec X, Vec Xdot, Vec F, void *ctx)
205 {
206   PetscScalar       *f;
207   const PetscScalar *x, *xdot;
208 
209   PetscFunctionBeginUser;
210   PetscCall(VecGetArrayRead(X, &x));
211   PetscCall(VecGetArrayRead(Xdot, &xdot));
212   PetscCall(VecGetArray(F, &f));
213   f[0] = xdot[0] - 77.27 * (x[1] + x[0] * (1. - 8.375e-6 * x[0] - x[1]));
214   f[1] = xdot[1] - 1 / 77.27 * (x[2] - (1. + x[0]) * x[1]);
215   f[2] = xdot[2] - 0.161 * (x[0] - x[2]);
216   PetscCall(VecRestoreArrayRead(X, &x));
217   PetscCall(VecRestoreArrayRead(Xdot, &xdot));
218   PetscCall(VecRestoreArray(F, &f));
219   PetscFunctionReturn(PETSC_SUCCESS);
220 }
221 
222 static PetscErrorCode OregoJacobian(TS ts, PetscReal t, Vec X, Vec Xdot, PetscReal a, Mat A, Mat B, void *ctx)
223 {
224   PetscInt           rowcol[] = {0, 1, 2};
225   PetscScalar        J[3][3];
226   const PetscScalar *x, *xdot;
227 
228   PetscFunctionBeginUser;
229   PetscCall(VecGetArrayRead(X, &x));
230   PetscCall(VecGetArrayRead(Xdot, &xdot));
231   J[0][0] = a - 77.27 * ((1. - 8.375e-6 * x[0] - x[1]) - 8.375e-6 * x[0]);
232   J[0][1] = -77.27 * (1. - x[0]);
233   J[0][2] = 0;
234   J[1][0] = 1. / 77.27 * x[1];
235   J[1][1] = a + 1. / 77.27 * (1. + x[0]);
236   J[1][2] = -1. / 77.27;
237   J[2][0] = -0.161;
238   J[2][1] = 0;
239   J[2][2] = a + 0.161;
240   PetscCall(MatSetValues(B, 3, rowcol, 3, rowcol, &J[0][0], INSERT_VALUES));
241   PetscCall(VecRestoreArrayRead(X, &x));
242   PetscCall(VecRestoreArrayRead(Xdot, &xdot));
243 
244   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
245   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
246   if (A != B) {
247     PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
248     PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
249   }
250   PetscFunctionReturn(PETSC_SUCCESS);
251 }
252 
253 static PetscErrorCode OregoSolution(PetscReal t, Vec X, void *ctx)
254 {
255   PetscScalar *x;
256 
257   PetscFunctionBeginUser;
258   PetscCheck(t == 0, PETSC_COMM_WORLD, PETSC_ERR_SUP, "not implemented");
259   PetscCall(VecGetArray(X, &x));
260   x[0] = 1;
261   x[1] = 2;
262   x[2] = 3;
263   PetscCall(VecRestoreArray(X, &x));
264   PetscFunctionReturn(PETSC_SUCCESS);
265 }
266 
267 static PetscErrorCode OregoCreate(Problem p)
268 {
269   PetscFunctionBeginUser;
270   p->destroy    = 0;
271   p->function   = &OregoFunction;
272   p->jacobian   = &OregoJacobian;
273   p->solution   = &OregoSolution;
274   p->final_time = 360;
275   p->n          = 3;
276   PetscFunctionReturn(PETSC_SUCCESS);
277 }
278 
279 /*
280    User-defined monitor for comparing to exact solutions when possible
281 */
282 typedef struct {
283   MPI_Comm comm;
284   Problem  problem;
285   Vec      x;
286 } MonitorCtx;
287 
288 static PetscErrorCode MonitorError(TS ts, PetscInt step, PetscReal t, Vec x, void *ctx)
289 {
290   MonitorCtx *mon = (MonitorCtx *)ctx;
291   PetscReal   h, nrm_x, nrm_exact, nrm_diff;
292 
293   PetscFunctionBeginUser;
294   if (!mon->problem->solution) PetscFunctionReturn(PETSC_SUCCESS);
295   PetscCall((*mon->problem->solution)(t, mon->x, mon->problem->data));
296   PetscCall(VecNorm(x, NORM_2, &nrm_x));
297   PetscCall(VecNorm(mon->x, NORM_2, &nrm_exact));
298   PetscCall(VecAYPX(mon->x, -1, x));
299   PetscCall(VecNorm(mon->x, NORM_2, &nrm_diff));
300   PetscCall(TSGetTimeStep(ts, &h));
301   if (step < 0) PetscCall(PetscPrintf(mon->comm, "Interpolated final solution "));
302   PetscCall(PetscPrintf(mon->comm, "step %4" PetscInt_FMT " t=%12.8e h=% 8.2e  |x|=%9.2e  |x_e|=%9.2e  |x-x_e|=%9.2e\n", step, (double)t, (double)h, (double)nrm_x, (double)nrm_exact, (double)nrm_diff));
303   PetscFunctionReturn(PETSC_SUCCESS);
304 }
305 
306 int main(int argc, char **argv)
307 {
308   PetscFunctionList plist = NULL;
309   char              pname[256];
310   TS                ts;   /* nonlinear solver */
311   Vec               x, r; /* solution, residual vectors */
312   Mat               A;    /* Jacobian matrix */
313   Problem           problem;
314   PetscBool         use_monitor = PETSC_FALSE;
315   PetscBool         use_result  = PETSC_FALSE;
316   PetscInt          steps, nonlinits, linits, snesfails, rejects;
317   PetscReal         ftime;
318   MonitorCtx        mon;
319   PetscMPIInt       size;
320 
321   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
322      Initialize program
323      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
324   PetscFunctionBeginUser;
325   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
326   PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
327   PetscCheck(size == 1, PETSC_COMM_WORLD, PETSC_ERR_WRONG_MPI_SIZE, "Only for sequential runs");
328 
329   /* Register the available problems */
330   PetscCall(PetscFunctionListAdd(&plist, "rober", &RoberCreate));
331   PetscCall(PetscFunctionListAdd(&plist, "ce", &CECreate));
332   PetscCall(PetscFunctionListAdd(&plist, "orego", &OregoCreate));
333   PetscCall(PetscStrncpy(pname, "ce", sizeof(pname)));
334 
335   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
336     Set runtime options
337     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
338   PetscOptionsBegin(PETSC_COMM_WORLD, NULL, "Timestepping benchmark options", "");
339   {
340     PetscCall(PetscOptionsFList("-problem_type", "Name of problem to run", "", plist, pname, pname, sizeof(pname), NULL));
341     use_monitor = PETSC_FALSE;
342     PetscCall(PetscOptionsBool("-monitor_error", "Display errors relative to exact solutions", "", use_monitor, &use_monitor, NULL));
343     PetscCall(PetscOptionsBool("-monitor_result", "Display result", "", use_result, &use_result, NULL));
344   }
345   PetscOptionsEnd();
346 
347   /* Create the new problem */
348   PetscCall(PetscNew(&problem));
349   problem->comm = MPI_COMM_WORLD;
350   {
351     PetscErrorCode (*pcreate)(Problem);
352 
353     PetscCall(PetscFunctionListFind(plist, pname, &pcreate));
354     PetscCheck(pcreate, PETSC_COMM_SELF, PETSC_ERR_ARG_UNKNOWN_TYPE, "No problem '%s'", pname);
355     PetscCall((*pcreate)(problem));
356   }
357 
358   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
359     Create necessary matrix and vectors
360     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
361   PetscCall(MatCreate(PETSC_COMM_WORLD, &A));
362   PetscCall(MatSetSizes(A, problem->n, problem->n, PETSC_DETERMINE, PETSC_DETERMINE));
363   PetscCall(MatSetFromOptions(A));
364   PetscCall(MatSetUp(A));
365 
366   PetscCall(MatCreateVecs(A, &x, NULL));
367   PetscCall(VecDuplicate(x, &r));
368 
369   mon.comm    = PETSC_COMM_WORLD;
370   mon.problem = problem;
371   PetscCall(VecDuplicate(x, &mon.x));
372 
373   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
374      Create timestepping solver context
375      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
376   PetscCall(TSCreate(PETSC_COMM_WORLD, &ts));
377   PetscCall(TSSetProblemType(ts, TS_NONLINEAR));
378   PetscCall(TSSetType(ts, TSROSW)); /* Rosenbrock-W */
379   PetscCall(TSSetIFunction(ts, NULL, problem->function, problem->data));
380   PetscCall(TSSetIJacobian(ts, A, A, problem->jacobian, problem->data));
381   PetscCall(TSSetMaxTime(ts, problem->final_time));
382   PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_STEPOVER));
383   PetscCall(TSSetMaxStepRejections(ts, 10));
384   PetscCall(TSSetMaxSNESFailures(ts, -1)); /* unlimited */
385   if (use_monitor) PetscCall(TSMonitorSet(ts, &MonitorError, &mon, NULL));
386 
387   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
388      Set initial conditions
389    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
390   PetscCall((*problem->solution)(0, x, problem->data));
391   PetscCall(TSSetTimeStep(ts, .001));
392   PetscCall(TSSetSolution(ts, x));
393 
394   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
395      Set runtime options
396    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
397   PetscCall(TSSetFromOptions(ts));
398 
399   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
400      Solve nonlinear system
401      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
402   PetscCall(TSSolve(ts, x));
403   PetscCall(TSGetSolveTime(ts, &ftime));
404   PetscCall(TSGetStepNumber(ts, &steps));
405   PetscCall(TSGetSNESFailures(ts, &snesfails));
406   PetscCall(TSGetStepRejections(ts, &rejects));
407   PetscCall(TSGetSNESIterations(ts, &nonlinits));
408   PetscCall(TSGetKSPIterations(ts, &linits));
409   if (use_result) {
410     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "steps %" PetscInt_FMT " (%" PetscInt_FMT " rejected, %" PetscInt_FMT " SNES fails), ftime %g, nonlinits %" PetscInt_FMT ", linits %" PetscInt_FMT "\n", steps, rejects, snesfails, (double)ftime, nonlinits, linits));
411   }
412 
413   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
414      Free work space.  All PETSc objects should be destroyed when they
415      are no longer needed.
416    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
417   PetscCall(MatDestroy(&A));
418   PetscCall(VecDestroy(&x));
419   PetscCall(VecDestroy(&r));
420   PetscCall(VecDestroy(&mon.x));
421   PetscCall(TSDestroy(&ts));
422   if (problem->destroy) PetscCall((*problem->destroy)(problem));
423   PetscCall(PetscFree(problem));
424   PetscCall(PetscFunctionListDestroy(&plist));
425 
426   PetscCall(PetscFinalize());
427   return 0;
428 }
429 
430 /*TEST
431 
432     test:
433       requires: !complex
434       args: -monitor_result -monitor_error -ts_atol 1e-2 -ts_rtol 1e-2 -ts_exact_final_time interpolate -ts_type arkimex
435 
436     test:
437       suffix: 2
438       requires: !single !complex
439       args: -monitor_result -ts_atol 1e-2 -ts_rtol 1e-2 -ts_max_time 15 -ts_type arkimex -ts_arkimex_type 2e -problem_type orego -ts_arkimex_initial_guess_extrapolate 0 -ts_adapt_time_step_increase_delay 4
440 
441     test:
442       suffix: 3
443       requires: !single !complex
444       args: -monitor_result -ts_atol 1e-2 -ts_rtol 1e-2 -ts_max_time 15 -ts_type arkimex -ts_arkimex_type 2e -problem_type orego -ts_arkimex_initial_guess_extrapolate 1
445 
446     test:
447       suffix: 4
448       output_file: output/empty.out
449 
450     test:
451       suffix: 5
452       args: -snes_lag_jacobian 20 -snes_lag_jacobian_persists
453       output_file: output/empty.out
454 
455 TEST*/
456