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 self.assertEqual(getrefcount(self._getCtx()), 2) 326 self.A.destroy() # XXX 327 self.A = None 328 PETSc.garbage_cleanup() 329 330 def testBasic(self): 331 self.assertEqual(getrefcount(self._getCtx()), 2) 332 ctx = self.A.getPythonContext() 333 self.assertTrue(self._getCtx() is ctx) 334 335 def testSetUp(self): 336 ctx = self.A.getPythonContext() 337 setupcalled = ctx.setupcalled 338 self.A.setUp() 339 self.assertEqual(setupcalled, ctx.setupcalled) 340 self.A.setPythonContext(ctx) 341 self.A.setUp() 342 self.assertEqual(setupcalled + 1, ctx.setupcalled) 343 344 def testZeroEntries(self): 345 f = lambda: self.A.zeroEntries() 346 self.assertRaises(Exception, f) 347 348 def testMult(self): 349 x, y = self.A.createVecs() 350 f = lambda: self.A.mult(x, y) 351 self.assertRaises(Exception, f) 352 353 def testMultTranspose(self): 354 x, y = self.A.createVecs() 355 f = lambda: self.A.multTranspose(x, y) 356 self.assertRaises(Exception, f) 357 358 def testGetDiagonal(self): 359 d = self.A.createVecLeft() 360 f = lambda: self.A.getDiagonal(d) 361 self.assertRaises(Exception, f) 362 363 def testSetDiagonal(self): 364 d = self.A.createVecLeft() 365 f = lambda: self.A.setDiagonal(d) 366 self.assertRaises(Exception, f) 367 368 def testDiagonalScale(self): 369 x, y = self.A.createVecs() 370 f = lambda: self.A.diagonalScale(x, y) 371 self.assertRaises(Exception, f) 372 373 def testDuplicate(self): 374 f1 = lambda: self.A.duplicate(True) 375 f2 = lambda: self.A.duplicate(False) 376 self.assertRaises(Exception, f1) 377 self.assertRaises(Exception, f2) 378 379 def testSetVecType(self): 380 self.A.setVecType('mpi') 381 self.assertTrue('mpi' == self.A.getVecType()) 382 383 def testH2Opus(self): 384 if not PETSc.Sys.hasExternalPackage('h2opus'): 385 return 386 if self.A.getComm().Get_size() > 1: 387 return 388 h = PETSc.Mat() 389 390 # need matrix vector and its transpose for norm estimation 391 AA = self.A.getPythonContext() 392 if not hasattr(AA, 'mult'): 393 return 394 AA.multTranspose = AA.mult 395 396 # without coordinates 397 h.createH2OpusFromMat(self.A, leafsize=2) 398 h.assemble() 399 h.destroy() 400 401 # with coordinates 402 coords = numpy.linspace( 403 (1, 2, 3), (10, 20, 30), self.A.getSize()[0], dtype=PETSc.RealType 404 ) 405 h.createH2OpusFromMat(self.A, coords, leafsize=2) 406 h.assemble() 407 408 # test API 409 h.H2OpusOrthogonalize() 410 h.H2OpusCompress(1.0e-1) 411 412 # Low-rank update 413 U = PETSc.Mat() 414 U.createDense([h.getSizes()[0], 3], comm=h.getComm()) 415 U.setUp() 416 U.setRandom() 417 418 he = PETSc.Mat() 419 h.convert('dense', he) 420 he.axpy(1.0, U.matTransposeMult(U)) 421 422 h.H2OpusLowRankUpdate(U) 423 self.assertTrue(he.equal(h)) 424 425 h.destroy() 426 427 del AA.multTranspose 428 429 def testGetType(self): 430 ctx = self.A.getPythonContext() 431 pytype = f'{ctx.__module__}.{type(ctx).__name__}' 432 self.assertTrue(self.A.getPythonType() == pytype) 433 434 435class TestScaledIdentity(TestMatrix): 436 PYCLS = 'ScaledIdentity' 437 438 def testMult(self): 439 s = self._getCtx().s 440 x, y = self.A.createVecs() 441 x.setRandom() 442 self.A.mult(x, y) 443 self.assertTrue(y.equal(s * x)) 444 445 def testMultTransposeSymmKnown(self): 446 s = self._getCtx().s 447 x, y = self.A.createVecs() 448 x.setRandom() 449 self.A.setOption(PETSc.Mat.Option.SYMMETRIC, True) 450 self.A.multTranspose(x, y) 451 self.assertTrue(y.equal(s * x)) 452 self.A.setOption(PETSc.Mat.Option.SYMMETRIC, False) 453 f = lambda: self.A.multTranspose(x, y) 454 self.assertRaises(Exception, f) 455 456 def testMultTransposeNewMeth(self): 457 s = self._getCtx().s 458 x, y = self.A.createVecs() 459 x.setRandom() 460 AA = self.A.getPythonContext() 461 AA.multTranspose = AA.mult 462 self.A.multTranspose(x, y) 463 del AA.multTranspose 464 self.assertTrue(y.equal(s * x)) 465 466 def testGetDiagonal(self): 467 s = self._getCtx().s 468 d = self.A.createVecLeft() 469 o = d.duplicate() 470 o.set(s) 471 self.A.getDiagonal(d) 472 self.assertTrue(o.equal(d)) 473 474 def testDuplicate(self): 475 B = self.A.duplicate() 476 self.assertTrue(B.getPythonContext().s == 2) 477 B = self.A.duplicate(False) 478 self.assertTrue(B.getPythonContext().s == 2) 479 B = self.A.duplicate(True) 480 self.assertTrue(B.getPythonContext().s == self.A.getPythonContext().s) 481 B = self.A.duplicate(PETSc.Mat.DuplicateOption.DO_NOT_COPY_VALUES) 482 self.assertTrue(B.getPythonContext().s == 2) 483 B = self.A.duplicate(PETSc.Mat.DuplicateOption.SHARE_NONZERO_PATTERN) 484 self.assertTrue(B.getPythonContext().s == 2) 485 B = self.A.duplicate(PETSc.Mat.DuplicateOption.COPY_VALUES) 486 self.assertTrue(B.getPythonContext().s == self.A.getPythonContext().s) 487 488 def testMatMat(self): 489 s = self._getCtx().s 490 R = PETSc.Random().create(self.COMM) 491 R.setFromOptions() 492 A = PETSc.Mat().create(self.COMM) 493 A.setSizes(self.A.getSizes()) 494 A.setType(PETSc.Mat.Type.AIJ) 495 A.setPreallocationNNZ(None) 496 A.setRandom(R) 497 B = PETSc.Mat().create(self.COMM) 498 B.setSizes(self.A.getSizes()) 499 B.setType(PETSc.Mat.Type.AIJ) 500 B.setPreallocationNNZ(None) 501 B.setRandom(R) 502 Id = PETSc.Mat().create(self.COMM) 503 Id.setSizes(self.A.getSizes()) 504 Id.setType(PETSc.Mat.Type.AIJ) 505 Id.setUp() 506 Id.assemble() 507 Id.shift(s) 508 509 self.assertTrue(self.A.matMult(A).equal(Id.matMult(A))) 510 self.assertTrue(A.matMult(self.A).equal(A.matMult(Id))) 511 if self.A.getComm().Get_size() == 1: 512 self.assertTrue(self.A.matTransposeMult(A).equal(Id.matTransposeMult(A))) 513 self.assertTrue(A.matTransposeMult(self.A).equal(A.matTransposeMult(Id))) 514 self.assertTrue(self.A.transposeMatMult(A).equal(Id.transposeMatMult(A))) 515 self.assertTrue(A.transposeMatMult(self.A).equal(A.transposeMatMult(Id))) 516 self.assertAlmostEqual((self.A.ptap(A) - Id.ptap(A)).norm(), 0.0, places=5) 517 self.assertAlmostEqual((A.ptap(self.A) - A.ptap(Id)).norm(), 0.0, places=5) 518 if self.A.getComm().Get_size() == 1: 519 self.assertAlmostEqual((self.A.rart(A) - Id.rart(A)).norm(), 0.0, places=5) 520 self.assertAlmostEqual((A.rart(self.A) - A.rart(Id)).norm(), 0.0, places=5) 521 self.assertAlmostEqual( 522 (self.A.matMatMult(A, B) - Id.matMatMult(A, B)).norm(), 0.0, places=5 523 ) 524 self.assertAlmostEqual( 525 (A.matMatMult(self.A, B) - A.matMatMult(Id, B)).norm(), 0.0, places=5 526 ) 527 self.assertAlmostEqual( 528 (A.matMatMult(B, self.A) - A.matMatMult(B, Id)).norm(), 0.0, places=5 529 ) 530 531 def testShift(self): 532 sold = self._getCtx().s 533 self.A.shift(-0.5) 534 s = self._getCtx().s 535 self.assertTrue(s == sold - 0.5) 536 537 def testScale(self): 538 sold = self._getCtx().s 539 self.A.scale(-0.5) 540 s = self._getCtx().s 541 self.assertTrue(s == sold * -0.5) 542 543 def testDiagonalMat(self): 544 s = self._getCtx().s 545 B = PETSc.Mat().createConstantDiagonal( 546 self.A.getSizes(), s, comm=self.A.getComm() 547 ) 548 self.assertTrue(self.A.equal(B)) 549 550 551class TestDiagonal(TestMatrix): 552 PYCLS = 'Diagonal' 553 CREATE_WITH_NONE = True 554 555 def setUp(self): 556 super().setUp() 557 D = self.A.createVecLeft() 558 s, e = D.getOwnershipRange() 559 for i in range(s, e): 560 D[i] = i + 1 561 D.assemble() 562 self.A.setDiagonal(D) 563 564 def testZeroEntries(self): 565 self.A.zeroEntries() 566 D = self._getCtx().D 567 self.assertEqual(D.norm(), 0) 568 569 def testMult(self): 570 x, y = self.A.createVecs() 571 x.set(1) 572 self.A.mult(x, y) 573 self.assertTrue(y.equal(self._getCtx().D)) 574 575 def testMultTransposeSymmKnown(self): 576 x, y = self.A.createVecs() 577 x.set(1) 578 self.A.setOption(PETSc.Mat.Option.SYMMETRIC, True) 579 self.A.multTranspose(x, y) 580 self.assertTrue(y.equal(self._getCtx().D)) 581 self.A.setOption(PETSc.Mat.Option.SYMMETRIC, False) 582 f = lambda: self.A.multTranspose(x, y) 583 self.assertRaises(Exception, f) 584 585 def testMultTransposeNewMeth(self): 586 x, y = self.A.createVecs() 587 x.set(1) 588 AA = self.A.getPythonContext() 589 AA.multTranspose = AA.mult 590 self.A.multTranspose(x, y) 591 del AA.multTranspose 592 self.assertTrue(y.equal(self._getCtx().D)) 593 594 def testDuplicate(self): 595 B = self.A.duplicate(False) 596 B = self.A.duplicate(True) 597 self.assertTrue(B.getPythonContext().D.equal(self.A.getPythonContext().D)) 598 599 def testGetDiagonal(self): 600 d = self.A.createVecLeft() 601 self.A.getDiagonal(d) 602 self.assertTrue(d.equal(self._getCtx().D)) 603 604 def testSetDiagonal(self): 605 d = self.A.createVecLeft() 606 d.setRandom() 607 self.A.setDiagonal(d) 608 self.assertTrue(d.equal(self._getCtx().D)) 609 610 def testDiagonalScale(self): 611 x, y = self.A.createVecs() 612 x.set(2) 613 y.set(3) 614 old = self._getCtx().D.copy() 615 self.A.diagonalScale(x, y) 616 D = self._getCtx().D 617 self.assertTrue(D.equal(old * 6)) 618 619 def testCreateTranspose(self): 620 A = self.A 621 A.setOption(PETSc.Mat.Option.SYMMETRIC, True) 622 AT = PETSc.Mat().createTranspose(A) 623 x, y = A.createVecs() 624 xt, yt = AT.createVecs() 625 # 626 y.setRandom() 627 A.multTranspose(y, x) 628 y.copy(xt) 629 AT.mult(xt, yt) 630 self.assertTrue(yt.equal(x)) 631 # 632 x.setRandom() 633 A.mult(x, y) 634 x.copy(yt) 635 AT.multTranspose(yt, xt) 636 self.assertTrue(xt.equal(y)) 637 del A 638 639 def testConvert(self): 640 self.assertTrue(self.A.convert(PETSc.Mat.Type.AIJ, PETSc.Mat()).equal(self.A)) 641 self.assertTrue(self.A.convert(PETSc.Mat.Type.BAIJ, PETSc.Mat()).equal(self.A)) 642 self.assertTrue(self.A.convert(PETSc.Mat.Type.SBAIJ, PETSc.Mat()).equal(self.A)) 643 self.assertTrue(self.A.convert(PETSc.Mat.Type.DENSE, PETSc.Mat()).equal(self.A)) 644 645 def testShift(self): 646 old = self._getCtx().D.copy() 647 self.A.shift(-0.5) 648 D = self._getCtx().D 649 self.assertTrue(D.equal(old - 0.5)) 650 651 def testScale(self): 652 old = self._getCtx().D.copy() 653 self.A.scale(-0.5) 654 D = self._getCtx().D 655 self.assertTrue(D.equal(-0.5 * old)) 656 657 def testDiagonalMat(self): 658 D = self._getCtx().D.copy() 659 B = PETSc.Mat().createDiagonal(D) 660 self.assertTrue(self.A.equal(B)) 661 662 663# -------------------------------------------------------------------- 664 665if __name__ == '__main__': 666 unittest.main() 667