xref: /petsc/src/binding/petsc4py/demo/legacy/taosolve/rosenbrock.py (revision 5a48edb989d3ea10d6aff6c0e26d581c18691deb)
1"""
2This example demonstrates the use of TAO for Python to solve an
3unconstrained minimization problem on a single processor.
4
5We minimize the extended Rosenbrock function::
6
7   sum_{i=0}^{n/2-1} ( alpha*(x_{2i+1}-x_{2i}^2)^2 + (1-x_{2i})^2 )
8"""
9
10try: range = xrange
11except NameError: pass
12
13# the two lines below are only
14# needed to build options database
15# from command line arguments
16import sys, petsc4py
17petsc4py.init(sys.argv)
18
19from petsc4py import PETSc
20
21
22class AppCtx(object):
23
24    """
25    Extended Rosenbrock function.
26    """
27
28    def __init__(self, n=2, alpha=99.0):
29        self.size  = int(n)
30        self.alpha = float(alpha)
31
32    def formObjective(self, tao, x):
33        #print 'AppCtx.formObjective()'
34        alpha = self.alpha
35        nn = self.size // 2
36        ff = 0.0
37        for i in range(nn):
38            t1 = x[2*i+1] - x[2*i] * x[2*i]
39            t2 = 1 - x[2*i];
40            ff += alpha*t1*t1 + t2*t2;
41        return ff
42
43    def formGradient(self, tao, x, G):
44        #print 'AppCtx.formGradient()'
45        alpha = self.alpha
46        nn = self.size // 2
47        G.zeroEntries()
48        for i in range(nn):
49            t1 = x[2*i+1] - x[2*i] * x[2*i]
50            t2 = 1 - x[2*i];
51            G[2*i]   = -4*alpha*t1*x[2*i] - 2*t2;
52            G[2*i+1] = 2*alpha*t1;
53
54    def formObjGrad(self, tao, x, G):
55        #print 'AppCtx.formObjGrad()'
56        alpha = self.alpha
57        nn = self.size // 2
58        ff = 0.0
59        G.zeroEntries()
60        for i in range(nn):
61            t1 = x[2*i+1] - x[2*i] * x[2*i]
62            t2 = 1 - x[2*i];
63            ff += alpha*t1*t1 + t2*t2;
64            G[2*i]   = -4*alpha*t1*x[2*i] - 2*t2;
65            G[2*i+1] = 2*alpha*t1;
66        return ff
67
68    def formHessian(self, tao, x, H, HP):
69        #print 'AppCtx.formHessian()'
70        alpha = self.alpha
71        nn = self.size // 2
72        idx = [0, 0]
73        v = [[0.0, 0.0],
74             [0.0, 0.0]]
75        H.zeroEntries()
76        for i in range(nn):
77            v[1][1] = 2*alpha
78            v[0][0] = -4*alpha*(x[2*i+1]-3*x[2*i]*x[2*i]) + 2
79            v[1][0] = v[0][1] = -4.0*alpha*x[2*i];
80            idx[0] = 2*i
81            idx[1] = 2*i+1
82            H[idx,idx] = v
83        H.assemble()
84
85# access PETSc options database
86OptDB = PETSc.Options()
87
88
89# create user application context
90# and configure user parameters
91user = AppCtx()
92user.size  = OptDB.getInt (    'n', user.size)
93user.alpha = OptDB.getReal('alpha', user.alpha)
94
95# create solution vector
96x = PETSc.Vec().create(PETSc.COMM_SELF)
97x.setSizes(user.size)
98x.setFromOptions()
99
100# create Hessian matrix
101H = PETSc.Mat().create(PETSc.COMM_SELF)
102H.setSizes([user.size, user.size])
103H.setFromOptions()
104H.setOption(PETSc.Mat.Option.SYMMETRIC, True)
105H.setUp()
106
107# pass the following to command line:
108#  $ ... -methods nm,lmvm,nls,ntr,cg,blmvm,tron
109# to try many methods
110methods = OptDB.getString('methods', '')
111methods = methods.split(',')
112for meth in methods:
113    # create TAO Solver
114    tao = PETSc.TAO().create(PETSc.COMM_SELF)
115    if meth: tao.setType(meth)
116    tao.setFromOptions()
117    # solve the problem
118    tao.setObjectiveGradient(user.formObjGrad)
119    tao.setObjective(user.formObjective)
120    tao.setGradient(user.formGradient)
121    tao.setHessian(user.formHessian, H)
122    #app.getKSP().getPC().setFromOptions()
123    x.set(0) # zero initial guess
124    tao.setSolution(x)
125    tao.solve()
126    tao.destroy()
127
128## # this is just for testing
129## x     = app.getSolution()
130## G     = app.getGradient()
131## H, HP = app.getHessian()
132
133## f = tao.computeObjective(x)
134## tao.computeGradient(x, G)
135## f = tao.computeObjectiveGradient(x, G)
136## tao.computeHessian(x, H, HP)
137