Actual source code: ex16opt_ic.cxx

  1: static char help[] = "Demonstrates automatic Jacobian generation using ADOL-C for an ODE-constrained optimization problem.\n\
  2: Input parameters include:\n\
  3:       -mu : stiffness parameter\n\n";

  5: /*
  6:    REQUIRES configuration of PETSc with option --download-adolc.

  8:    For documentation on ADOL-C, see
  9:      $PETSC_ARCH/externalpackages/ADOL-C-2.6.0/ADOL-C/doc/adolc-manual.pdf
 10: */
 11: /* ------------------------------------------------------------------------
 12:   See ex16opt_ic for a description of the problem being solved.
 13:   ------------------------------------------------------------------------- */
 14: #include <petsctao.h>
 15: #include <petscts.h>
 16: #include <petscmat.h>
 17: #include "adolc-utils/drivers.cxx"
 18: #include <adolc/adolc.h>

 20: typedef struct _n_User *User;
 21: struct _n_User {
 22:   PetscReal mu;
 23:   PetscReal next_output;
 24:   PetscInt  steps;

 26:   /* Sensitivity analysis support */
 27:   PetscReal ftime, x_ob[2];
 28:   Mat       A;            /* Jacobian matrix */
 29:   Vec       x, lambda[2]; /* adjoint variables */

 31:   /* Automatic differentiation support */
 32:   AdolcCtx *adctx;
 33: };

 35: PetscErrorCode FormFunctionGradient(Tao, Vec, PetscReal *, Vec, void *);

 37: /*
 38:   'Passive' RHS function, used in residual evaluations during the time integration.
 39: */
 40: static PetscErrorCode RHSFunctionPassive(TS ts, PetscReal t, Vec X, Vec F, void *ctx)
 41: {
 42:   User               user = (User)ctx;
 43:   PetscScalar       *f;
 44:   const PetscScalar *x;

 46:   PetscFunctionBeginUser;
 47:   PetscCall(VecGetArrayRead(X, &x));
 48:   PetscCall(VecGetArray(F, &f));
 49:   f[0] = x[1];
 50:   f[1] = user->mu * (1. - x[0] * x[0]) * x[1] - x[0];
 51:   PetscCall(VecRestoreArrayRead(X, &x));
 52:   PetscCall(VecRestoreArray(F, &f));
 53:   PetscFunctionReturn(PETSC_SUCCESS);
 54: }

 56: /*
 57:   Trace RHS to mark on tape 1 the dependence of f upon x. This tape is used in generating the
 58:   Jacobian transform.
 59: */
 60: static PetscErrorCode RHSFunctionActive(TS ts, PetscReal t, Vec X, Vec F, void *ctx)
 61: {
 62:   User               user = (User)ctx;
 63:   PetscReal          mu   = user->mu;
 64:   PetscScalar       *f;
 65:   const PetscScalar *x;

 67:   adouble f_a[2]; /* adouble for dependent variables */
 68:   adouble x_a[2]; /* adouble for independent variables */

 70:   PetscFunctionBeginUser;
 71:   PetscCall(VecGetArrayRead(X, &x));
 72:   PetscCall(VecGetArray(F, &f));

 74:   trace_on(1); /* Start of active section */
 75:   x_a[0] <<= x[0];
 76:   x_a[1] <<= x[1]; /* Mark as independent */
 77:   f_a[0] = x_a[1];
 78:   f_a[1] = mu * (1. - x_a[0] * x_a[0]) * x_a[1] - x_a[0];
 79:   f_a[0] >>= f[0];
 80:   f_a[1] >>= f[1]; /* Mark as dependent */
 81:   trace_off(1);    /* End of active section */

 83:   PetscCall(VecRestoreArrayRead(X, &x));
 84:   PetscCall(VecRestoreArray(F, &f));
 85:   PetscFunctionReturn(PETSC_SUCCESS);
 86: }

 88: /*
 89:   Compute the Jacobian w.r.t. x using PETSc-ADOL-C driver.
 90: */
 91: static PetscErrorCode RHSJacobian(TS ts, PetscReal t, Vec X, Mat A, Mat B, void *ctx)
 92: {
 93:   User               user = (User)ctx;
 94:   const PetscScalar *x;

 96:   PetscFunctionBeginUser;
 97:   PetscCall(VecGetArrayRead(X, &x));
 98:   PetscCall(PetscAdolcComputeRHSJacobian(1, A, x, user->adctx));
 99:   PetscCall(VecRestoreArrayRead(X, &x));
100:   PetscFunctionReturn(PETSC_SUCCESS);
101: }

103: /*
104:   Monitor timesteps and use interpolation to output at integer multiples of 0.1
105: */
106: static PetscErrorCode Monitor(TS ts, PetscInt step, PetscReal t, Vec X, void *ctx)
107: {
108:   const PetscScalar *x;
109:   PetscReal          tfinal, dt, tprev;
110:   User               user = (User)ctx;

112:   PetscFunctionBeginUser;
113:   PetscCall(TSGetTimeStep(ts, &dt));
114:   PetscCall(TSGetMaxTime(ts, &tfinal));
115:   PetscCall(TSGetPrevTime(ts, &tprev));
116:   PetscCall(VecGetArrayRead(X, &x));
117:   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "[%.1f] %" PetscInt_FMT " TS %.6f (dt = %.6f) X % 12.6e % 12.6e\n", (double)user->next_output, step, (double)t, (double)dt, (double)PetscRealPart(x[0]), (double)PetscRealPart(x[1])));
118:   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "t %.6f (tprev = %.6f) \n", (double)t, (double)tprev));
119:   PetscCall(VecGetArrayRead(X, &x));
120:   PetscFunctionReturn(PETSC_SUCCESS);
121: }

123: int main(int argc, char **argv)
124: {
125:   TS             ts = NULL; /* nonlinear solver */
126:   Vec            ic, r;
127:   PetscBool      monitor = PETSC_FALSE;
128:   PetscScalar   *x_ptr;
129:   PetscMPIInt    size;
130:   struct _n_User user;
131:   AdolcCtx      *adctx;
132:   Tao            tao;
133:   KSP            ksp;
134:   PC             pc;

136:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
137:      Initialize program
138:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
139:   PetscFunctionBeginUser;
140:   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
141:   PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
142:   PetscCheck(size == 1, PETSC_COMM_WORLD, PETSC_ERR_WRONG_MPI_SIZE, "This is a uniprocessor example only!");

144:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
145:     Set runtime options and create AdolcCtx
146:     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
147:   PetscCall(PetscNew(&adctx));
148:   user.mu          = 1.0;
149:   user.next_output = 0.0;
150:   user.steps       = 0;
151:   user.ftime       = 0.5;
152:   adctx->m         = 2;
153:   adctx->n         = 2;
154:   adctx->p         = 2;
155:   user.adctx       = adctx;

157:   PetscCall(PetscOptionsGetReal(NULL, NULL, "-mu", &user.mu, NULL));
158:   PetscCall(PetscOptionsGetBool(NULL, NULL, "-monitor", &monitor, NULL));

160:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
161:     Create necessary matrix and vectors, solve same ODE on every process
162:     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
163:   PetscCall(MatCreate(PETSC_COMM_WORLD, &user.A));
164:   PetscCall(MatSetSizes(user.A, PETSC_DECIDE, PETSC_DECIDE, 2, 2));
165:   PetscCall(MatSetFromOptions(user.A));
166:   PetscCall(MatSetUp(user.A));
167:   PetscCall(MatCreateVecs(user.A, &user.x, NULL));

169:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
170:      Set initial conditions
171:    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
172:   PetscCall(VecGetArray(user.x, &x_ptr));
173:   x_ptr[0] = 2.0;
174:   x_ptr[1] = 0.66666654321;
175:   PetscCall(VecRestoreArray(user.x, &x_ptr));

177:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
178:      Trace just once on each tape and put zeros on Jacobian diagonal
179:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
180:   PetscCall(VecDuplicate(user.x, &r));
181:   PetscCall(RHSFunctionActive(ts, 0., user.x, r, &user));
182:   PetscCall(VecSet(r, 0));
183:   PetscCall(MatDiagonalSet(user.A, r, INSERT_VALUES));
184:   PetscCall(VecDestroy(&r));

186:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
187:      Create timestepping solver context
188:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
189:   PetscCall(TSCreate(PETSC_COMM_WORLD, &ts));
190:   PetscCall(TSSetType(ts, TSRK));
191:   PetscCall(TSSetRHSFunction(ts, NULL, RHSFunctionPassive, &user));
192:   PetscCall(TSSetRHSJacobian(ts, user.A, user.A, RHSJacobian, &user));
193:   PetscCall(TSSetMaxTime(ts, user.ftime));
194:   PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP));
195:   if (monitor) PetscCall(TSMonitorSet(ts, Monitor, &user, NULL));

197:   PetscCall(TSSetTime(ts, 0.0));
198:   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "mu %g, steps %" PetscInt_FMT ", ftime %g\n", (double)user.mu, user.steps, (double)user.ftime));

200:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
201:     Save trajectory of solution so that TSAdjointSolve() may be used
202:    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
203:   PetscCall(TSSetSaveTrajectory(ts));

205:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
206:      Set runtime options
207:    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
208:   PetscCall(TSSetFromOptions(ts));

210:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
211:      Solve nonlinear system
212:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
213:   PetscCall(TSSolve(ts, user.x));
214:   PetscCall(TSGetSolveTime(ts, &user.ftime));
215:   PetscCall(TSGetStepNumber(ts, &user.steps));
216:   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "mu %g, steps %" PetscInt_FMT ", ftime %g\n", (double)user.mu, user.steps, (double)user.ftime));

218:   PetscCall(VecGetArray(user.x, &x_ptr));
219:   user.x_ob[0] = x_ptr[0];
220:   user.x_ob[1] = x_ptr[1];
221:   PetscCall(VecRestoreArray(user.x, &x_ptr));

223:   PetscCall(MatCreateVecs(user.A, &user.lambda[0], NULL));

225:   /* Create TAO solver and set desired solution method */
226:   PetscCall(TaoCreate(PETSC_COMM_WORLD, &tao));
227:   PetscCall(TaoSetType(tao, TAOCG));

229:   /* Set initial solution guess */
230:   PetscCall(MatCreateVecs(user.A, &ic, NULL));
231:   PetscCall(VecGetArray(ic, &x_ptr));
232:   x_ptr[0] = 2.1;
233:   x_ptr[1] = 0.7;
234:   PetscCall(VecRestoreArray(ic, &x_ptr));

236:   PetscCall(TaoSetSolution(tao, ic));

238:   /* Set routine for function and gradient evaluation */
239:   PetscCall(TaoSetObjectiveAndGradient(tao, NULL, FormFunctionGradient, (void *)&user));

241:   /* Check for any TAO command line options */
242:   PetscCall(TaoSetFromOptions(tao));
243:   PetscCall(TaoGetKSP(tao, &ksp));
244:   if (ksp) {
245:     PetscCall(KSPGetPC(ksp, &pc));
246:     PetscCall(PCSetType(pc, PCNONE));
247:   }

249:   PetscCall(TaoSetTolerances(tao, 1e-10, PETSC_CURRENT, PETSC_CURRENT));

251:   /* SOLVE THE APPLICATION */
252:   PetscCall(TaoSolve(tao));

254:   /* Free TAO data structures */
255:   PetscCall(TaoDestroy(&tao));

257:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
258:      Free work space.  All PETSc objects should be destroyed when they
259:      are no longer needed.
260:    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
261:   PetscCall(MatDestroy(&user.A));
262:   PetscCall(VecDestroy(&user.x));
263:   PetscCall(VecDestroy(&user.lambda[0]));
264:   PetscCall(TSDestroy(&ts));
265:   PetscCall(VecDestroy(&ic));
266:   PetscCall(PetscFree(adctx));
267:   PetscCall(PetscFinalize());
268:   return 0;
269: }

271: /* ------------------------------------------------------------------ */
272: /*
273:    FormFunctionGradient - Evaluates the function and corresponding gradient.

275:    Input Parameters:
276:    tao - the Tao context
277:    X   - the input vector
278:    ptr - optional user-defined context, as set by TaoSetObjectiveAndGradient()

280:    Output Parameters:
281:    f   - the newly evaluated function
282:    G   - the newly evaluated gradient
283: */
284: PetscErrorCode FormFunctionGradient(Tao tao, Vec IC, PetscReal *f, Vec G, void *ctx)
285: {
286:   User         user = (User)ctx;
287:   TS           ts;
288:   PetscScalar *x_ptr, *y_ptr;

290:   PetscFunctionBeginUser;
291:   PetscCall(VecCopy(IC, user->x));

293:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
294:      Create timestepping solver context
295:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
296:   PetscCall(TSCreate(PETSC_COMM_WORLD, &ts));
297:   PetscCall(TSSetType(ts, TSRK));
298:   PetscCall(TSSetRHSFunction(ts, NULL, RHSFunctionPassive, user));
299:   /*   Set RHS Jacobian  for the adjoint integration */
300:   PetscCall(TSSetRHSJacobian(ts, user->A, user->A, RHSJacobian, user));

302:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
303:      Set time
304:    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
305:   PetscCall(TSSetTime(ts, 0.0));
306:   PetscCall(TSSetTimeStep(ts, .001));
307:   PetscCall(TSSetMaxTime(ts, 0.5));
308:   PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP));

310:   PetscCall(TSSetTolerances(ts, 1e-7, NULL, 1e-7, NULL));

312:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
313:     Save trajectory of solution so that TSAdjointSolve() may be used
314:    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
315:   PetscCall(TSSetSaveTrajectory(ts));

317:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
318:      Set runtime options
319:    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
320:   PetscCall(TSSetFromOptions(ts));

322:   PetscCall(TSSolve(ts, user->x));
323:   PetscCall(TSGetSolveTime(ts, &user->ftime));
324:   PetscCall(TSGetStepNumber(ts, &user->steps));
325:   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "mu %.6f, steps %" PetscInt_FMT ", ftime %g\n", (double)user->mu, user->steps, (double)user->ftime));

327:   PetscCall(VecGetArray(user->x, &x_ptr));
328:   *f = (x_ptr[0] - user->x_ob[0]) * (x_ptr[0] - user->x_ob[0]) + (x_ptr[1] - user->x_ob[1]) * (x_ptr[1] - user->x_ob[1]);
329:   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Observed value y_ob=[%f; %f], ODE solution y=[%f;%f], Cost function f=%f\n", (double)user->x_ob[0], (double)user->x_ob[1], (double)x_ptr[0], (double)x_ptr[1], (double)(*f)));

331:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
332:      Adjoint model starts here
333:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
334:   /*   Redet initial conditions for the adjoint integration */
335:   PetscCall(VecGetArray(user->lambda[0], &y_ptr));
336:   y_ptr[0] = 2. * (x_ptr[0] - user->x_ob[0]);
337:   y_ptr[1] = 2. * (x_ptr[1] - user->x_ob[1]);
338:   PetscCall(VecRestoreArray(user->lambda[0], &y_ptr));
339:   PetscCall(VecRestoreArray(user->x, &x_ptr));
340:   PetscCall(TSSetCostGradients(ts, 1, user->lambda, NULL));

342:   PetscCall(TSAdjointSolve(ts));

344:   PetscCall(VecCopy(user->lambda[0], G));

346:   PetscCall(TSDestroy(&ts));
347:   PetscFunctionReturn(PETSC_SUCCESS);
348: }

350: /*TEST

352:   build:
353:     requires: double !complex adolc

355:   test:
356:     suffix: 1
357:     args: -ts_rhs_jacobian_test_mult_transpose FALSE -tao_max_it 2 -ts_rhs_jacobian_test_mult FALSE
358:     output_file: output/ex16opt_ic_1.out

360: TEST*/