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 315class TestIdentity(TestMatrix): 316 317 PYCLS = 'Identity' 318 319 def testMult(self): 320 x, y = self.A.createVecs() 321 x.setRandom() 322 self.A.mult(x,y) 323 self.assertTrue(y.equal(x)) 324 325 def testMultTransposeSymmKnown(self): 326 x, y = self.A.createVecs() 327 x.setRandom() 328 self.A.setOption(PETSc.Mat.Option.SYMMETRIC, True) 329 self.A.multTranspose(x,y) 330 self.assertTrue(y.equal(x)) 331 self.A.setOption(PETSc.Mat.Option.SYMMETRIC, False) 332 f = lambda : self.A.multTranspose(x, y) 333 self.assertRaises(Exception, f) 334 335 def testMultTransposeNewMeth(self): 336 x, y = self.A.createVecs() 337 x.setRandom() 338 AA = self.A.getPythonContext() 339 AA.multTranspose = AA.mult 340 self.A.multTranspose(x,y) 341 del AA.multTranspose 342 self.assertTrue(y.equal(x)) 343 344 def testGetDiagonal(self): 345 d = self.A.createVecLeft() 346 o = d.duplicate() 347 o.set(1) 348 self.A.getDiagonal(d) 349 self.assertTrue(o.equal(d)) 350 351 def testMatMat(self): 352 R = PETSc.Random().create(self.COMM) 353 R.setFromOptions() 354 A = PETSc.Mat().create(self.COMM) 355 A.setSizes(self.A.getSizes()) 356 A.setType(PETSc.Mat.Type.AIJ) 357 A.setUp() 358 A.setRandom(R) 359 B = PETSc.Mat().create(self.COMM) 360 B.setSizes(self.A.getSizes()) 361 B.setType(PETSc.Mat.Type.AIJ) 362 B.setUp() 363 B.setRandom(R) 364 I = PETSc.Mat().create(self.COMM) 365 I.setSizes(self.A.getSizes()) 366 I.setType(PETSc.Mat.Type.AIJ) 367 I.setUp() 368 I.assemble() 369 I.shift(1.) 370 371 self.assertTrue(self.A.matMult(A).equal(I.matMult(A))) 372 self.assertTrue(A.matMult(self.A).equal(A.matMult(I))) 373 if self.A.getComm().Get_size() == 1: 374 self.assertTrue(self.A.matTransposeMult(A).equal(I.matTransposeMult(A))) 375 self.assertTrue(A.matTransposeMult(self.A).equal(A.matTransposeMult(I))) 376 self.assertTrue(self.A.transposeMatMult(A).equal(I.transposeMatMult(A))) 377 self.assertTrue(A.transposeMatMult(self.A).equal(A.transposeMatMult(I))) 378 self.assertAlmostEqual((self.A.ptap(A) - I.ptap(A)).norm(), 0.0, places=5) 379 self.assertAlmostEqual((A.ptap(self.A) - A.ptap(I)).norm(), 0.0, places=5) 380 if self.A.getComm().Get_size() == 1: 381 self.assertAlmostEqual((self.A.rart(A) - I.rart(A)).norm(), 0.0, places=5) 382 self.assertAlmostEqual((A.rart(self.A) - A.rart(I)).norm(), 0.0, places=5) 383 self.assertAlmostEqual((self.A.matMatMult(A,B)-I.matMatMult(A,B)).norm(), 0.0, places=5) 384 self.assertAlmostEqual((A.matMatMult(self.A,B)-A.matMatMult(I,B)).norm(), 0.0, places=5) 385 self.assertAlmostEqual((A.matMatMult(B,self.A)-A.matMatMult(B,I)).norm(), 0.0, places=5) 386 387 def testH2Opus(self): 388 if not PETSc.Sys.hasExternalPackage("h2opus"): 389 return 390 if self.A.getComm().Get_size() > 1: 391 return 392 h = PETSc.Mat() 393 394 # need transpose operation for norm estimation 395 AA = self.A.getPythonContext() 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((1,2,3),(10,20,30),self.A.getSize()[0],dtype=PETSc.RealType) 405 h.createH2OpusFromMat(self.A,coords,leafsize=2) 406 h.assemble() 407 h.destroy() 408 409 del AA.multTranspose 410 411class TestDiagonal(TestMatrix): 412 413 PYCLS = 'Diagonal' 414 415 def setUp(self): 416 super(TestDiagonal, self).setUp() 417 D = self.A.createVecLeft() 418 s, e = D.getOwnershipRange() 419 for i in range(s, e): 420 D[i] = i+1 421 D.assemble() 422 self.A.setDiagonal(D) 423 424 def testZeroEntries(self): 425 self.A.zeroEntries() 426 D = self._getCtx().D 427 self.assertEqual(D.norm(), 0) 428 429 def testMult(self): 430 x, y = self.A.createVecs() 431 x.set(1) 432 self.A.mult(x,y) 433 self.assertTrue(y.equal(self._getCtx().D)) 434 435 def testMultTransposeSymmKnown(self): 436 x, y = self.A.createVecs() 437 x.set(1) 438 self.A.setOption(PETSc.Mat.Option.SYMMETRIC, True) 439 self.A.multTranspose(x,y) 440 self.assertTrue(y.equal(self._getCtx().D)) 441 self.A.setOption(PETSc.Mat.Option.SYMMETRIC, False) 442 f = lambda : self.A.multTranspose(x, y) 443 self.assertRaises(Exception, f) 444 445 def testMultTransposeNewMeth(self): 446 x, y = self.A.createVecs() 447 x.set(1) 448 AA = self.A.getPythonContext() 449 AA.multTranspose = AA.mult 450 self.A.multTranspose(x,y) 451 del AA.multTranspose 452 self.assertTrue(y.equal(self._getCtx().D)) 453 454 def testGetDiagonal(self): 455 d = self.A.createVecLeft() 456 self.A.getDiagonal(d) 457 self.assertTrue(d.equal(self._getCtx().D)) 458 459 def testSetDiagonal(self): 460 d = self.A.createVecLeft() 461 d.setRandom() 462 self.A.setDiagonal(d) 463 self.assertTrue(d.equal(self._getCtx().D)) 464 465 def testDiagonalScale(self): 466 x, y = self.A.createVecs() 467 x.set(2) 468 y.set(3) 469 old = self._getCtx().D.copy() 470 self.A.diagonalScale(x, y) 471 D = self._getCtx().D 472 self.assertTrue(D.equal(old*6)) 473 474 def testCreateTranspose(self): 475 A = self.A 476 A.setOption(PETSc.Mat.Option.SYMMETRIC, True) 477 AT = PETSc.Mat().createTranspose(A) 478 x, y = A.createVecs() 479 xt, yt = AT.createVecs() 480 # 481 y.setRandom() 482 A.multTranspose(y, x) 483 y.copy(xt) 484 AT.mult(xt, yt) 485 self.assertTrue(yt.equal(x)) 486 # 487 x.setRandom() 488 A.mult(x, y) 489 x.copy(yt) 490 AT.multTranspose(yt, xt) 491 self.assertTrue(xt.equal(y)) 492 del A 493 494 def testConvert(self): 495 self.assertTrue(self.A.convert(PETSc.Mat.Type.AIJ,PETSc.Mat()).equal(self.A)) 496 self.assertTrue(self.A.convert(PETSc.Mat.Type.BAIJ,PETSc.Mat()).equal(self.A)) 497 self.assertTrue(self.A.convert(PETSc.Mat.Type.SBAIJ,PETSc.Mat()).equal(self.A)) 498 self.assertTrue(self.A.convert(PETSc.Mat.Type.DENSE,PETSc.Mat()).equal(self.A)) 499 500# -------------------------------------------------------------------- 501 502if __name__ == '__main__': 503 unittest.main() 504