1from petsc4py import PETSc 2import unittest 3 4# -------------------------------------------------------------------- 5 6 7class BaseTestDA: 8 COMM = PETSc.COMM_WORLD 9 SIZES = None 10 BOUNDARY = None 11 DOF = 1 12 STENCIL = PETSc.DMDA.StencilType.STAR 13 SWIDTH = 1 14 15 def setUp(self): 16 self.da = PETSc.DMDA().create( 17 dim=len(self.SIZES), 18 dof=self.DOF, 19 sizes=self.SIZES, 20 boundary_type=self.BOUNDARY, 21 stencil_type=self.STENCIL, 22 stencil_width=self.SWIDTH, 23 comm=self.COMM, 24 ) 25 26 def tearDown(self): 27 self.da = None 28 PETSc.garbage_cleanup() 29 30 def testGetInfo(self): 31 dim = self.da.getDim() 32 dof = self.da.getDof() 33 sizes = self.da.getSizes() 34 boundary = self.da.getBoundaryType() 35 stencil_type = self.da.getStencilType() 36 stencil_width = self.da.getStencilWidth() 37 self.assertEqual(dim, len(self.SIZES)) 38 self.assertEqual(dof, self.DOF) 39 self.assertEqual(sizes, tuple(self.SIZES)) 40 self.assertEqual(boundary, self.BOUNDARY or (0,) * dim) 41 self.assertEqual(stencil_type, self.STENCIL) 42 self.assertEqual(stencil_width, self.SWIDTH) 43 44 def testRangesCorners(self): 45 dim = self.da.getDim() 46 ranges = self.da.getRanges() 47 starts, lsizes = self.da.getCorners() 48 self.assertEqual(dim, len(ranges)) 49 self.assertEqual(dim, len(starts)) 50 self.assertEqual(dim, len(lsizes)) 51 for i in range(dim): 52 s, e = ranges[i] 53 self.assertEqual(s, starts[i]) 54 self.assertEqual(e - s, lsizes[i]) 55 56 def testGhostRangesCorners(self): 57 dim = self.da.getDim() 58 ranges = self.da.getGhostRanges() 59 starts, lsizes = self.da.getGhostCorners() 60 self.assertEqual(dim, len(ranges)) 61 self.assertEqual(dim, len(starts)) 62 self.assertEqual(dim, len(lsizes)) 63 for i in range(dim): 64 s, e = ranges[i] 65 self.assertEqual(s, starts[i]) 66 self.assertEqual(e - s, lsizes[i]) 67 68 def testOwnershipRanges(self): 69 ownership_ranges = self.da.getOwnershipRanges() 70 procsizes = self.da.getProcSizes() 71 self.assertEqual(len(procsizes), len(ownership_ranges)) 72 for i, m in enumerate(procsizes): 73 self.assertEqual(m, len(ownership_ranges[i])) 74 75 def testFieldName(self): 76 for i in range(self.da.getDof()): 77 self.da.setFieldName(i, 'field%d' % i) 78 for i in range(self.da.getDof()): 79 name = self.da.getFieldName(i) 80 self.assertEqual(name, 'field%d' % i) 81 82 def testCoordinates(self): 83 self.da.setUniformCoordinates(0, 1, 0, 1, 0, 1) 84 # 85 c = self.da.getCoordinates() 86 self.da.setCoordinates(c) 87 c.destroy() 88 cda = self.da.getCoordinateDM() 89 cda.destroy() 90 # 91 c = self.da.getCoordinates() 92 self.da.setCoordinates(c) 93 c.destroy() 94 gc = self.da.getCoordinatesLocal() 95 gc.destroy() 96 97 def testCreateVecMat(self): 98 vn = self.da.createNaturalVec() 99 vg = self.da.createGlobalVec() 100 vl = self.da.createLocalVec() 101 mat = self.da.createMat() 102 self.assertTrue(mat.getType() in ('aij', 'seqaij', 'mpiaij')) 103 vn.set(1.0) 104 self.da.naturalToGlobal(vn, vg) 105 self.assertEqual(vg.max()[1], 1.0) 106 self.assertEqual(vg.min()[1], 1.0) 107 self.da.globalToLocal(vg, vl) 108 self.assertEqual(vl.max()[1], 1.0) 109 self.assertTrue(vl.min()[1] in (1.0, 0.0)) 110 vn.set(0.0) 111 self.da.globalToNatural(vg, vn) 112 self.assertEqual(vn.max()[1], 1.0) 113 self.assertEqual(vn.min()[1], 1.0) 114 vl2 = self.da.createLocalVec() 115 self.da.localToLocal(vl, vl2) 116 self.assertEqual(vl2.max()[1], 1.0) 117 self.assertTrue(vl2.min()[1] in (1.0, 0.0)) 118 NONE = PETSc.DM.BoundaryType.NONE 119 btype = self.da.boundary_type 120 psize = self.da.proc_sizes 121 for b, p in zip(btype, psize): 122 if b != NONE and p == 1: 123 return 124 vg2 = self.da.createGlobalVec() 125 self.da.localToGlobal(vl2, vg2) 126 127 def testGetVec(self): 128 vg = self.da.getGlobalVec() 129 vl = self.da.getLocalVec() 130 try: 131 vg.set(1.0) 132 self.assertEqual(vg.max()[1], 1.0) 133 self.assertEqual(vg.min()[1], 1.0) 134 self.da.globalToLocal(vg, vl) 135 self.assertEqual(vl.max()[1], 1.0) 136 self.assertTrue(vl.min()[1] in (1.0, 0.0)) 137 vl.set(2.0) 138 NONE = PETSc.DM.BoundaryType.NONE 139 btype = self.da.boundary_type 140 psize = self.da.proc_sizes 141 for b, p in zip(btype, psize): 142 if b != NONE and p == 1: 143 return 144 self.da.localToGlobal(vl, vg) 145 self.assertEqual(vg.max()[1], 2.0) 146 self.assertTrue(vg.min()[1] in (2.0, 0.0)) 147 finally: 148 self.da.restoreGlobalVec(vg) 149 self.da.restoreLocalVec(vl) 150 self.assertFalse(vg) 151 self.assertFalse(vl) 152 153 name = 'abcd' 154 155 vg = self.da.getGlobalVec(name) 156 vg.set(4.0) 157 self.da.restoreGlobalVec(vg, name) 158 self.assertFalse(vg) 159 vg = self.da.getGlobalVec() 160 vg.setRandom() 161 self.da.restoreGlobalVec(vg) 162 vg = self.da.getGlobalVec(name) 163 vg.shift(-4.0) 164 self.assertEqual(vg.max()[1], 0.0) 165 self.da.restoreGlobalVec(vg, name) 166 167 vl = self.da.getLocalVec(name) 168 vl.set(4.0) 169 self.da.restoreLocalVec(vl, name) 170 self.assertFalse(vl) 171 vl = self.da.getLocalVec() 172 vl.setRandom() 173 self.da.restoreLocalVec(vl) 174 vl = self.da.getLocalVec(name) 175 vl.shift(-4.0) 176 self.assertEqual(vl.max()[1], 0.0) 177 self.da.restoreLocalVec(vl, name) 178 179 def testGetArray(self): 180 for r in [True, False]: 181 vg = self.da.getGlobalVec() 182 vl = self.da.getLocalVec() 183 184 ag = self.da.getVecArray(vg, readonly=r) 185 al = self.da.getVecArray(vl, readonly=r) 186 187 # test reading lower-left and upper-right corners 188 ranges = self.da.getRanges() 189 ranges = list(zip(*ranges)) 190 dim = len(ranges[0]) 191 if dim == 1: 192 _ = ag[ranges[0][0] : ranges[1][0]] 193 _ = al[ranges[0][0] : ranges[1][0]] 194 elif dim == 2: 195 _ = ag[ranges[0][0] : ranges[1][0], ranges[0][1] : ranges[1][1]] 196 _ = al[ranges[0][0] : ranges[1][0], ranges[0][1] : ranges[1][1]] 197 elif dim == 3: 198 _ = ag[ 199 ranges[0][0] : ranges[1][0], 200 ranges[0][1] : ranges[1][1], 201 ranges[0][2] : ranges[1][2], 202 ] 203 _ = al[ 204 ranges[0][0] : ranges[1][0], 205 ranges[0][1] : ranges[1][1], 206 ranges[0][2] : ranges[1][2], 207 ] 208 209 # test writing 210 if not r: 211 ag[:] = 1.0 212 al[:] = 1.0 213 else: 214 with self.assertRaises(ValueError): 215 ag[:] = 1.0 216 with self.assertRaises(ValueError): 217 al[:] = 1.0 218 219 self.da.restoreGlobalVec(vg) 220 self.da.restoreLocalVec(vl) 221 222 def testGetOther(self): 223 _ = self.da.getAO() 224 _ = self.da.getLGMap() 225 _, _ = self.da.getScatter() 226 227 def testRefineCoarsen(self): 228 da = self.da 229 rda = da.refine() 230 self.assertEqual(da.getDim(), rda.getDim()) 231 self.assertEqual(da.getDof(), rda.getDof()) 232 if da.dim != 1: 233 self.assertEqual(da.getStencilType(), rda.getStencilType()) 234 self.assertEqual(da.getStencilWidth(), rda.getStencilWidth()) 235 cda = rda.coarsen() 236 self.assertEqual(rda.getDim(), cda.getDim()) 237 self.assertEqual(rda.getDof(), cda.getDof()) 238 for n1, n2 in zip(self.da.getSizes(), cda.getSizes()): 239 self.assertTrue(abs(n1 - n2) <= 1) 240 241 def testCoarsenRefine(self): 242 if PETSc.COMM_WORLD.getSize() > 6: 243 return 244 da = self.da 245 cda = self.da.coarsen() 246 self.assertEqual(da.getDim(), cda.getDim()) 247 self.assertEqual(da.getDof(), cda.getDof()) 248 if da.dim != 1: 249 self.assertEqual(da.getStencilType(), cda.getStencilType()) 250 self.assertEqual(da.getStencilWidth(), cda.getStencilWidth()) 251 rda = cda.refine() 252 for n1, n2 in zip(self.da.getSizes(), rda.getSizes()): 253 self.assertTrue(abs(n1 - n2) <= 1) 254 255 def testRefineHierarchy(self): 256 levels = self.da.refineHierarchy(2) 257 self.assertTrue(isinstance(levels, list)) 258 self.assertEqual(len(levels), 2) 259 for item in levels: 260 self.assertTrue(isinstance(item, PETSc.DM)) 261 262 def testCoarsenHierarchy(self): 263 if PETSc.COMM_WORLD.getSize() > 6: 264 return 265 levels = self.da.coarsenHierarchy(2) 266 self.assertTrue(isinstance(levels, list)) 267 self.assertEqual(len(levels), 2) 268 for item in levels: 269 self.assertTrue(isinstance(item, PETSc.DM)) 270 271 def testCreateInterpolation(self): 272 da = self.da 273 if da.dim == 1: 274 return 275 rda = da.refine() 276 _, _ = da.createInterpolation(rda) 277 278 def testCreateInjection(self): 279 if PETSc.COMM_WORLD.getSize() > 6: 280 return 281 da = self.da 282 if da.dim == 1: 283 return 284 rda = da.refine() 285 _ = da.createInjection(rda) 286 287 def testzeroRowsColumnsStencil(self): 288 da = self.da 289 A = da.createMatrix() 290 x = da.createGlobalVector() 291 x.set(2.0) 292 A.setDiagonal(x) 293 diag1 = x.duplicate() 294 A.getDiagonal(diag1) 295 if self.SIZES != 2: # only coded test for 2D case 296 return 297 istart, iend, jstart, jend = da.getRanges() 298 self.assertTrue(x.equal(diag1)) 299 zeroidx = [] 300 for i in range(istart, iend): 301 for j in range(jstart, jend): 302 row = PETSc.Mat.Stencil() 303 row.index = (i, j) 304 zeroidx = zeroidx + [row] 305 diag2 = x.duplicate() 306 diag2.set(1.0) 307 A.zeroRowsColumnsStencil(zeroidx, 1.0, x, diag2) 308 ans = x.duplicate() 309 ans.set(2.0) 310 self.assertTrue(ans.equal(diag2)) 311 312 313MIRROR = PETSc.DMDA.BoundaryType.MIRROR 314GHOSTED = PETSc.DMDA.BoundaryType.GHOSTED 315PERIODIC = PETSc.DMDA.BoundaryType.PERIODIC 316TWIST = PETSc.DMDA.BoundaryType.TWIST 317 318SCALE = 4 319 320 321class BaseTestDA_1D(BaseTestDA): 322 SIZES = [100 * SCALE] 323 324 325class BaseTestDA_2D(BaseTestDA): 326 SIZES = [9 * SCALE, 11 * SCALE] 327 328 329class BaseTestDA_3D(BaseTestDA): 330 SIZES = [6 * SCALE, 7 * SCALE, 8 * SCALE] 331 332 333# -------------------------------------------------------------------- 334 335 336class TestDA_1D(BaseTestDA_1D, unittest.TestCase): 337 pass 338 339 340class TestDA_1D_W0(TestDA_1D): 341 SWIDTH = 0 342 343 344class TestDA_1D_W2(TestDA_1D): 345 SWIDTH = 2 346 347 348class TestDA_2D(BaseTestDA_2D, unittest.TestCase): 349 pass 350 351 352class TestDA_2D_W0(TestDA_2D): 353 SWIDTH = 0 354 355 356class TestDA_2D_W0_N2(TestDA_2D): 357 DOF = 2 358 SWIDTH = 0 359 360 361class TestDA_2D_W2(TestDA_2D): 362 SWIDTH = 2 363 364 365class TestDA_2D_W2_N2(TestDA_2D): 366 DOF = 2 367 SWIDTH = 2 368 369 370class TestDA_2D_PXY(TestDA_2D): 371 SIZES = [13 * SCALE, 17 * SCALE] 372 DOF = 2 373 SWIDTH = 5 374 BOUNDARY = (PERIODIC,) * 2 375 376 377class TestDA_2D_GXY(TestDA_2D): 378 SIZES = [13 * SCALE, 17 * SCALE] 379 DOF = 2 380 SWIDTH = 5 381 BOUNDARY = (GHOSTED,) * 2 382 383 384class TestDA_2D_TXY(TestDA_2D): 385 SIZES = [13 * SCALE, 17 * SCALE] 386 DOF = 2 387 SWIDTH = 5 388 BOUNDARY = (TWIST,) * 2 389 390 391class TestDA_3D(BaseTestDA_3D, unittest.TestCase): 392 pass 393 394 395class TestDA_3D_W0(TestDA_3D): 396 SWIDTH = 0 397 398 399class TestDA_3D_W0_N2(TestDA_3D): 400 DOF = 2 401 SWIDTH = 0 402 403 404class TestDA_3D_W2(TestDA_3D): 405 SWIDTH = 2 406 407 408class TestDA_3D_W2_N2(TestDA_3D): 409 DOF = 2 410 SWIDTH = 2 411 412 413class TestDA_3D_PXYZ(TestDA_3D): 414 SIZES = [11 * SCALE, 13 * SCALE, 17 * SCALE] 415 DOF = 2 416 SWIDTH = 3 417 BOUNDARY = (PERIODIC,) * 3 418 419 420class TestDA_3D_GXYZ(TestDA_3D): 421 SIZES = [11 * SCALE, 13 * SCALE, 17 * SCALE] 422 DOF = 2 423 SWIDTH = 3 424 BOUNDARY = (GHOSTED,) * 3 425 426 427class TestDA_3D_TXYZ(TestDA_3D): 428 SIZES = [11 * SCALE, 13 * SCALE, 17 * SCALE] 429 DOF = 2 430 SWIDTH = 3 431 BOUNDARY = (TWIST,) * 3 432 433 434# -------------------------------------------------------------------- 435 436DIM = ( 437 1, 438 2, 439 3, 440) 441DOF = ( 442 None, 443 1, 444 2, 445 3, 446 4, 447 5, 448) 449BOUNDARY_TYPE = ( 450 None, 451 'none', 452 (0,) * 3, 453 0, 454 'ghosted', 455 (GHOSTED,) * 3, 456 GHOSTED, 457 'periodic', 458 (PERIODIC,) * 3, 459 PERIODIC, 460 'twist', 461 (TWIST,) * 3, 462 TWIST, 463) 464STENCIL_TYPE = (None, 'star', 'box') 465STENCIL_WIDTH = (None, 0, 1, 2, 3) 466 467 468DIM = (1, 2, 3) 469DOF = (None, 2, 5) 470BOUNDARY_TYPE = (None, 'none', 'periodic', 'ghosted', 'twist') 471STENCIL_TYPE = (None, 'box') 472STENCIL_WIDTH = (None, 1, 2) 473 474 475class TestDACreate(unittest.TestCase): 476 pass 477 478 479counter = 0 480for dim in DIM: 481 for dof in DOF: 482 for boundary in BOUNDARY_TYPE: 483 if isinstance(boundary, tuple): 484 boundary = boundary[:dim] 485 for stencil in STENCIL_TYPE: 486 for width in STENCIL_WIDTH: 487 kargs = { 488 'sizes': [8 * SCALE] * dim, 489 'dim': dim, 490 'dof': dof, 491 'boundary_type': boundary, 492 'stencil_type': stencil, 493 'stencil_width': width, 494 } 495 496 def testCreate(self, kargs=kargs): 497 kargs = dict(kargs) 498 da = PETSc.DMDA().create(**kargs) 499 da.destroy() 500 501 setattr(TestDACreate, 'testCreate%04d' % counter, testCreate) 502 del testCreate, kargs 503 counter += 1 504del counter, dim, dof, boundary, stencil, width 505 506 507class TestDADuplicate(unittest.TestCase): 508 pass 509 510 511counter = 0 512for dim in DIM: 513 for dof in DOF: 514 for boundary in BOUNDARY_TYPE: 515 if isinstance(boundary, tuple): 516 boundary = boundary[:dim] 517 for stencil in STENCIL_TYPE: 518 for width in STENCIL_WIDTH: 519 kargs = { 520 'dim': dim, 521 'dof': dof, 522 'boundary_type': boundary, 523 'stencil_type': stencil, 524 'stencil_width': width, 525 } 526 527 def testDuplicate(self, kargs=kargs): 528 kargs = dict(kargs) 529 dim = kargs.pop('dim') 530 dof = kargs['dof'] 531 boundary = kargs['boundary_type'] 532 stencil = kargs['stencil_type'] 533 width = kargs['stencil_width'] 534 da = PETSc.DMDA().create([8 * SCALE] * dim) 535 newda = da.duplicate(**kargs) 536 self.assertEqual(newda.dim, da.dim) 537 self.assertEqual(newda.sizes, da.sizes) 538 self.assertEqual(newda.proc_sizes, da.proc_sizes) 539 self.assertEqual(newda.ranges, da.ranges) 540 self.assertEqual(newda.corners, da.corners) 541 if ( 542 newda.boundary_type == da.boundary_type 543 and newda.stencil_width == da.stencil_width 544 ): 545 self.assertEqual(newda.ghost_ranges, da.ghost_ranges) 546 self.assertEqual(newda.ghost_corners, da.ghost_corners) 547 if dof is None: 548 dof = da.dof 549 if boundary is None: 550 boundary = da.boundary_type 551 elif boundary == 'none': 552 boundary = (0,) * dim 553 elif boundary == 'mirror': 554 boundary = (MIRROR,) * dim 555 elif boundary == 'ghosted': 556 boundary = (GHOSTED,) * dim 557 elif boundary == 'periodic': 558 boundary = (PERIODIC,) * dim 559 elif boundary == 'twist': 560 boundary = (TWIST,) * dim 561 elif isinstance(boundary, int): 562 boundary = (boundary,) * dim 563 if stencil is None: 564 stencil = da.stencil[0] 565 if width is None: 566 width = da.stencil_width 567 self.assertEqual(newda.dof, dof) 568 self.assertEqual(newda.boundary_type, boundary) 569 if dim == 1: 570 self.assertEqual(newda.stencil, (stencil, width)) 571 newda.destroy() 572 da.destroy() 573 574 setattr( 575 TestDADuplicate, 'testDuplicate%04d' % counter, testDuplicate 576 ) 577 del testDuplicate, kargs 578 counter += 1 579del counter, dim, dof, boundary, stencil, width 580 581# -------------------------------------------------------------------- 582 583if PETSc.COMM_WORLD.getSize() > 1: 584 del TestDA_1D_W0 585 del TestDA_2D_W0, TestDA_2D_W0_N2 586 del TestDA_3D_W0, TestDA_3D_W0_N2 587 588# -------------------------------------------------------------------- 589 590if __name__ == '__main__': 591 unittest.main() 592 593# -------------------------------------------------------------------- 594