1 #include <petsctao.h> 2 #include <petscts.h> 3 4 typedef struct _n_aircraft *Aircraft; 5 struct _n_aircraft { 6 TS ts, quadts; 7 Vec V, W; /* control variables V and W */ 8 PetscInt nsteps; /* number of time steps */ 9 PetscReal ftime; 10 Mat A, H; 11 Mat Jacp, DRDU, DRDP; 12 Vec U, Lambda[1], Mup[1], Lambda2[1], Mup2[1], Dir; 13 Vec rhshp1[1], rhshp2[1], rhshp3[1], rhshp4[1], inthp1[1], inthp2[1], inthp3[1], inthp4[1]; 14 PetscReal lv, lw; 15 PetscBool mf, eh; 16 }; 17 18 PetscErrorCode FormObjFunctionGradient(Tao, Vec, PetscReal *, Vec, void *); 19 PetscErrorCode FormObjHessian(Tao, Vec, Mat, Mat, void *); 20 PetscErrorCode ComputeObjHessianWithSOA(Vec, PetscScalar[], Aircraft); 21 PetscErrorCode MatrixFreeObjHessian(Tao, Vec, Mat, Mat, void *); 22 PetscErrorCode MyMatMult(Mat, Vec, Vec); 23 24 static PetscErrorCode RHSFunction(TS ts, PetscReal t, Vec U, Vec F, void *ctx) 25 { 26 Aircraft actx = (Aircraft)ctx; 27 const PetscScalar *u, *v, *w; 28 PetscScalar *f; 29 PetscInt step; 30 31 PetscFunctionBeginUser; 32 PetscCall(TSGetStepNumber(ts, &step)); 33 PetscCall(VecGetArrayRead(U, &u)); 34 PetscCall(VecGetArrayRead(actx->V, &v)); 35 PetscCall(VecGetArrayRead(actx->W, &w)); 36 PetscCall(VecGetArray(F, &f)); 37 f[0] = v[step] * PetscCosReal(w[step]); 38 f[1] = v[step] * PetscSinReal(w[step]); 39 PetscCall(VecRestoreArrayRead(U, &u)); 40 PetscCall(VecRestoreArrayRead(actx->V, &v)); 41 PetscCall(VecRestoreArrayRead(actx->W, &w)); 42 PetscCall(VecRestoreArray(F, &f)); 43 PetscFunctionReturn(PETSC_SUCCESS); 44 } 45 46 static PetscErrorCode RHSJacobianP(TS ts, PetscReal t, Vec U, Mat A, void *ctx) 47 { 48 Aircraft actx = (Aircraft)ctx; 49 const PetscScalar *u, *v, *w; 50 PetscInt step, rows[2] = {0, 1}, rowcol[2]; 51 PetscScalar Jp[2][2]; 52 53 PetscFunctionBeginUser; 54 PetscCall(MatZeroEntries(A)); 55 PetscCall(TSGetStepNumber(ts, &step)); 56 PetscCall(VecGetArrayRead(U, &u)); 57 PetscCall(VecGetArrayRead(actx->V, &v)); 58 PetscCall(VecGetArrayRead(actx->W, &w)); 59 60 Jp[0][0] = PetscCosReal(w[step]); 61 Jp[0][1] = -v[step] * PetscSinReal(w[step]); 62 Jp[1][0] = PetscSinReal(w[step]); 63 Jp[1][1] = v[step] * PetscCosReal(w[step]); 64 65 PetscCall(VecRestoreArrayRead(U, &u)); 66 PetscCall(VecRestoreArrayRead(actx->V, &v)); 67 PetscCall(VecRestoreArrayRead(actx->W, &w)); 68 69 rowcol[0] = 2 * step; 70 rowcol[1] = 2 * step + 1; 71 PetscCall(MatSetValues(A, 2, rows, 2, rowcol, &Jp[0][0], INSERT_VALUES)); 72 73 PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY)); 74 PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY)); 75 PetscFunctionReturn(PETSC_SUCCESS); 76 } 77 78 static PetscErrorCode RHSHessianProductUU(TS ts, PetscReal t, Vec U, Vec *Vl, Vec Vr, Vec *VHV, void *ctx) 79 { 80 PetscFunctionBeginUser; 81 PetscFunctionReturn(PETSC_SUCCESS); 82 } 83 84 static PetscErrorCode RHSHessianProductUP(TS ts, PetscReal t, Vec U, Vec *Vl, Vec Vr, Vec *VHV, void *ctx) 85 { 86 PetscFunctionBeginUser; 87 PetscFunctionReturn(PETSC_SUCCESS); 88 } 89 90 static PetscErrorCode RHSHessianProductPU(TS ts, PetscReal t, Vec U, Vec *Vl, Vec Vr, Vec *VHV, void *ctx) 91 { 92 PetscFunctionBeginUser; 93 PetscFunctionReturn(PETSC_SUCCESS); 94 } 95 96 static PetscErrorCode RHSHessianProductPP(TS ts, PetscReal t, Vec U, Vec *Vl, Vec Vr, Vec *VHV, void *ctx) 97 { 98 Aircraft actx = (Aircraft)ctx; 99 const PetscScalar *v, *w, *vl, *vr, *u; 100 PetscScalar *vhv; 101 PetscScalar dJpdP[2][2][2] = {{{0}}}; 102 PetscInt step, i, j, k; 103 104 PetscFunctionBeginUser; 105 PetscCall(TSGetStepNumber(ts, &step)); 106 PetscCall(VecGetArrayRead(U, &u)); 107 PetscCall(VecGetArrayRead(actx->V, &v)); 108 PetscCall(VecGetArrayRead(actx->W, &w)); 109 PetscCall(VecGetArrayRead(Vl[0], &vl)); 110 PetscCall(VecGetArrayRead(Vr, &vr)); 111 PetscCall(VecSet(VHV[0], 0.0)); 112 PetscCall(VecGetArray(VHV[0], &vhv)); 113 114 dJpdP[0][0][1] = -PetscSinReal(w[step]); 115 dJpdP[0][1][0] = -PetscSinReal(w[step]); 116 dJpdP[0][1][1] = -v[step] * PetscCosReal(w[step]); 117 dJpdP[1][0][1] = PetscCosReal(w[step]); 118 dJpdP[1][1][0] = PetscCosReal(w[step]); 119 dJpdP[1][1][1] = -v[step] * PetscSinReal(w[step]); 120 121 for (j = 0; j < 2; j++) { 122 vhv[2 * step + j] = 0; 123 for (k = 0; k < 2; k++) 124 for (i = 0; i < 2; i++) vhv[2 * step + j] += vl[i] * dJpdP[i][j][k] * vr[2 * step + k]; 125 } 126 PetscCall(VecRestoreArrayRead(U, &u)); 127 PetscCall(VecRestoreArrayRead(Vl[0], &vl)); 128 PetscCall(VecRestoreArrayRead(Vr, &vr)); 129 PetscCall(VecRestoreArray(VHV[0], &vhv)); 130 PetscFunctionReturn(PETSC_SUCCESS); 131 } 132 133 /* Vl in NULL,updates to VHV must be added */ 134 static PetscErrorCode IntegrandHessianProductUU(TS ts, PetscReal t, Vec U, Vec *Vl, Vec Vr, Vec *VHV, void *ctx) 135 { 136 Aircraft actx = (Aircraft)ctx; 137 const PetscScalar *v, *w, *vr, *u; 138 PetscScalar *vhv; 139 PetscScalar dRudU[2][2] = {{0}}; 140 PetscInt step, j, k; 141 142 PetscFunctionBeginUser; 143 PetscCall(TSGetStepNumber(ts, &step)); 144 PetscCall(VecGetArrayRead(U, &u)); 145 PetscCall(VecGetArrayRead(actx->V, &v)); 146 PetscCall(VecGetArrayRead(actx->W, &w)); 147 PetscCall(VecGetArrayRead(Vr, &vr)); 148 PetscCall(VecGetArray(VHV[0], &vhv)); 149 150 dRudU[0][0] = 2.0; 151 dRudU[1][1] = 2.0; 152 153 for (j = 0; j < 2; j++) { 154 vhv[j] = 0; 155 for (k = 0; k < 2; k++) vhv[j] += dRudU[j][k] * vr[k]; 156 } 157 PetscCall(VecRestoreArrayRead(U, &u)); 158 PetscCall(VecRestoreArrayRead(Vr, &vr)); 159 PetscCall(VecRestoreArray(VHV[0], &vhv)); 160 PetscFunctionReturn(PETSC_SUCCESS); 161 } 162 163 static PetscErrorCode IntegrandHessianProductUP(TS ts, PetscReal t, Vec U, Vec *Vl, Vec Vr, Vec *VHV, void *ctx) 164 { 165 PetscFunctionBeginUser; 166 PetscFunctionReturn(PETSC_SUCCESS); 167 } 168 169 static PetscErrorCode IntegrandHessianProductPU(TS ts, PetscReal t, Vec U, Vec *Vl, Vec Vr, Vec *VHV, void *ctx) 170 { 171 PetscFunctionBeginUser; 172 PetscFunctionReturn(PETSC_SUCCESS); 173 } 174 175 static PetscErrorCode IntegrandHessianProductPP(TS ts, PetscReal t, Vec U, Vec *Vl, Vec Vr, Vec *VHV, void *ctx) 176 { 177 PetscFunctionBeginUser; 178 PetscFunctionReturn(PETSC_SUCCESS); 179 } 180 181 static PetscErrorCode CostIntegrand(TS ts, PetscReal t, Vec U, Vec R, void *ctx) 182 { 183 Aircraft actx = (Aircraft)ctx; 184 PetscScalar *r; 185 PetscReal dx, dy; 186 const PetscScalar *u; 187 188 PetscFunctionBegin; 189 PetscCall(VecGetArrayRead(U, &u)); 190 PetscCall(VecGetArray(R, &r)); 191 dx = u[0] - actx->lv * t * PetscCosReal(actx->lw); 192 dy = u[1] - actx->lv * t * PetscSinReal(actx->lw); 193 r[0] = dx * dx + dy * dy; 194 PetscCall(VecRestoreArray(R, &r)); 195 PetscCall(VecRestoreArrayRead(U, &u)); 196 PetscFunctionReturn(PETSC_SUCCESS); 197 } 198 199 static PetscErrorCode DRDUJacobianTranspose(TS ts, PetscReal t, Vec U, Mat DRDU, Mat B, void *ctx) 200 { 201 Aircraft actx = (Aircraft)ctx; 202 PetscScalar drdu[2][1]; 203 const PetscScalar *u; 204 PetscReal dx, dy; 205 PetscInt row[] = {0, 1}, col[] = {0}; 206 207 PetscFunctionBegin; 208 PetscCall(VecGetArrayRead(U, &u)); 209 dx = u[0] - actx->lv * t * PetscCosReal(actx->lw); 210 dy = u[1] - actx->lv * t * PetscSinReal(actx->lw); 211 drdu[0][0] = 2. * dx; 212 drdu[1][0] = 2. * dy; 213 PetscCall(MatSetValues(DRDU, 2, row, 1, col, &drdu[0][0], INSERT_VALUES)); 214 PetscCall(VecRestoreArrayRead(U, &u)); 215 PetscCall(MatAssemblyBegin(DRDU, MAT_FINAL_ASSEMBLY)); 216 PetscCall(MatAssemblyEnd(DRDU, MAT_FINAL_ASSEMBLY)); 217 PetscFunctionReturn(PETSC_SUCCESS); 218 } 219 220 static PetscErrorCode DRDPJacobianTranspose(TS ts, PetscReal t, Vec U, Mat DRDP, void *ctx) 221 { 222 PetscFunctionBegin; 223 PetscCall(MatZeroEntries(DRDP)); 224 PetscFunctionReturn(PETSC_SUCCESS); 225 } 226 227 int main(int argc, char **argv) 228 { 229 Vec P, PL, PU; 230 struct _n_aircraft aircraft; 231 PetscMPIInt size; 232 Tao tao; 233 KSP ksp; 234 PC pc; 235 PetscScalar *u, *p; 236 PetscInt i; 237 238 /* Initialize program */ 239 PetscFunctionBeginUser; 240 PetscCall(PetscInitialize(&argc, &argv, NULL, NULL)); 241 PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size)); 242 PetscCheck(size == 1, PETSC_COMM_WORLD, PETSC_ERR_WRONG_MPI_SIZE, "This is a uniprocessor example only!"); 243 244 /* Parameter settings */ 245 aircraft.ftime = 1.; /* time interval in hour */ 246 aircraft.nsteps = 10; /* number of steps */ 247 aircraft.lv = 2.0; /* leader speed in kmph */ 248 aircraft.lw = PETSC_PI / 4.; /* leader heading angle */ 249 250 PetscCall(PetscOptionsGetReal(NULL, NULL, "-ftime", &aircraft.ftime, NULL)); 251 PetscCall(PetscOptionsGetInt(NULL, NULL, "-nsteps", &aircraft.nsteps, NULL)); 252 PetscCall(PetscOptionsHasName(NULL, NULL, "-matrixfree", &aircraft.mf)); 253 PetscCall(PetscOptionsHasName(NULL, NULL, "-exacthessian", &aircraft.eh)); 254 255 /* Create TAO solver and set desired solution method */ 256 PetscCall(TaoCreate(PETSC_COMM_WORLD, &tao)); 257 PetscCall(TaoSetType(tao, TAOBQNLS)); 258 259 /* Create necessary matrix and vectors, solve same ODE on every process */ 260 PetscCall(MatCreate(PETSC_COMM_WORLD, &aircraft.A)); 261 PetscCall(MatSetSizes(aircraft.A, PETSC_DECIDE, PETSC_DECIDE, 2, 2)); 262 PetscCall(MatSetFromOptions(aircraft.A)); 263 PetscCall(MatSetUp(aircraft.A)); 264 /* this is to set explicit zeros along the diagonal of the matrix */ 265 PetscCall(MatAssemblyBegin(aircraft.A, MAT_FINAL_ASSEMBLY)); 266 PetscCall(MatAssemblyEnd(aircraft.A, MAT_FINAL_ASSEMBLY)); 267 PetscCall(MatShift(aircraft.A, 1)); 268 PetscCall(MatShift(aircraft.A, -1)); 269 270 PetscCall(MatCreate(PETSC_COMM_WORLD, &aircraft.Jacp)); 271 PetscCall(MatSetSizes(aircraft.Jacp, PETSC_DECIDE, PETSC_DECIDE, 2, 2 * aircraft.nsteps)); 272 PetscCall(MatSetFromOptions(aircraft.Jacp)); 273 PetscCall(MatSetUp(aircraft.Jacp)); 274 PetscCall(MatSetOption(aircraft.Jacp, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_FALSE)); 275 PetscCall(MatCreateDense(PETSC_COMM_WORLD, PETSC_DECIDE, PETSC_DECIDE, 2 * aircraft.nsteps, 1, NULL, &aircraft.DRDP)); 276 PetscCall(MatSetUp(aircraft.DRDP)); 277 PetscCall(MatCreateDense(PETSC_COMM_WORLD, PETSC_DECIDE, PETSC_DECIDE, 2, 1, NULL, &aircraft.DRDU)); 278 PetscCall(MatSetUp(aircraft.DRDU)); 279 280 /* Create timestepping solver context */ 281 PetscCall(TSCreate(PETSC_COMM_WORLD, &aircraft.ts)); 282 PetscCall(TSSetType(aircraft.ts, TSRK)); 283 PetscCall(TSSetRHSFunction(aircraft.ts, NULL, RHSFunction, &aircraft)); 284 PetscCall(TSSetRHSJacobian(aircraft.ts, aircraft.A, aircraft.A, TSComputeRHSJacobianConstant, &aircraft)); 285 PetscCall(TSSetRHSJacobianP(aircraft.ts, aircraft.Jacp, RHSJacobianP, &aircraft)); 286 PetscCall(TSSetExactFinalTime(aircraft.ts, TS_EXACTFINALTIME_MATCHSTEP)); 287 PetscCall(TSSetEquationType(aircraft.ts, TS_EQ_ODE_EXPLICIT)); /* less Jacobian evaluations when adjoint BEuler is used, otherwise no effect */ 288 289 /* Set initial conditions */ 290 PetscCall(MatCreateVecs(aircraft.A, &aircraft.U, NULL)); 291 PetscCall(TSSetSolution(aircraft.ts, aircraft.U)); 292 PetscCall(VecGetArray(aircraft.U, &u)); 293 u[0] = 1.5; 294 u[1] = 0; 295 PetscCall(VecRestoreArray(aircraft.U, &u)); 296 PetscCall(VecCreate(PETSC_COMM_WORLD, &aircraft.V)); 297 PetscCall(VecSetSizes(aircraft.V, PETSC_DECIDE, aircraft.nsteps)); 298 PetscCall(VecSetUp(aircraft.V)); 299 PetscCall(VecDuplicate(aircraft.V, &aircraft.W)); 300 PetscCall(VecSet(aircraft.V, 1.)); 301 PetscCall(VecSet(aircraft.W, PETSC_PI / 4.)); 302 303 /* Save trajectory of solution so that TSAdjointSolve() may be used */ 304 PetscCall(TSSetSaveTrajectory(aircraft.ts)); 305 306 /* Set sensitivity context */ 307 PetscCall(TSCreateQuadratureTS(aircraft.ts, PETSC_FALSE, &aircraft.quadts)); 308 PetscCall(TSSetRHSFunction(aircraft.quadts, NULL, (TSRHSFunctionFn *)CostIntegrand, &aircraft)); 309 PetscCall(TSSetRHSJacobian(aircraft.quadts, aircraft.DRDU, aircraft.DRDU, (TSRHSJacobian)DRDUJacobianTranspose, &aircraft)); 310 PetscCall(TSSetRHSJacobianP(aircraft.quadts, aircraft.DRDP, (TSRHSJacobianPFn *)DRDPJacobianTranspose, &aircraft)); 311 PetscCall(MatCreateVecs(aircraft.A, &aircraft.Lambda[0], NULL)); 312 PetscCall(MatCreateVecs(aircraft.Jacp, &aircraft.Mup[0], NULL)); 313 if (aircraft.eh) { 314 PetscCall(MatCreateVecs(aircraft.A, &aircraft.rhshp1[0], NULL)); 315 PetscCall(MatCreateVecs(aircraft.A, &aircraft.rhshp2[0], NULL)); 316 PetscCall(MatCreateVecs(aircraft.Jacp, &aircraft.rhshp3[0], NULL)); 317 PetscCall(MatCreateVecs(aircraft.Jacp, &aircraft.rhshp4[0], NULL)); 318 PetscCall(MatCreateVecs(aircraft.DRDU, &aircraft.inthp1[0], NULL)); 319 PetscCall(MatCreateVecs(aircraft.DRDU, &aircraft.inthp2[0], NULL)); 320 PetscCall(MatCreateVecs(aircraft.DRDP, &aircraft.inthp3[0], NULL)); 321 PetscCall(MatCreateVecs(aircraft.DRDP, &aircraft.inthp4[0], NULL)); 322 PetscCall(MatCreateVecs(aircraft.Jacp, &aircraft.Dir, NULL)); 323 PetscCall(TSSetRHSHessianProduct(aircraft.ts, aircraft.rhshp1, RHSHessianProductUU, aircraft.rhshp2, RHSHessianProductUP, aircraft.rhshp3, RHSHessianProductPU, aircraft.rhshp4, RHSHessianProductPP, &aircraft)); 324 PetscCall(TSSetRHSHessianProduct(aircraft.quadts, aircraft.inthp1, IntegrandHessianProductUU, aircraft.inthp2, IntegrandHessianProductUP, aircraft.inthp3, IntegrandHessianProductPU, aircraft.inthp4, IntegrandHessianProductPP, &aircraft)); 325 PetscCall(MatCreateVecs(aircraft.A, &aircraft.Lambda2[0], NULL)); 326 PetscCall(MatCreateVecs(aircraft.Jacp, &aircraft.Mup2[0], NULL)); 327 } 328 PetscCall(TSSetFromOptions(aircraft.ts)); 329 PetscCall(TSSetMaxTime(aircraft.ts, aircraft.ftime)); 330 PetscCall(TSSetTimeStep(aircraft.ts, aircraft.ftime / aircraft.nsteps)); 331 332 /* Set initial solution guess */ 333 PetscCall(MatCreateVecs(aircraft.Jacp, &P, NULL)); 334 PetscCall(VecGetArray(P, &p)); 335 for (i = 0; i < aircraft.nsteps; i++) { 336 p[2 * i] = 2.0; 337 p[2 * i + 1] = PETSC_PI / 2.0; 338 } 339 PetscCall(VecRestoreArray(P, &p)); 340 PetscCall(VecDuplicate(P, &PU)); 341 PetscCall(VecDuplicate(P, &PL)); 342 PetscCall(VecGetArray(PU, &p)); 343 for (i = 0; i < aircraft.nsteps; i++) { 344 p[2 * i] = 2.0; 345 p[2 * i + 1] = PETSC_PI; 346 } 347 PetscCall(VecRestoreArray(PU, &p)); 348 PetscCall(VecGetArray(PL, &p)); 349 for (i = 0; i < aircraft.nsteps; i++) { 350 p[2 * i] = 0.0; 351 p[2 * i + 1] = -PETSC_PI; 352 } 353 PetscCall(VecRestoreArray(PL, &p)); 354 355 PetscCall(TaoSetSolution(tao, P)); 356 PetscCall(TaoSetVariableBounds(tao, PL, PU)); 357 /* Set routine for function and gradient evaluation */ 358 PetscCall(TaoSetObjectiveAndGradient(tao, NULL, FormObjFunctionGradient, (void *)&aircraft)); 359 360 if (aircraft.eh) { 361 if (aircraft.mf) { 362 PetscCall(MatCreateShell(PETSC_COMM_WORLD, PETSC_DECIDE, PETSC_DECIDE, 2 * aircraft.nsteps, 2 * aircraft.nsteps, (void *)&aircraft, &aircraft.H)); 363 PetscCall(MatShellSetOperation(aircraft.H, MATOP_MULT, (PetscErrorCodeFn *)MyMatMult)); 364 PetscCall(MatSetOption(aircraft.H, MAT_SYMMETRIC, PETSC_TRUE)); 365 PetscCall(TaoSetHessian(tao, aircraft.H, aircraft.H, MatrixFreeObjHessian, (void *)&aircraft)); 366 } else { 367 PetscCall(MatCreateDense(MPI_COMM_WORLD, PETSC_DETERMINE, PETSC_DETERMINE, 2 * aircraft.nsteps, 2 * aircraft.nsteps, NULL, &aircraft.H)); 368 PetscCall(MatSetOption(aircraft.H, MAT_SYMMETRIC, PETSC_TRUE)); 369 PetscCall(TaoSetHessian(tao, aircraft.H, aircraft.H, FormObjHessian, (void *)&aircraft)); 370 } 371 } 372 373 /* Check for any TAO command line options */ 374 PetscCall(TaoGetKSP(tao, &ksp)); 375 if (ksp) { 376 PetscCall(KSPGetPC(ksp, &pc)); 377 PetscCall(PCSetType(pc, PCNONE)); 378 } 379 PetscCall(TaoSetFromOptions(tao)); 380 381 PetscCall(TaoSolve(tao)); 382 PetscCall(VecView(P, PETSC_VIEWER_STDOUT_WORLD)); 383 384 /* Free TAO data structures */ 385 PetscCall(TaoDestroy(&tao)); 386 387 /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 388 Free work space. All PETSc objects should be destroyed when they 389 are no longer needed. 390 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */ 391 PetscCall(TSDestroy(&aircraft.ts)); 392 PetscCall(MatDestroy(&aircraft.A)); 393 PetscCall(VecDestroy(&aircraft.U)); 394 PetscCall(VecDestroy(&aircraft.V)); 395 PetscCall(VecDestroy(&aircraft.W)); 396 PetscCall(VecDestroy(&P)); 397 PetscCall(VecDestroy(&PU)); 398 PetscCall(VecDestroy(&PL)); 399 PetscCall(MatDestroy(&aircraft.Jacp)); 400 PetscCall(MatDestroy(&aircraft.DRDU)); 401 PetscCall(MatDestroy(&aircraft.DRDP)); 402 PetscCall(VecDestroy(&aircraft.Lambda[0])); 403 PetscCall(VecDestroy(&aircraft.Mup[0])); 404 PetscCall(VecDestroy(&P)); 405 if (aircraft.eh) { 406 PetscCall(VecDestroy(&aircraft.Lambda2[0])); 407 PetscCall(VecDestroy(&aircraft.Mup2[0])); 408 PetscCall(VecDestroy(&aircraft.Dir)); 409 PetscCall(VecDestroy(&aircraft.rhshp1[0])); 410 PetscCall(VecDestroy(&aircraft.rhshp2[0])); 411 PetscCall(VecDestroy(&aircraft.rhshp3[0])); 412 PetscCall(VecDestroy(&aircraft.rhshp4[0])); 413 PetscCall(VecDestroy(&aircraft.inthp1[0])); 414 PetscCall(VecDestroy(&aircraft.inthp2[0])); 415 PetscCall(VecDestroy(&aircraft.inthp3[0])); 416 PetscCall(VecDestroy(&aircraft.inthp4[0])); 417 PetscCall(MatDestroy(&aircraft.H)); 418 } 419 PetscCall(PetscFinalize()); 420 return 0; 421 } 422 423 /* 424 FormObjFunctionGradient - Evaluates the function and corresponding gradient. 425 426 Input Parameters: 427 tao - the Tao context 428 P - the input vector 429 ctx - optional aircraft-defined context, as set by TaoSetObjectiveAndGradient() 430 431 Output Parameters: 432 f - the newly evaluated function 433 G - the newly evaluated gradient 434 */ 435 PetscErrorCode FormObjFunctionGradient(Tao tao, Vec P, PetscReal *f, Vec G, void *ctx) 436 { 437 Aircraft actx = (Aircraft)ctx; 438 TS ts = actx->ts; 439 Vec Q; 440 const PetscScalar *p, *q; 441 PetscScalar *u, *v, *w; 442 PetscInt i; 443 444 PetscFunctionBeginUser; 445 PetscCall(VecGetArrayRead(P, &p)); 446 PetscCall(VecGetArray(actx->V, &v)); 447 PetscCall(VecGetArray(actx->W, &w)); 448 for (i = 0; i < actx->nsteps; i++) { 449 v[i] = p[2 * i]; 450 w[i] = p[2 * i + 1]; 451 } 452 PetscCall(VecRestoreArrayRead(P, &p)); 453 PetscCall(VecRestoreArray(actx->V, &v)); 454 PetscCall(VecRestoreArray(actx->W, &w)); 455 456 PetscCall(TSSetTime(ts, 0.0)); 457 PetscCall(TSSetStepNumber(ts, 0)); 458 PetscCall(TSSetFromOptions(ts)); 459 PetscCall(TSSetTimeStep(ts, actx->ftime / actx->nsteps)); 460 461 /* reinitialize system state */ 462 PetscCall(VecGetArray(actx->U, &u)); 463 u[0] = 2.0; 464 u[1] = 0; 465 PetscCall(VecRestoreArray(actx->U, &u)); 466 467 /* reinitialize the integral value */ 468 PetscCall(TSGetCostIntegral(ts, &Q)); 469 PetscCall(VecSet(Q, 0.0)); 470 471 PetscCall(TSSolve(ts, actx->U)); 472 473 /* Reset initial conditions for the adjoint integration */ 474 PetscCall(VecSet(actx->Lambda[0], 0.0)); 475 PetscCall(VecSet(actx->Mup[0], 0.0)); 476 PetscCall(TSSetCostGradients(ts, 1, actx->Lambda, actx->Mup)); 477 478 PetscCall(TSAdjointSolve(ts)); 479 PetscCall(VecCopy(actx->Mup[0], G)); 480 PetscCall(TSGetCostIntegral(ts, &Q)); 481 PetscCall(VecGetArrayRead(Q, &q)); 482 *f = q[0]; 483 PetscCall(VecRestoreArrayRead(Q, &q)); 484 PetscFunctionReturn(PETSC_SUCCESS); 485 } 486 487 PetscErrorCode FormObjHessian(Tao tao, Vec P, Mat H, Mat Hpre, void *ctx) 488 { 489 Aircraft actx = (Aircraft)ctx; 490 const PetscScalar *p; 491 PetscScalar *harr, *v, *w, one = 1.0; 492 PetscInt ind[1]; 493 PetscInt *cols, i; 494 Vec Dir; 495 496 PetscFunctionBeginUser; 497 /* set up control parameters */ 498 PetscCall(VecGetArrayRead(P, &p)); 499 PetscCall(VecGetArray(actx->V, &v)); 500 PetscCall(VecGetArray(actx->W, &w)); 501 for (i = 0; i < actx->nsteps; i++) { 502 v[i] = p[2 * i]; 503 w[i] = p[2 * i + 1]; 504 } 505 PetscCall(VecRestoreArrayRead(P, &p)); 506 PetscCall(VecRestoreArray(actx->V, &v)); 507 PetscCall(VecRestoreArray(actx->W, &w)); 508 509 PetscCall(PetscMalloc1(2 * actx->nsteps, &harr)); 510 PetscCall(PetscMalloc1(2 * actx->nsteps, &cols)); 511 for (i = 0; i < 2 * actx->nsteps; i++) cols[i] = i; 512 PetscCall(VecDuplicate(P, &Dir)); 513 for (i = 0; i < 2 * actx->nsteps; i++) { 514 ind[0] = i; 515 PetscCall(VecSet(Dir, 0.0)); 516 PetscCall(VecSetValues(Dir, 1, ind, &one, INSERT_VALUES)); 517 PetscCall(VecAssemblyBegin(Dir)); 518 PetscCall(VecAssemblyEnd(Dir)); 519 PetscCall(ComputeObjHessianWithSOA(Dir, harr, actx)); 520 PetscCall(MatSetValues(H, 1, ind, 2 * actx->nsteps, cols, harr, INSERT_VALUES)); 521 PetscCall(MatAssemblyBegin(H, MAT_FINAL_ASSEMBLY)); 522 PetscCall(MatAssemblyEnd(H, MAT_FINAL_ASSEMBLY)); 523 if (H != Hpre) { 524 PetscCall(MatAssemblyBegin(Hpre, MAT_FINAL_ASSEMBLY)); 525 PetscCall(MatAssemblyEnd(Hpre, MAT_FINAL_ASSEMBLY)); 526 } 527 } 528 PetscCall(PetscFree(cols)); 529 PetscCall(PetscFree(harr)); 530 PetscCall(VecDestroy(&Dir)); 531 PetscFunctionReturn(PETSC_SUCCESS); 532 } 533 534 PetscErrorCode MatrixFreeObjHessian(Tao tao, Vec P, Mat H, Mat Hpre, void *ctx) 535 { 536 Aircraft actx = (Aircraft)ctx; 537 PetscScalar *v, *w; 538 const PetscScalar *p; 539 PetscInt i; 540 541 PetscFunctionBegin; 542 PetscCall(VecGetArrayRead(P, &p)); 543 PetscCall(VecGetArray(actx->V, &v)); 544 PetscCall(VecGetArray(actx->W, &w)); 545 for (i = 0; i < actx->nsteps; i++) { 546 v[i] = p[2 * i]; 547 w[i] = p[2 * i + 1]; 548 } 549 PetscCall(VecRestoreArrayRead(P, &p)); 550 PetscCall(VecRestoreArray(actx->V, &v)); 551 PetscCall(VecRestoreArray(actx->W, &w)); 552 PetscFunctionReturn(PETSC_SUCCESS); 553 } 554 555 PetscErrorCode MyMatMult(Mat H_shell, Vec X, Vec Y) 556 { 557 PetscScalar *y; 558 void *ptr; 559 560 PetscFunctionBegin; 561 PetscCall(MatShellGetContext(H_shell, &ptr)); 562 PetscCall(VecGetArray(Y, &y)); 563 PetscCall(ComputeObjHessianWithSOA(X, y, (Aircraft)ptr)); 564 PetscCall(VecRestoreArray(Y, &y)); 565 PetscFunctionReturn(PETSC_SUCCESS); 566 } 567 568 PetscErrorCode ComputeObjHessianWithSOA(Vec Dir, PetscScalar arr[], Aircraft actx) 569 { 570 TS ts = actx->ts; 571 const PetscScalar *z_ptr; 572 PetscScalar *u; 573 Vec Q; 574 PetscInt i; 575 576 PetscFunctionBeginUser; 577 /* Reset TSAdjoint so that AdjointSetUp will be called again */ 578 PetscCall(TSAdjointReset(ts)); 579 580 PetscCall(TSSetTime(ts, 0.0)); 581 PetscCall(TSSetStepNumber(ts, 0)); 582 PetscCall(TSSetFromOptions(ts)); 583 PetscCall(TSSetTimeStep(ts, actx->ftime / actx->nsteps)); 584 PetscCall(TSSetCostHessianProducts(actx->ts, 1, actx->Lambda2, actx->Mup2, Dir)); 585 586 /* reinitialize system state */ 587 PetscCall(VecGetArray(actx->U, &u)); 588 u[0] = 2.0; 589 u[1] = 0; 590 PetscCall(VecRestoreArray(actx->U, &u)); 591 592 /* reinitialize the integral value */ 593 PetscCall(TSGetCostIntegral(ts, &Q)); 594 PetscCall(VecSet(Q, 0.0)); 595 596 /* initialize tlm variable */ 597 PetscCall(MatZeroEntries(actx->Jacp)); 598 PetscCall(TSAdjointSetForward(ts, actx->Jacp)); 599 600 PetscCall(TSSolve(ts, actx->U)); 601 602 /* Set terminal conditions for first- and second-order adjonts */ 603 PetscCall(VecSet(actx->Lambda[0], 0.0)); 604 PetscCall(VecSet(actx->Mup[0], 0.0)); 605 PetscCall(VecSet(actx->Lambda2[0], 0.0)); 606 PetscCall(VecSet(actx->Mup2[0], 0.0)); 607 PetscCall(TSSetCostGradients(ts, 1, actx->Lambda, actx->Mup)); 608 609 PetscCall(TSGetCostIntegral(ts, &Q)); 610 611 /* Reset initial conditions for the adjoint integration */ 612 PetscCall(TSAdjointSolve(ts)); 613 614 /* initial condition does not depend on p, so that lambda is not needed to assemble G */ 615 PetscCall(VecGetArrayRead(actx->Mup2[0], &z_ptr)); 616 for (i = 0; i < 2 * actx->nsteps; i++) arr[i] = z_ptr[i]; 617 PetscCall(VecRestoreArrayRead(actx->Mup2[0], &z_ptr)); 618 619 /* Disable second-order adjoint mode */ 620 PetscCall(TSAdjointReset(ts)); 621 PetscCall(TSAdjointResetForward(ts)); 622 PetscFunctionReturn(PETSC_SUCCESS); 623 } 624 625 /*TEST 626 build: 627 requires: !complex !single 628 629 test: 630 args: -ts_adapt_type none -ts_type rk -ts_rk_type 3 -viewer_binary_skip_info -tao_monitor -tao_gatol 1e-7 631 632 test: 633 suffix: 2 634 args: -ts_adapt_type none -ts_type rk -ts_rk_type 3 -viewer_binary_skip_info -tao_monitor -tao_view -tao_type bntr -tao_bnk_pc_type none -exacthessian 635 636 test: 637 suffix: 3 638 args: -ts_adapt_type none -ts_type rk -ts_rk_type 3 -viewer_binary_skip_info -tao_monitor -tao_view -tao_type bntr -tao_bnk_pc_type none -exacthessian -matrixfree 639 TEST*/ 640