xref: /petsc/src/tao/unconstrained/tutorials/rosenbrock2.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, TAOLMVM));
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(TaoLMVMRecycle(tao, PETSC_TRUE));
80   reason = TAO_CONTINUE_ITERATING;
81   while (reason != TAO_CONVERGED_GATOL) {
82     PetscCall(TaoSolve(tao));
83     PetscCall(TaoGetConvergedReason(tao, &reason));
84     PetscCall(TaoGetIterationNumber(tao, &its));
85     recycled_its += its;
86     PetscCall(PetscPrintf(PETSC_COMM_SELF, "-----------------------\n"));
87   }
88 
89   /* Disable recycling and solve again! */
90   PetscCall(TaoSetMaximumIterations(tao, 100));
91   PetscCall(TaoLMVMRecycle(tao, PETSC_FALSE));
92   PetscCall(VecSet(x, zero));
93   PetscCall(TaoSolve(tao));
94   PetscCall(TaoGetConvergedReason(tao, &reason));
95   PetscCheck(reason == TAO_CONVERGED_GATOL, PETSC_COMM_SELF, PETSC_ERR_NOT_CONVERGED, "Solution failed to converge!");
96   PetscCall(TaoGetIterationNumber(tao, &oneshot_its));
97   PetscCall(PetscPrintf(PETSC_COMM_SELF, "-----------------------\n"));
98   PetscCall(PetscPrintf(PETSC_COMM_SELF, "recycled its: %" PetscInt_FMT " | oneshot its: %" PetscInt_FMT "\n", recycled_its, oneshot_its));
99   PetscCheck(recycled_its == oneshot_its, PETSC_COMM_SELF, PETSC_ERR_NOT_CONVERGED, "LMVM recycling does not work!");
100 
101   PetscCall(TaoDestroy(&tao));
102   PetscCall(VecDestroy(&x));
103   PetscCall(MatDestroy(&H));
104 
105   PetscCall(PetscFinalize());
106   return 0;
107 }
108 
109 /* -------------------------------------------------------------------- */
110 /*
111     FormFunctionGradient - Evaluates the function, f(X), and gradient, G(X).
112 
113     Input Parameters:
114 .   tao  - the Tao context
115 .   X    - input vector
116 .   ptr  - optional user-defined context, as set by TaoSetFunctionGradient()
117 
118     Output Parameters:
119 .   G - vector containing the newly evaluated gradient
120 .   f - function value
121 
122     Note:
123     Some optimization methods ask for the function and the gradient evaluation
124     at the same time.  Evaluating both at once may be more efficient that
125     evaluating each separately.
126 */
127 PetscErrorCode FormFunctionGradient(Tao tao, Vec X, PetscReal *f, Vec G, void *ptr) {
128   AppCtx            *user = (AppCtx *)ptr;
129   PetscInt           i, nn = user->n / 2;
130   PetscReal          ff = 0, t1, t2, alpha = user->alpha;
131   PetscScalar       *g;
132   const PetscScalar *x;
133 
134   PetscFunctionBeginUser;
135   /* Get pointers to vector data */
136   PetscCall(VecGetArrayRead(X, &x));
137   PetscCall(VecGetArray(G, &g));
138 
139   /* Compute G(X) */
140   if (user->chained) {
141     g[0] = 0;
142     for (i = 0; i < user->n - 1; i++) {
143       t1 = x[i + 1] - x[i] * x[i];
144       ff += PetscSqr(1 - x[i]) + alpha * t1 * t1;
145       g[i] += -2 * (1 - x[i]) + 2 * alpha * t1 * (-2 * x[i]);
146       g[i + 1] = 2 * alpha * t1;
147     }
148   } else {
149     for (i = 0; i < nn; i++) {
150       t1 = x[2 * i + 1] - x[2 * i] * x[2 * i];
151       t2 = 1 - x[2 * i];
152       ff += alpha * t1 * t1 + t2 * t2;
153       g[2 * i]     = -4 * alpha * t1 * x[2 * i] - 2.0 * t2;
154       g[2 * i + 1] = 2 * alpha * t1;
155     }
156   }
157 
158   /* Restore vectors */
159   PetscCall(VecRestoreArrayRead(X, &x));
160   PetscCall(VecRestoreArray(G, &g));
161   *f = ff;
162 
163   PetscCall(PetscLogFlops(15.0 * nn));
164   PetscFunctionReturn(0);
165 }
166 
167 /* ------------------------------------------------------------------- */
168 /*
169    FormHessian - Evaluates Hessian matrix.
170 
171    Input Parameters:
172 .  tao   - the Tao context
173 .  x     - input vector
174 .  ptr   - optional user-defined context, as set by TaoSetHessian()
175 
176    Output Parameters:
177 .  H     - Hessian matrix
178 
179    Note:  Providing the Hessian may not be necessary.  Only some solvers
180    require this matrix.
181 */
182 PetscErrorCode FormHessian(Tao tao, Vec X, Mat H, Mat Hpre, void *ptr) {
183   AppCtx            *user = (AppCtx *)ptr;
184   PetscInt           i, ind[2];
185   PetscReal          alpha = user->alpha;
186   PetscReal          v[2][2];
187   const PetscScalar *x;
188   PetscBool          assembled;
189 
190   PetscFunctionBeginUser;
191   /* Zero existing matrix entries */
192   PetscCall(MatAssembled(H, &assembled));
193   if (assembled) PetscCall(MatZeroEntries(H));
194 
195   /* Get a pointer to vector data */
196   PetscCall(VecGetArrayRead(X, &x));
197 
198   /* Compute H(X) entries */
199   if (user->chained) {
200     PetscCall(MatZeroEntries(H));
201     for (i = 0; i < user->n - 1; i++) {
202       PetscScalar t1 = x[i + 1] - x[i] * x[i];
203       v[0][0]        = 2 + 2 * alpha * (t1 * (-2) - 2 * x[i]);
204       v[0][1]        = 2 * alpha * (-2 * x[i]);
205       v[1][0]        = 2 * alpha * (-2 * x[i]);
206       v[1][1]        = 2 * alpha * t1;
207       ind[0]         = i;
208       ind[1]         = i + 1;
209       PetscCall(MatSetValues(H, 2, ind, 2, ind, v[0], ADD_VALUES));
210     }
211   } else {
212     for (i = 0; i < user->n / 2; i++) {
213       v[1][1] = 2 * alpha;
214       v[0][0] = -4 * alpha * (x[2 * i + 1] - 3 * x[2 * i] * x[2 * i]) + 2;
215       v[1][0] = v[0][1] = -4.0 * alpha * x[2 * i];
216       ind[0]            = 2 * i;
217       ind[1]            = 2 * i + 1;
218       PetscCall(MatSetValues(H, 2, ind, 2, ind, v[0], INSERT_VALUES));
219     }
220   }
221   PetscCall(VecRestoreArrayRead(X, &x));
222 
223   /* Assemble matrix */
224   PetscCall(MatAssemblyBegin(H, MAT_FINAL_ASSEMBLY));
225   PetscCall(MatAssemblyEnd(H, MAT_FINAL_ASSEMBLY));
226   PetscCall(PetscLogFlops(9.0 * user->n / 2.0));
227   PetscFunctionReturn(0);
228 }
229 
230 /*TEST
231 
232    build:
233       requires: !complex
234 
235    test:
236       args: -tao_type lmvm -tao_monitor
237       requires: !single
238 
239 TEST*/
240