xref: /petsc/src/tao/unconstrained/tutorials/rosenbrock2.c (revision c3b5f7ba6bc5ce25a01a67bb37ba5d34b02bbbd7)
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 {
30   PetscReal          zero=0.0;
31   Vec                x;                     /* solution vector */
32   Mat                H;
33   Tao                tao;                   /* Tao solver context */
34   PetscBool          flg, test_lmvm = PETSC_FALSE;
35   PetscMPIInt        size;                  /* number of processes running */
36   AppCtx             user;                  /* user-defined application context */
37   TaoConvergedReason reason;
38   PetscInt           its, recycled_its=0, oneshot_its=0;
39 
40   /* Initialize TAO and PETSc */
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; user.alpha = 99.0; user.chained = PETSC_FALSE;
47   /* Check for command line arguments to override defaults */
48   PetscCall(PetscOptionsGetInt(NULL,NULL,"-n",&user.n,&flg));
49   PetscCall(PetscOptionsGetReal(NULL,NULL,"-alpha",&user.alpha,&flg));
50   PetscCall(PetscOptionsGetBool(NULL,NULL,"-chained",&user.chained,&flg));
51   PetscCall(PetscOptionsGetBool(NULL,NULL,"-test_lmvm",&test_lmvm,&flg));
52 
53   /* Allocate vectors for the solution and gradient */
54   PetscCall(VecCreateSeq(PETSC_COMM_SELF,user.n,&x));
55   PetscCall(MatCreateSeqBAIJ(PETSC_COMM_SELF,2,user.n,user.n,1,NULL,&H));
56 
57   /* The TAO code begins here */
58 
59   /* Create TAO solver with desired solution method */
60   PetscCall(TaoCreate(PETSC_COMM_SELF,&tao));
61   PetscCall(TaoSetType(tao,TAOLMVM));
62 
63   /* Set solution vec and an initial guess */
64   PetscCall(VecSet(x, zero));
65   PetscCall(TaoSetSolution(tao,x));
66 
67   /* Set routines for function, gradient, hessian evaluation */
68   PetscCall(TaoSetObjectiveAndGradient(tao,NULL,FormFunctionGradient,&user));
69   PetscCall(TaoSetHessian(tao,H,H,FormHessian,&user));
70 
71   /* Check for TAO command line options */
72   PetscCall(TaoSetFromOptions(tao));
73 
74   /* Solve the problem */
75   PetscCall(TaoSetTolerances(tao, 1.e-5, 0.0, 0.0));
76   PetscCall(TaoSetMaximumIterations(tao, 5));
77   PetscCall(TaoLMVMRecycle(tao, PETSC_TRUE));
78   reason = TAO_CONTINUE_ITERATING;
79   while (reason != TAO_CONVERGED_GATOL) {
80     PetscCall(TaoSolve(tao));
81     PetscCall(TaoGetConvergedReason(tao, &reason));
82     PetscCall(TaoGetIterationNumber(tao, &its));
83     recycled_its += its;
84     PetscCall(PetscPrintf(PETSC_COMM_SELF, "-----------------------\n"));
85   }
86 
87   /* Disable recycling and solve again! */
88   PetscCall(TaoSetMaximumIterations(tao, 100));
89   PetscCall(TaoLMVMRecycle(tao, PETSC_FALSE));
90   PetscCall(VecSet(x, zero));
91   PetscCall(TaoSolve(tao));
92   PetscCall(TaoGetConvergedReason(tao, &reason));
93   PetscCheck(reason == TAO_CONVERGED_GATOL,PETSC_COMM_SELF, PETSC_ERR_NOT_CONVERGED, "Solution failed to converge!");
94   PetscCall(TaoGetIterationNumber(tao, &oneshot_its));
95   PetscCall(PetscPrintf(PETSC_COMM_SELF, "-----------------------\n"));
96   PetscCall(PetscPrintf(PETSC_COMM_SELF, "recycled its: %D | oneshot its: %D\n", recycled_its, oneshot_its));
97   PetscCheck(recycled_its == oneshot_its,PETSC_COMM_SELF, PETSC_ERR_NOT_CONVERGED, "LMVM recycling does not work!");
98 
99   PetscCall(TaoDestroy(&tao));
100   PetscCall(VecDestroy(&x));
101   PetscCall(MatDestroy(&H));
102 
103   PetscCall(PetscFinalize());
104   return 0;
105 }
106 
107 /* -------------------------------------------------------------------- */
108 /*
109     FormFunctionGradient - Evaluates the function, f(X), and gradient, G(X).
110 
111     Input Parameters:
112 .   tao  - the Tao context
113 .   X    - input vector
114 .   ptr  - optional user-defined context, as set by TaoSetFunctionGradient()
115 
116     Output Parameters:
117 .   G - vector containing the newly evaluated gradient
118 .   f - function value
119 
120     Note:
121     Some optimization methods ask for the function and the gradient evaluation
122     at the same time.  Evaluating both at once may be more efficient that
123     evaluating each separately.
124 */
125 PetscErrorCode FormFunctionGradient(Tao tao,Vec X,PetscReal *f, Vec G,void *ptr)
126 {
127   AppCtx            *user = (AppCtx *) ptr;
128   PetscInt          i,nn=user->n/2;
129   PetscReal         ff=0,t1,t2,alpha=user->alpha;
130   PetscScalar       *g;
131   const PetscScalar *x;
132 
133   PetscFunctionBeginUser;
134   /* Get pointers to vector data */
135   PetscCall(VecGetArrayRead(X,&x));
136   PetscCall(VecGetArray(G,&g));
137 
138   /* Compute G(X) */
139   if (user->chained) {
140     g[0] = 0;
141     for (i=0; i<user->n-1; i++) {
142       t1 = x[i+1] - x[i]*x[i];
143       ff += PetscSqr(1 - x[i]) + alpha*t1*t1;
144       g[i] += -2*(1 - x[i]) + 2*alpha*t1*(-2*x[i]);
145       g[i+1] = 2*alpha*t1;
146     }
147   } else {
148     for (i=0; i<nn; i++) {
149       t1 = x[2*i+1]-x[2*i]*x[2*i]; 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   }
155 
156   /* Restore vectors */
157   PetscCall(VecRestoreArrayRead(X,&x));
158   PetscCall(VecRestoreArray(G,&g));
159   *f   = ff;
160 
161   PetscCall(PetscLogFlops(15.0*nn));
162   PetscFunctionReturn(0);
163 }
164 
165 /* ------------------------------------------------------------------- */
166 /*
167    FormHessian - Evaluates Hessian matrix.
168 
169    Input Parameters:
170 .  tao   - the Tao context
171 .  x     - input vector
172 .  ptr   - optional user-defined context, as set by TaoSetHessian()
173 
174    Output Parameters:
175 .  H     - Hessian matrix
176 
177    Note:  Providing the Hessian may not be necessary.  Only some solvers
178    require this matrix.
179 */
180 PetscErrorCode FormHessian(Tao tao,Vec X,Mat H, Mat Hpre, void *ptr)
181 {
182   AppCtx            *user = (AppCtx*)ptr;
183   PetscInt          i, ind[2];
184   PetscReal         alpha=user->alpha;
185   PetscReal         v[2][2];
186   const PetscScalar *x;
187   PetscBool         assembled;
188 
189   PetscFunctionBeginUser;
190   /* Zero existing matrix entries */
191   PetscCall(MatAssembled(H,&assembled));
192   if (assembled) PetscCall(MatZeroEntries(H));
193 
194   /* Get a pointer to vector data */
195   PetscCall(VecGetArrayRead(X,&x));
196 
197   /* Compute H(X) entries */
198   if (user->chained) {
199     PetscCall(MatZeroEntries(H));
200     for (i=0; i<user->n-1; i++) {
201       PetscScalar t1 = x[i+1] - x[i]*x[i];
202       v[0][0] = 2 + 2*alpha*(t1*(-2) - 2*x[i]);
203       v[0][1] = 2*alpha*(-2*x[i]);
204       v[1][0] = 2*alpha*(-2*x[i]);
205       v[1][1] = 2*alpha*t1;
206       ind[0] = i; ind[1] = i+1;
207       PetscCall(MatSetValues(H,2,ind,2,ind,v[0],ADD_VALUES));
208     }
209   } else {
210     for (i=0; i<user->n/2; i++) {
211       v[1][1] = 2*alpha;
212       v[0][0] = -4*alpha*(x[2*i+1]-3*x[2*i]*x[2*i]) + 2;
213       v[1][0] = v[0][1] = -4.0*alpha*x[2*i];
214       ind[0]=2*i; ind[1]=2*i+1;
215       PetscCall(MatSetValues(H,2,ind,2,ind,v[0],INSERT_VALUES));
216     }
217   }
218   PetscCall(VecRestoreArrayRead(X,&x));
219 
220   /* Assemble matrix */
221   PetscCall(MatAssemblyBegin(H,MAT_FINAL_ASSEMBLY));
222   PetscCall(MatAssemblyEnd(H,MAT_FINAL_ASSEMBLY));
223   PetscCall(PetscLogFlops(9.0*user->n/2.0));
224   PetscFunctionReturn(0);
225 }
226 
227 /*TEST
228 
229    build:
230       requires: !complex
231 
232    test:
233       args: -tao_type lmvm -tao_monitor
234       requires: !single
235 
236 TEST*/
237