xref: /petsc/src/tao/unconstrained/tutorials/rosenbrock3.c (revision 58d68138c660dfb4e9f5b03334792cd4f2ffd7cc)
1 /* Program usage: mpiexec -n 1 rosenbrock2 [-help] [all TAO options] */
2 
3 /*  Include "petsctao.h" so we can use TAO solvers.  */
4 #include <petsctao.h>
5 
6 static char help[] = "This example demonstrates use of the TAO package to \n\
7 solve an unconstrained minimization problem on a single processor.  We \n\
8 minimize the extended Rosenbrock function: \n\
9    sum_{i=0}^{n/2-1} (alpha*(x_{2i+1}-x_{2i}^2)^2 + (1-x_{2i})^2) \n\
10 or the chained Rosenbrock function:\n\
11    sum_{i=0}^{n-1} alpha*(x_{i+1} - x_i^2)^2 + (1 - x_i)^2\n";
12 
13 /*
14    User-defined application context - contains data needed by the
15    application-provided call-back routines that evaluate the function,
16    gradient, and hessian.
17 */
18 typedef struct {
19   PetscInt  n;     /* dimension */
20   PetscReal alpha; /* condition parameter */
21   PetscBool chained;
22 } AppCtx;
23 
24 /* -------------- User-defined routines ---------- */
25 PetscErrorCode FormFunctionGradient(Tao, Vec, PetscReal *, Vec, void *);
26 PetscErrorCode FormHessian(Tao, Vec, Mat, Mat, void *);
27 
28 int main(int argc, char **argv) {
29   PetscReal          zero = 0.0;
30   Vec                x; /* solution vector */
31   Mat                H;
32   Tao                tao; /* Tao solver context */
33   PetscBool          flg, test_lmvm = PETSC_FALSE;
34   PetscMPIInt        size; /* number of processes running */
35   AppCtx             user; /* user-defined application context */
36   TaoConvergedReason reason;
37   PetscInt           its, recycled_its = 0, oneshot_its = 0;
38 
39   /* Initialize TAO and PETSc */
40   PetscFunctionBeginUser;
41   PetscCall(PetscInitialize(&argc, &argv, (char *)0, help));
42   PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
43   PetscCheck(size == 1, PETSC_COMM_WORLD, PETSC_ERR_WRONG_MPI_SIZE, "Incorrect number of processors");
44 
45   /* Initialize problem parameters */
46   user.n       = 2;
47   user.alpha   = 99.0;
48   user.chained = PETSC_FALSE;
49   /* Check for command line arguments to override defaults */
50   PetscCall(PetscOptionsGetInt(NULL, NULL, "-n", &user.n, &flg));
51   PetscCall(PetscOptionsGetReal(NULL, NULL, "-alpha", &user.alpha, &flg));
52   PetscCall(PetscOptionsGetBool(NULL, NULL, "-chained", &user.chained, &flg));
53   PetscCall(PetscOptionsGetBool(NULL, NULL, "-test_lmvm", &test_lmvm, &flg));
54 
55   /* Allocate vectors for the solution and gradient */
56   PetscCall(VecCreateSeq(PETSC_COMM_SELF, user.n, &x));
57   PetscCall(MatCreateSeqBAIJ(PETSC_COMM_SELF, 2, user.n, user.n, 1, NULL, &H));
58 
59   /* The TAO code begins here */
60 
61   /* Create TAO solver with desired solution method */
62   PetscCall(TaoCreate(PETSC_COMM_SELF, &tao));
63   PetscCall(TaoSetType(tao, TAOBQNLS));
64 
65   /* Set solution vec and an initial guess */
66   PetscCall(VecSet(x, zero));
67   PetscCall(TaoSetSolution(tao, x));
68 
69   /* Set routines for function, gradient, hessian evaluation */
70   PetscCall(TaoSetObjectiveAndGradient(tao, NULL, FormFunctionGradient, &user));
71   PetscCall(TaoSetHessian(tao, H, H, FormHessian, &user));
72 
73   /* Check for TAO command line options */
74   PetscCall(TaoSetFromOptions(tao));
75 
76   /* Solve the problem */
77   PetscCall(TaoSetTolerances(tao, 1.e-5, 0.0, 0.0));
78   PetscCall(TaoSetMaximumIterations(tao, 5));
79   PetscCall(TaoSetRecycleHistory(tao, PETSC_TRUE));
80   reason = TAO_CONTINUE_ITERATING;
81   flg    = PETSC_FALSE;
82   PetscCall(TaoGetRecycleHistory(tao, &flg));
83   if (flg) PetscCall(PetscPrintf(PETSC_COMM_SELF, "Recycle: enabled\n"));
84   while (reason != TAO_CONVERGED_GATOL) {
85     PetscCall(TaoSolve(tao));
86     PetscCall(TaoGetConvergedReason(tao, &reason));
87     PetscCall(TaoGetIterationNumber(tao, &its));
88     recycled_its += its;
89     PetscCall(PetscPrintf(PETSC_COMM_SELF, "-----------------------\n"));
90   }
91 
92   /* Disable recycling and solve again! */
93   PetscCall(TaoSetMaximumIterations(tao, 100));
94   PetscCall(TaoSetRecycleHistory(tao, PETSC_FALSE));
95   PetscCall(VecSet(x, zero));
96   PetscCall(TaoGetRecycleHistory(tao, &flg));
97   if (!flg) PetscCall(PetscPrintf(PETSC_COMM_SELF, "Recycle: disabled\n"));
98   PetscCall(TaoSolve(tao));
99   PetscCall(TaoGetConvergedReason(tao, &reason));
100   PetscCheck(reason == TAO_CONVERGED_GATOL, PETSC_COMM_SELF, PETSC_ERR_NOT_CONVERGED, "Solution failed to converge!");
101   PetscCall(TaoGetIterationNumber(tao, &oneshot_its));
102   PetscCall(PetscPrintf(PETSC_COMM_SELF, "-----------------------\n"));
103   PetscCall(PetscPrintf(PETSC_COMM_SELF, "recycled its: %" PetscInt_FMT " | oneshot its: %" PetscInt_FMT "\n", recycled_its, oneshot_its));
104   PetscCheck(recycled_its == oneshot_its, PETSC_COMM_SELF, PETSC_ERR_NOT_CONVERGED, "Recycled solution does not match oneshot solution!");
105 
106   PetscCall(TaoDestroy(&tao));
107   PetscCall(VecDestroy(&x));
108   PetscCall(MatDestroy(&H));
109 
110   PetscCall(PetscFinalize());
111   return 0;
112 }
113 
114 /* -------------------------------------------------------------------- */
115 /*
116     FormFunctionGradient - Evaluates the function, f(X), and gradient, G(X).
117 
118     Input Parameters:
119 .   tao  - the Tao context
120 .   X    - input vector
121 .   ptr  - optional user-defined context, as set by TaoSetFunctionGradient()
122 
123     Output Parameters:
124 .   G - vector containing the newly evaluated gradient
125 .   f - function value
126 
127     Note:
128     Some optimization methods ask for the function and the gradient evaluation
129     at the same time.  Evaluating both at once may be more efficient than
130     evaluating each separately.
131 */
132 PetscErrorCode FormFunctionGradient(Tao tao, Vec X, PetscReal *f, Vec G, void *ptr) {
133   AppCtx            *user = (AppCtx *)ptr;
134   PetscInt           i, nn = user->n / 2;
135   PetscReal          ff = 0, t1, t2, alpha = user->alpha;
136   PetscScalar       *g;
137   const PetscScalar *x;
138 
139   PetscFunctionBeginUser;
140   /* Get pointers to vector data */
141   PetscCall(VecGetArrayRead(X, &x));
142   PetscCall(VecGetArrayWrite(G, &g));
143 
144   /* Compute G(X) */
145   if (user->chained) {
146     g[0] = 0;
147     for (i = 0; i < user->n - 1; i++) {
148       t1 = x[i + 1] - x[i] * x[i];
149       ff += PetscSqr(1 - x[i]) + alpha * t1 * t1;
150       g[i] += -2 * (1 - x[i]) + 2 * alpha * t1 * (-2 * x[i]);
151       g[i + 1] = 2 * alpha * t1;
152     }
153   } else {
154     for (i = 0; i < nn; i++) {
155       t1 = x[2 * i + 1] - x[2 * i] * x[2 * i];
156       t2 = 1 - x[2 * i];
157       ff += alpha * t1 * t1 + t2 * t2;
158       g[2 * i]     = -4 * alpha * t1 * x[2 * i] - 2.0 * t2;
159       g[2 * i + 1] = 2 * alpha * t1;
160     }
161   }
162 
163   /* Restore vectors */
164   PetscCall(VecRestoreArrayRead(X, &x));
165   PetscCall(VecRestoreArrayWrite(G, &g));
166   *f = ff;
167 
168   PetscCall(PetscLogFlops(15.0 * nn));
169   PetscFunctionReturn(0);
170 }
171 
172 /* ------------------------------------------------------------------- */
173 /*
174    FormHessian - Evaluates Hessian matrix.
175 
176    Input Parameters:
177 .  tao   - the Tao context
178 .  x     - input vector
179 .  ptr   - optional user-defined context, as set by TaoSetHessian()
180 
181    Output Parameters:
182 .  H     - Hessian matrix
183 
184    Note:  Providing the Hessian may not be necessary.  Only some solvers
185    require this matrix.
186 */
187 PetscErrorCode FormHessian(Tao tao, Vec X, Mat H, Mat Hpre, void *ptr) {
188   AppCtx            *user = (AppCtx *)ptr;
189   PetscInt           i, ind[2];
190   PetscReal          alpha = user->alpha;
191   PetscReal          v[2][2];
192   const PetscScalar *x;
193   PetscBool          assembled;
194 
195   PetscFunctionBeginUser;
196   /* Zero existing matrix entries */
197   PetscCall(MatAssembled(H, &assembled));
198   if (assembled || user->chained) PetscCall(MatZeroEntries(H));
199 
200   /* Get a pointer to vector data */
201   PetscCall(VecGetArrayRead(X, &x));
202 
203   /* Compute H(X) entries */
204   if (user->chained) {
205     for (i = 0; i < user->n - 1; i++) {
206       PetscScalar t1 = x[i + 1] - x[i] * x[i];
207       v[0][0]        = 2 + 2 * alpha * (t1 * (-2) - 2 * x[i]);
208       v[0][1]        = 2 * alpha * (-2 * x[i]);
209       v[1][0]        = 2 * alpha * (-2 * x[i]);
210       v[1][1]        = 2 * alpha * t1;
211       ind[0]         = i;
212       ind[1]         = i + 1;
213       PetscCall(MatSetValues(H, 2, ind, 2, ind, v[0], ADD_VALUES));
214     }
215   } else {
216     for (i = 0; i < user->n / 2; i++) {
217       v[1][1] = 2 * alpha;
218       v[0][0] = -4 * alpha * (x[2 * i + 1] - 3 * x[2 * i] * x[2 * i]) + 2;
219       v[1][0] = v[0][1] = -4.0 * alpha * x[2 * i];
220       ind[0]            = 2 * i;
221       ind[1]            = 2 * i + 1;
222       PetscCall(MatSetValues(H, 2, ind, 2, ind, v[0], INSERT_VALUES));
223     }
224   }
225   PetscCall(VecRestoreArrayRead(X, &x));
226 
227   /* Assemble matrix */
228   PetscCall(MatAssemblyBegin(H, MAT_FINAL_ASSEMBLY));
229   PetscCall(MatAssemblyEnd(H, MAT_FINAL_ASSEMBLY));
230   PetscCall(PetscLogFlops(9.0 * user->n / 2.0));
231   PetscFunctionReturn(0);
232 }
233 
234 /*TEST
235 
236    build:
237       requires: !complex
238 
239    test:
240       args: -tao_type bqnls -tao_monitor
241       requires: !single
242 
243 TEST*/
244