1from petsc4py import PETSc 2import unittest, numpy 3from sys import getrefcount 4# -------------------------------------------------------------------- 5 6class Matrix(object): 7 8 def __init__(self): 9 pass 10 11 def create(self, mat): 12 pass 13 14 def destroy(self, mat): 15 pass 16 17class ScaledIdentity(Matrix): 18 19 s = 2.0 20 21 def scale(self, mat, s): 22 self.s *= s 23 24 def shift(self, mat, s): 25 self.s += s 26 27 def mult(self, mat, x, y): 28 x.copy(y) 29 y.scale(self.s) 30 31 def duplicate(self, mat, op): 32 dmat = PETSc.Mat() 33 dctx = ScaledIdentity() 34 dmat.createPython(mat.getSizes(), dctx, comm=mat.getComm()) 35 if op == PETSc.Mat.DuplicateOption.COPY_VALUES: 36 dctx.s = self.s 37 dmat.setUp() 38 return dmat 39 40 def getDiagonal(self, mat, vd): 41 vd.set(self.s) 42 43 def productSetFromOptions(self, mat, producttype, A, B, C): 44 return True 45 46 def productSymbolic(self, mat, product, producttype, A, B, C): 47 if producttype == 'AB': 48 if mat is A: # product = identity * B 49 product.setType(B.getType()) 50 product.setSizes(B.getSizes()) 51 product.setUp() 52 product.assemble() 53 B.copy(product) 54 elif mat is B: # product = A * identity 55 product.setType(A.getType()) 56 product.setSizes(A.getSizes()) 57 product.setUp() 58 product.assemble() 59 A.copy(product) 60 else: 61 raise RuntimeError('wrong configuration') 62 elif producttype == 'AtB': 63 if mat is A: # product = identity^T * B 64 product.setType(B.getType()) 65 product.setSizes(B.getSizes()) 66 product.setUp() 67 product.assemble() 68 B.copy(product) 69 elif mat is B: # product = A^T * identity 70 tmp = PETSc.Mat() 71 A.transpose(tmp) 72 product.setType(tmp.getType()) 73 product.setSizes(tmp.getSizes()) 74 product.setUp() 75 product.assemble() 76 tmp.copy(product) 77 else: 78 raise RuntimeError('wrong configuration') 79 elif producttype == 'ABt': 80 if mat is A: # product = identity * B^T 81 tmp = PETSc.Mat() 82 B.transpose(tmp) 83 product.setType(tmp.getType()) 84 product.setSizes(tmp.getSizes()) 85 product.setUp() 86 product.assemble() 87 tmp.copy(product) 88 elif mat is B: # product = A * identity^T 89 product.setType(A.getType()) 90 product.setSizes(A.getSizes()) 91 product.setUp() 92 product.assemble() 93 A.copy(product) 94 else: 95 raise RuntimeError('wrong configuration') 96 elif producttype == 'PtAP': 97 if mat is A: # product = P^T * identity * P 98 self.tmp = PETSc.Mat() 99 B.transposeMatMult(B, self.tmp) 100 product.setType(self.tmp.getType()) 101 product.setSizes(self.tmp.getSizes()) 102 product.setUp() 103 product.assemble() 104 self.tmp.copy(product) 105 elif mat is B: # product = identity^T * A * identity 106 product.setType(A.getType()) 107 product.setSizes(A.getSizes()) 108 product.setUp() 109 product.assemble() 110 A.copy(product) 111 else: 112 raise RuntimeError('wrong configuration') 113 elif producttype == 'RARt': 114 if mat is A: # product = R * identity * R^t 115 self.tmp = PETSc.Mat() 116 B.matTransposeMult(B, self.tmp) 117 product.setType(self.tmp.getType()) 118 product.setSizes(self.tmp.getSizes()) 119 product.setUp() 120 product.assemble() 121 self.tmp.copy(product) 122 elif mat is B: # product = identity * A * identity^T 123 product.setType(A.getType()) 124 product.setSizes(A.getSizes()) 125 product.setUp() 126 product.assemble() 127 A.copy(product) 128 else: 129 raise RuntimeError('wrong configuration') 130 elif producttype == 'ABC': 131 if mat is A: # product = identity * B * C 132 self.tmp = PETSc.Mat() 133 B.matMult(C, self.tmp) 134 product.setType(self.tmp.getType()) 135 product.setSizes(self.tmp.getSizes()) 136 product.setUp() 137 product.assemble() 138 self.tmp.copy(product) 139 elif mat is B: # product = A * identity * C 140 self.tmp = PETSc.Mat() 141 A.matMult(C, self.tmp) 142 product.setType(self.tmp.getType()) 143 product.setSizes(self.tmp.getSizes()) 144 product.setUp() 145 product.assemble() 146 self.tmp.copy(product) 147 elif mat is C: # product = A * B * identity 148 self.tmp = PETSc.Mat() 149 A.matMult(B, self.tmp) 150 product.setType(self.tmp.getType()) 151 product.setSizes(self.tmp.getSizes()) 152 product.setUp() 153 product.assemble() 154 self.tmp.copy(product) 155 else: 156 raise RuntimeError('wrong configuration') 157 else: 158 raise RuntimeError('Product {} not implemented'.format(producttype)) 159 product.zeroEntries() 160 161 def productNumeric(self, mat, product, producttype, A, B, C): 162 if producttype == 'AB': 163 if mat is A: # product = identity * B 164 B.copy(product, structure=True) 165 elif mat is B: # product = A * identity 166 A.copy(product, structure=True) 167 else: 168 raise RuntimeError('wrong configuration') 169 product.scale(self.s) 170 elif producttype == 'AtB': 171 if mat is A: # product = identity^T * B 172 B.copy(product, structure=True) 173 elif mat is B: # product = A^T * identity 174 A.transpose(product) 175 else: 176 raise RuntimeError('wrong configuration') 177 product.scale(self.s) 178 elif producttype == 'ABt': 179 if mat is A: # product = identity * B^T 180 B.transpose(product) 181 elif mat is B: # product = A * identity^T 182 A.copy(product, structure=True) 183 else: 184 raise RuntimeError('wrong configuration') 185 product.scale(self.s) 186 elif producttype == 'PtAP': 187 if mat is A: # product = P^T * identity * P 188 B.transposeMatMult(B, self.tmp) 189 self.tmp.copy(product, structure=True) 190 product.scale(self.s) 191 elif mat is B: # product = identity^T * A * identity 192 A.copy(product, structure=True) 193 product.scale(self.s**2) 194 else: 195 raise RuntimeError('wrong configuration') 196 elif producttype == 'RARt': 197 if mat is A: # product = R * identity * R^t 198 B.matTransposeMult(B, self.tmp) 199 self.tmp.copy(product, structure=True) 200 product.scale(self.s) 201 elif mat is B: # product = identity * A * identity^T 202 A.copy(product, structure=True) 203 product.scale(self.s**2) 204 else: 205 raise RuntimeError('wrong configuration') 206 elif producttype == 'ABC': 207 if mat is A: # product = identity * B * C 208 B.matMult(C, self.tmp) 209 self.tmp.copy(product, structure=True) 210 elif mat is B: # product = A * identity * C 211 A.matMult(C, self.tmp) 212 self.tmp.copy(product, structure=True) 213 elif mat is C: # product = A * B * identity 214 A.matMult(B, self.tmp) 215 self.tmp.copy(product, structure=True) 216 else: 217 raise RuntimeError('wrong configuration') 218 product.scale(self.s) 219 else: 220 raise RuntimeError('Product {} not implemented'.format(producttype)) 221 222class Diagonal(Matrix): 223 224 def create(self, mat): 225 super(Diagonal,self).create(mat) 226 mat.setUp() 227 self.D = mat.createVecLeft() 228 229 def destroy(self, mat): 230 self.D.destroy() 231 super(Diagonal,self).destroy(mat) 232 233 def scale(self, mat, a): 234 self.D.scale(a) 235 236 def shift(self, mat, a): 237 self.D.shift(a) 238 239 def zeroEntries(self, mat): 240 self.D.zeroEntries() 241 242 def mult(self, mat, x, y): 243 y.pointwiseMult(x, self.D) 244 245 def duplicate(self, mat, op): 246 dmat = PETSc.Mat() 247 dctx = Diagonal() 248 dmat.createPython(mat.getSizes(), dctx, comm=mat.getComm()) 249 dctx.D = self.D.duplicate() 250 if op == PETSc.Mat.DuplicateOption.COPY_VALUES: 251 self.D.copy(dctx.D) 252 dmat.setUp() 253 return dmat 254 255 def getDiagonal(self, mat, vd): 256 self.D.copy(vd) 257 258 def setDiagonal(self, mat, vd, im): 259 if isinstance (im, bool): 260 addv = im 261 if addv: 262 self.D.axpy(1, vd) 263 else: 264 vd.copy(self.D) 265 elif im == PETSc.InsertMode.INSERT_VALUES: 266 vd.copy(self.D) 267 elif im == PETSc.InsertMode.ADD_VALUES: 268 self.D.axpy(1, vd) 269 else: 270 raise ValueError('wrong InsertMode %d'% im) 271 272 def diagonalScale(self, mat, vl, vr): 273 if vl: self.D.pointwiseMult(self.D, vl) 274 if vr: self.D.pointwiseMult(self.D, vr) 275 276# -------------------------------------------------------------------- 277 278class TestMatrix(unittest.TestCase): 279 280 COMM = PETSc.COMM_WORLD 281 PYMOD = __name__ 282 PYCLS = 'Matrix' 283 284 def _getCtx(self): 285 return self.A.getPythonContext() 286 287 def setUp(self): 288 N = self.N = 10 289 self.A = PETSc.Mat() 290 if 0: # command line way 291 self.A.create(self.COMM) 292 self.A.setSizes([N,N]) 293 self.A.setType('python') 294 OptDB = PETSc.Options(self.A) 295 OptDB['mat_python_type'] = '%s.%s' % (self.PYMOD,self.PYCLS) 296 self.A.setFromOptions() 297 self.A.setUp() 298 del OptDB['mat_python_type'] 299 self.assertTrue(self._getCtx() is not None) 300 else: # python way 301 context = globals()[self.PYCLS]() 302 self.A.createPython([N,N], context, comm=self.COMM) 303 self.A.setUp() 304 self.assertTrue(self._getCtx() is context) 305 self.assertEqual(getrefcount(context), 3) 306 del context 307 self.assertEqual(getrefcount(self._getCtx()), 2) 308 309 def tearDown(self): 310 ctx = self.A.getPythonContext() 311 self.assertEqual(getrefcount(ctx), 3) 312 self.A.destroy() # XXX 313 self.A = None 314 self.assertEqual(getrefcount(ctx), 2) 315 #import gc,pprint; pprint.pprint(gc.get_referrers(ctx)) 316 317 def testBasic(self): 318 ctx = self.A.getPythonContext() 319 self.assertTrue(self._getCtx() is ctx) 320 self.assertEqual(getrefcount(ctx), 3) 321 322 def testZeroEntries(self): 323 f = lambda : self.A.zeroEntries() 324 self.assertRaises(Exception, f) 325 326 def testMult(self): 327 x, y = self.A.createVecs() 328 f = lambda : self.A.mult(x, y) 329 self.assertRaises(Exception, f) 330 331 def testMultTranspose(self): 332 x, y = self.A.createVecs() 333 f = lambda : self.A.multTranspose(x, y) 334 self.assertRaises(Exception, f) 335 336 def testGetDiagonal(self): 337 d = self.A.createVecLeft() 338 f = lambda : self.A.getDiagonal(d) 339 self.assertRaises(Exception, f) 340 341 def testSetDiagonal(self): 342 d = self.A.createVecLeft() 343 f = lambda : self.A.setDiagonal(d) 344 self.assertRaises(Exception, f) 345 346 def testDiagonalScale(self): 347 x, y = self.A.createVecs() 348 f = lambda : self.A.diagonalScale(x, y) 349 self.assertRaises(Exception, f) 350 351 def testDuplicate(self): 352 f1 = lambda : self.A.duplicate(x, True) 353 f2 = lambda : self.A.duplicate(x, False) 354 self.assertRaises(Exception, f1) 355 self.assertRaises(Exception, f2) 356 357 def testSetVecType(self): 358 self.A.setVecType('mpi') 359 self.assertTrue('mpi' == self.A.getVecType()) 360 361class TestScaledIdentity(TestMatrix): 362 363 PYCLS = 'ScaledIdentity' 364 365 def testMult(self): 366 s = self._getCtx().s 367 x, y = self.A.createVecs() 368 x.setRandom() 369 self.A.mult(x,y) 370 self.assertTrue(y.equal(s*x)) 371 372 def testMultTransposeSymmKnown(self): 373 s = self._getCtx().s 374 x, y = self.A.createVecs() 375 x.setRandom() 376 self.A.setOption(PETSc.Mat.Option.SYMMETRIC, True) 377 self.A.multTranspose(x,y) 378 self.assertTrue(y.equal(s*x)) 379 self.A.setOption(PETSc.Mat.Option.SYMMETRIC, False) 380 f = lambda : self.A.multTranspose(x, y) 381 self.assertRaises(Exception, f) 382 383 def testMultTransposeNewMeth(self): 384 s = self._getCtx().s 385 x, y = self.A.createVecs() 386 x.setRandom() 387 AA = self.A.getPythonContext() 388 AA.multTranspose = AA.mult 389 self.A.multTranspose(x,y) 390 del AA.multTranspose 391 self.assertTrue(y.equal(s*x)) 392 393 def testGetDiagonal(self): 394 s = self._getCtx().s 395 d = self.A.createVecLeft() 396 o = d.duplicate() 397 o.set(s) 398 self.A.getDiagonal(d) 399 self.assertTrue(o.equal(d)) 400 401 def testDuplicate(self): 402 B = self.A.duplicate(False) 403 self.assertTrue(B.getPythonContext().s == 2) 404 B = self.A.duplicate(True) 405 self.assertTrue(B.getPythonContext().s == self.A.getPythonContext().s) 406 407 def testMatMat(self): 408 s = self._getCtx().s 409 R = PETSc.Random().create(self.COMM) 410 R.setFromOptions() 411 A = PETSc.Mat().create(self.COMM) 412 A.setSizes(self.A.getSizes()) 413 A.setType(PETSc.Mat.Type.AIJ) 414 A.setUp() 415 A.setRandom(R) 416 B = PETSc.Mat().create(self.COMM) 417 B.setSizes(self.A.getSizes()) 418 B.setType(PETSc.Mat.Type.AIJ) 419 B.setUp() 420 B.setRandom(R) 421 I = PETSc.Mat().create(self.COMM) 422 I.setSizes(self.A.getSizes()) 423 I.setType(PETSc.Mat.Type.AIJ) 424 I.setUp() 425 I.assemble() 426 I.shift(s) 427 428 self.assertTrue(self.A.matMult(A).equal(I.matMult(A))) 429 self.assertTrue(A.matMult(self.A).equal(A.matMult(I))) 430 if self.A.getComm().Get_size() == 1: 431 self.assertTrue(self.A.matTransposeMult(A).equal(I.matTransposeMult(A))) 432 self.assertTrue(A.matTransposeMult(self.A).equal(A.matTransposeMult(I))) 433 self.assertTrue(self.A.transposeMatMult(A).equal(I.transposeMatMult(A))) 434 self.assertTrue(A.transposeMatMult(self.A).equal(A.transposeMatMult(I))) 435 self.assertAlmostEqual((self.A.ptap(A) - I.ptap(A)).norm(), 0.0, places=5) 436 self.assertAlmostEqual((A.ptap(self.A) - A.ptap(I)).norm(), 0.0, places=5) 437 if self.A.getComm().Get_size() == 1: 438 self.assertAlmostEqual((self.A.rart(A) - I.rart(A)).norm(), 0.0, places=5) 439 self.assertAlmostEqual((A.rart(self.A) - A.rart(I)).norm(), 0.0, places=5) 440 self.assertAlmostEqual((self.A.matMatMult(A,B)-I.matMatMult(A,B)).norm(), 0.0, places=5) 441 self.assertAlmostEqual((A.matMatMult(self.A,B)-A.matMatMult(I,B)).norm(), 0.0, places=5) 442 self.assertAlmostEqual((A.matMatMult(B,self.A)-A.matMatMult(B,I)).norm(), 0.0, places=5) 443 444 def testH2Opus(self): 445 if not PETSc.Sys.hasExternalPackage("h2opus"): 446 return 447 if self.A.getComm().Get_size() > 1: 448 return 449 h = PETSc.Mat() 450 451 # need transpose operation for norm estimation 452 AA = self.A.getPythonContext() 453 AA.multTranspose = AA.mult 454 455 # without coordinates 456 h.createH2OpusFromMat(self.A,leafsize=2) 457 h.assemble() 458 h.destroy() 459 460 # with coordinates 461 coords = numpy.linspace((1,2,3),(10,20,30),self.A.getSize()[0],dtype=PETSc.RealType) 462 h.createH2OpusFromMat(self.A,coords,leafsize=2) 463 h.assemble() 464 h.destroy() 465 466 del AA.multTranspose 467 468 def testShift(self): 469 sold = self._getCtx().s 470 self.A.shift(-0.5) 471 s = self._getCtx().s 472 self.assertTrue(s == sold - 0.5) 473 474 def testScale(self): 475 sold = self._getCtx().s 476 self.A.scale(-0.5) 477 s = self._getCtx().s 478 self.assertTrue(s == sold * -0.5) 479 480class TestDiagonal(TestMatrix): 481 482 PYCLS = 'Diagonal' 483 484 def setUp(self): 485 super(TestDiagonal, self).setUp() 486 D = self.A.createVecLeft() 487 s, e = D.getOwnershipRange() 488 for i in range(s, e): 489 D[i] = i+1 490 D.assemble() 491 self.A.setDiagonal(D) 492 493 def testZeroEntries(self): 494 self.A.zeroEntries() 495 D = self._getCtx().D 496 self.assertEqual(D.norm(), 0) 497 498 def testMult(self): 499 x, y = self.A.createVecs() 500 x.set(1) 501 self.A.mult(x,y) 502 self.assertTrue(y.equal(self._getCtx().D)) 503 504 def testMultTransposeSymmKnown(self): 505 x, y = self.A.createVecs() 506 x.set(1) 507 self.A.setOption(PETSc.Mat.Option.SYMMETRIC, True) 508 self.A.multTranspose(x,y) 509 self.assertTrue(y.equal(self._getCtx().D)) 510 self.A.setOption(PETSc.Mat.Option.SYMMETRIC, False) 511 f = lambda : self.A.multTranspose(x, y) 512 self.assertRaises(Exception, f) 513 514 def testMultTransposeNewMeth(self): 515 x, y = self.A.createVecs() 516 x.set(1) 517 AA = self.A.getPythonContext() 518 AA.multTranspose = AA.mult 519 self.A.multTranspose(x,y) 520 del AA.multTranspose 521 self.assertTrue(y.equal(self._getCtx().D)) 522 523 def testDuplicate(self): 524 B = self.A.duplicate(False) 525 B = self.A.duplicate(True) 526 self.assertTrue(B.getPythonContext().D.equal(self.A.getPythonContext().D)) 527 528 def testGetDiagonal(self): 529 d = self.A.createVecLeft() 530 self.A.getDiagonal(d) 531 self.assertTrue(d.equal(self._getCtx().D)) 532 533 def testSetDiagonal(self): 534 d = self.A.createVecLeft() 535 d.setRandom() 536 self.A.setDiagonal(d) 537 self.assertTrue(d.equal(self._getCtx().D)) 538 539 def testDiagonalScale(self): 540 x, y = self.A.createVecs() 541 x.set(2) 542 y.set(3) 543 old = self._getCtx().D.copy() 544 self.A.diagonalScale(x, y) 545 D = self._getCtx().D 546 self.assertTrue(D.equal(old*6)) 547 548 def testCreateTranspose(self): 549 A = self.A 550 A.setOption(PETSc.Mat.Option.SYMMETRIC, True) 551 AT = PETSc.Mat().createTranspose(A) 552 x, y = A.createVecs() 553 xt, yt = AT.createVecs() 554 # 555 y.setRandom() 556 A.multTranspose(y, x) 557 y.copy(xt) 558 AT.mult(xt, yt) 559 self.assertTrue(yt.equal(x)) 560 # 561 x.setRandom() 562 A.mult(x, y) 563 x.copy(yt) 564 AT.multTranspose(yt, xt) 565 self.assertTrue(xt.equal(y)) 566 del A 567 568 def testConvert(self): 569 self.assertTrue(self.A.convert(PETSc.Mat.Type.AIJ,PETSc.Mat()).equal(self.A)) 570 self.assertTrue(self.A.convert(PETSc.Mat.Type.BAIJ,PETSc.Mat()).equal(self.A)) 571 self.assertTrue(self.A.convert(PETSc.Mat.Type.SBAIJ,PETSc.Mat()).equal(self.A)) 572 self.assertTrue(self.A.convert(PETSc.Mat.Type.DENSE,PETSc.Mat()).equal(self.A)) 573 574 def testShift(self): 575 old = self._getCtx().D.copy() 576 self.A.shift(-0.5) 577 D = self._getCtx().D 578 self.assertTrue(D.equal(old-0.5)) 579 580 def testScale(self): 581 old = self._getCtx().D.copy() 582 self.A.scale(-0.5) 583 D = self._getCtx().D 584 self.assertTrue(D.equal(-0.5*old)) 585 586 587# -------------------------------------------------------------------- 588 589if __name__ == '__main__': 590 unittest.main() 591