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