Date: 2025-03-29T05:58:51+00:00

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*/