Actual source code: ex16adj.cxx

  1: static char help[] = "Demonstrates automatic Jacobian generation using ADOL-C for an adjoint sensitivity analysis of the van der Pol equation.\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 ex16adj for a description of the problem being solved.
 13:   ------------------------------------------------------------------------- */

 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:   PetscReal tprev;

 26:   /* Automatic differentiation support */
 27:   AdolcCtx *adctx;
 28: };

 30: /*
 31:   'Passive' RHS function, used in residual evaluations during the time integration.
 32: */
 33: static PetscErrorCode RHSFunctionPassive(TS ts, PetscReal t, Vec X, Vec F, void *ctx)
 34: {
 35:   User               user = (User)ctx;
 36:   PetscScalar       *f;
 37:   const PetscScalar *x;

 39:   PetscFunctionBeginUser;
 40:   PetscCall(VecGetArrayRead(X, &x));
 41:   PetscCall(VecGetArray(F, &f));
 42:   f[0] = x[1];
 43:   f[1] = user->mu * (1. - x[0] * x[0]) * x[1] - x[0];
 44:   PetscCall(VecRestoreArrayRead(X, &x));
 45:   PetscCall(VecRestoreArray(F, &f));
 46:   PetscFunctionReturn(PETSC_SUCCESS);
 47: }

 49: /*
 50:   Trace RHS to mark on tape 1 the dependence of f upon x. This tape is used in generating the
 51:   Jacobian transform.
 52: */
 53: static PetscErrorCode RHSFunctionActive(TS ts, PetscReal t, Vec X, Vec F, void *ctx)
 54: {
 55:   User               user = (User)ctx;
 56:   PetscScalar       *f;
 57:   const PetscScalar *x;

 59:   adouble f_a[2]; /* 'active' double for dependent variables */
 60:   adouble x_a[2]; /* 'active' double for independent variables */

 62:   PetscFunctionBeginUser;
 63:   PetscCall(VecGetArrayRead(X, &x));
 64:   PetscCall(VecGetArray(F, &f));

 66:   /* Start of active section */
 67:   trace_on(1);
 68:   x_a[0] <<= x[0];
 69:   x_a[1] <<= x[1]; /* Mark independence */
 70:   f_a[0] = x_a[1];
 71:   f_a[1] = user->mu * (1. - x_a[0] * x_a[0]) * x_a[1] - x_a[0];
 72:   f_a[0] >>= f[0];
 73:   f_a[1] >>= f[1]; /* Mark dependence */
 74:   trace_off();
 75:   /* End of active section */

 77:   PetscCall(VecRestoreArrayRead(X, &x));
 78:   PetscCall(VecRestoreArray(F, &f));
 79:   PetscFunctionReturn(PETSC_SUCCESS);
 80: }

 82: /*
 83:   Trace RHS again to mark on tape 2 the dependence of f upon the parameter mu. This tape is used in
 84:   generating JacobianP.
 85: */
 86: static PetscErrorCode RHSFunctionActiveP(TS ts, PetscReal t, Vec X, Vec F, void *ctx)
 87: {
 88:   User               user = (User)ctx;
 89:   PetscScalar       *f;
 90:   const PetscScalar *x;

 92:   adouble f_a[2];       /* 'active' double for dependent variables */
 93:   adouble x_a[2], mu_a; /* 'active' double for independent variables */

 95:   PetscFunctionBeginUser;
 96:   PetscCall(VecGetArrayRead(X, &x));
 97:   PetscCall(VecGetArray(F, &f));

 99:   /* Start of active section */
100:   trace_on(3);
101:   x_a[0] <<= x[0];
102:   x_a[1] <<= x[1];
103:   mu_a <<= user->mu; /* Mark independence */
104:   f_a[0] = x_a[1];
105:   f_a[1] = mu_a * (1. - x_a[0] * x_a[0]) * x_a[1] - x_a[0];
106:   f_a[0] >>= f[0];
107:   f_a[1] >>= f[1]; /* Mark dependence */
108:   trace_off();
109:   /* End of active section */

111:   PetscCall(VecRestoreArrayRead(X, &x));
112:   PetscCall(VecRestoreArray(F, &f));
113:   PetscFunctionReturn(PETSC_SUCCESS);
114: }

116: /*
117:   Compute the Jacobian w.r.t. x using PETSc-ADOL-C driver for explicit TS.
118: */
119: static PetscErrorCode RHSJacobian(TS ts, PetscReal t, Vec X, Mat A, Mat B, void *ctx)
120: {
121:   User               user = (User)ctx;
122:   const PetscScalar *x;

124:   PetscFunctionBeginUser;
125:   PetscCall(VecGetArrayRead(X, &x));
126:   PetscCall(PetscAdolcComputeRHSJacobian(1, A, x, user->adctx));
127:   PetscCall(VecRestoreArrayRead(X, &x));
128:   PetscFunctionReturn(PETSC_SUCCESS);
129: }

131: /*
132:   Compute the Jacobian w.r.t. mu using PETSc-ADOL-C driver for explicit TS.
133: */
134: static PetscErrorCode RHSJacobianP(TS ts, PetscReal t, Vec X, Mat A, void *ctx)
135: {
136:   User               user = (User)ctx;
137:   const PetscScalar *x;

139:   PetscFunctionBeginUser;
140:   PetscCall(VecGetArrayRead(X, &x));
141:   PetscCall(PetscAdolcComputeRHSJacobianP(3, A, x, &user->mu, user->adctx));
142:   PetscCall(VecRestoreArrayRead(X, &x));
143:   PetscFunctionReturn(PETSC_SUCCESS);
144: }

146: /*
147:   Monitor timesteps and use interpolation to output at integer multiples of 0.1
148: */
149: static PetscErrorCode Monitor(TS ts, PetscInt step, PetscReal t, Vec X, void *ctx)
150: {
151:   const PetscScalar *x;
152:   PetscReal          tfinal, dt, tprev;
153:   User               user = (User)ctx;

155:   PetscFunctionBeginUser;
156:   PetscCall(TSGetTimeStep(ts, &dt));
157:   PetscCall(TSGetMaxTime(ts, &tfinal));
158:   PetscCall(TSGetPrevTime(ts, &tprev));
159:   PetscCall(VecGetArrayRead(X, &x));
160:   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])));
161:   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "t %.6f (tprev = %.6f) \n", (double)t, (double)tprev));
162:   PetscCall(VecRestoreArrayRead(X, &x));
163:   PetscFunctionReturn(PETSC_SUCCESS);
164: }

166: int main(int argc, char **argv)
167: {
168:   TS             ts;   /* nonlinear solver */
169:   Vec            x;    /* solution, residual vectors */
170:   Mat            A;    /* Jacobian matrix */
171:   Mat            Jacp; /* JacobianP matrix */
172:   PetscInt       steps;
173:   PetscReal      ftime   = 0.5;
174:   PetscBool      monitor = PETSC_FALSE;
175:   PetscScalar   *x_ptr;
176:   PetscMPIInt    size;
177:   struct _n_User user;
178:   AdolcCtx      *adctx;
179:   Vec            lambda[2], mu[2], r;

181:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
182:      Initialize program
183:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
184:   PetscFunctionBeginUser;
185:   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
186:   PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
187:   PetscCheck(size == 1, PETSC_COMM_WORLD, PETSC_ERR_WRONG_MPI_SIZE, "This is a uniprocessor example only!");

189:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
190:     Set runtime options and create AdolcCtx
191:     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
192:   PetscCall(PetscNew(&adctx));
193:   user.mu           = 1;
194:   user.next_output  = 0.0;
195:   adctx->m          = 2;
196:   adctx->n          = 2;
197:   adctx->p          = 2;
198:   adctx->num_params = 1;
199:   user.adctx        = adctx;

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

204:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
205:     Create necessary matrix and vectors, solve same ODE on every process
206:     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
207:   PetscCall(MatCreate(PETSC_COMM_WORLD, &A));
208:   PetscCall(MatSetSizes(A, PETSC_DECIDE, PETSC_DECIDE, 2, 2));
209:   PetscCall(MatSetFromOptions(A));
210:   PetscCall(MatSetUp(A));
211:   PetscCall(MatCreateVecs(A, &x, NULL));

213:   PetscCall(MatCreate(PETSC_COMM_WORLD, &Jacp));
214:   PetscCall(MatSetSizes(Jacp, PETSC_DECIDE, PETSC_DECIDE, 2, 1));
215:   PetscCall(MatSetFromOptions(Jacp));
216:   PetscCall(MatSetUp(Jacp));

218:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
219:      Create timestepping solver context
220:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
221:   PetscCall(TSCreate(PETSC_COMM_WORLD, &ts));
222:   PetscCall(TSSetType(ts, TSRK));
223:   PetscCall(TSSetRHSFunction(ts, NULL, RHSFunctionPassive, &user));

225:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
226:      Set initial conditions
227:    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
228:   PetscCall(VecGetArray(x, &x_ptr));
229:   x_ptr[0] = 2;
230:   x_ptr[1] = 0.66666654321;
231:   PetscCall(VecRestoreArray(x, &x_ptr));

233:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
234:      Trace just once on each tape and put zeros on Jacobian diagonal
235:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
236:   PetscCall(VecDuplicate(x, &r));
237:   PetscCall(RHSFunctionActive(ts, 0., x, r, &user));
238:   PetscCall(RHSFunctionActiveP(ts, 0., x, r, &user));
239:   PetscCall(VecSet(r, 0));
240:   PetscCall(MatDiagonalSet(A, r, INSERT_VALUES));
241:   PetscCall(VecDestroy(&r));

243:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
244:      Set RHS Jacobian for the adjoint integration
245:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
246:   PetscCall(TSSetRHSJacobian(ts, A, A, RHSJacobian, &user));
247:   PetscCall(TSSetMaxTime(ts, ftime));
248:   PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP));
249:   if (monitor) PetscCall(TSMonitorSet(ts, Monitor, &user, NULL));
250:   PetscCall(TSSetTimeStep(ts, .001));

252:   /*
253:     Have the TS save its trajectory so that TSAdjointSolve() may be used
254:   */
255:   PetscCall(TSSetSaveTrajectory(ts));

257:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
258:      Set runtime options
259:    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
260:   PetscCall(TSSetFromOptions(ts));

262:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
263:      Solve nonlinear system
264:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
265:   PetscCall(TSSolve(ts, x));
266:   PetscCall(TSGetSolveTime(ts, &ftime));
267:   PetscCall(TSGetStepNumber(ts, &steps));
268:   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "mu %g, steps %" PetscInt_FMT ", ftime %g\n", (double)user.mu, steps, (double)ftime));
269:   PetscCall(VecView(x, PETSC_VIEWER_STDOUT_WORLD));

271:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
272:      Start the Adjoint model
273:      - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
274:   PetscCall(MatCreateVecs(A, &lambda[0], NULL));
275:   PetscCall(MatCreateVecs(A, &lambda[1], NULL));
276:   /*   Reset initial conditions for the adjoint integration */
277:   PetscCall(VecGetArray(lambda[0], &x_ptr));
278:   x_ptr[0] = 1.0;
279:   x_ptr[1] = 0.0;
280:   PetscCall(VecRestoreArray(lambda[0], &x_ptr));
281:   PetscCall(VecGetArray(lambda[1], &x_ptr));
282:   x_ptr[0] = 0.0;
283:   x_ptr[1] = 1.0;
284:   PetscCall(VecRestoreArray(lambda[1], &x_ptr));

286:   PetscCall(MatCreateVecs(Jacp, &mu[0], NULL));
287:   PetscCall(MatCreateVecs(Jacp, &mu[1], NULL));
288:   PetscCall(VecGetArray(mu[0], &x_ptr));
289:   x_ptr[0] = 0.0;
290:   PetscCall(VecRestoreArray(mu[0], &x_ptr));
291:   PetscCall(VecGetArray(mu[1], &x_ptr));
292:   x_ptr[0] = 0.0;
293:   PetscCall(VecRestoreArray(mu[1], &x_ptr));
294:   PetscCall(TSSetCostGradients(ts, 2, lambda, mu));

296:   /*   Set RHS JacobianP */
297:   PetscCall(TSSetRHSJacobianP(ts, Jacp, RHSJacobianP, &user));

299:   PetscCall(TSAdjointSolve(ts));

301:   PetscCall(VecView(lambda[0], PETSC_VIEWER_STDOUT_WORLD));
302:   PetscCall(VecView(lambda[1], PETSC_VIEWER_STDOUT_WORLD));
303:   PetscCall(VecView(mu[0], PETSC_VIEWER_STDOUT_WORLD));
304:   PetscCall(VecView(mu[1], PETSC_VIEWER_STDOUT_WORLD));

306:   /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
307:      Free work space.  All PETSc objects should be destroyed when they
308:      are no longer needed.
309:    - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
310:   PetscCall(MatDestroy(&A));
311:   PetscCall(MatDestroy(&Jacp));
312:   PetscCall(VecDestroy(&x));
313:   PetscCall(VecDestroy(&lambda[0]));
314:   PetscCall(VecDestroy(&lambda[1]));
315:   PetscCall(VecDestroy(&mu[0]));
316:   PetscCall(VecDestroy(&mu[1]));
317:   PetscCall(TSDestroy(&ts));
318:   PetscCall(PetscFree(adctx));
319:   PetscCall(PetscFinalize());
320:   return 0;
321: }

323: /*TEST

325:   build:
326:     requires: double !complex adolc

328:   test:
329:     suffix: 1
330:     args: -ts_max_steps 10 -ts_monitor -ts_adjoint_monitor
331:     output_file: output/ex16adj_1.out

333:   test:
334:     suffix: 2
335:     args: -ts_max_steps 10 -ts_monitor -ts_adjoint_monitor -mu 5
336:     output_file: output/ex16adj_2.out

338:   test:
339:     suffix: 3
340:     args: -ts_max_steps 10 -monitor
341:     output_file: output/ex16adj_3.out

343:   test:
344:     suffix: 4
345:     args: -ts_max_steps 10 -monitor -mu 5
346:     output_file: output/ex16adj_4.out

348: TEST*/