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