1# -------------------------------------------------------------------- 2 3from petsc4py import PETSc 4import unittest 5 6# -------------------------------------------------------------------- 7class Objective: 8 def __call__(self, tao, x): 9 return (x[0] - 2.0)**2 + (x[1] - 2.0)**2 - 2.0*(x[0] + x[1]) 10 11class Gradient: 12 def __call__(self, tao, x, g): 13 g[0] = 2.0*(x[0] - 2.0) - 2.0 14 g[1] = 2.0*(x[1] - 2.0) - 2.0 15 g.assemble() 16 17class EqConstraints: 18 def __call__(self, tao, x, c): 19 c[0] = x[0]**2 + x[1] - 2.0 20 c.assemble() 21 22class EqJacobian: 23 def __call__(self, tao, x, J, P): 24 P[0,0] = 2.0*x[0] 25 P[0,1] = 1.0 26 P.assemble() 27 if J != P: J.assemble() 28 29class BaseTestTAO(object): 30 31 COMM = None 32 33 def setUp(self): 34 self.tao = PETSc.TAO().create(comm=self.COMM) 35 36 def tearDown(self): 37 self.tao = None 38 39 def testSetRoutinesToNone(self): 40 tao = self.tao 41 objective, gradient, objgrad = None, None, None 42 constraint, varbounds = None, None 43 hessian, jacobian = None, None 44 tao.setObjective(objective) 45 tao.setGradient(gradient,None) 46 tao.setVariableBounds(varbounds) 47 tao.setObjectiveGradient(objgrad,None) 48 tao.setConstraints(constraint) 49 tao.setHessian(hessian) 50 tao.setJacobian(jacobian) 51 52 def testGetVecsAndMats(self): 53 tao = self.tao 54 x = tao.getSolution() 55 (g, _) = tao.getGradient() 56 l, u = tao.getVariableBounds() 57 r = None#tao.getConstraintVec() 58 H, HP = None,None#tao.getHessianMat() 59 J, JP = None,None#tao.getJacobianMat() 60 for o in [x, g, r, l, u ,H, HP, J, JP,]: 61 self.assertFalse(o) 62 63 def testGetKSP(self): 64 ksp = self.tao.getKSP() 65 self.assertFalse(ksp) 66 67 def testEqualityConstraints(self): 68 if self.tao.getComm().Get_size() > 1: 69 return 70 tao = self.tao 71 72 x = PETSc.Vec().create(tao.getComm()) 73 x.setType('standard') 74 x.setSizes(2) 75 c = PETSc.Vec().create(tao.getComm()) 76 c.setSizes(1) 77 c.setType(x.getType()) 78 J = PETSc.Mat().create(tao.getComm()) 79 J.setSizes([1, 2]) 80 J.setType(PETSc.Mat.Type.DENSE) 81 J.setUp() 82 83 tao.setObjective(Objective()) 84 tao.setGradient(Gradient(),None) 85 tao.setEqualityConstraints(EqConstraints(),c) 86 tao.setJacobianEquality(EqJacobian(),J,J) 87 tao.setSolution(x) 88 tao.setType(PETSc.TAO.Type.ALMM) 89 tao.setTolerances(gatol=1.e-4) 90 tao.setFromOptions() 91 tao.solve() 92 self.assertAlmostEqual(abs(x[0]**2 + x[1] - 2.0), 0.0, places=4) 93 94# -------------------------------------------------------------------- 95 96class TestTAOSelf(BaseTestTAO, unittest.TestCase): 97 COMM = PETSc.COMM_SELF 98 99class TestTAOWorld(BaseTestTAO, unittest.TestCase): 100 COMM = PETSc.COMM_WORLD 101 102# -------------------------------------------------------------------- 103 104import numpy 105if numpy.iscomplexobj(PETSc.ScalarType()): 106 del BaseTestTAO 107 del TestTAOSelf 108 del TestTAOWorld 109 110if __name__ == '__main__': 111 unittest.main() 112