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