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(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 482 def testMatMat(self): 483 s = self._getCtx().s 484 R = PETSc.Random().create(self.COMM) 485 R.setFromOptions() 486 A = PETSc.Mat().create(self.COMM) 487 A.setSizes(self.A.getSizes()) 488 A.setType(PETSc.Mat.Type.AIJ) 489 A.setPreallocationNNZ(None) 490 A.setRandom(R) 491 B = PETSc.Mat().create(self.COMM) 492 B.setSizes(self.A.getSizes()) 493 B.setType(PETSc.Mat.Type.AIJ) 494 B.setPreallocationNNZ(None) 495 B.setRandom(R) 496 Id = PETSc.Mat().create(self.COMM) 497 Id.setSizes(self.A.getSizes()) 498 Id.setType(PETSc.Mat.Type.AIJ) 499 Id.setUp() 500 Id.assemble() 501 Id.shift(s) 502 503 self.assertTrue(self.A.matMult(A).equal(Id.matMult(A))) 504 self.assertTrue(A.matMult(self.A).equal(A.matMult(Id))) 505 if self.A.getComm().Get_size() == 1: 506 self.assertTrue(self.A.matTransposeMult(A).equal(Id.matTransposeMult(A))) 507 self.assertTrue(A.matTransposeMult(self.A).equal(A.matTransposeMult(Id))) 508 self.assertTrue(self.A.transposeMatMult(A).equal(Id.transposeMatMult(A))) 509 self.assertTrue(A.transposeMatMult(self.A).equal(A.transposeMatMult(Id))) 510 self.assertAlmostEqual((self.A.ptap(A) - Id.ptap(A)).norm(), 0.0, places=5) 511 self.assertAlmostEqual((A.ptap(self.A) - A.ptap(Id)).norm(), 0.0, places=5) 512 if self.A.getComm().Get_size() == 1: 513 self.assertAlmostEqual((self.A.rart(A) - Id.rart(A)).norm(), 0.0, places=5) 514 self.assertAlmostEqual((A.rart(self.A) - A.rart(Id)).norm(), 0.0, places=5) 515 self.assertAlmostEqual( 516 (self.A.matMatMult(A, B) - Id.matMatMult(A, B)).norm(), 0.0, places=5 517 ) 518 self.assertAlmostEqual( 519 (A.matMatMult(self.A, B) - A.matMatMult(Id, B)).norm(), 0.0, places=5 520 ) 521 self.assertAlmostEqual( 522 (A.matMatMult(B, self.A) - A.matMatMult(B, Id)).norm(), 0.0, places=5 523 ) 524 525 def testShift(self): 526 sold = self._getCtx().s 527 self.A.shift(-0.5) 528 s = self._getCtx().s 529 self.assertTrue(s == sold - 0.5) 530 531 def testScale(self): 532 sold = self._getCtx().s 533 self.A.scale(-0.5) 534 s = self._getCtx().s 535 self.assertTrue(s == sold * -0.5) 536 537 def testDiagonalMat(self): 538 s = self._getCtx().s 539 B = PETSc.Mat().createConstantDiagonal( 540 self.A.getSizes(), s, comm=self.A.getComm() 541 ) 542 self.assertTrue(self.A.equal(B)) 543 544 545class TestDiagonal(TestMatrix): 546 PYCLS = 'Diagonal' 547 CREATE_WITH_NONE = True 548 549 def setUp(self): 550 super().setUp() 551 D = self.A.createVecLeft() 552 s, e = D.getOwnershipRange() 553 for i in range(s, e): 554 D[i] = i + 1 555 D.assemble() 556 self.A.setDiagonal(D) 557 558 def testZeroEntries(self): 559 self.A.zeroEntries() 560 D = self._getCtx().D 561 self.assertEqual(D.norm(), 0) 562 563 def testMult(self): 564 x, y = self.A.createVecs() 565 x.set(1) 566 self.A.mult(x, y) 567 self.assertTrue(y.equal(self._getCtx().D)) 568 569 def testMultTransposeSymmKnown(self): 570 x, y = self.A.createVecs() 571 x.set(1) 572 self.A.setOption(PETSc.Mat.Option.SYMMETRIC, True) 573 self.A.multTranspose(x, y) 574 self.assertTrue(y.equal(self._getCtx().D)) 575 self.A.setOption(PETSc.Mat.Option.SYMMETRIC, False) 576 f = lambda: self.A.multTranspose(x, y) 577 self.assertRaises(Exception, f) 578 579 def testMultTransposeNewMeth(self): 580 x, y = self.A.createVecs() 581 x.set(1) 582 AA = self.A.getPythonContext() 583 AA.multTranspose = AA.mult 584 self.A.multTranspose(x, y) 585 del AA.multTranspose 586 self.assertTrue(y.equal(self._getCtx().D)) 587 588 def testDuplicate(self): 589 B = self.A.duplicate(False) 590 B = self.A.duplicate(True) 591 self.assertTrue(B.getPythonContext().D.equal(self.A.getPythonContext().D)) 592 593 def testGetDiagonal(self): 594 d = self.A.createVecLeft() 595 self.A.getDiagonal(d) 596 self.assertTrue(d.equal(self._getCtx().D)) 597 598 def testSetDiagonal(self): 599 d = self.A.createVecLeft() 600 d.setRandom() 601 self.A.setDiagonal(d) 602 self.assertTrue(d.equal(self._getCtx().D)) 603 604 def testDiagonalScale(self): 605 x, y = self.A.createVecs() 606 x.set(2) 607 y.set(3) 608 old = self._getCtx().D.copy() 609 self.A.diagonalScale(x, y) 610 D = self._getCtx().D 611 self.assertTrue(D.equal(old * 6)) 612 613 def testCreateTranspose(self): 614 A = self.A 615 A.setOption(PETSc.Mat.Option.SYMMETRIC, True) 616 AT = PETSc.Mat().createTranspose(A) 617 x, y = A.createVecs() 618 xt, yt = AT.createVecs() 619 # 620 y.setRandom() 621 A.multTranspose(y, x) 622 y.copy(xt) 623 AT.mult(xt, yt) 624 self.assertTrue(yt.equal(x)) 625 # 626 x.setRandom() 627 A.mult(x, y) 628 x.copy(yt) 629 AT.multTranspose(yt, xt) 630 self.assertTrue(xt.equal(y)) 631 del A 632 633 def testConvert(self): 634 self.assertTrue(self.A.convert(PETSc.Mat.Type.AIJ, PETSc.Mat()).equal(self.A)) 635 self.assertTrue(self.A.convert(PETSc.Mat.Type.BAIJ, PETSc.Mat()).equal(self.A)) 636 self.assertTrue(self.A.convert(PETSc.Mat.Type.SBAIJ, PETSc.Mat()).equal(self.A)) 637 self.assertTrue(self.A.convert(PETSc.Mat.Type.DENSE, PETSc.Mat()).equal(self.A)) 638 639 def testShift(self): 640 old = self._getCtx().D.copy() 641 self.A.shift(-0.5) 642 D = self._getCtx().D 643 self.assertTrue(D.equal(old - 0.5)) 644 645 def testScale(self): 646 old = self._getCtx().D.copy() 647 self.A.scale(-0.5) 648 D = self._getCtx().D 649 self.assertTrue(D.equal(-0.5 * old)) 650 651 def testDiagonalMat(self): 652 D = self._getCtx().D.copy() 653 B = PETSc.Mat().createDiagonal(D) 654 self.assertTrue(self.A.equal(B)) 655 656 657# -------------------------------------------------------------------- 658 659if __name__ == '__main__': 660 unittest.main() 661