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