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