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 131 def testInequlityConstraints(self): 132 if self.tao.getComm().Get_size() > 1: 133 return 134 tao = self.tao 135 136 x = PETSc.Vec().create(tao.getComm()) 137 x.setType('standard') 138 x.setSizes(2) 139 c = PETSc.Vec().create(tao.getComm()) 140 c.setSizes(1) 141 c.setType(x.getType()) 142 J = PETSc.Mat().create(tao.getComm()) 143 J.setSizes([1, 2]) 144 J.setType(PETSc.Mat.Type.DENSE) 145 J.setUp() 146 147 tao.setObjective(Objective()) 148 tao.setGradient(Gradient(), None) 149 tao.setInequalityConstraints(InEqConstraints(), c) 150 tao.setJacobianInequality(InEqJacobian(), J, J) 151 tao.setSolution(x) 152 tao.setType(PETSc.TAO.Type.ALMM) 153 tao.setALMMType(PETSc.TAO.ALMMType.CLASSIC) 154 tao.setTolerances(gatol=1.0e-4) 155 tao.setFromOptions() 156 tao.solve() 157 158 self.assertTrue(tao.getALMMType() == PETSc.TAO.ALMMType.CLASSIC) 159 self.assertTrue(x[1] - x[0] ** 2 >= -1.0e-4) 160 self.assertAlmostEqual(x[0], 0.5 + sqrt(7) / 2, places=4) 161 self.assertAlmostEqual(x[1], 2 + sqrt(7) / 2, places=4) 162 163 def testBNCG(self): 164 if self.tao.getComm().Get_size() > 1: 165 return 166 tao = self.tao 167 168 x = PETSc.Vec().create(tao.getComm()) 169 x.setType('standard') 170 x.setSizes(2) 171 xl = PETSc.Vec().create(tao.getComm()) 172 xl.setType('standard') 173 xl.setSizes(2) 174 xl.set(0.0) 175 xu = PETSc.Vec().create(tao.getComm()) 176 xu.setType('standard') 177 xu.setSizes(2) 178 xu.set(2.0) 179 tao.setVariableBounds((xl, xu)) 180 tao.setObjective(Objective()) 181 tao.setGradient(Gradient(), None) 182 tao.setSolution(x) 183 tao.setType(PETSc.TAO.Type.BNCG) 184 tao.setTolerances(gatol=1.0e-4) 185 ls = tao.getLineSearch() 186 ls.setType(PETSc.TAOLineSearch.Type.UNIT) 187 tao.setFromOptions() 188 tao.solve() 189 self.assertAlmostEqual(x[0], 2.0, places=4) 190 self.assertAlmostEqual(x[1], 2.0, places=4) 191 192 def templateBQNLS(self, lmvm_setup): 193 if self.tao.getComm().Get_size() > 1: 194 return 195 tao = self.tao 196 197 x = PETSc.Vec().create(tao.getComm()) 198 x.setType('standard') 199 x.setSizes(2) 200 xl = PETSc.Vec().create(tao.getComm()) 201 xl.setType('standard') 202 xl.setSizes(2) 203 xl.set(0.0) 204 xu = PETSc.Vec().create(tao.getComm()) 205 xu.setType('standard') 206 xu.setSizes(2) 207 xu.set(2.0) 208 tao.setVariableBounds((xl, xu)) 209 tao.setObjective(Objective()) 210 tao.setGradient(Gradient(), None) 211 tao.setSolution(x) 212 tao.setType(PETSc.TAO.Type.BQNLS) 213 tao.setTolerances(gatol=1.0e-4) 214 215 H = PETSc.Mat() 216 if lmvm_setup == 'dense' or lmvm_setup == 'ksp': 217 H.createDense((2, 2), comm=tao.getComm()) 218 H[0, 0] = 2 219 H[0, 1] = 0 220 H[1, 0] = 0 221 H[1, 1] = 2 222 H.assemble() 223 elif lmvm_setup == 'diagonal': 224 H_vec = PETSc.Vec().createSeq(2) 225 H_vec[0] = 2 226 H_vec[1] = 2 227 H_vec.assemble() 228 H.createDiagonal(H_vec) 229 H.assemble() 230 231 if lmvm_setup == 'dense' or lmvm_setup == 'diagonal': 232 tao.getLMVMMat().setLMVMJ0(H) 233 elif lmvm_setup == 'ksp': 234 lmvm_ksp = PETSc.KSP().create(tao.getComm()) 235 lmvm_ksp.setType(PETSc.KSP.Type.CG) 236 lmvm_ksp.setOperators(H) 237 tao.getLMVMMat().setLMVMJ0KSP(lmvm_ksp) 238 239 tao.setFromOptions() 240 tao.solve() 241 self.assertEqual(tao.getIterationNumber(), 1) 242 self.assertAlmostEqual(x[0], 2.0, places=4) 243 self.assertAlmostEqual(x[1], 2.0, places=4) 244 245 if lmvm_setup == 'dense' or lmvm_setup == 'diagonal': 246 self.assertTrue(tao.getLMVMMat().getLMVMJ0().equal(H)) 247 elif lmvm_setup == 'ksp': 248 self.assertTrue( 249 tao.getLMVMMat().getLMVMJ0KSP().getType() == PETSc.KSP.Type.CG 250 ) 251 252 def testBQNLS_dense(self): 253 self.templateBQNLS('dense') 254 255 def testBQNLS_ksp(self): 256 self.templateBQNLS('ksp') 257 258 def testBQNLS_diagonal(self): 259 self.templateBQNLS('diagonal') 260 261 262# -------------------------------------------------------------------- 263 264 265class TestTAOSelf(BaseTestTAO, unittest.TestCase): 266 COMM = PETSc.COMM_SELF 267 268 269class TestTAOWorld(BaseTestTAO, unittest.TestCase): 270 COMM = PETSc.COMM_WORLD 271 272 273# -------------------------------------------------------------------- 274 275 276if numpy.iscomplexobj(PETSc.ScalarType()): 277 del BaseTestTAO 278 del TestTAOSelf 279 del TestTAOWorld 280 281if __name__ == '__main__': 282 unittest.main() 283