1# -------------------------------------------------------------------- 2 3from math import sqrt 4from petsc4py import PETSc 5import unittest 6import numpy 7 8 9# -------------------------------------------------------------------- 10class Objective: 11 def __call__(self, tao, x): 12 return (x[0] - 2.0) ** 2 + (x[1] - 2.0) ** 2 - 2.0 * (x[0] + x[1]) 13 14 15class Gradient: 16 def __call__(self, tao, x, g): 17 g[0] = 2.0 * (x[0] - 2.0) - 2.0 18 g[1] = 2.0 * (x[1] - 2.0) - 2.0 19 g.assemble() 20 21 22class EqConstraints: 23 def __call__(self, tao, x, c): 24 c[0] = x[0] ** 2 + x[1] - 2.0 25 c.assemble() 26 27 28class EqJacobian: 29 def __call__(self, tao, x, J, P): 30 P[0, 0] = 2.0 * x[0] 31 P[0, 1] = 1.0 32 P.assemble() 33 if J != P: 34 J.assemble() 35 36 37class InEqConstraints: 38 def __call__(self, tao, x, c): 39 c[0] = x[1] - x[0] ** 2 40 c.assemble() 41 42 43class InEqJacobian: 44 def __call__(self, tao, x, J, P): 45 P[0, 0] = -2.0 * x[0] 46 P[0, 1] = 1.0 47 P.assemble() 48 if J != P: 49 J.assemble() 50 51 52class BaseTestTAO: 53 COMM = None 54 55 def setUp(self): 56 self.tao = PETSc.TAO().create(comm=self.COMM) 57 58 def tearDown(self): 59 self.tao = None 60 PETSc.garbage_cleanup() 61 62 def testSetRoutinesToNone(self): 63 tao = self.tao 64 objective, gradient, objgrad = None, None, None 65 constraint, varbounds = None, None 66 hessian, jacobian = None, None 67 tao.setObjective(objective) 68 tao.setGradient(gradient, None) 69 tao.setVariableBounds(varbounds) 70 tao.setObjectiveGradient(objgrad, None) 71 tao.setConstraints(constraint) 72 tao.setHessian(hessian) 73 tao.setJacobian(jacobian) 74 75 def testGetVecsAndMats(self): 76 tao = self.tao 77 x = tao.getSolution() 78 (g, _) = tao.getGradient() 79 low, up = tao.getVariableBounds() 80 r = None # tao.getConstraintVec() 81 H, HP = None, None # tao.getHessianMat() 82 J, JP = None, None # tao.getJacobianMat() 83 for o in [ 84 x, 85 g, 86 r, 87 low, 88 up, 89 H, 90 HP, 91 J, 92 JP, 93 ]: 94 self.assertFalse(o) 95 96 def testGetKSP(self): 97 ksp = self.tao.getKSP() 98 self.assertFalse(ksp) 99 100 def testEqualityConstraints(self): 101 if self.tao.getComm().Get_size() > 1: 102 return 103 tao = self.tao 104 105 x = PETSc.Vec().create(tao.getComm()) 106 x.setType('standard') 107 x.setSizes(2) 108 c = PETSc.Vec().create(tao.getComm()) 109 c.setSizes(1) 110 c.setType(x.getType()) 111 J = PETSc.Mat().create(tao.getComm()) 112 J.setSizes([1, 2]) 113 J.setType(PETSc.Mat.Type.DENSE) 114 J.setUp() 115 116 tao.setObjective(Objective()) 117 tao.setGradient(Gradient(), None) 118 tao.setEqualityConstraints(EqConstraints(), c) 119 tao.setJacobianEquality(EqJacobian(), J, J) 120 tao.setSolution(x) 121 tao.setType(PETSc.TAO.Type.ALMM) 122 tao.setALMMType(PETSc.TAO.ALMMType.PHR) 123 tao.setTolerances(gatol=1.0e-4) 124 tao.setFromOptions() 125 tao.solve() 126 self.assertTrue(tao.getALMMType() == PETSc.TAO.ALMMType.PHR) 127 self.assertAlmostEqual(abs(x[0] ** 2 + x[1] - 2.0), 0.0, places=4) 128 self.assertAlmostEqual(x[0], 0.7351392590499015014254200465, places=4) 129 self.assertAlmostEqual(x[1], 1.4595702698035618134357683666, places=4) 130 self.assertTrue(tao.getObjective() is not None) 131 132 def testInequlityConstraints(self): 133 if self.tao.getComm().Get_size() > 1: 134 return 135 tao = self.tao 136 137 x = PETSc.Vec().create(tao.getComm()) 138 x.setType('standard') 139 x.setSizes(2) 140 c = PETSc.Vec().create(tao.getComm()) 141 c.setSizes(1) 142 c.setType(x.getType()) 143 J = PETSc.Mat().create(tao.getComm()) 144 J.setSizes([1, 2]) 145 J.setType(PETSc.Mat.Type.DENSE) 146 J.setUp() 147 148 tao.setObjective(Objective()) 149 tao.setGradient(Gradient(), None) 150 tao.setInequalityConstraints(InEqConstraints(), c) 151 tao.setJacobianInequality(InEqJacobian(), J, J) 152 tao.setSolution(x) 153 tao.setType(PETSc.TAO.Type.ALMM) 154 tao.setALMMType(PETSc.TAO.ALMMType.CLASSIC) 155 tao.setTolerances(gatol=1.0e-4) 156 tao.setFromOptions() 157 tao.solve() 158 159 self.assertTrue(tao.getALMMType() == PETSc.TAO.ALMMType.CLASSIC) 160 self.assertTrue(x[1] - x[0] ** 2 >= -1.0e-4) 161 self.assertAlmostEqual(x[0], 0.5 + sqrt(7) / 2, places=4) 162 self.assertAlmostEqual(x[1], 2 + sqrt(7) / 2, places=4) 163 164 def testBNCG(self): 165 if self.tao.getComm().Get_size() > 1: 166 return 167 tao = self.tao 168 169 x = PETSc.Vec().create(tao.getComm()) 170 x.setType('standard') 171 x.setSizes(2) 172 xl = PETSc.Vec().create(tao.getComm()) 173 xl.setType('standard') 174 xl.setSizes(2) 175 xl.set(0.0) 176 xu = PETSc.Vec().create(tao.getComm()) 177 xu.setType('standard') 178 xu.setSizes(2) 179 xu.set(2.0) 180 tao.setVariableBounds((xl, xu)) 181 tao.setObjective(Objective()) 182 tao.setGradient(Gradient(), None) 183 tao.setSolution(x) 184 tao.setType(PETSc.TAO.Type.BNCG) 185 tao.setTolerances(gatol=1.0e-4) 186 ls = tao.getLineSearch() 187 ls.setType(PETSc.TAOLineSearch.Type.UNIT) 188 tao.setFromOptions() 189 tao.solve() 190 self.assertAlmostEqual(x[0], 2.0, places=4) 191 self.assertAlmostEqual(x[1], 2.0, places=4) 192 193 def templateBQNLS(self, lmvm_setup): 194 if self.tao.getComm().Get_size() > 1: 195 return 196 tao = self.tao 197 198 x = PETSc.Vec().create(tao.getComm()) 199 x.setType('standard') 200 x.setSizes(2) 201 xl = PETSc.Vec().create(tao.getComm()) 202 xl.setType('standard') 203 xl.setSizes(2) 204 xl.set(0.0) 205 xu = PETSc.Vec().create(tao.getComm()) 206 xu.setType('standard') 207 xu.setSizes(2) 208 xu.set(2.0) 209 tao.setVariableBounds((xl, xu)) 210 tao.setObjective(Objective()) 211 tao.setGradient(Gradient(), None) 212 tao.setSolution(x) 213 tao.setType(PETSc.TAO.Type.BQNLS) 214 tao.setTolerances(gatol=1.0e-4) 215 216 H = PETSc.Mat() 217 if lmvm_setup == 'dense' or lmvm_setup == 'ksp': 218 H.createDense((2, 2), comm=tao.getComm()) 219 H[0, 0] = 2 220 H[0, 1] = 0 221 H[1, 0] = 0 222 H[1, 1] = 2 223 H.assemble() 224 elif lmvm_setup == 'diagonal': 225 H_vec = PETSc.Vec().createSeq(2) 226 H_vec[0] = 2 227 H_vec[1] = 2 228 H_vec.assemble() 229 H.createDiagonal(H_vec) 230 H.assemble() 231 232 if lmvm_setup == 'dense' or lmvm_setup == 'diagonal': 233 tao.getLMVMMat().setLMVMJ0(H) 234 elif lmvm_setup == 'ksp': 235 lmvm_ksp = PETSc.KSP().create(tao.getComm()) 236 lmvm_ksp.setType(PETSc.KSP.Type.CG) 237 lmvm_ksp.setOperators(H) 238 tao.getLMVMMat().setLMVMJ0KSP(lmvm_ksp) 239 240 tao.setFromOptions() 241 tao.solve() 242 if lmvm_setup == 'dense': 243 self.assertEqual(tao.getIterationNumber(), 1) 244 self.assertAlmostEqual(x[0], 2.0, places=4) 245 self.assertAlmostEqual(x[1], 2.0, places=4) 246 247 if lmvm_setup == 'dense' or lmvm_setup == 'diagonal': 248 self.assertTrue(tao.getLMVMMat().getLMVMJ0().equal(H)) 249 elif lmvm_setup == 'ksp': 250 self.assertTrue( 251 tao.getLMVMMat().getLMVMJ0KSP().getType() == PETSc.KSP.Type.CG 252 ) 253 254 def testBQNLS_dense(self): 255 self.templateBQNLS('dense') 256 257 def testBQNLS_ksp(self): 258 self.templateBQNLS('ksp') 259 260 def testBQNLS_diagonal(self): 261 self.templateBQNLS('diagonal') 262 263 264# -------------------------------------------------------------------- 265 266 267class TestTAOSelf(BaseTestTAO, unittest.TestCase): 268 COMM = PETSc.COMM_SELF 269 270 271class TestTAOWorld(BaseTestTAO, unittest.TestCase): 272 COMM = PETSc.COMM_WORLD 273 274 275# -------------------------------------------------------------------- 276 277 278if numpy.iscomplexobj(PETSc.ScalarType()): 279 del BaseTestTAO 280 del TestTAOSelf 281 del TestTAOWorld 282 283if __name__ == '__main__': 284 unittest.main() 285