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