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