xref: /petsc/src/ts/tutorials/ex24.c (revision 9371c9d470a9602b6d10a8bf50c9b2280a79e45a)
1 static char help[] = "Pseudotransient continuation to solve a many-variable system that comes from the 2 variable Rosenbrock function + trivial.\n\n";
2 
3 #include <petscts.h>
4 
5 static PetscErrorCode FormIJacobian(TS, PetscReal, Vec, Vec, PetscReal, Mat, Mat, void *);
6 static PetscErrorCode FormIFunction(TS, PetscReal, Vec, Vec, Vec, void *);
7 static PetscErrorCode MonitorObjective(TS, PetscInt, PetscReal, Vec, void *);
8 
9 typedef struct {
10   PetscInt  n;
11   PetscBool monitor_short;
12 } Ctx;
13 
14 int main(int argc, char **argv) {
15   TS           ts; /* time integration context */
16   Vec          X;  /* solution, residual vectors */
17   Mat          J;  /* Jacobian matrix */
18   PetscScalar *x;
19   PetscReal    ftime;
20   PetscInt     i, steps, nits, lits;
21   PetscBool    view_final;
22   Ctx          ctx;
23 
24   PetscFunctionBeginUser;
25   PetscCall(PetscInitialize(&argc, &argv, (char *)0, help));
26   ctx.n = 3;
27   PetscCall(PetscOptionsGetInt(NULL, NULL, "-n", &ctx.n, NULL));
28   PetscCheck(ctx.n >= 2, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "The dimension specified with -n must be at least 2");
29 
30   view_final = PETSC_FALSE;
31   PetscCall(PetscOptionsGetBool(NULL, NULL, "-view_final", &view_final, NULL));
32 
33   ctx.monitor_short = PETSC_FALSE;
34   PetscCall(PetscOptionsGetBool(NULL, NULL, "-monitor_short", &ctx.monitor_short, NULL));
35 
36   /*
37      Create Jacobian matrix data structure and state vector
38   */
39   PetscCall(MatCreate(PETSC_COMM_WORLD, &J));
40   PetscCall(MatSetSizes(J, PETSC_DECIDE, PETSC_DECIDE, ctx.n, ctx.n));
41   PetscCall(MatSetFromOptions(J));
42   PetscCall(MatSetUp(J));
43   PetscCall(MatCreateVecs(J, &X, NULL));
44 
45   /* Create time integration context */
46   PetscCall(TSCreate(PETSC_COMM_WORLD, &ts));
47   PetscCall(TSSetType(ts, TSPSEUDO));
48   PetscCall(TSSetIFunction(ts, NULL, FormIFunction, &ctx));
49   PetscCall(TSSetIJacobian(ts, J, J, FormIJacobian, &ctx));
50   PetscCall(TSSetMaxSteps(ts, 1000));
51   PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_STEPOVER));
52   PetscCall(TSSetTimeStep(ts, 1e-3));
53   PetscCall(TSMonitorSet(ts, MonitorObjective, &ctx, NULL));
54 
55   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
56      Customize time integrator; set runtime options
57    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
58   PetscCall(TSSetFromOptions(ts));
59 
60   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
61      Evaluate initial guess; then solve nonlinear system
62    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
63   PetscCall(VecSet(X, 0.0));
64   PetscCall(VecGetArray(X, &x));
65 #if 1
66   x[0] = 5.;
67   x[1] = -5.;
68   for (i = 2; i < ctx.n; i++) x[i] = 5.;
69 #else
70   x[0] = 1.0;
71   x[1] = 15.0;
72   for (i = 2; i < ctx.n; i++) x[i] = 10.0;
73 #endif
74   PetscCall(VecRestoreArray(X, &x));
75 
76   PetscCall(TSSolve(ts, X));
77   PetscCall(TSGetSolveTime(ts, &ftime));
78   PetscCall(TSGetStepNumber(ts, &steps));
79   PetscCall(TSGetSNESIterations(ts, &nits));
80   PetscCall(TSGetKSPIterations(ts, &lits));
81   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Time integrator took (%" PetscInt_FMT ",%" PetscInt_FMT ",%" PetscInt_FMT ") iterations to reach final time %g\n", steps, nits, lits, (double)ftime));
82   if (view_final) PetscCall(VecView(X, PETSC_VIEWER_STDOUT_WORLD));
83 
84   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
85      Free work space.  All PETSc objects should be destroyed when they
86      are no longer needed.
87    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
88 
89   PetscCall(VecDestroy(&X));
90   PetscCall(MatDestroy(&J));
91   PetscCall(TSDestroy(&ts));
92   PetscCall(PetscFinalize());
93   return 0;
94 }
95 
96 static PetscErrorCode MonitorObjective(TS ts, PetscInt step, PetscReal t, Vec X, void *ictx) {
97   Ctx               *ctx = (Ctx *)ictx;
98   const PetscScalar *x;
99   PetscScalar        f;
100   PetscReal          dt, gnorm;
101   PetscInt           i, snesit, linit;
102   SNES               snes;
103   Vec                Xdot, F;
104 
105   PetscFunctionBeginUser;
106   /* Compute objective functional */
107   PetscCall(VecGetArrayRead(X, &x));
108   f = 0;
109   for (i = 0; i < ctx->n - 1; i++) f += PetscSqr(1. - x[i]) + 100. * PetscSqr(x[i + 1] - PetscSqr(x[i]));
110   PetscCall(VecRestoreArrayRead(X, &x));
111 
112   /* Compute norm of gradient */
113   PetscCall(VecDuplicate(X, &Xdot));
114   PetscCall(VecDuplicate(X, &F));
115   PetscCall(VecZeroEntries(Xdot));
116   PetscCall(FormIFunction(ts, t, X, Xdot, F, ictx));
117   PetscCall(VecNorm(F, NORM_2, &gnorm));
118   PetscCall(VecDestroy(&Xdot));
119   PetscCall(VecDestroy(&F));
120 
121   PetscCall(TSGetTimeStep(ts, &dt));
122   PetscCall(TSGetSNES(ts, &snes));
123   PetscCall(SNESGetIterationNumber(snes, &snesit));
124   PetscCall(SNESGetLinearSolveIterations(snes, &linit));
125   PetscCall(PetscPrintf(PETSC_COMM_WORLD, (ctx->monitor_short ? "%3" PetscInt_FMT " t=%10.1e  dt=%10.1e  f=%10.1e  df=%10.1e  it=(%2" PetscInt_FMT ",%3" PetscInt_FMT ")\n" : "%3" PetscInt_FMT " t=%10.4e  dt=%10.4e  f=%10.4e  df=%10.4e  it=(%2" PetscInt_FMT ",%3" PetscInt_FMT ")\n"), step, (double)t, (double)dt, (double)PetscRealPart(f), (double)gnorm, snesit, linit));
126   PetscFunctionReturn(0);
127 }
128 
129 /* ------------------------------------------------------------------- */
130 /*
131    FormIFunction - Evaluates nonlinear function, F(X,Xdot) = Xdot + grad(objective(X))
132 
133    Input Parameters:
134 +  ts   - the TS context
135 .  t - time
136 .  X    - input vector
137 .  Xdot - time derivative
138 -  ctx  - optional user-defined context
139 
140    Output Parameter:
141 .  F - function vector
142  */
143 static PetscErrorCode FormIFunction(TS ts, PetscReal t, Vec X, Vec Xdot, Vec F, void *ictx) {
144   const PetscScalar *x;
145   PetscScalar       *f;
146   PetscInt           i;
147   Ctx               *ctx = (Ctx *)ictx;
148 
149   PetscFunctionBeginUser;
150   /*
151     Get pointers to vector data.
152     - For default PETSc vectors, VecGetArray() returns a pointer to
153     the data array.  Otherwise, the routine is implementation dependent.
154     - You MUST call VecRestoreArray() when you no longer need access to
155     the array.
156   */
157   PetscCall(VecGetArrayRead(X, &x));
158   PetscCall(VecZeroEntries(F));
159   PetscCall(VecGetArray(F, &f));
160 
161   /* Compute gradient of objective */
162   for (i = 0; i < ctx->n - 1; i++) {
163     PetscScalar a, a0, a1;
164     a  = x[i + 1] - PetscSqr(x[i]);
165     a0 = -2. * x[i];
166     a1 = 1.;
167     f[i] += -2. * (1. - x[i]) + 200. * a * a0;
168     f[i + 1] += 200. * a * a1;
169   }
170   /* Restore vectors */
171   PetscCall(VecRestoreArrayRead(X, &x));
172   PetscCall(VecRestoreArray(F, &f));
173   PetscCall(VecAXPY(F, 1.0, Xdot));
174   PetscFunctionReturn(0);
175 }
176 /* ------------------------------------------------------------------- */
177 /*
178    FormIJacobian - Evaluates Jacobian matrix.
179 
180    Input Parameters:
181 +  ts - the TS context
182 .  t - pseudo-time
183 .  X - input vector
184 .  Xdot - time derivative
185 .  shift - multiplier for mass matrix
186 .  dummy - user-defined context
187 
188    Output Parameters:
189 .  J - Jacobian matrix
190 .  B - optionally different preconditioning matrix
191 .  flag - flag indicating matrix structure
192 */
193 static PetscErrorCode FormIJacobian(TS ts, PetscReal t, Vec X, Vec Xdot, PetscReal shift, Mat J, Mat B, void *ictx) {
194   const PetscScalar *x;
195   PetscInt           i;
196   Ctx               *ctx = (Ctx *)ictx;
197 
198   PetscFunctionBeginUser;
199   PetscCall(MatZeroEntries(B));
200   /*
201      Get pointer to vector data
202   */
203   PetscCall(VecGetArrayRead(X, &x));
204 
205   /*
206      Compute Jacobian entries and insert into matrix.
207   */
208   for (i = 0; i < ctx->n - 1; i++) {
209     PetscInt    rowcol[2];
210     PetscScalar v[2][2], a, a0, a1, a00, a01, a10, a11;
211     rowcol[0] = i;
212     rowcol[1] = i + 1;
213     a         = x[i + 1] - PetscSqr(x[i]);
214     a0        = -2. * x[i];
215     a00       = -2.;
216     a01       = 0.;
217     a1        = 1.;
218     a10       = 0.;
219     a11       = 0.;
220     v[0][0]   = 2. + 200. * (a * a00 + a0 * a0);
221     v[0][1]   = 200. * (a * a01 + a1 * a0);
222     v[1][0]   = 200. * (a * a10 + a0 * a1);
223     v[1][1]   = 200. * (a * a11 + a1 * a1);
224     PetscCall(MatSetValues(B, 2, rowcol, 2, rowcol, &v[0][0], ADD_VALUES));
225   }
226   for (i = 0; i < ctx->n; i++) { PetscCall(MatSetValue(B, i, i, (PetscScalar)shift, ADD_VALUES)); }
227 
228   PetscCall(VecRestoreArrayRead(X, &x));
229 
230   /*
231      Assemble matrix
232   */
233   PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
234   PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
235   if (J != B) {
236     PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY));
237     PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY));
238   }
239   PetscFunctionReturn(0);
240 }
241 
242 /*TEST
243 
244     test:
245       requires: !single
246 
247     test:
248       args: -pc_type lu -ts_dt 1e-5 -ts_max_time 1e5 -n 50 -monitor_short -snes_max_it 5 -snes_type newtonls -ts_max_snes_failures -1
249       requires: !single
250       suffix: 2
251 
252 TEST*/
253