1 static char help[] = "Solves the van der Pol equation.\n\
2 Input parameters include:\n";
3
4 /* ------------------------------------------------------------------------
5
6 Notes:
7 This code demonstrates how to solve a DAE-constrained optimization problem with TAO, TSAdjoint and TS.
8 The nonlinear problem is written in a DAE equivalent form.
9 The objective is to minimize the difference between observation and model prediction by finding an optimal value for parameter \mu.
10 The gradient is computed with the discrete adjoint of an implicit theta method, see ex20adj.c for details.
11 ------------------------------------------------------------------------- */
12 #include <petsctao.h>
13 #include <petscts.h>
14
15 typedef struct _n_User *User;
16 struct _n_User {
17 TS ts;
18 PetscReal mu;
19 PetscReal next_output;
20
21 /* Sensitivity analysis support */
22 PetscReal ftime;
23 Mat A; /* Jacobian matrix */
24 Mat Jacp; /* JacobianP matrix */
25 Mat H; /* Hessian matrix for optimization */
26 Vec U, Lambda[1], Mup[1]; /* adjoint variables */
27 Vec Lambda2[1], Mup2[1]; /* second-order adjoint variables */
28 Vec Ihp1[1]; /* working space for Hessian evaluations */
29 Vec Ihp2[1]; /* working space for Hessian evaluations */
30 Vec Ihp3[1]; /* working space for Hessian evaluations */
31 Vec Ihp4[1]; /* working space for Hessian evaluations */
32 Vec Dir; /* direction vector */
33 PetscReal ob[2]; /* observation used by the cost function */
34 PetscBool implicitform; /* implicit ODE? */
35 };
36
37 PetscErrorCode FormFunctionGradient(Tao, Vec, PetscReal *, Vec, void *);
38 PetscErrorCode FormHessian(Tao, Vec, Mat, Mat, void *);
39 PetscErrorCode Adjoint2(Vec, PetscScalar[], User);
40
41 /* ----------------------- Explicit form of the ODE -------------------- */
42
RHSFunction(TS ts,PetscReal t,Vec U,Vec F,PetscCtx ctx)43 static PetscErrorCode RHSFunction(TS ts, PetscReal t, Vec U, Vec F, PetscCtx ctx)
44 {
45 User user = (User)ctx;
46 PetscScalar *f;
47 const PetscScalar *u;
48
49 PetscFunctionBeginUser;
50 PetscCall(VecGetArrayRead(U, &u));
51 PetscCall(VecGetArray(F, &f));
52 f[0] = u[1];
53 f[1] = user->mu * ((1. - u[0] * u[0]) * u[1] - u[0]);
54 PetscCall(VecRestoreArrayRead(U, &u));
55 PetscCall(VecRestoreArray(F, &f));
56 PetscFunctionReturn(PETSC_SUCCESS);
57 }
58
RHSJacobian(TS ts,PetscReal t,Vec U,Mat A,Mat B,PetscCtx ctx)59 static PetscErrorCode RHSJacobian(TS ts, PetscReal t, Vec U, Mat A, Mat B, PetscCtx ctx)
60 {
61 User user = (User)ctx;
62 PetscReal mu = user->mu;
63 PetscInt rowcol[] = {0, 1};
64 PetscScalar J[2][2];
65 const PetscScalar *u;
66
67 PetscFunctionBeginUser;
68 PetscCall(VecGetArrayRead(U, &u));
69
70 J[0][0] = 0;
71 J[1][0] = -mu * (2.0 * u[1] * u[0] + 1.);
72 J[0][1] = 1.0;
73 J[1][1] = mu * (1.0 - u[0] * u[0]);
74 PetscCall(MatSetValues(A, 2, rowcol, 2, rowcol, &J[0][0], INSERT_VALUES));
75 PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
76 PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
77 if (B && A != B) {
78 PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
79 PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
80 }
81
82 PetscCall(VecRestoreArrayRead(U, &u));
83 PetscFunctionReturn(PETSC_SUCCESS);
84 }
85
RHSHessianProductUU(TS ts,PetscReal t,Vec U,Vec * Vl,Vec Vr,Vec * VHV,PetscCtx ctx)86 static PetscErrorCode RHSHessianProductUU(TS ts, PetscReal t, Vec U, Vec *Vl, Vec Vr, Vec *VHV, PetscCtx ctx)
87 {
88 const PetscScalar *vl, *vr, *u;
89 PetscScalar *vhv;
90 PetscScalar dJdU[2][2][2] = {{{0}}};
91 PetscInt i, j, k;
92 User user = (User)ctx;
93
94 PetscFunctionBeginUser;
95 PetscCall(VecGetArrayRead(U, &u));
96 PetscCall(VecGetArrayRead(Vl[0], &vl));
97 PetscCall(VecGetArrayRead(Vr, &vr));
98 PetscCall(VecGetArray(VHV[0], &vhv));
99
100 dJdU[1][0][0] = -2. * user->mu * u[1];
101 dJdU[1][1][0] = -2. * user->mu * u[0];
102 dJdU[1][0][1] = -2. * user->mu * u[0];
103 for (j = 0; j < 2; j++) {
104 vhv[j] = 0;
105 for (k = 0; k < 2; k++)
106 for (i = 0; i < 2; i++) vhv[j] += vl[i] * dJdU[i][j][k] * vr[k];
107 }
108
109 PetscCall(VecRestoreArrayRead(U, &u));
110 PetscCall(VecRestoreArrayRead(Vl[0], &vl));
111 PetscCall(VecRestoreArrayRead(Vr, &vr));
112 PetscCall(VecRestoreArray(VHV[0], &vhv));
113 PetscFunctionReturn(PETSC_SUCCESS);
114 }
115
RHSHessianProductUP(TS ts,PetscReal t,Vec U,Vec * Vl,Vec Vr,Vec * VHV,PetscCtx ctx)116 static PetscErrorCode RHSHessianProductUP(TS ts, PetscReal t, Vec U, Vec *Vl, Vec Vr, Vec *VHV, PetscCtx ctx)
117 {
118 const PetscScalar *vl, *vr, *u;
119 PetscScalar *vhv;
120 PetscScalar dJdP[2][2][1] = {{{0}}};
121 PetscInt i, j, k;
122
123 PetscFunctionBeginUser;
124 PetscCall(VecGetArrayRead(U, &u));
125 PetscCall(VecGetArrayRead(Vl[0], &vl));
126 PetscCall(VecGetArrayRead(Vr, &vr));
127 PetscCall(VecGetArray(VHV[0], &vhv));
128
129 dJdP[1][0][0] = -(1. + 2. * u[0] * u[1]);
130 dJdP[1][1][0] = 1. - u[0] * u[0];
131 for (j = 0; j < 2; j++) {
132 vhv[j] = 0;
133 for (k = 0; k < 1; k++)
134 for (i = 0; i < 2; i++) vhv[j] += vl[i] * dJdP[i][j][k] * vr[k];
135 }
136
137 PetscCall(VecRestoreArrayRead(U, &u));
138 PetscCall(VecRestoreArrayRead(Vl[0], &vl));
139 PetscCall(VecRestoreArrayRead(Vr, &vr));
140 PetscCall(VecRestoreArray(VHV[0], &vhv));
141 PetscFunctionReturn(PETSC_SUCCESS);
142 }
143
RHSHessianProductPU(TS ts,PetscReal t,Vec U,Vec * Vl,Vec Vr,Vec * VHV,PetscCtx ctx)144 static PetscErrorCode RHSHessianProductPU(TS ts, PetscReal t, Vec U, Vec *Vl, Vec Vr, Vec *VHV, PetscCtx ctx)
145 {
146 const PetscScalar *vl, *vr, *u;
147 PetscScalar *vhv;
148 PetscScalar dJdU[2][1][2] = {{{0}}};
149 PetscInt i, j, k;
150
151 PetscFunctionBeginUser;
152 PetscCall(VecGetArrayRead(U, &u));
153 PetscCall(VecGetArrayRead(Vl[0], &vl));
154 PetscCall(VecGetArrayRead(Vr, &vr));
155 PetscCall(VecGetArray(VHV[0], &vhv));
156
157 dJdU[1][0][0] = -1. - 2. * u[1] * u[0];
158 dJdU[1][0][1] = 1. - u[0] * u[0];
159 for (j = 0; j < 1; j++) {
160 vhv[j] = 0;
161 for (k = 0; k < 2; k++)
162 for (i = 0; i < 2; i++) vhv[j] += vl[i] * dJdU[i][j][k] * vr[k];
163 }
164
165 PetscCall(VecRestoreArrayRead(U, &u));
166 PetscCall(VecRestoreArrayRead(Vl[0], &vl));
167 PetscCall(VecRestoreArrayRead(Vr, &vr));
168 PetscCall(VecRestoreArray(VHV[0], &vhv));
169 PetscFunctionReturn(PETSC_SUCCESS);
170 }
171
RHSHessianProductPP(TS ts,PetscReal t,Vec U,Vec * Vl,Vec Vr,Vec * VHV,PetscCtx ctx)172 static PetscErrorCode RHSHessianProductPP(TS ts, PetscReal t, Vec U, Vec *Vl, Vec Vr, Vec *VHV, PetscCtx ctx)
173 {
174 PetscFunctionBeginUser;
175 PetscFunctionReturn(PETSC_SUCCESS);
176 }
177
178 /* ----------------------- Implicit form of the ODE -------------------- */
179
IFunction(TS ts,PetscReal t,Vec U,Vec Udot,Vec F,PetscCtx ctx)180 static PetscErrorCode IFunction(TS ts, PetscReal t, Vec U, Vec Udot, Vec F, PetscCtx ctx)
181 {
182 User user = (User)ctx;
183 PetscScalar *f;
184 const PetscScalar *u, *udot;
185
186 PetscFunctionBeginUser;
187 PetscCall(VecGetArrayRead(U, &u));
188 PetscCall(VecGetArrayRead(Udot, &udot));
189 PetscCall(VecGetArray(F, &f));
190
191 f[0] = udot[0] - u[1];
192 f[1] = udot[1] - user->mu * ((1.0 - u[0] * u[0]) * u[1] - u[0]);
193
194 PetscCall(VecRestoreArrayRead(U, &u));
195 PetscCall(VecRestoreArrayRead(Udot, &udot));
196 PetscCall(VecRestoreArray(F, &f));
197 PetscFunctionReturn(PETSC_SUCCESS);
198 }
199
IJacobian(TS ts,PetscReal t,Vec U,Vec Udot,PetscReal a,Mat A,Mat B,PetscCtx ctx)200 static PetscErrorCode IJacobian(TS ts, PetscReal t, Vec U, Vec Udot, PetscReal a, Mat A, Mat B, PetscCtx ctx)
201 {
202 User user = (User)ctx;
203 PetscInt rowcol[] = {0, 1};
204 PetscScalar J[2][2];
205 const PetscScalar *u;
206
207 PetscFunctionBeginUser;
208 PetscCall(VecGetArrayRead(U, &u));
209
210 J[0][0] = a;
211 J[0][1] = -1.0;
212 J[1][0] = user->mu * (1.0 + 2.0 * u[0] * u[1]);
213 J[1][1] = a - user->mu * (1.0 - u[0] * u[0]);
214 PetscCall(MatSetValues(B, 2, rowcol, 2, rowcol, &J[0][0], INSERT_VALUES));
215 PetscCall(VecRestoreArrayRead(U, &u));
216 PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
217 PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
218 if (A != B) {
219 PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
220 PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
221 }
222
223 PetscCall(VecRestoreArrayRead(U, &u));
224 PetscFunctionReturn(PETSC_SUCCESS);
225 }
226
RHSJacobianP(TS ts,PetscReal t,Vec U,Mat A,PetscCtx ctx)227 static PetscErrorCode RHSJacobianP(TS ts, PetscReal t, Vec U, Mat A, PetscCtx ctx)
228 {
229 PetscInt row[] = {0, 1}, col[] = {0};
230 PetscScalar J[2][1];
231 const PetscScalar *u;
232
233 PetscFunctionBeginUser;
234 PetscCall(VecGetArrayRead(U, &u));
235
236 J[0][0] = 0;
237 J[1][0] = (1. - u[0] * u[0]) * u[1] - u[0];
238 PetscCall(MatSetValues(A, 2, row, 1, col, &J[0][0], INSERT_VALUES));
239 PetscCall(VecRestoreArrayRead(U, &u));
240 PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
241 PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
242
243 PetscCall(VecRestoreArrayRead(U, &u));
244 PetscFunctionReturn(PETSC_SUCCESS);
245 }
246
IHessianProductUU(TS ts,PetscReal t,Vec U,Vec * Vl,Vec Vr,Vec * VHV,PetscCtx ctx)247 static PetscErrorCode IHessianProductUU(TS ts, PetscReal t, Vec U, Vec *Vl, Vec Vr, Vec *VHV, PetscCtx ctx)
248 {
249 const PetscScalar *vl, *vr, *u;
250 PetscScalar *vhv;
251 PetscScalar dJdU[2][2][2] = {{{0}}};
252 PetscInt i, j, k;
253 User user = (User)ctx;
254
255 PetscFunctionBeginUser;
256 PetscCall(VecGetArrayRead(U, &u));
257 PetscCall(VecGetArrayRead(Vl[0], &vl));
258 PetscCall(VecGetArrayRead(Vr, &vr));
259 PetscCall(VecGetArray(VHV[0], &vhv));
260
261 dJdU[1][0][0] = 2. * user->mu * u[1];
262 dJdU[1][1][0] = 2. * user->mu * u[0];
263 dJdU[1][0][1] = 2. * user->mu * u[0];
264 for (j = 0; j < 2; j++) {
265 vhv[j] = 0;
266 for (k = 0; k < 2; k++)
267 for (i = 0; i < 2; i++) vhv[j] += vl[i] * dJdU[i][j][k] * vr[k];
268 }
269
270 PetscCall(VecRestoreArrayRead(U, &u));
271 PetscCall(VecRestoreArrayRead(Vl[0], &vl));
272 PetscCall(VecRestoreArrayRead(Vr, &vr));
273 PetscCall(VecRestoreArray(VHV[0], &vhv));
274 PetscFunctionReturn(PETSC_SUCCESS);
275 }
276
IHessianProductUP(TS ts,PetscReal t,Vec U,Vec * Vl,Vec Vr,Vec * VHV,PetscCtx ctx)277 static PetscErrorCode IHessianProductUP(TS ts, PetscReal t, Vec U, Vec *Vl, Vec Vr, Vec *VHV, PetscCtx ctx)
278 {
279 const PetscScalar *vl, *vr, *u;
280 PetscScalar *vhv;
281 PetscScalar dJdP[2][2][1] = {{{0}}};
282 PetscInt i, j, k;
283
284 PetscFunctionBeginUser;
285 PetscCall(VecGetArrayRead(U, &u));
286 PetscCall(VecGetArrayRead(Vl[0], &vl));
287 PetscCall(VecGetArrayRead(Vr, &vr));
288 PetscCall(VecGetArray(VHV[0], &vhv));
289
290 dJdP[1][0][0] = 1. + 2. * u[0] * u[1];
291 dJdP[1][1][0] = u[0] * u[0] - 1.;
292 for (j = 0; j < 2; j++) {
293 vhv[j] = 0;
294 for (k = 0; k < 1; k++)
295 for (i = 0; i < 2; i++) vhv[j] += vl[i] * dJdP[i][j][k] * vr[k];
296 }
297
298 PetscCall(VecRestoreArrayRead(U, &u));
299 PetscCall(VecRestoreArrayRead(Vl[0], &vl));
300 PetscCall(VecRestoreArrayRead(Vr, &vr));
301 PetscCall(VecRestoreArray(VHV[0], &vhv));
302 PetscFunctionReturn(PETSC_SUCCESS);
303 }
304
IHessianProductPU(TS ts,PetscReal t,Vec U,Vec * Vl,Vec Vr,Vec * VHV,PetscCtx ctx)305 static PetscErrorCode IHessianProductPU(TS ts, PetscReal t, Vec U, Vec *Vl, Vec Vr, Vec *VHV, PetscCtx ctx)
306 {
307 const PetscScalar *vl, *vr, *u;
308 PetscScalar *vhv;
309 PetscScalar dJdU[2][1][2] = {{{0}}};
310 PetscInt i, j, k;
311
312 PetscFunctionBeginUser;
313 PetscCall(VecGetArrayRead(U, &u));
314 PetscCall(VecGetArrayRead(Vl[0], &vl));
315 PetscCall(VecGetArrayRead(Vr, &vr));
316 PetscCall(VecGetArray(VHV[0], &vhv));
317
318 dJdU[1][0][0] = 1. + 2. * u[1] * u[0];
319 dJdU[1][0][1] = u[0] * u[0] - 1.;
320 for (j = 0; j < 1; j++) {
321 vhv[j] = 0;
322 for (k = 0; k < 2; k++)
323 for (i = 0; i < 2; i++) vhv[j] += vl[i] * dJdU[i][j][k] * vr[k];
324 }
325
326 PetscCall(VecRestoreArrayRead(U, &u));
327 PetscCall(VecRestoreArrayRead(Vl[0], &vl));
328 PetscCall(VecRestoreArrayRead(Vr, &vr));
329 PetscCall(VecRestoreArray(VHV[0], &vhv));
330 PetscFunctionReturn(PETSC_SUCCESS);
331 }
332
IHessianProductPP(TS ts,PetscReal t,Vec U,Vec * Vl,Vec Vr,Vec * VHV,PetscCtx ctx)333 static PetscErrorCode IHessianProductPP(TS ts, PetscReal t, Vec U, Vec *Vl, Vec Vr, Vec *VHV, PetscCtx ctx)
334 {
335 PetscFunctionBeginUser;
336 PetscFunctionReturn(PETSC_SUCCESS);
337 }
338
339 /* Monitor timesteps and use interpolation to output at integer multiples of 0.1 */
Monitor(TS ts,PetscInt step,PetscReal t,Vec X,PetscCtx ctx)340 static PetscErrorCode Monitor(TS ts, PetscInt step, PetscReal t, Vec X, PetscCtx ctx)
341 {
342 const PetscScalar *x;
343 PetscReal tfinal, dt;
344 User user = (User)ctx;
345 Vec interpolatedX;
346
347 PetscFunctionBeginUser;
348 PetscCall(TSGetTimeStep(ts, &dt));
349 PetscCall(TSGetMaxTime(ts, &tfinal));
350
351 while (user->next_output <= t && user->next_output <= tfinal) {
352 PetscCall(VecDuplicate(X, &interpolatedX));
353 PetscCall(TSInterpolate(ts, user->next_output, interpolatedX));
354 PetscCall(VecGetArrayRead(interpolatedX, &x));
355 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "[%g] %" PetscInt_FMT " TS %g (dt = %g) X %g %g\n", (double)user->next_output, step, (double)t, (double)dt, (double)PetscRealPart(x[0]), (double)PetscRealPart(x[1])));
356 PetscCall(VecRestoreArrayRead(interpolatedX, &x));
357 PetscCall(VecDestroy(&interpolatedX));
358 user->next_output += PetscRealConstant(0.1);
359 }
360 PetscFunctionReturn(PETSC_SUCCESS);
361 }
362
main(int argc,char ** argv)363 int main(int argc, char **argv)
364 {
365 Vec P;
366 PetscBool monitor = PETSC_FALSE;
367 PetscScalar *x_ptr;
368 const PetscScalar *y_ptr;
369 PetscMPIInt size;
370 struct _n_User user;
371 Tao tao;
372 KSP ksp;
373 PC pc;
374
375 /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
376 Initialize program
377 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
378 PetscFunctionBeginUser;
379 PetscCall(PetscInitialize(&argc, &argv, NULL, help));
380 PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
381 PetscCheck(size == 1, PETSC_COMM_WORLD, PETSC_ERR_WRONG_MPI_SIZE, "This is a uniprocessor example only!");
382
383 /* Create TAO solver and set desired solution method */
384 PetscCall(TaoCreate(PETSC_COMM_WORLD, &tao));
385 PetscCall(TaoSetType(tao, TAOBQNLS));
386
387 /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
388 Set runtime options
389 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
390 user.next_output = 0.0;
391 user.mu = PetscRealConstant(1.0e3);
392 user.ftime = PetscRealConstant(0.5);
393 user.implicitform = PETSC_TRUE;
394
395 PetscCall(PetscOptionsGetBool(NULL, NULL, "-monitor", &monitor, NULL));
396 PetscCall(PetscOptionsGetReal(NULL, NULL, "-mu", &user.mu, NULL));
397 PetscCall(PetscOptionsGetBool(NULL, NULL, "-implicitform", &user.implicitform, NULL));
398
399 /* Create necessary matrix and vectors, solve same ODE on every process */
400 PetscCall(MatCreate(PETSC_COMM_WORLD, &user.A));
401 PetscCall(MatSetSizes(user.A, PETSC_DECIDE, PETSC_DECIDE, 2, 2));
402 PetscCall(MatSetFromOptions(user.A));
403 PetscCall(MatSetUp(user.A));
404 PetscCall(MatCreateVecs(user.A, &user.U, NULL));
405 PetscCall(MatCreateVecs(user.A, &user.Lambda[0], NULL));
406 PetscCall(MatCreateVecs(user.A, &user.Lambda2[0], NULL));
407 PetscCall(MatCreateVecs(user.A, &user.Ihp1[0], NULL));
408 PetscCall(MatCreateVecs(user.A, &user.Ihp2[0], NULL));
409
410 PetscCall(MatCreate(PETSC_COMM_WORLD, &user.Jacp));
411 PetscCall(MatSetSizes(user.Jacp, PETSC_DECIDE, PETSC_DECIDE, 2, 1));
412 PetscCall(MatSetFromOptions(user.Jacp));
413 PetscCall(MatSetUp(user.Jacp));
414 PetscCall(MatCreateVecs(user.Jacp, &user.Dir, NULL));
415 PetscCall(MatCreateVecs(user.Jacp, &user.Mup[0], NULL));
416 PetscCall(MatCreateVecs(user.Jacp, &user.Mup2[0], NULL));
417 PetscCall(MatCreateVecs(user.Jacp, &user.Ihp3[0], NULL));
418 PetscCall(MatCreateVecs(user.Jacp, &user.Ihp4[0], NULL));
419
420 /* Create timestepping solver context */
421 PetscCall(TSCreate(PETSC_COMM_WORLD, &user.ts));
422 PetscCall(TSSetEquationType(user.ts, TS_EQ_ODE_EXPLICIT)); /* less Jacobian evaluations when adjoint BEuler is used, otherwise no effect */
423 if (user.implicitform) {
424 PetscCall(TSSetIFunction(user.ts, NULL, IFunction, &user));
425 PetscCall(TSSetIJacobian(user.ts, user.A, user.A, IJacobian, &user));
426 PetscCall(TSSetType(user.ts, TSCN));
427 } else {
428 PetscCall(TSSetRHSFunction(user.ts, NULL, RHSFunction, &user));
429 PetscCall(TSSetRHSJacobian(user.ts, user.A, user.A, RHSJacobian, &user));
430 PetscCall(TSSetType(user.ts, TSRK));
431 }
432 PetscCall(TSSetRHSJacobianP(user.ts, user.Jacp, RHSJacobianP, &user));
433 PetscCall(TSSetMaxTime(user.ts, user.ftime));
434 PetscCall(TSSetExactFinalTime(user.ts, TS_EXACTFINALTIME_MATCHSTEP));
435
436 if (monitor) PetscCall(TSMonitorSet(user.ts, Monitor, &user, NULL));
437
438 /* Set ODE initial conditions */
439 PetscCall(VecGetArray(user.U, &x_ptr));
440 x_ptr[0] = 2.0;
441 x_ptr[1] = -2.0 / 3.0 + 10.0 / (81.0 * user.mu) - 292.0 / (2187.0 * user.mu * user.mu);
442 PetscCall(VecRestoreArray(user.U, &x_ptr));
443 PetscCall(TSSetTimeStep(user.ts, PetscRealConstant(0.001)));
444
445 /* Set runtime options */
446 PetscCall(TSSetFromOptions(user.ts));
447
448 PetscCall(TSSolve(user.ts, user.U));
449 PetscCall(VecGetArrayRead(user.U, &y_ptr));
450 user.ob[0] = y_ptr[0];
451 user.ob[1] = y_ptr[1];
452 PetscCall(VecRestoreArrayRead(user.U, &y_ptr));
453
454 /* Save trajectory of solution so that TSAdjointSolve() may be used.
455 Skip checkpointing for the first TSSolve since no adjoint run follows it.
456 */
457 PetscCall(TSSetSaveTrajectory(user.ts));
458
459 /* Optimization starts */
460 PetscCall(MatCreate(PETSC_COMM_WORLD, &user.H));
461 PetscCall(MatSetSizes(user.H, PETSC_DECIDE, PETSC_DECIDE, 1, 1));
462 PetscCall(MatSetUp(user.H)); /* Hessian should be symmetric. Do we need to do MatSetOption(user.H,MAT_SYMMETRIC,PETSC_TRUE) ? */
463
464 /* Set initial solution guess */
465 PetscCall(MatCreateVecs(user.Jacp, &P, NULL));
466 PetscCall(VecGetArray(P, &x_ptr));
467 x_ptr[0] = PetscRealConstant(1.2);
468 PetscCall(VecRestoreArray(P, &x_ptr));
469 PetscCall(TaoSetSolution(tao, P));
470
471 /* Set routine for function and gradient evaluation */
472 PetscCall(TaoSetObjectiveAndGradient(tao, NULL, FormFunctionGradient, (void *)&user));
473 PetscCall(TaoSetHessian(tao, user.H, user.H, FormHessian, (void *)&user));
474
475 /* Check for any TAO command line options */
476 PetscCall(TaoGetKSP(tao, &ksp));
477 if (ksp) {
478 PetscCall(KSPGetPC(ksp, &pc));
479 PetscCall(PCSetType(pc, PCNONE));
480 }
481 PetscCall(TaoSetFromOptions(tao));
482
483 PetscCall(TaoSolve(tao));
484
485 PetscCall(VecView(P, PETSC_VIEWER_STDOUT_WORLD));
486 /* Free TAO data structures */
487 PetscCall(TaoDestroy(&tao));
488
489 /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
490 Free work space. All PETSc objects should be destroyed when they
491 are no longer needed.
492 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
493 PetscCall(MatDestroy(&user.H));
494 PetscCall(MatDestroy(&user.A));
495 PetscCall(VecDestroy(&user.U));
496 PetscCall(MatDestroy(&user.Jacp));
497 PetscCall(VecDestroy(&user.Lambda[0]));
498 PetscCall(VecDestroy(&user.Mup[0]));
499 PetscCall(VecDestroy(&user.Lambda2[0]));
500 PetscCall(VecDestroy(&user.Mup2[0]));
501 PetscCall(VecDestroy(&user.Ihp1[0]));
502 PetscCall(VecDestroy(&user.Ihp2[0]));
503 PetscCall(VecDestroy(&user.Ihp3[0]));
504 PetscCall(VecDestroy(&user.Ihp4[0]));
505 PetscCall(VecDestroy(&user.Dir));
506 PetscCall(TSDestroy(&user.ts));
507 PetscCall(VecDestroy(&P));
508 PetscCall(PetscFinalize());
509 return 0;
510 }
511
512 /* ------------------------------------------------------------------ */
513 /*
514 FormFunctionGradient - Evaluates the function and corresponding gradient.
515
516 Input Parameters:
517 tao - the Tao context
518 X - the input vector
519 ptr - optional user-defined context, as set by TaoSetObjectiveAndGradient()
520
521 Output Parameters:
522 f - the newly evaluated function
523 G - the newly evaluated gradient
524 */
FormFunctionGradient(Tao tao,Vec P,PetscReal * f,Vec G,PetscCtx ctx)525 PetscErrorCode FormFunctionGradient(Tao tao, Vec P, PetscReal *f, Vec G, PetscCtx ctx)
526 {
527 User user_ptr = (User)ctx;
528 TS ts = user_ptr->ts;
529 PetscScalar *x_ptr, *g;
530 const PetscScalar *y_ptr;
531
532 PetscFunctionBeginUser;
533 PetscCall(VecGetArrayRead(P, &y_ptr));
534 user_ptr->mu = y_ptr[0];
535 PetscCall(VecRestoreArrayRead(P, &y_ptr));
536
537 PetscCall(TSSetTime(ts, 0.0));
538 PetscCall(TSSetStepNumber(ts, 0));
539 PetscCall(TSSetTimeStep(ts, PetscRealConstant(0.001))); /* can be overwritten by command line options */
540 PetscCall(TSSetFromOptions(ts));
541 PetscCall(VecGetArray(user_ptr->U, &x_ptr));
542 x_ptr[0] = 2.0;
543 x_ptr[1] = -2.0 / 3.0 + 10.0 / (81.0 * user_ptr->mu) - 292.0 / (2187.0 * user_ptr->mu * user_ptr->mu);
544 PetscCall(VecRestoreArray(user_ptr->U, &x_ptr));
545
546 PetscCall(TSSolve(ts, user_ptr->U));
547
548 PetscCall(VecGetArrayRead(user_ptr->U, &y_ptr));
549 *f = (y_ptr[0] - user_ptr->ob[0]) * (y_ptr[0] - user_ptr->ob[0]) + (y_ptr[1] - user_ptr->ob[1]) * (y_ptr[1] - user_ptr->ob[1]);
550
551 /* Reset initial conditions for the adjoint integration */
552 PetscCall(VecGetArray(user_ptr->Lambda[0], &x_ptr));
553 x_ptr[0] = 2. * (y_ptr[0] - user_ptr->ob[0]);
554 x_ptr[1] = 2. * (y_ptr[1] - user_ptr->ob[1]);
555 PetscCall(VecRestoreArrayRead(user_ptr->U, &y_ptr));
556 PetscCall(VecRestoreArray(user_ptr->Lambda[0], &x_ptr));
557
558 PetscCall(VecGetArray(user_ptr->Mup[0], &x_ptr));
559 x_ptr[0] = 0.0;
560 PetscCall(VecRestoreArray(user_ptr->Mup[0], &x_ptr));
561 PetscCall(TSSetCostGradients(ts, 1, user_ptr->Lambda, user_ptr->Mup));
562
563 PetscCall(TSAdjointSolve(ts));
564
565 PetscCall(VecGetArray(user_ptr->Mup[0], &x_ptr));
566 PetscCall(VecGetArrayRead(user_ptr->Lambda[0], &y_ptr));
567 PetscCall(VecGetArray(G, &g));
568 g[0] = y_ptr[1] * (-10.0 / (81.0 * user_ptr->mu * user_ptr->mu) + 2.0 * 292.0 / (2187.0 * user_ptr->mu * user_ptr->mu * user_ptr->mu)) + x_ptr[0];
569 PetscCall(VecRestoreArray(user_ptr->Mup[0], &x_ptr));
570 PetscCall(VecRestoreArrayRead(user_ptr->Lambda[0], &y_ptr));
571 PetscCall(VecRestoreArray(G, &g));
572 PetscFunctionReturn(PETSC_SUCCESS);
573 }
574
FormHessian(Tao tao,Vec P,Mat H,Mat Hpre,PetscCtx ctx)575 PetscErrorCode FormHessian(Tao tao, Vec P, Mat H, Mat Hpre, PetscCtx ctx)
576 {
577 User user_ptr = (User)ctx;
578 PetscScalar harr[1];
579 const PetscInt rows[1] = {0};
580 PetscInt col = 0;
581
582 PetscFunctionBeginUser;
583 PetscCall(Adjoint2(P, harr, user_ptr));
584 PetscCall(MatSetValues(H, 1, rows, 1, &col, harr, INSERT_VALUES));
585
586 PetscCall(MatAssemblyBegin(H, MAT_FINAL_ASSEMBLY));
587 PetscCall(MatAssemblyEnd(H, MAT_FINAL_ASSEMBLY));
588 if (H != Hpre) {
589 PetscCall(MatAssemblyBegin(Hpre, MAT_FINAL_ASSEMBLY));
590 PetscCall(MatAssemblyEnd(Hpre, MAT_FINAL_ASSEMBLY));
591 }
592 PetscFunctionReturn(PETSC_SUCCESS);
593 }
594
Adjoint2(Vec P,PetscScalar arr[],User ctx)595 PetscErrorCode Adjoint2(Vec P, PetscScalar arr[], User ctx)
596 {
597 TS ts = ctx->ts;
598 const PetscScalar *z_ptr;
599 PetscScalar *x_ptr, *y_ptr, dzdp, dzdp2;
600 Mat tlmsen;
601
602 PetscFunctionBeginUser;
603 /* Reset TSAdjoint so that AdjointSetUp will be called again */
604 PetscCall(TSAdjointReset(ts));
605
606 /* The directional vector should be 1 since it is one-dimensional */
607 PetscCall(VecGetArray(ctx->Dir, &x_ptr));
608 x_ptr[0] = 1.;
609 PetscCall(VecRestoreArray(ctx->Dir, &x_ptr));
610
611 PetscCall(VecGetArrayRead(P, &z_ptr));
612 ctx->mu = z_ptr[0];
613 PetscCall(VecRestoreArrayRead(P, &z_ptr));
614
615 dzdp = -10.0 / (81.0 * ctx->mu * ctx->mu) + 2.0 * 292.0 / (2187.0 * ctx->mu * ctx->mu * ctx->mu);
616 dzdp2 = 2. * 10.0 / (81.0 * ctx->mu * ctx->mu * ctx->mu) - 3.0 * 2.0 * 292.0 / (2187.0 * ctx->mu * ctx->mu * ctx->mu * ctx->mu);
617
618 PetscCall(TSSetTime(ts, 0.0));
619 PetscCall(TSSetStepNumber(ts, 0));
620 PetscCall(TSSetTimeStep(ts, PetscRealConstant(0.001))); /* can be overwritten by command line options */
621 PetscCall(TSSetFromOptions(ts));
622 PetscCall(TSSetCostHessianProducts(ts, 1, ctx->Lambda2, ctx->Mup2, ctx->Dir));
623
624 PetscCall(MatZeroEntries(ctx->Jacp));
625 PetscCall(MatSetValue(ctx->Jacp, 1, 0, dzdp, INSERT_VALUES));
626 PetscCall(MatAssemblyBegin(ctx->Jacp, MAT_FINAL_ASSEMBLY));
627 PetscCall(MatAssemblyEnd(ctx->Jacp, MAT_FINAL_ASSEMBLY));
628
629 PetscCall(TSAdjointSetForward(ts, ctx->Jacp));
630 PetscCall(VecGetArray(ctx->U, &y_ptr));
631 y_ptr[0] = 2.0;
632 y_ptr[1] = -2.0 / 3.0 + 10.0 / (81.0 * ctx->mu) - 292.0 / (2187.0 * ctx->mu * ctx->mu);
633 PetscCall(VecRestoreArray(ctx->U, &y_ptr));
634 PetscCall(TSSolve(ts, ctx->U));
635
636 /* Set terminal conditions for first- and second-order adjonts */
637 PetscCall(VecGetArrayRead(ctx->U, &z_ptr));
638 PetscCall(VecGetArray(ctx->Lambda[0], &y_ptr));
639 y_ptr[0] = 2. * (z_ptr[0] - ctx->ob[0]);
640 y_ptr[1] = 2. * (z_ptr[1] - ctx->ob[1]);
641 PetscCall(VecRestoreArray(ctx->Lambda[0], &y_ptr));
642 PetscCall(VecRestoreArrayRead(ctx->U, &z_ptr));
643 PetscCall(VecGetArray(ctx->Mup[0], &y_ptr));
644 y_ptr[0] = 0.0;
645 PetscCall(VecRestoreArray(ctx->Mup[0], &y_ptr));
646 PetscCall(TSForwardGetSensitivities(ts, NULL, &tlmsen));
647 PetscCall(MatDenseGetColumn(tlmsen, 0, &x_ptr));
648 PetscCall(VecGetArray(ctx->Lambda2[0], &y_ptr));
649 y_ptr[0] = 2. * x_ptr[0];
650 y_ptr[1] = 2. * x_ptr[1];
651 PetscCall(VecRestoreArray(ctx->Lambda2[0], &y_ptr));
652 PetscCall(VecGetArray(ctx->Mup2[0], &y_ptr));
653 y_ptr[0] = 0.0;
654 PetscCall(VecRestoreArray(ctx->Mup2[0], &y_ptr));
655 PetscCall(MatDenseRestoreColumn(tlmsen, &x_ptr));
656 PetscCall(TSSetCostGradients(ts, 1, ctx->Lambda, ctx->Mup));
657 if (ctx->implicitform) {
658 PetscCall(TSSetIHessianProduct(ts, ctx->Ihp1, IHessianProductUU, ctx->Ihp2, IHessianProductUP, ctx->Ihp3, IHessianProductPU, ctx->Ihp4, IHessianProductPP, ctx));
659 } else {
660 PetscCall(TSSetRHSHessianProduct(ts, ctx->Ihp1, RHSHessianProductUU, ctx->Ihp2, RHSHessianProductUP, ctx->Ihp3, RHSHessianProductPU, ctx->Ihp4, RHSHessianProductPP, ctx));
661 }
662 PetscCall(TSAdjointSolve(ts));
663
664 PetscCall(VecGetArray(ctx->Lambda[0], &x_ptr));
665 PetscCall(VecGetArray(ctx->Lambda2[0], &y_ptr));
666 PetscCall(VecGetArrayRead(ctx->Mup2[0], &z_ptr));
667
668 arr[0] = x_ptr[1] * dzdp2 + y_ptr[1] * dzdp2 + z_ptr[0];
669
670 PetscCall(VecRestoreArray(ctx->Lambda2[0], &x_ptr));
671 PetscCall(VecRestoreArray(ctx->Lambda2[0], &y_ptr));
672 PetscCall(VecRestoreArrayRead(ctx->Mup2[0], &z_ptr));
673
674 /* Disable second-order adjoint mode */
675 PetscCall(TSAdjointReset(ts));
676 PetscCall(TSAdjointResetForward(ts));
677 PetscFunctionReturn(PETSC_SUCCESS);
678 }
679
680 /*TEST
681 build:
682 requires: !complex !single
683 test:
684 args: -implicitform 0 -ts_type rk -ts_adapt_type none -mu 10 -ts_time_step 0.1 -viewer_binary_skip_info -tao_monitor -tao_view
685 output_file: output/ex20opt_p_1.out
686
687 test:
688 suffix: 2
689 args: -implicitform 0 -ts_type rk -ts_adapt_type none -mu 10 -ts_time_step 0.01 -viewer_binary_skip_info -tao_monitor -tao_type bntr -tao_bnk_pc_type none
690 output_file: output/ex20opt_p_2.out
691
692 test:
693 suffix: 3
694 args: -ts_type cn -ts_adapt_type none -mu 100 -ts_time_step 0.01 -viewer_binary_skip_info -tao_monitor -tao_view
695 output_file: output/ex20opt_p_3.out
696
697 test:
698 suffix: 4
699 args: -ts_type cn -ts_adapt_type none -mu 100 -ts_time_step 0.01 -viewer_binary_skip_info -tao_monitor -tao_type bntr -tao_bnk_pc_type none
700 output_file: output/ex20opt_p_4.out
701
702 TEST*/
703