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 151 def testGetOther(self): 152 _ = self.da.getAO() 153 _ = self.da.getLGMap() 154 _, _ = self.da.getScatter() 155 156 def testRefineCoarsen(self): 157 da = self.da 158 rda = da.refine() 159 self.assertEqual(da.getDim(), rda.getDim()) 160 self.assertEqual(da.getDof(), rda.getDof()) 161 if da.dim != 1: 162 self.assertEqual(da.getStencilType(), rda.getStencilType()) 163 self.assertEqual(da.getStencilWidth(), rda.getStencilWidth()) 164 cda = rda.coarsen() 165 self.assertEqual(rda.getDim(), cda.getDim()) 166 self.assertEqual(rda.getDof(), cda.getDof()) 167 for n1, n2 in zip(self.da.getSizes(), cda.getSizes()): 168 self.assertTrue(abs(n1 - n2) <= 1) 169 170 def testCoarsenRefine(self): 171 if PETSc.COMM_WORLD.getSize() > 6: 172 return 173 da = self.da 174 cda = self.da.coarsen() 175 self.assertEqual(da.getDim(), cda.getDim()) 176 self.assertEqual(da.getDof(), cda.getDof()) 177 if da.dim != 1: 178 self.assertEqual(da.getStencilType(), cda.getStencilType()) 179 self.assertEqual(da.getStencilWidth(), cda.getStencilWidth()) 180 rda = cda.refine() 181 for n1, n2 in zip(self.da.getSizes(), rda.getSizes()): 182 self.assertTrue(abs(n1 - n2) <= 1) 183 184 def testRefineHierarchy(self): 185 levels = self.da.refineHierarchy(2) 186 self.assertTrue(isinstance(levels, list)) 187 self.assertEqual(len(levels), 2) 188 for item in levels: 189 self.assertTrue(isinstance(item, PETSc.DM)) 190 191 def testCoarsenHierarchy(self): 192 if PETSc.COMM_WORLD.getSize() > 6: 193 return 194 levels = self.da.coarsenHierarchy(2) 195 self.assertTrue(isinstance(levels, list)) 196 self.assertEqual(len(levels), 2) 197 for item in levels: 198 self.assertTrue(isinstance(item, PETSc.DM)) 199 200 def testCreateInterpolation(self): 201 da = self.da 202 if da.dim == 1: 203 return 204 rda = da.refine() 205 _, _ = da.createInterpolation(rda) 206 207 def testCreateInjection(self): 208 if PETSc.COMM_WORLD.getSize() > 6: 209 return 210 da = self.da 211 if da.dim == 1: 212 return 213 rda = da.refine() 214 _ = da.createInjection(rda) 215 216 def testzeroRowsColumnsStencil(self): 217 da = self.da 218 A = da.createMatrix() 219 x = da.createGlobalVector() 220 x.set(2.0) 221 A.setDiagonal(x) 222 diag1 = x.duplicate() 223 A.getDiagonal(diag1) 224 if self.SIZES != 2: # only coded test for 2D case 225 return 226 istart, iend, jstart, jend = da.getRanges() 227 self.assertTrue(x.equal(diag1)) 228 zeroidx = [] 229 for i in range(istart, iend): 230 for j in range(jstart, jend): 231 row = PETSc.Mat.Stencil() 232 row.index = (i, j) 233 zeroidx = zeroidx + [row] 234 diag2 = x.duplicate() 235 diag2.set(1.0) 236 A.zeroRowsColumnsStencil(zeroidx, 1.0, x, diag2) 237 ans = x.duplicate() 238 ans.set(2.0) 239 self.assertTrue(ans.equal(diag2)) 240 241 242MIRROR = PETSc.DMDA.BoundaryType.MIRROR 243GHOSTED = PETSc.DMDA.BoundaryType.GHOSTED 244PERIODIC = PETSc.DMDA.BoundaryType.PERIODIC 245TWIST = PETSc.DMDA.BoundaryType.TWIST 246 247SCALE = 4 248 249 250class BaseTestDA_1D(BaseTestDA): 251 SIZES = [100 * SCALE] 252 253 254class BaseTestDA_2D(BaseTestDA): 255 SIZES = [9 * SCALE, 11 * SCALE] 256 257 258class BaseTestDA_3D(BaseTestDA): 259 SIZES = [6 * SCALE, 7 * SCALE, 8 * SCALE] 260 261 262# -------------------------------------------------------------------- 263 264 265class TestDA_1D(BaseTestDA_1D, unittest.TestCase): 266 pass 267 268 269class TestDA_1D_W0(TestDA_1D): 270 SWIDTH = 0 271 272 273class TestDA_1D_W2(TestDA_1D): 274 SWIDTH = 2 275 276 277class TestDA_2D(BaseTestDA_2D, unittest.TestCase): 278 pass 279 280 281class TestDA_2D_W0(TestDA_2D): 282 SWIDTH = 0 283 284 285class TestDA_2D_W0_N2(TestDA_2D): 286 DOF = 2 287 SWIDTH = 0 288 289 290class TestDA_2D_W2(TestDA_2D): 291 SWIDTH = 2 292 293 294class TestDA_2D_W2_N2(TestDA_2D): 295 DOF = 2 296 SWIDTH = 2 297 298 299class TestDA_2D_PXY(TestDA_2D): 300 SIZES = [13 * SCALE, 17 * SCALE] 301 DOF = 2 302 SWIDTH = 5 303 BOUNDARY = (PERIODIC,) * 2 304 305 306class TestDA_2D_GXY(TestDA_2D): 307 SIZES = [13 * SCALE, 17 * SCALE] 308 DOF = 2 309 SWIDTH = 5 310 BOUNDARY = (GHOSTED,) * 2 311 312 313class TestDA_2D_TXY(TestDA_2D): 314 SIZES = [13 * SCALE, 17 * SCALE] 315 DOF = 2 316 SWIDTH = 5 317 BOUNDARY = (TWIST,) * 2 318 319 320class TestDA_3D(BaseTestDA_3D, unittest.TestCase): 321 pass 322 323 324class TestDA_3D_W0(TestDA_3D): 325 SWIDTH = 0 326 327 328class TestDA_3D_W0_N2(TestDA_3D): 329 DOF = 2 330 SWIDTH = 0 331 332 333class TestDA_3D_W2(TestDA_3D): 334 SWIDTH = 2 335 336 337class TestDA_3D_W2_N2(TestDA_3D): 338 DOF = 2 339 SWIDTH = 2 340 341 342class TestDA_3D_PXYZ(TestDA_3D): 343 SIZES = [11 * SCALE, 13 * SCALE, 17 * SCALE] 344 DOF = 2 345 SWIDTH = 3 346 BOUNDARY = (PERIODIC,) * 3 347 348 349class TestDA_3D_GXYZ(TestDA_3D): 350 SIZES = [11 * SCALE, 13 * SCALE, 17 * SCALE] 351 DOF = 2 352 SWIDTH = 3 353 BOUNDARY = (GHOSTED,) * 3 354 355 356class TestDA_3D_TXYZ(TestDA_3D): 357 SIZES = [11 * SCALE, 13 * SCALE, 17 * SCALE] 358 DOF = 2 359 SWIDTH = 3 360 BOUNDARY = (TWIST,) * 3 361 362 363# -------------------------------------------------------------------- 364 365DIM = ( 366 1, 367 2, 368 3, 369) 370DOF = ( 371 None, 372 1, 373 2, 374 3, 375 4, 376 5, 377) 378BOUNDARY_TYPE = ( 379 None, 380 'none', 381 (0,) * 3, 382 0, 383 'ghosted', 384 (GHOSTED,) * 3, 385 GHOSTED, 386 'periodic', 387 (PERIODIC,) * 3, 388 PERIODIC, 389 'twist', 390 (TWIST,) * 3, 391 TWIST, 392) 393STENCIL_TYPE = (None, 'star', 'box') 394STENCIL_WIDTH = (None, 0, 1, 2, 3) 395 396 397DIM = (1, 2, 3) 398DOF = (None, 2, 5) 399BOUNDARY_TYPE = (None, 'none', 'periodic', 'ghosted', 'twist') 400STENCIL_TYPE = (None, 'box') 401STENCIL_WIDTH = (None, 1, 2) 402 403 404class TestDACreate(unittest.TestCase): 405 pass 406 407 408counter = 0 409for dim in DIM: 410 for dof in DOF: 411 for boundary in BOUNDARY_TYPE: 412 if isinstance(boundary, tuple): 413 boundary = boundary[:dim] 414 for stencil in STENCIL_TYPE: 415 for width in STENCIL_WIDTH: 416 kargs = { 417 'sizes': [8 * SCALE] * dim, 418 'dim': dim, 419 'dof': dof, 420 'boundary_type': boundary, 421 'stencil_type': stencil, 422 'stencil_width': width, 423 } 424 425 def testCreate(self, kargs=kargs): 426 kargs = dict(kargs) 427 da = PETSc.DMDA().create(**kargs) 428 da.destroy() 429 430 setattr(TestDACreate, 'testCreate%04d' % counter, testCreate) 431 del testCreate, kargs 432 counter += 1 433del counter, dim, dof, boundary, stencil, width 434 435 436class TestDADuplicate(unittest.TestCase): 437 pass 438 439 440counter = 0 441for dim in DIM: 442 for dof in DOF: 443 for boundary in BOUNDARY_TYPE: 444 if isinstance(boundary, tuple): 445 boundary = boundary[:dim] 446 for stencil in STENCIL_TYPE: 447 for width in STENCIL_WIDTH: 448 kargs = { 449 'dim': dim, 450 'dof': dof, 451 'boundary_type': boundary, 452 'stencil_type': stencil, 453 'stencil_width': width, 454 } 455 456 def testDuplicate(self, kargs=kargs): 457 kargs = dict(kargs) 458 dim = kargs.pop('dim') 459 dof = kargs['dof'] 460 boundary = kargs['boundary_type'] 461 stencil = kargs['stencil_type'] 462 width = kargs['stencil_width'] 463 da = PETSc.DMDA().create([8 * SCALE] * dim) 464 newda = da.duplicate(**kargs) 465 self.assertEqual(newda.dim, da.dim) 466 self.assertEqual(newda.sizes, da.sizes) 467 self.assertEqual(newda.proc_sizes, da.proc_sizes) 468 self.assertEqual(newda.ranges, da.ranges) 469 self.assertEqual(newda.corners, da.corners) 470 if ( 471 newda.boundary_type == da.boundary_type 472 and newda.stencil_width == da.stencil_width 473 ): 474 self.assertEqual(newda.ghost_ranges, da.ghost_ranges) 475 self.assertEqual(newda.ghost_corners, da.ghost_corners) 476 if dof is None: 477 dof = da.dof 478 if boundary is None: 479 boundary = da.boundary_type 480 elif boundary == 'none': 481 boundary = (0,) * dim 482 elif boundary == 'mirror': 483 boundary = (MIRROR,) * dim 484 elif boundary == 'ghosted': 485 boundary = (GHOSTED,) * dim 486 elif boundary == 'periodic': 487 boundary = (PERIODIC,) * dim 488 elif boundary == 'twist': 489 boundary = (TWIST,) * dim 490 elif isinstance(boundary, int): 491 boundary = (boundary,) * dim 492 if stencil is None: 493 stencil = da.stencil[0] 494 if width is None: 495 width = da.stencil_width 496 self.assertEqual(newda.dof, dof) 497 self.assertEqual(newda.boundary_type, boundary) 498 if dim == 1: 499 self.assertEqual(newda.stencil, (stencil, width)) 500 newda.destroy() 501 da.destroy() 502 503 setattr( 504 TestDADuplicate, 'testDuplicate%04d' % counter, testDuplicate 505 ) 506 del testDuplicate, kargs 507 counter += 1 508del counter, dim, dof, boundary, stencil, width 509 510# -------------------------------------------------------------------- 511 512if PETSc.COMM_WORLD.getSize() > 1: 513 del TestDA_1D_W0 514 del TestDA_2D_W0, TestDA_2D_W0_N2 515 del TestDA_3D_W0, TestDA_3D_W0_N2 516 517# -------------------------------------------------------------------- 518 519if __name__ == '__main__': 520 unittest.main() 521 522# -------------------------------------------------------------------- 523