1# -------------------------------------------------------------------- 2 3from petsc4py import PETSc 4import unittest 5import numpy 6 7 8# -------------------------------------------------------------------- 9class Objective: 10 def __call__(self, tao, x): 11 return (x[0] - 2.0) ** 2 + (x[1] - 2.0) ** 2 - 2.0 * (x[0] + x[1]) 12 13 14class Gradient: 15 def __call__(self, tao, x, g): 16 g[0] = 2.0 * (x[0] - 2.0) - 2.0 17 g[1] = 2.0 * (x[1] - 2.0) - 2.0 18 g.assemble() 19 20 21class EqConstraints: 22 def __call__(self, tao, x, c): 23 c[0] = x[0] ** 2 + x[1] - 2.0 24 c.assemble() 25 26 27class EqJacobian: 28 def __call__(self, tao, x, J, P): 29 P[0, 0] = 2.0 * x[0] 30 P[0, 1] = 1.0 31 P.assemble() 32 if J != P: 33 J.assemble() 34 35 36class BaseTestTAO: 37 COMM = None 38 39 def setUp(self): 40 self.tao = PETSc.TAO().create(comm=self.COMM) 41 42 def tearDown(self): 43 self.tao = None 44 PETSc.garbage_cleanup() 45 46 def testSetRoutinesToNone(self): 47 tao = self.tao 48 objective, gradient, objgrad = None, None, None 49 constraint, varbounds = None, None 50 hessian, jacobian = None, None 51 tao.setObjective(objective) 52 tao.setGradient(gradient, None) 53 tao.setVariableBounds(varbounds) 54 tao.setObjectiveGradient(objgrad, None) 55 tao.setConstraints(constraint) 56 tao.setHessian(hessian) 57 tao.setJacobian(jacobian) 58 59 def testGetVecsAndMats(self): 60 tao = self.tao 61 x = tao.getSolution() 62 (g, _) = tao.getGradient() 63 low, up = tao.getVariableBounds() 64 r = None # tao.getConstraintVec() 65 H, HP = None, None # tao.getHessianMat() 66 J, JP = None, None # tao.getJacobianMat() 67 for o in [ 68 x, 69 g, 70 r, 71 low, 72 up, 73 H, 74 HP, 75 J, 76 JP, 77 ]: 78 self.assertFalse(o) 79 80 def testGetKSP(self): 81 ksp = self.tao.getKSP() 82 self.assertFalse(ksp) 83 84 def testEqualityConstraints(self): 85 if self.tao.getComm().Get_size() > 1: 86 return 87 tao = self.tao 88 89 x = PETSc.Vec().create(tao.getComm()) 90 x.setType('standard') 91 x.setSizes(2) 92 c = PETSc.Vec().create(tao.getComm()) 93 c.setSizes(1) 94 c.setType(x.getType()) 95 J = PETSc.Mat().create(tao.getComm()) 96 J.setSizes([1, 2]) 97 J.setType(PETSc.Mat.Type.DENSE) 98 J.setUp() 99 100 tao.setObjective(Objective()) 101 tao.setGradient(Gradient(), None) 102 tao.setEqualityConstraints(EqConstraints(), c) 103 tao.setJacobianEquality(EqJacobian(), J, J) 104 tao.setSolution(x) 105 tao.setType(PETSc.TAO.Type.ALMM) 106 tao.setTolerances(gatol=1.0e-4) 107 tao.setFromOptions() 108 tao.solve() 109 self.assertAlmostEqual(abs(x[0] ** 2 + x[1] - 2.0), 0.0, places=4) 110 111 def testBNCG(self): 112 if self.tao.getComm().Get_size() > 1: 113 return 114 tao = self.tao 115 116 x = PETSc.Vec().create(tao.getComm()) 117 x.setType('standard') 118 x.setSizes(2) 119 xl = PETSc.Vec().create(tao.getComm()) 120 xl.setType('standard') 121 xl.setSizes(2) 122 xl.set(0.0) 123 xu = PETSc.Vec().create(tao.getComm()) 124 xu.setType('standard') 125 xu.setSizes(2) 126 xu.set(2.0) 127 tao.setVariableBounds((xl, xu)) 128 tao.setObjective(Objective()) 129 tao.setGradient(Gradient(), None) 130 tao.setSolution(x) 131 tao.setType(PETSc.TAO.Type.BNCG) 132 tao.setTolerances(gatol=1.0e-4) 133 ls = tao.getLineSearch() 134 ls.setType(PETSc.TAOLineSearch.Type.UNIT) 135 tao.setFromOptions() 136 tao.solve() 137 self.assertAlmostEqual(x[0], 2.0, places=4) 138 self.assertAlmostEqual(x[1], 2.0, places=4) 139 140 141# -------------------------------------------------------------------- 142 143 144class TestTAOSelf(BaseTestTAO, unittest.TestCase): 145 COMM = PETSc.COMM_SELF 146 147 148class TestTAOWorld(BaseTestTAO, unittest.TestCase): 149 COMM = PETSc.COMM_WORLD 150 151 152# -------------------------------------------------------------------- 153 154 155if numpy.iscomplexobj(PETSc.ScalarType()): 156 del BaseTestTAO 157 del TestTAOSelf 158 del TestTAOWorld 159 160if __name__ == '__main__': 161 unittest.main() 162