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 PETSc.garbage_cleanup() 39 40 def testSetRoutinesToNone(self): 41 tao = self.tao 42 objective, gradient, objgrad = None, None, None 43 constraint, varbounds = None, None 44 hessian, jacobian = None, None 45 tao.setObjective(objective) 46 tao.setGradient(gradient,None) 47 tao.setVariableBounds(varbounds) 48 tao.setObjectiveGradient(objgrad,None) 49 tao.setConstraints(constraint) 50 tao.setHessian(hessian) 51 tao.setJacobian(jacobian) 52 53 def testGetVecsAndMats(self): 54 tao = self.tao 55 x = tao.getSolution() 56 (g, _) = tao.getGradient() 57 l, u = tao.getVariableBounds() 58 r = None#tao.getConstraintVec() 59 H, HP = None,None#tao.getHessianMat() 60 J, JP = None,None#tao.getJacobianMat() 61 for o in [x, g, r, l, u ,H, HP, J, JP,]: 62 self.assertFalse(o) 63 64 def testGetKSP(self): 65 ksp = self.tao.getKSP() 66 self.assertFalse(ksp) 67 68 def testEqualityConstraints(self): 69 if self.tao.getComm().Get_size() > 1: 70 return 71 tao = self.tao 72 73 x = PETSc.Vec().create(tao.getComm()) 74 x.setType('standard') 75 x.setSizes(2) 76 c = PETSc.Vec().create(tao.getComm()) 77 c.setSizes(1) 78 c.setType(x.getType()) 79 J = PETSc.Mat().create(tao.getComm()) 80 J.setSizes([1, 2]) 81 J.setType(PETSc.Mat.Type.DENSE) 82 J.setUp() 83 84 tao.setObjective(Objective()) 85 tao.setGradient(Gradient(),None) 86 tao.setEqualityConstraints(EqConstraints(),c) 87 tao.setJacobianEquality(EqJacobian(),J,J) 88 tao.setSolution(x) 89 tao.setType(PETSc.TAO.Type.ALMM) 90 tao.setTolerances(gatol=1.e-4) 91 tao.setFromOptions() 92 tao.solve() 93 self.assertAlmostEqual(abs(x[0]**2 + x[1] - 2.0), 0.0, places=4) 94 95# -------------------------------------------------------------------- 96 97class TestTAOSelf(BaseTestTAO, unittest.TestCase): 98 COMM = PETSc.COMM_SELF 99 100class TestTAOWorld(BaseTestTAO, unittest.TestCase): 101 COMM = PETSc.COMM_WORLD 102 103# -------------------------------------------------------------------- 104 105import numpy 106if numpy.iscomplexobj(PETSc.ScalarType()): 107 del BaseTestTAO 108 del TestTAOSelf 109 del TestTAOWorld 110 111if __name__ == '__main__': 112 unittest.main() 113