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 InEqConstraints: 37 def __call__(self, tao, x, c): 38 c[0] = x[1] - x[0] ** 2 39 c.assemble() 40 41 42class InEqJacobian: 43 def __call__(self, tao, x, J, P): 44 P[0, 0] = -2.0 * x[0] 45 P[0, 1] = 1.0 46 P.assemble() 47 if J != P: 48 J.assemble() 49 50 51class BaseTestTAO: 52 COMM = None 53 54 def setUp(self): 55 self.tao = PETSc.TAO().create(comm=self.COMM) 56 57 def tearDown(self): 58 self.tao = None 59 PETSc.garbage_cleanup() 60 61 def testSetRoutinesToNone(self): 62 tao = self.tao 63 objective, gradient, objgrad = None, None, None 64 constraint, varbounds = None, None 65 hessian, jacobian = None, None 66 tao.setObjective(objective) 67 tao.setGradient(gradient, None) 68 tao.setVariableBounds(varbounds) 69 tao.setObjectiveGradient(objgrad, None) 70 tao.setConstraints(constraint) 71 tao.setHessian(hessian) 72 tao.setJacobian(jacobian) 73 74 def testGetVecsAndMats(self): 75 tao = self.tao 76 x = tao.getSolution() 77 (g, _) = tao.getGradient() 78 low, up = tao.getVariableBounds() 79 r = None # tao.getConstraintVec() 80 H, HP = None, None # tao.getHessianMat() 81 J, JP = None, None # tao.getJacobianMat() 82 for o in [ 83 x, 84 g, 85 r, 86 low, 87 up, 88 H, 89 HP, 90 J, 91 JP, 92 ]: 93 self.assertFalse(o) 94 95 def testGetKSP(self): 96 ksp = self.tao.getKSP() 97 self.assertFalse(ksp) 98 99 def testEqualityConstraints(self): 100 if self.tao.getComm().Get_size() > 1: 101 return 102 tao = self.tao 103 104 x = PETSc.Vec().create(tao.getComm()) 105 x.setType('standard') 106 x.setSizes(2) 107 c = PETSc.Vec().create(tao.getComm()) 108 c.setSizes(1) 109 c.setType(x.getType()) 110 J = PETSc.Mat().create(tao.getComm()) 111 J.setSizes([1, 2]) 112 J.setType(PETSc.Mat.Type.DENSE) 113 J.setUp() 114 115 tao.setObjective(Objective()) 116 tao.setGradient(Gradient(), None) 117 tao.setEqualityConstraints(EqConstraints(), c) 118 tao.setJacobianEquality(EqJacobian(), J, J) 119 tao.setSolution(x) 120 tao.setType(PETSc.TAO.Type.ALMM) 121 tao.setTolerances(gatol=1.0e-4) 122 tao.setFromOptions() 123 tao.solve() 124 self.assertAlmostEqual(abs(x[0] ** 2 + x[1] - 2.0), 0.0, places=4) 125 126 def testInequlityConstraints(self): 127 if self.tao.getComm().Get_size() > 1: 128 return 129 tao = self.tao 130 131 x = PETSc.Vec().create(tao.getComm()) 132 x.setType('standard') 133 x.setSizes(2) 134 c = PETSc.Vec().create(tao.getComm()) 135 c.setSizes(1) 136 c.setType(x.getType()) 137 J = PETSc.Mat().create(tao.getComm()) 138 J.setSizes([1, 2]) 139 J.setType(PETSc.Mat.Type.DENSE) 140 J.setUp() 141 142 tao.setObjective(Objective()) 143 tao.setGradient(Gradient(), None) 144 tao.setInequalityConstraints(InEqConstraints(), c) 145 tao.setJacobianInequality(InEqJacobian(), J, J) 146 tao.setSolution(x) 147 tao.setType(PETSc.TAO.Type.ALMM) 148 tao.setTolerances(gatol=1.0e-4) 149 tao.setFromOptions() 150 tao.solve() 151 self.assertTrue(x[1] - x[0] ** 2 >= -1.0e-4) 152 153 def testBNCG(self): 154 if self.tao.getComm().Get_size() > 1: 155 return 156 tao = self.tao 157 158 x = PETSc.Vec().create(tao.getComm()) 159 x.setType('standard') 160 x.setSizes(2) 161 xl = PETSc.Vec().create(tao.getComm()) 162 xl.setType('standard') 163 xl.setSizes(2) 164 xl.set(0.0) 165 xu = PETSc.Vec().create(tao.getComm()) 166 xu.setType('standard') 167 xu.setSizes(2) 168 xu.set(2.0) 169 tao.setVariableBounds((xl, xu)) 170 tao.setObjective(Objective()) 171 tao.setGradient(Gradient(), None) 172 tao.setSolution(x) 173 tao.setType(PETSc.TAO.Type.BNCG) 174 tao.setTolerances(gatol=1.0e-4) 175 ls = tao.getLineSearch() 176 ls.setType(PETSc.TAOLineSearch.Type.UNIT) 177 tao.setFromOptions() 178 tao.solve() 179 self.assertAlmostEqual(x[0], 2.0, places=4) 180 self.assertAlmostEqual(x[1], 2.0, places=4) 181 182 183# -------------------------------------------------------------------- 184 185 186class TestTAOSelf(BaseTestTAO, unittest.TestCase): 187 COMM = PETSc.COMM_SELF 188 189 190class TestTAOWorld(BaseTestTAO, unittest.TestCase): 191 COMM = PETSc.COMM_WORLD 192 193 194# -------------------------------------------------------------------- 195 196 197if numpy.iscomplexobj(PETSc.ScalarType()): 198 del BaseTestTAO 199 del TestTAOSelf 200 del TestTAOWorld 201 202if __name__ == '__main__': 203 unittest.main() 204