1from petsc4py import PETSc 2import unittest 3 4# -------------------------------------------------------------------- 5 6class BaseTestDA(object): 7 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(dim=len(self.SIZES), 17 dof=self.DOF, 18 sizes=self.SIZES, 19 boundary_type=self.BOUNDARY, 20 stencil_type=self.STENCIL, 21 stencil_width=self.SWIDTH, 22 comm=self.COMM) 23 24 def tearDown(self): 25 self.da = None 26 27 def testGetInfo(self): 28 dim = self.da.getDim() 29 dof = self.da.getDof() 30 sizes = self.da.getSizes() 31 psizes = self.da.getProcSizes() 32 boundary = self.da.getBoundaryType() 33 stencil_type = self.da.getStencilType() 34 stencil_width = self.da.getStencilWidth() 35 self.assertEqual(dim, len(self.SIZES)) 36 self.assertEqual(dof, self.DOF) 37 self.assertEqual(sizes, tuple(self.SIZES)) 38 self.assertEqual(boundary, self.BOUNDARY or (0,)*dim) 39 self.assertEqual(stencil_type, self.STENCIL) 40 self.assertEqual(stencil_width, self.SWIDTH) 41 42 def testRangesCorners(self): 43 dim = self.da.getDim() 44 ranges = self.da.getRanges() 45 starts, lsizes = self.da.getCorners() 46 self.assertEqual(dim, len(ranges)) 47 self.assertEqual(dim, len(starts)) 48 self.assertEqual(dim, len(lsizes)) 49 for i in range(dim): 50 s, e = ranges[i] 51 self.assertEqual(s, starts[i]) 52 self.assertEqual(e-s, lsizes[i]) 53 54 def testGhostRangesCorners(self): 55 dim = self.da.getDim() 56 ranges = self.da.getGhostRanges() 57 starts, lsizes = self.da.getGhostCorners() 58 self.assertEqual(dim, len(ranges)) 59 self.assertEqual(dim, len(starts)) 60 self.assertEqual(dim, len(lsizes)) 61 for i in range(dim): 62 s, e = ranges[i] 63 self.assertEqual(s, starts[i]) 64 self.assertEqual(e-s, lsizes[i]) 65 66 def testOwnershipRanges(self): 67 dim = self.da.getDim() 68 ownership_ranges = self.da.getOwnershipRanges() 69 procsizes = self.da.getProcSizes() 70 self.assertEqual(len(procsizes), len(ownership_ranges)) 71 for i,m in enumerate(procsizes): 72 self.assertEqual(m, len(ownership_ranges[i])) 73 74 def testFieldName(self): 75 for i in range(self.da.getDof()): 76 self.da.setFieldName(i, "field%d" % i) 77 for i in range(self.da.getDof()): 78 name = self.da.getFieldName(i) 79 self.assertEqual(name, "field%d" % i) 80 81 def testCoordinates(self): 82 self.da.setUniformCoordinates(0,1,0,1,0,1) 83 # 84 c = self.da.getCoordinates() 85 self.da.setCoordinates(c) 86 c.destroy() 87 cda = self.da.getCoordinateDM() 88 cda.destroy() 89 # 90 c = self.da.getCoordinates() 91 self.da.setCoordinates(c) 92 c.destroy() 93 gc = self.da.getCoordinatesLocal() 94 gc.destroy() 95 96 def testCreateVecMat(self): 97 vn = self.da.createNaturalVec() 98 vg = self.da.createGlobalVec() 99 vl = self.da.createLocalVec() 100 mat = self.da.createMat() 101 self.assertTrue(mat.getType() in ('aij', 'seqaij', 'mpiaij')) 102 vn.set(1.0) 103 self.da.naturalToGlobal(vn,vg) 104 self.assertEqual(vg.max()[1], 1.0) 105 self.assertEqual(vg.min()[1], 1.0) 106 self.da.globalToLocal(vg,vl) 107 self.assertEqual(vl.max()[1], 1.0) 108 self.assertTrue (vl.min()[1] in (1.0, 0.0)) 109 vn.set(0.0) 110 self.da.globalToNatural(vg,vn) 111 self.assertEqual(vn.max()[1], 1.0) 112 self.assertEqual(vn.min()[1], 1.0) 113 vl2 = self.da.createLocalVec() 114 self.da.localToLocal(vl,vl2) 115 self.assertEqual(vl2.max()[1], 1.0) 116 self.assertTrue (vl2.min()[1] in (1.0, 0.0)) 117 NONE = PETSc.DM.BoundaryType.NONE 118 s = self.da.stencil_width 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: return 123 vg2 = self.da.createGlobalVec() 124 self.da.localToGlobal(vl2,vg2) 125 126 def testGetVec(self): 127 vg = self.da.getGlobalVec() 128 vl = self.da.getLocalVec() 129 try: 130 vg.set(1.0) 131 self.assertEqual(vg.max()[1], 1.0) 132 self.assertEqual(vg.min()[1], 1.0) 133 self.da.globalToLocal(vg,vl) 134 self.assertEqual(vl.max()[1], 1.0) 135 self.assertTrue (vl.min()[1] in (1.0, 0.0)) 136 vl.set(2.0) 137 NONE = PETSc.DM.BoundaryType.NONE 138 s = self.da.stencil_width 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: return 143 self.da.localToGlobal(vl,vg) 144 self.assertEqual(vg.max()[1], 2.0) 145 self.assertTrue (vg.min()[1] in (2.0, 0.0)) 146 finally: 147 self.da.restoreGlobalVec(vg) 148 self.da.restoreLocalVec(vl) 149 150 def testGetOther(self): 151 ao = self.da.getAO() 152 lgmap = self.da.getLGMap() 153 l2g, g2l = self.da.getScatter() 154 155 def testRefineCoarsen(self): 156 da = self.da 157 rda = da.refine() 158 self.assertEqual(da.getDim(), rda.getDim()) 159 self.assertEqual(da.getDof(), rda.getDof()) 160 if da.dim != 1: 161 self.assertEqual(da.getStencilType(), rda.getStencilType()) 162 self.assertEqual(da.getStencilWidth(), rda.getStencilWidth()) 163 cda = rda.coarsen() 164 self.assertEqual(rda.getDim(), cda.getDim()) 165 self.assertEqual(rda.getDof(), cda.getDof()) 166 for n1, n2 in zip(self.da.getSizes(), cda.getSizes()): 167 self.assertTrue(abs(n1-n2)<=1) 168 169 def testCoarsenRefine(self): 170 da = self.da 171 cda = self.da.coarsen() 172 self.assertEqual(da.getDim(), cda.getDim()) 173 self.assertEqual(da.getDof(), cda.getDof()) 174 if da.dim != 1: 175 self.assertEqual(da.getStencilType(), cda.getStencilType()) 176 self.assertEqual(da.getStencilWidth(), cda.getStencilWidth()) 177 rda = cda.refine() 178 for n1, n2 in zip(self.da.getSizes(), rda.getSizes()): 179 self.assertTrue(abs(n1-n2)<=1) 180 181 def testRefineHierarchy(self): 182 levels = self.da.refineHierarchy(2) 183 self.assertTrue(isinstance(levels, list)) 184 self.assertEqual(len(levels), 2) 185 for item in levels: 186 self.assertTrue(isinstance(item, PETSc.DM)) 187 188 def testCoarsenHierarchy(self): 189 levels = self.da.coarsenHierarchy(2) 190 self.assertTrue(isinstance(levels, list)) 191 self.assertEqual(len(levels), 2) 192 for item in levels: 193 self.assertTrue(isinstance(item, PETSc.DM)) 194 195 def testCreateInterpolation(self): 196 da = self.da 197 if da.dim == 1: return 198 rda = da.refine() 199 mat, vec = da.createInterpolation(rda) 200 201 def testCreateInjection(self): 202 da = self.da 203 if da.dim == 1: return 204 rda = da.refine() 205 scatter = da.createInjection(rda) 206 207 def testzeroRowsColumnsStencil(self): 208 da = self.da 209 A = da.createMatrix() 210 x = da.createGlobalVector() 211 x.set(2.0) 212 A.setDiagonal(x) 213 diag1 = x.duplicate() 214 A.getDiagonal(diag1) 215 if self.SIZES != 2: #only coded test for 2D case 216 return 217 istart,iend, jstart, jend = da.getRanges() 218 self.assertTrue(x.equal(diag1)) 219 zeroidx = [] 220 for i in range(istart,iend): 221 for j in range(jstart,jend): 222 row = PETSc.Mat.Stencil() 223 row.index = (i,j) 224 zeroidx = zeroidx + [row] 225 diag2 = x.duplicate() 226 diag2.set(1.0) 227 A.zeroRowsColumnsStencil(zeroidx, 1.0, x, diag2) 228 ans = x.duplicate() 229 ans.set(2.0) 230 self.assertTrue(ans.equal(diag2)) 231 232 233MIRROR = PETSc.DMDA.BoundaryType.MIRROR 234GHOSTED = PETSc.DMDA.BoundaryType.GHOSTED 235PERIODIC = PETSc.DMDA.BoundaryType.PERIODIC 236TWIST = PETSc.DMDA.BoundaryType.TWIST 237 238SCALE = 4 239 240class BaseTestDA_1D(BaseTestDA): 241 SIZES = [100*SCALE] 242 243class BaseTestDA_2D(BaseTestDA): 244 SIZES = [9*SCALE,11*SCALE] 245 246class BaseTestDA_3D(BaseTestDA): 247 SIZES = [6*SCALE,7*SCALE,8*SCALE] 248 249# -------------------------------------------------------------------- 250 251class TestDA_1D(BaseTestDA_1D, unittest.TestCase): 252 pass 253class TestDA_1D_W0(TestDA_1D): 254 SWIDTH = 0 255class TestDA_1D_W2(TestDA_1D): 256 SWIDTH = 2 257 258class TestDA_2D(BaseTestDA_2D, unittest.TestCase): 259 pass 260class TestDA_2D_W0(TestDA_2D): 261 SWIDTH = 0 262class TestDA_2D_W0_N2(TestDA_2D): 263 DOF = 2 264 SWIDTH = 0 265class TestDA_2D_W2(TestDA_2D): 266 SWIDTH = 2 267class TestDA_2D_W2_N2(TestDA_2D): 268 DOF = 2 269 SWIDTH = 2 270class TestDA_2D_PXY(TestDA_2D): 271 SIZES = [13*SCALE,17*SCALE] 272 DOF = 2 273 SWIDTH = 5 274 BOUNDARY = (PERIODIC,)*2 275class TestDA_2D_GXY(TestDA_2D): 276 SIZES = [13*SCALE,17*SCALE] 277 DOF = 2 278 SWIDTH = 5 279 BOUNDARY = (GHOSTED,)*2 280class TestDA_2D_TXY(TestDA_2D): 281 SIZES = [13*SCALE,17*SCALE] 282 DOF = 2 283 SWIDTH = 5 284 BOUNDARY = (TWIST,)*2 285 286class TestDA_3D(BaseTestDA_3D, unittest.TestCase): 287 pass 288class TestDA_3D_W0(TestDA_3D): 289 SWIDTH = 0 290class TestDA_3D_W0_N2(TestDA_3D): 291 DOF = 2 292 SWIDTH = 0 293class TestDA_3D_W2(TestDA_3D): 294 SWIDTH = 2 295class TestDA_3D_W2_N2(TestDA_3D): 296 DOF = 2 297 SWIDTH = 2 298class TestDA_3D_PXYZ(TestDA_3D): 299 SIZES = [11*SCALE,13*SCALE,17*SCALE] 300 DOF = 2 301 SWIDTH = 3 302 BOUNDARY = (PERIODIC,)*3 303class TestDA_3D_GXYZ(TestDA_3D): 304 SIZES = [11*SCALE,13*SCALE,17*SCALE] 305 DOF = 2 306 SWIDTH = 3 307 BOUNDARY = (GHOSTED,)*3 308class TestDA_3D_TXYZ(TestDA_3D): 309 SIZES = [11*SCALE,13*SCALE,17*SCALE] 310 DOF = 2 311 SWIDTH = 3 312 BOUNDARY = (TWIST,)*3 313 314# -------------------------------------------------------------------- 315 316DIM = (1,2,3,) 317DOF = (None,1,2,3,4,5,) 318BOUNDARY_TYPE = ( 319 None, 320 "none", (0,)*3, 0, 321 "ghosted", (GHOSTED,)*3, GHOSTED, 322 "periodic", (PERIODIC,)*3, PERIODIC, 323 "twist", (TWIST,)*3, TWIST, 324 ) 325STENCIL_TYPE = (None,"star","box") 326STENCIL_WIDTH = (None,0,1,2,3) 327 328 329DIM = (1,2,3) 330DOF = (None,2,5) 331BOUNDARY_TYPE = (None,"none","periodic","ghosted","twist") 332STENCIL_TYPE = (None,"box") 333STENCIL_WIDTH = (None,1,2) 334 335class TestDACreate(unittest.TestCase): 336 pass 337counter = 0 338for dim in DIM: 339 for dof in DOF: 340 for boundary in BOUNDARY_TYPE: 341 if isinstance(boundary, tuple): 342 boundary = boundary[:dim] 343 for stencil in STENCIL_TYPE: 344 for width in STENCIL_WIDTH: 345 kargs = dict(sizes=[8*SCALE]*dim, 346 dim=dim, dof=dof, 347 boundary_type=boundary, 348 stencil_type=stencil, 349 stencil_width=width) 350 def testCreate(self, kargs=kargs): 351 kargs = dict(kargs) 352 da = PETSc.DMDA().create(**kargs) 353 da.destroy() 354 setattr(TestDACreate, 355 "testCreate%04d"%counter, 356 testCreate) 357 del testCreate, kargs 358 counter += 1 359del counter, dim, dof, boundary, stencil, width 360 361class TestDADuplicate(unittest.TestCase): 362 pass 363counter = 0 364for dim in DIM: 365 for dof in DOF: 366 for boundary in BOUNDARY_TYPE: 367 if isinstance(boundary, tuple): 368 boundary = boundary[:dim] 369 for stencil in STENCIL_TYPE: 370 for width in STENCIL_WIDTH: 371 kargs = dict(dim=dim, dof=dof, 372 boundary_type=boundary, 373 stencil_type=stencil, 374 stencil_width=width) 375 def testDuplicate(self, kargs=kargs): 376 kargs = dict(kargs) 377 dim = kargs.pop('dim') 378 dof = kargs['dof'] 379 boundary = kargs['boundary_type'] 380 stencil = kargs['stencil_type'] 381 width = kargs['stencil_width'] 382 da = PETSc.DMDA().create([8*SCALE]*dim) 383 newda = da.duplicate(**kargs) 384 self.assertEqual(newda.dim, da.dim) 385 self.assertEqual(newda.sizes, da.sizes) 386 self.assertEqual(newda.proc_sizes, 387 da.proc_sizes) 388 self.assertEqual(newda.ranges, da.ranges) 389 self.assertEqual(newda.corners, da.corners) 390 if (newda.boundary_type == da.boundary_type 391 and 392 newda.stencil_width == da.stencil_width): 393 self.assertEqual(newda.ghost_ranges, 394 da.ghost_ranges) 395 self.assertEqual(newda.ghost_corners, 396 da.ghost_corners) 397 if dof is None: 398 dof = da.dof 399 if boundary is None: 400 boundary = da.boundary_type 401 elif boundary == "none": 402 boundary = (0,) * dim 403 elif boundary == "mirror": 404 boundary = (MIRROR,) * dim 405 elif boundary == "ghosted": 406 boundary = (GHOSTED,) * dim 407 elif boundary == "periodic": 408 boundary = (PERIODIC,) * dim 409 elif boundary == "twist": 410 boundary = (TWIST,) * dim 411 elif isinstance(boundary, int): 412 boundary = (boundary,) * dim 413 if stencil is None: 414 stencil = da.stencil[0] 415 if width is None: 416 width = da.stencil_width 417 self.assertEqual(newda.dof, dof) 418 self.assertEqual(newda.boundary_type, 419 boundary) 420 if dim == 1: 421 self.assertEqual(newda.stencil, 422 (stencil, width)) 423 newda.destroy() 424 da.destroy() 425 setattr(TestDADuplicate, 426 "testDuplicate%04d"%counter, 427 testDuplicate) 428 del testDuplicate, kargs 429 counter += 1 430del counter, dim, dof, boundary, stencil, width 431 432# -------------------------------------------------------------------- 433 434if PETSc.COMM_WORLD.getSize() > 1: 435 del TestDA_1D_W0 436 del TestDA_2D_W0, TestDA_2D_W0_N2 437 del TestDA_3D_W0, TestDA_3D_W0_N2 438 439# -------------------------------------------------------------------- 440 441if __name__ == '__main__': 442 unittest.main() 443 444# -------------------------------------------------------------------- 445