Actual source code: ex20opt_p.c

      1: static char help[] = "Solves the van der Pol equation.\n\
      2: Input parameters include:\n";

      4: /* ------------------------------------------------------------------------

      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>

     15: typedef struct _n_User *User;
     16: struct _n_User {
     17:   TS        ts;
     18:   PetscReal mu;
     19:   PetscReal next_output;

     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: };

     37: PetscErrorCode FormFunctionGradient(Tao, Vec, PetscReal *, Vec, void *);
     38: PetscErrorCode FormHessian(Tao, Vec, Mat, Mat, void *);
     39: PetscErrorCode Adjoint2(Vec, PetscScalar[], User);

     41: /* ----------------------- Explicit form of the ODE  -------------------- */

     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;

     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: }

     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;

     67:   PetscFunctionBeginUser;
     68:   PetscCall(VecGetArrayRead(U, &u));

     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:   }

     82:   PetscCall(VecRestoreArrayRead(U, &u));
     83:   PetscFunctionReturn(PETSC_SUCCESS);
     84: }

     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;

     94:   PetscFunctionBeginUser;
     95:   PetscCall(VecGetArrayRead(U, &u));
     96:   PetscCall(VecGetArrayRead(Vl[0], &vl));
     97:   PetscCall(VecGetArrayRead(Vr, &vr));
     98:   PetscCall(VecGetArray(VHV[0], &vhv));

    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:   }

    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: }

    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;

    123:   PetscFunctionBeginUser;
    124:   PetscCall(VecGetArrayRead(U, &u));
    125:   PetscCall(VecGetArrayRead(Vl[0], &vl));
    126:   PetscCall(VecGetArrayRead(Vr, &vr));
    127:   PetscCall(VecGetArray(VHV[0], &vhv));

    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:   }

    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: }

    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;

    151:   PetscFunctionBeginUser;
    152:   PetscCall(VecGetArrayRead(U, &u));
    153:   PetscCall(VecGetArrayRead(Vl[0], &vl));
    154:   PetscCall(VecGetArrayRead(Vr, &vr));
    155:   PetscCall(VecGetArray(VHV[0], &vhv));

    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:   }

    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: }

    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: }

    178: /* ----------------------- Implicit form of the ODE  -------------------- */

    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;

    186:   PetscFunctionBeginUser;
    187:   PetscCall(VecGetArrayRead(U, &u));
    188:   PetscCall(VecGetArrayRead(Udot, &udot));
    189:   PetscCall(VecGetArray(F, &f));

    191:   f[0] = udot[0] - u[1];
    192:   f[1] = udot[1] - user->mu * ((1.0 - u[0] * u[0]) * u[1] - u[0]);

    194:   PetscCall(VecRestoreArrayRead(U, &u));
    195:   PetscCall(VecRestoreArrayRead(Udot, &udot));
    196:   PetscCall(VecRestoreArray(F, &f));
    197:   PetscFunctionReturn(PETSC_SUCCESS);
    198: }

    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;

    207:   PetscFunctionBeginUser;
    208:   PetscCall(VecGetArrayRead(U, &u));

    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:   }

    223:   PetscCall(VecRestoreArrayRead(U, &u));
    224:   PetscFunctionReturn(PETSC_SUCCESS);
    225: }

    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;

    233:   PetscFunctionBeginUser;
    234:   PetscCall(VecGetArrayRead(U, &u));

    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));

    243:   PetscCall(VecRestoreArrayRead(U, &u));
    244:   PetscFunctionReturn(PETSC_SUCCESS);
    245: }

    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;

    255:   PetscFunctionBeginUser;
    256:   PetscCall(VecGetArrayRead(U, &u));
    257:   PetscCall(VecGetArrayRead(Vl[0], &vl));
    258:   PetscCall(VecGetArrayRead(Vr, &vr));
    259:   PetscCall(VecGetArray(VHV[0], &vhv));

    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:   }

    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: }

    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;

    284:   PetscFunctionBeginUser;
    285:   PetscCall(VecGetArrayRead(U, &u));
    286:   PetscCall(VecGetArrayRead(Vl[0], &vl));
    287:   PetscCall(VecGetArrayRead(Vr, &vr));
    288:   PetscCall(VecGetArray(VHV[0], &vhv));

    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:   }

    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: }

    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;

    312:   PetscFunctionBeginUser;
    313:   PetscCall(VecGetArrayRead(U, &u));
    314:   PetscCall(VecGetArrayRead(Vl[0], &vl));
    315:   PetscCall(VecGetArrayRead(Vr, &vr));
    316:   PetscCall(VecGetArray(VHV[0], &vhv));

    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:   }

    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: }

    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: }

    339: /* Monitor timesteps and use interpolation to output at integer multiples of 0.1 */
    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;

    347:   PetscFunctionBeginUser;
    348:   PetscCall(TSGetTimeStep(ts, &dt));
    349:   PetscCall(TSGetMaxTime(ts, &tfinal));

    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: }

    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;

    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!");

    383:   /* Create TAO solver and set desired solution method */
    384:   PetscCall(TaoCreate(PETSC_COMM_WORLD, &tao));
    385:   PetscCall(TaoSetType(tao, TAOBQNLS));

    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;

    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));

    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));

    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));

    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));

    436:   if (monitor) PetscCall(TSMonitorSet(user.ts, Monitor, &user, NULL));

    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)));

    445:   /* Set runtime options */
    446:   PetscCall(TSSetFromOptions(user.ts));

    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));

    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));

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

    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));

    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));

    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));

    483:   PetscCall(TaoSolve(tao));

    485:   PetscCall(VecView(P, PETSC_VIEWER_STDOUT_WORLD));
    486:   /* Free TAO data structures */
    487:   PetscCall(TaoDestroy(&tao));

    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: }

    512: /* ------------------------------------------------------------------ */
    513: /*
    514:    FormFunctionGradient - Evaluates the function and corresponding gradient.

    516:    Input Parameters:
    517:    tao - the Tao context
    518:    X   - the input vector
    519:    ptr - optional user-defined context, as set by TaoSetObjectiveAndGradient()

    521:    Output Parameters:
    522:    f   - the newly evaluated function
    523:    G   - the newly evaluated gradient
    524: */
    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;

    532:   PetscFunctionBeginUser;
    533:   PetscCall(VecGetArrayRead(P, &y_ptr));
    534:   user_ptr->mu = y_ptr[0];
    535:   PetscCall(VecRestoreArrayRead(P, &y_ptr));

    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));

    546:   PetscCall(TSSolve(ts, user_ptr->U));

    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]);

    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));

    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));

    563:   PetscCall(TSAdjointSolve(ts));

    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: }

    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;

    582:   PetscFunctionBeginUser;
    583:   PetscCall(Adjoint2(P, harr, user_ptr));
    584:   PetscCall(MatSetValues(H, 1, rows, 1, &col, harr, INSERT_VALUES));

    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: }

    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;

    602:   PetscFunctionBeginUser;
    603:   /* Reset TSAdjoint so that AdjointSetUp will be called again */
    604:   PetscCall(TSAdjointReset(ts));

    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));

    611:   PetscCall(VecGetArrayRead(P, &z_ptr));
    612:   ctx->mu = z_ptr[0];
    613:   PetscCall(VecRestoreArrayRead(P, &z_ptr));

    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);

    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));

    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));

    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));

    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));

    664:   PetscCall(VecGetArray(ctx->Lambda[0], &x_ptr));
    665:   PetscCall(VecGetArray(ctx->Lambda2[0], &y_ptr));
    666:   PetscCall(VecGetArrayRead(ctx->Mup2[0], &z_ptr));

    668:   arr[0] = x_ptr[1] * dzdp2 + y_ptr[1] * dzdp2 + z_ptr[0];

    670:   PetscCall(VecRestoreArray(ctx->Lambda2[0], &x_ptr));
    671:   PetscCall(VecRestoreArray(ctx->Lambda2[0], &y_ptr));
    672:   PetscCall(VecRestoreArrayRead(ctx->Mup2[0], &z_ptr));

    674:   /* Disable second-order adjoint mode */
    675:   PetscCall(TSAdjointReset(ts));
    676:   PetscCall(TSAdjointResetForward(ts));
    677:   PetscFunctionReturn(PETSC_SUCCESS);
    678: }

    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

    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

    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

    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

    702: TEST*/