1from petsc4py import PETSc 2import unittest, numpy 3from sys import getrefcount 4# -------------------------------------------------------------------- 5 6class Matrix(object): 7 8 def __init__(self): 9 pass 10 11 def create(self, mat): 12 pass 13 14 def destroy(self, mat): 15 pass 16 17class Identity(Matrix): 18 19 def mult(self, mat, x, y): 20 x.copy(y) 21 22 def getDiagonal(self, mat, vd): 23 vd.set(1) 24 25 def productSetFromOptions(self, mat, producttype, A, B, C): 26 return True 27 28 def productSymbolic(self, mat, product, producttype, A, B, C): 29 if producttype == 'AB': 30 if mat is A: # product = identity * B 31 product.setType(B.getType()) 32 product.setSizes(B.getSizes()) 33 product.setUp() 34 product.assemble() 35 B.copy(product) 36 elif mat is B: # product = A * identity 37 product.setType(A.getType()) 38 product.setSizes(A.getSizes()) 39 product.setUp() 40 product.assemble() 41 A.copy(product) 42 else: 43 raise RuntimeError('wrong configuration') 44 elif producttype == 'AtB': 45 if mat is A: # product = identity^T * B 46 product.setType(B.getType()) 47 product.setSizes(B.getSizes()) 48 product.setUp() 49 product.assemble() 50 B.copy(product) 51 elif mat is B: # product = A^T * identity 52 tmp = PETSc.Mat() 53 A.transpose(tmp) 54 product.setType(tmp.getType()) 55 product.setSizes(tmp.getSizes()) 56 product.setUp() 57 product.assemble() 58 tmp.copy(product) 59 else: 60 raise RuntimeError('wrong configuration') 61 elif producttype == 'ABt': 62 if mat is A: # product = identity * B^T 63 tmp = PETSc.Mat() 64 B.transpose(tmp) 65 product.setType(tmp.getType()) 66 product.setSizes(tmp.getSizes()) 67 product.setUp() 68 product.assemble() 69 tmp.copy(product) 70 elif mat is B: # product = A * identity^T 71 product.setType(A.getType()) 72 product.setSizes(A.getSizes()) 73 product.setUp() 74 product.assemble() 75 A.copy(product) 76 else: 77 raise RuntimeError('wrong configuration') 78 elif producttype == 'PtAP': 79 if mat is A: # product = P^T * identity * P 80 self.tmp = PETSc.Mat() 81 B.transposeMatMult(B, self.tmp) 82 product.setType(self.tmp.getType()) 83 product.setSizes(self.tmp.getSizes()) 84 product.setUp() 85 product.assemble() 86 self.tmp.copy(product) 87 elif mat is B: # product = identity^T * A * identity 88 product.setType(A.getType()) 89 product.setSizes(A.getSizes()) 90 product.setUp() 91 product.assemble() 92 A.copy(product) 93 else: 94 raise RuntimeError('wrong configuration') 95 elif producttype == 'RARt': 96 if mat is A: # product = R * identity * R^t 97 self.tmp = PETSc.Mat() 98 B.matTransposeMult(B, self.tmp) 99 product.setType(self.tmp.getType()) 100 product.setSizes(self.tmp.getSizes()) 101 product.setUp() 102 product.assemble() 103 self.tmp.copy(product) 104 elif mat is B: # product = identity * A * identity^T 105 product.setType(A.getType()) 106 product.setSizes(A.getSizes()) 107 product.setUp() 108 product.assemble() 109 A.copy(product) 110 else: 111 raise RuntimeError('wrong configuration') 112 elif producttype == 'ABC': 113 if mat is A: # product = identity * B * C 114 self.tmp = PETSc.Mat() 115 B.matMult(C, self.tmp) 116 product.setType(self.tmp.getType()) 117 product.setSizes(self.tmp.getSizes()) 118 product.setUp() 119 product.assemble() 120 self.tmp.copy(product) 121 elif mat is B: # product = A * identity * C 122 self.tmp = PETSc.Mat() 123 A.matMult(C, self.tmp) 124 product.setType(self.tmp.getType()) 125 product.setSizes(self.tmp.getSizes()) 126 product.setUp() 127 product.assemble() 128 self.tmp.copy(product) 129 elif mat is C: # product = A * B * identity 130 self.tmp = PETSc.Mat() 131 A.matMult(B, self.tmp) 132 product.setType(self.tmp.getType()) 133 product.setSizes(self.tmp.getSizes()) 134 product.setUp() 135 product.assemble() 136 self.tmp.copy(product) 137 else: 138 raise RuntimeError('wrong configuration') 139 else: 140 raise RuntimeError('Product {} not implemented'.format(producttype)) 141 product.zeroEntries() 142 143 def productNumeric(self, mat, product, producttype, A, B, C): 144 if producttype == 'AB': 145 if mat is A: # product = identity * B 146 B.copy(product, structure=True) 147 elif mat is B: # product = A * identity 148 A.copy(product, structure=True) 149 else: 150 raise RuntimeError('wrong configuration') 151 elif producttype == 'AtB': 152 if mat is A: # product = identity^T * B 153 B.copy(product, structure=True) 154 elif mat is B: # product = A^T * identity 155 A.transpose(product) 156 else: 157 raise RuntimeError('wrong configuration') 158 elif producttype == 'ABt': 159 if mat is A: # product = identity * B^T 160 B.transpose(product) 161 elif mat is B: # product = A * identity^T 162 A.copy(product, structure=True) 163 else: 164 raise RuntimeError('wrong configuration') 165 elif producttype == 'PtAP': 166 if mat is A: # product = P^T * identity * P 167 B.transposeMatMult(B, self.tmp) 168 self.tmp.copy(product, structure=True) 169 elif mat is B: # product = identity^T * A * identity 170 A.copy(product, structure=True) 171 else: 172 raise RuntimeError('wrong configuration') 173 elif producttype == 'RARt': 174 if mat is A: # product = R * identity * R^t 175 B.matTransposeMult(B, self.tmp) 176 self.tmp.copy(product, structure=True) 177 elif mat is B: # product = identity * A * identity^T 178 A.copy(product, structure=True) 179 else: 180 raise RuntimeError('wrong configuration') 181 elif producttype == 'ABC': 182 if mat is A: # product = identity * B * C 183 B.matMult(C, self.tmp) 184 self.tmp.copy(product, structure=True) 185 elif mat is B: # product = A * identity * C 186 A.matMult(C, self.tmp) 187 self.tmp.copy(product, structure=True) 188 elif mat is C: # product = A * B * identity 189 A.matMult(B, self.tmp) 190 self.tmp.copy(product, structure=True) 191 else: 192 raise RuntimeError('wrong configuration') 193 else: 194 raise RuntimeError('Product {} not implemented'.format(producttype)) 195 196class Diagonal(Matrix): 197 198 def create(self, mat): 199 super(Diagonal,self).create(mat) 200 mat.setUp() 201 self.D = mat.createVecLeft() 202 203 def destroy(self, mat): 204 self.D.destroy() 205 super(Diagonal,self).destroy(mat) 206 207 def scale(self, mat, a): 208 self.D.scale(a) 209 210 def shift(self, mat, a): 211 self.D.shift(a) 212 213 def zeroEntries(self, mat): 214 self.D.zeroEntries() 215 216 def mult(self, mat, x, y): 217 y.pointwiseMult(x, self.D) 218 219 def getDiagonal(self, mat, vd): 220 self.D.copy(vd) 221 222 def setDiagonal(self, mat, vd, im): 223 if isinstance (im, bool): 224 addv = im 225 if addv: 226 self.D.axpy(1, vd) 227 else: 228 vd.copy(self.D) 229 elif im == PETSc.InsertMode.INSERT_VALUES: 230 vd.copy(self.D) 231 elif im == PETSc.InsertMode.ADD_VALUES: 232 self.D.axpy(1, vd) 233 else: 234 raise ValueError('wrong InsertMode %d'% im) 235 236 def diagonalScale(self, mat, vl, vr): 237 if vl: self.D.pointwiseMult(self.D, vl) 238 if vr: self.D.pointwiseMult(self.D, vr) 239 240# -------------------------------------------------------------------- 241 242class TestMatrix(unittest.TestCase): 243 244 COMM = PETSc.COMM_WORLD 245 PYMOD = __name__ 246 PYCLS = 'Matrix' 247 248 def _getCtx(self): 249 return self.A.getPythonContext() 250 251 def setUp(self): 252 N = self.N = 10 253 self.A = PETSc.Mat() 254 if 0: # command line way 255 self.A.create(self.COMM) 256 self.A.setSizes([N,N]) 257 self.A.setType('python') 258 OptDB = PETSc.Options(self.A) 259 OptDB['mat_python_type'] = '%s.%s' % (self.PYMOD,self.PYCLS) 260 self.A.setFromOptions() 261 self.A.setUp() 262 del OptDB['mat_python_type'] 263 self.assertTrue(self._getCtx() is not None) 264 else: # python way 265 context = globals()[self.PYCLS]() 266 self.A.createPython([N,N], context, comm=self.COMM) 267 self.A.setUp() 268 self.assertTrue(self._getCtx() is context) 269 self.assertEqual(getrefcount(context), 3) 270 del context 271 self.assertEqual(getrefcount(self._getCtx()), 2) 272 273 def tearDown(self): 274 ctx = self.A.getPythonContext() 275 self.assertEqual(getrefcount(ctx), 3) 276 self.A.destroy() # XXX 277 self.A = None 278 self.assertEqual(getrefcount(ctx), 2) 279 #import gc,pprint; pprint.pprint(gc.get_referrers(ctx)) 280 281 def testBasic(self): 282 ctx = self.A.getPythonContext() 283 self.assertTrue(self._getCtx() is ctx) 284 self.assertEqual(getrefcount(ctx), 3) 285 286 def testZeroEntries(self): 287 f = lambda : self.A.zeroEntries() 288 self.assertRaises(Exception, f) 289 290 def testMult(self): 291 x, y = self.A.createVecs() 292 f = lambda : self.A.mult(x, y) 293 self.assertRaises(Exception, f) 294 295 def testMultTranspose(self): 296 x, y = self.A.createVecs() 297 f = lambda : self.A.multTranspose(x, y) 298 self.assertRaises(Exception, f) 299 300 def testGetDiagonal(self): 301 d = self.A.createVecLeft() 302 f = lambda : self.A.getDiagonal(d) 303 self.assertRaises(Exception, f) 304 305 def testSetDiagonal(self): 306 d = self.A.createVecLeft() 307 f = lambda : self.A.setDiagonal(d) 308 self.assertRaises(Exception, f) 309 310 def testDiagonalScale(self): 311 x, y = self.A.createVecs() 312 f = lambda : self.A.diagonalScale(x, y) 313 self.assertRaises(Exception, f) 314 315 316class TestIdentity(TestMatrix): 317 318 PYCLS = 'Identity' 319 320 def testMult(self): 321 x, y = self.A.createVecs() 322 x.setRandom() 323 self.A.mult(x,y) 324 self.assertTrue(y.equal(x)) 325 326 def testMultTransposeSymmKnown(self): 327 x, y = self.A.createVecs() 328 x.setRandom() 329 self.A.setOption(PETSc.Mat.Option.SYMMETRIC, True) 330 self.A.multTranspose(x,y) 331 self.assertTrue(y.equal(x)) 332 self.A.setOption(PETSc.Mat.Option.SYMMETRIC, False) 333 f = lambda : self.A.multTranspose(x, y) 334 self.assertRaises(Exception, f) 335 336 def testMultTransposeNewMeth(self): 337 x, y = self.A.createVecs() 338 x.setRandom() 339 AA = self.A.getPythonContext() 340 AA.multTranspose = AA.mult 341 self.A.multTranspose(x,y) 342 del AA.multTranspose 343 self.assertTrue(y.equal(x)) 344 345 def testGetDiagonal(self): 346 d = self.A.createVecLeft() 347 o = d.duplicate() 348 o.set(1) 349 self.A.getDiagonal(d) 350 self.assertTrue(o.equal(d)) 351 352 def testMatMat(self): 353 R = PETSc.Random().create(self.COMM) 354 R.setFromOptions() 355 A = PETSc.Mat().create(self.COMM) 356 A.setSizes(self.A.getSizes()) 357 A.setType(PETSc.Mat.Type.AIJ) 358 A.setUp() 359 A.setRandom(R) 360 B = PETSc.Mat().create(self.COMM) 361 B.setSizes(self.A.getSizes()) 362 B.setType(PETSc.Mat.Type.AIJ) 363 B.setUp() 364 B.setRandom(R) 365 I = PETSc.Mat().create(self.COMM) 366 I.setSizes(self.A.getSizes()) 367 I.setType(PETSc.Mat.Type.AIJ) 368 I.setUp() 369 I.assemble() 370 I.shift(1.) 371 372 self.assertTrue(self.A.matMult(A).equal(I.matMult(A))) 373 self.assertTrue(A.matMult(self.A).equal(A.matMult(I))) 374 if self.A.getComm().Get_size() == 1: 375 self.assertTrue(self.A.matTransposeMult(A).equal(I.matTransposeMult(A))) 376 self.assertTrue(A.matTransposeMult(self.A).equal(A.matTransposeMult(I))) 377 self.assertTrue(self.A.transposeMatMult(A).equal(I.transposeMatMult(A))) 378 self.assertTrue(A.transposeMatMult(self.A).equal(A.transposeMatMult(I))) 379 self.assertAlmostEqual((self.A.ptap(A) - I.ptap(A)).norm(), 0.0, places=5) 380 self.assertAlmostEqual((A.ptap(self.A) - A.ptap(I)).norm(), 0.0, places=5) 381 if self.A.getComm().Get_size() == 1: 382 self.assertAlmostEqual((self.A.rart(A) - I.rart(A)).norm(), 0.0, places=5) 383 self.assertAlmostEqual((A.rart(self.A) - A.rart(I)).norm(), 0.0, places=5) 384 self.assertAlmostEqual((self.A.matMatMult(A,B)-I.matMatMult(A,B)).norm(), 0.0, places=5) 385 self.assertAlmostEqual((A.matMatMult(self.A,B)-A.matMatMult(I,B)).norm(), 0.0, places=5) 386 self.assertAlmostEqual((A.matMatMult(B,self.A)-A.matMatMult(B,I)).norm(), 0.0, places=5) 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 transpose operation for norm estimation 396 AA = self.A.getPythonContext() 397 AA.multTranspose = AA.mult 398 399 # without coordinates 400 h.createH2OpusFromMat(self.A,leafsize=2) 401 h.assemble() 402 h.destroy() 403 404 # with coordinates 405 coords = numpy.linspace((1,2,3),(10,20,30),self.A.getSize()[0],dtype=PETSc.RealType) 406 h.createH2OpusFromMat(self.A,coords,leafsize=2) 407 h.assemble() 408 h.destroy() 409 410 del AA.multTranspose 411 412class TestDiagonal(TestMatrix): 413 414 PYCLS = 'Diagonal' 415 416 def setUp(self): 417 super(TestDiagonal, self).setUp() 418 D = self.A.createVecLeft() 419 s, e = D.getOwnershipRange() 420 for i in range(s, e): 421 D[i] = i+1 422 D.assemble() 423 self.A.setDiagonal(D) 424 425 def testZeroEntries(self): 426 self.A.zeroEntries() 427 D = self._getCtx().D 428 self.assertEqual(D.norm(), 0) 429 430 def testMult(self): 431 x, y = self.A.createVecs() 432 x.set(1) 433 self.A.mult(x,y) 434 self.assertTrue(y.equal(self._getCtx().D)) 435 436 def testMultTransposeSymmKnown(self): 437 x, y = self.A.createVecs() 438 x.set(1) 439 self.A.setOption(PETSc.Mat.Option.SYMMETRIC, True) 440 self.A.multTranspose(x,y) 441 self.assertTrue(y.equal(self._getCtx().D)) 442 self.A.setOption(PETSc.Mat.Option.SYMMETRIC, False) 443 f = lambda : self.A.multTranspose(x, y) 444 self.assertRaises(Exception, f) 445 446 def testMultTransposeNewMeth(self): 447 x, y = self.A.createVecs() 448 x.set(1) 449 AA = self.A.getPythonContext() 450 AA.multTranspose = AA.mult 451 self.A.multTranspose(x,y) 452 del AA.multTranspose 453 self.assertTrue(y.equal(self._getCtx().D)) 454 455 def testGetDiagonal(self): 456 d = self.A.createVecLeft() 457 self.A.getDiagonal(d) 458 self.assertTrue(d.equal(self._getCtx().D)) 459 460 def testSetDiagonal(self): 461 d = self.A.createVecLeft() 462 d.setRandom() 463 self.A.setDiagonal(d) 464 self.assertTrue(d.equal(self._getCtx().D)) 465 466 def testDiagonalScale(self): 467 x, y = self.A.createVecs() 468 x.set(2) 469 y.set(3) 470 old = self._getCtx().D.copy() 471 self.A.diagonalScale(x, y) 472 D = self._getCtx().D 473 self.assertTrue(D.equal(old*6)) 474 475 def testCreateTranspose(self): 476 A = self.A 477 A.setOption(PETSc.Mat.Option.SYMMETRIC, True) 478 AT = PETSc.Mat().createTranspose(A) 479 x, y = A.createVecs() 480 xt, yt = AT.createVecs() 481 # 482 y.setRandom() 483 A.multTranspose(y, x) 484 y.copy(xt) 485 AT.mult(xt, yt) 486 self.assertTrue(yt.equal(x)) 487 # 488 x.setRandom() 489 A.mult(x, y) 490 x.copy(yt) 491 AT.multTranspose(yt, xt) 492 self.assertTrue(xt.equal(y)) 493 del A 494 495 def testConvert(self): 496 self.assertTrue(self.A.convert(PETSc.Mat.Type.AIJ,PETSc.Mat()).equal(self.A)) 497 self.assertTrue(self.A.convert(PETSc.Mat.Type.BAIJ,PETSc.Mat()).equal(self.A)) 498 self.assertTrue(self.A.convert(PETSc.Mat.Type.SBAIJ,PETSc.Mat()).equal(self.A)) 499 self.assertTrue(self.A.convert(PETSc.Mat.Type.DENSE,PETSc.Mat()).equal(self.A)) 500 501# -------------------------------------------------------------------- 502 503if __name__ == '__main__': 504 unittest.main() 505