Actual source code: rosenbrock1.h

      1: #pragma once
      2: // Common data structures for rosenbrock1.c and rosenbrock1_taoterm.c

      4: #include <petsctao.h>

      6: /* User-defined application context

      8:    contains data needed by the application-provided call-back routines that
      9:    evaluate the function, gradient, and hessian.
     10: */
     11: typedef struct {
     12:   MPI_Comm  comm;
     13:   PetscInt  n;         /* dimension */
     14:   PetscReal alpha;     /* condition parameter */
     15:   PetscBool chained;   /* chained vs. unchained Rosenbrock function */
     16:   PetscBool test;      /* run tests in AppCtxFinalize() */
     17:   PetscBool jacobi_pc; /* Create Jacobi Hpre */
     18:   PetscBool use_fd;    /* Use finite difference for grad and hess */
     19: } AppCtx;

     21: static PetscErrorCode AppCtxInitialize(MPI_Comm, AppCtx *); /* process options */
     22: static PetscErrorCode AppCtxFinalize(AppCtx *, Tao);        /* clean up and optionally run tests */
     23: static PetscErrorCode AppCtxCreateSolution(AppCtx *, Vec *);
     24: static PetscErrorCode AppCtxCreateHessianMatrices(AppCtx *, Mat *, Mat *);

     26: static PetscErrorCode AppCtxInitialize(MPI_Comm comm, AppCtx *usr)
     27: {
     28:   PetscBool flg;

     30:   PetscFunctionBeginUser;
     31:   usr->comm      = comm;
     32:   usr->n         = 2;
     33:   usr->alpha     = 99.0;
     34:   usr->chained   = PETSC_FALSE;
     35:   usr->test      = PETSC_FALSE;
     36:   usr->jacobi_pc = PETSC_FALSE;
     37:   usr->use_fd    = PETSC_FALSE;
     38:   /* Check for command line arguments to override defaults */
     39:   PetscCall(PetscOptionsGetInt(NULL, NULL, "-n", &usr->n, &flg));
     40:   PetscCall(PetscOptionsGetReal(NULL, NULL, "-alpha", &usr->alpha, &flg));
     41:   PetscCall(PetscOptionsGetBool(NULL, NULL, "-chained", &usr->chained, &flg));
     42:   PetscCall(PetscOptionsGetBool(NULL, NULL, "-test_lmvm", &usr->test, &flg));
     43:   PetscCall(PetscOptionsGetBool(NULL, NULL, "-jacobi_pc", &usr->jacobi_pc, &flg));
     44:   PetscCall(PetscOptionsGetBool(NULL, NULL, "-use_fd", &usr->use_fd, &flg));
     45:   PetscFunctionReturn(PETSC_SUCCESS);
     46: }

     48: static PetscErrorCode AppCtxFinalize(AppCtx *usr, Tao tao)
     49: {
     50:   PetscFunctionBeginUser;
     51:   if (usr->test) {
     52:     /* Test the LMVM matrix */
     53:     KSP       ksp;
     54:     PC        pc;
     55:     Mat       M;
     56:     Vec       in, out, out2;
     57:     PetscReal mult_solve_dist;
     58:     Vec       x;

     60:     PetscCall(TaoGetSolution(tao, &x));
     61:     PetscCall(TaoGetKSP(tao, &ksp));
     62:     PetscCall(KSPGetPC(ksp, &pc));
     63:     PetscCall(PCLMVMGetMatLMVM(pc, &M));
     64:     PetscCall(VecDuplicate(x, &in));
     65:     PetscCall(VecDuplicate(x, &out));
     66:     PetscCall(VecDuplicate(x, &out2));
     67:     PetscCall(VecSet(in, 1.0));
     68:     PetscCall(MatMult(M, in, out));
     69:     PetscCall(MatSolve(M, out, out2));
     70:     PetscCall(VecAXPY(out2, -1.0, in));
     71:     PetscCall(VecNorm(out2, NORM_2, &mult_solve_dist));
     72:     if (mult_solve_dist < 1.e-11) {
     73:       PetscCall(PetscPrintf(PetscObjectComm((PetscObject)tao), "Inverse error of LMVM MatMult and MatSolve: < 1.e-11\n"));
     74:     } else if (mult_solve_dist < 1.e-6) {
     75:       PetscCall(PetscPrintf(PetscObjectComm((PetscObject)tao), "Inverse error of LMVM MatMult and MatSolve: < 1.e-6\n"));
     76:     } else {
     77:       PetscCall(PetscPrintf(PetscObjectComm((PetscObject)tao), "Inverse error of LMVM MatMult and MatSolve is not small: %e\n", (double)mult_solve_dist));
     78:     }
     79:     PetscCall(VecDestroy(&in));
     80:     PetscCall(VecDestroy(&out));
     81:     PetscCall(VecDestroy(&out2));
     82:   }
     83:   PetscFunctionReturn(PETSC_SUCCESS);
     84: }

     86: static PetscErrorCode AppCtxCreateSolution(AppCtx *usr, Vec *solution)
     87: {
     88:   PetscFunctionBeginUser;
     89:   PetscCall(VecCreateSeq(PETSC_COMM_SELF, usr->n, solution));
     90:   PetscFunctionReturn(PETSC_SUCCESS);
     91: }

     93: static PETSC_UNUSED PetscErrorCode AppCtxCreateHessianMatrices(AppCtx *usr, Mat *H, Mat *Hpre)
     94: {
     95:   Mat hessian;

     97:   PetscFunctionBeginUser;
     98:   PetscCall(MatCreateSeqBAIJ(PETSC_COMM_SELF, 2, usr->n, usr->n, 1, NULL, &hessian));
     99:   if (H) {
    100:     PetscCall(PetscObjectReference((PetscObject)hessian));
    101:     *H = hessian;
    102:   }
    103:   if (Hpre) {
    104:     if (usr->jacobi_pc) PetscCall(MatCreateSeqBAIJ(PETSC_COMM_SELF, 1, usr->n, usr->n, 1, NULL, Hpre));
    105:     else {
    106:       PetscCall(PetscObjectReference((PetscObject)hessian));
    107:       *Hpre = hessian;
    108:     }
    109:   }
    110:   PetscCall(MatDestroy(&hessian));
    111:   PetscFunctionReturn(PETSC_SUCCESS);
    112: }

    114: /*
    115:   AppCtxFormFunctionGradient - Evaluates the function, f(X), and gradient, G(X).

    117:   Input Parameters:
    118: + user - AppCtx
    119: - X    - input vector

    121:   Output Parameters:
    122: + f - function value
    123: - G - vector containing the newly evaluated gradient
    124: */
    125: static PetscErrorCode AppCtxFormFunctionGradient(AppCtx *user, Vec X, PetscReal *f, Vec G)
    126: {
    127:   PetscInt           nn = user->n / 2;
    128:   PetscReal          ff = 0, t1, t2, alpha = user->alpha;
    129:   PetscScalar       *g;
    130:   const PetscScalar *x;

    132:   PetscFunctionBeginUser;
    133:   /* Get pointers to vector data */
    134:   PetscCall(VecGetArrayRead(X, &x));
    135:   PetscCall(VecGetArray(G, &g));

    137:   /* Compute G(X) */
    138:   if (user->chained) {
    139:     g[0] = 0;
    140:     for (PetscInt i = 0; i < user->n - 1; i++) {
    141:       t1 = x[i + 1] - x[i] * x[i];
    142:       ff += PetscSqr(1 - x[i]) + alpha * t1 * t1;
    143:       g[i] += -2 * (1 - x[i]) + 2 * alpha * t1 * (-2 * x[i]);
    144:       g[i + 1] = 2 * alpha * t1;
    145:     }
    146:   } else {
    147:     for (PetscInt i = 0; i < nn; i++) {
    148:       t1 = x[2 * i + 1] - x[2 * i] * x[2 * i];
    149:       t2 = 1 - x[2 * i];
    150:       ff += alpha * t1 * t1 + t2 * t2;
    151:       g[2 * i]     = -4 * alpha * t1 * x[2 * i] - 2.0 * t2;
    152:       g[2 * i + 1] = 2 * alpha * t1;
    153:     }
    154:   }

    156:   /* Restore vectors */
    157:   PetscCall(VecRestoreArrayRead(X, &x));
    158:   PetscCall(VecRestoreArray(G, &g));
    159:   *f = ff;

    161:   PetscCall(PetscLogFlops(15.0 * nn));
    162:   PetscFunctionReturn(PETSC_SUCCESS);
    163: }

    165: /*
    166:   AppCtxFormHessian - Evaluates Hessian matrix.

    168:   Input Parameters:
    169: + user - the context
    170: - X    - input vector

    172:   Output Parameters:
    173: . H - Hessian matrix
    174: */
    175: static PetscErrorCode AppCtxFormHessian(AppCtx *user, Vec X, Mat H)
    176: {
    177:   PetscInt           ind[2];
    178:   PetscReal          alpha = user->alpha;
    179:   PetscReal          v[2][2];
    180:   const PetscScalar *x;
    181:   PetscBool          assembled;

    183:   PetscFunctionBeginUser;
    184:   /* Zero existing matrix entries */
    185:   PetscCall(MatAssembled(H, &assembled));
    186:   if (assembled) PetscCall(MatZeroEntries(H));

    188:   /* Get a pointer to vector data */
    189:   PetscCall(VecGetArrayRead(X, &x));

    191:   /* Compute H(X) entries */
    192:   if (user->chained) {
    193:     PetscCall(MatZeroEntries(H));
    194:     for (PetscInt i = 0; i < user->n - 1; i++) {
    195:       PetscScalar t1 = x[i + 1] - x[i] * x[i];
    196:       v[0][0]        = 2 + 2 * alpha * (t1 * (-2) + 4 * x[i] * x[i]);
    197:       v[0][1]        = 2 * alpha * (-2 * x[i]);
    198:       v[1][0]        = 2 * alpha * (-2 * x[i]);
    199:       v[1][1]        = 2 * alpha;
    200:       ind[0]         = i;
    201:       ind[1]         = i + 1;
    202:       PetscCall(MatSetValues(H, 2, ind, 2, ind, v[0], ADD_VALUES));
    203:     }
    204:   } else {
    205:     for (PetscInt i = 0; i < user->n / 2; i++) {
    206:       v[1][1] = 2 * alpha;
    207:       v[0][0] = -4 * alpha * (x[2 * i + 1] - 3 * x[2 * i] * x[2 * i]) + 2;
    208:       v[1][0] = v[0][1] = -4.0 * alpha * x[2 * i];
    209:       ind[0]            = 2 * i;
    210:       ind[1]            = 2 * i + 1;
    211:       PetscCall(MatSetValues(H, 2, ind, 2, ind, v[0], INSERT_VALUES));
    212:     }
    213:   }
    214:   PetscCall(VecRestoreArrayRead(X, &x));

    216:   /* Assemble matrix */
    217:   PetscCall(MatAssemblyBegin(H, MAT_FINAL_ASSEMBLY));
    218:   PetscCall(MatAssemblyEnd(H, MAT_FINAL_ASSEMBLY));
    219:   PetscCall(PetscLogFlops(9.0 * user->n / 2.0));
    220:   PetscFunctionReturn(PETSC_SUCCESS);
    221: }