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 PETSc.garbage_cleanup() 27 28 def testGetInfo(self): 29 dim = self.da.getDim() 30 dof = self.da.getDof() 31 sizes = self.da.getSizes() 32 psizes = self.da.getProcSizes() 33 boundary = self.da.getBoundaryType() 34 stencil_type = self.da.getStencilType() 35 stencil_width = self.da.getStencilWidth() 36 self.assertEqual(dim, len(self.SIZES)) 37 self.assertEqual(dof, self.DOF) 38 self.assertEqual(sizes, tuple(self.SIZES)) 39 self.assertEqual(boundary, self.BOUNDARY or (0,)*dim) 40 self.assertEqual(stencil_type, self.STENCIL) 41 self.assertEqual(stencil_width, self.SWIDTH) 42 43 def testRangesCorners(self): 44 dim = self.da.getDim() 45 ranges = self.da.getRanges() 46 starts, lsizes = self.da.getCorners() 47 self.assertEqual(dim, len(ranges)) 48 self.assertEqual(dim, len(starts)) 49 self.assertEqual(dim, len(lsizes)) 50 for i in range(dim): 51 s, e = ranges[i] 52 self.assertEqual(s, starts[i]) 53 self.assertEqual(e-s, lsizes[i]) 54 55 def testGhostRangesCorners(self): 56 dim = self.da.getDim() 57 ranges = self.da.getGhostRanges() 58 starts, lsizes = self.da.getGhostCorners() 59 self.assertEqual(dim, len(ranges)) 60 self.assertEqual(dim, len(starts)) 61 self.assertEqual(dim, len(lsizes)) 62 for i in range(dim): 63 s, e = ranges[i] 64 self.assertEqual(s, starts[i]) 65 self.assertEqual(e-s, lsizes[i]) 66 67 def testOwnershipRanges(self): 68 dim = self.da.getDim() 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 s = self.da.stencil_width 120 btype = self.da.boundary_type 121 psize = self.da.proc_sizes 122 for b, p in zip(btype, psize): 123 if b != NONE and p == 1: 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 s = self.da.stencil_width 140 btype = self.da.boundary_type 141 psize = self.da.proc_sizes 142 for b, p in zip(btype, psize): 143 if b != NONE and p == 1: 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 ao = self.da.getAO() 153 lgmap = self.da.getLGMap() 154 l2g, g2l = 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 da = self.da 172 cda = self.da.coarsen() 173 self.assertEqual(da.getDim(), cda.getDim()) 174 self.assertEqual(da.getDof(), cda.getDof()) 175 if da.dim != 1: 176 self.assertEqual(da.getStencilType(), cda.getStencilType()) 177 self.assertEqual(da.getStencilWidth(), cda.getStencilWidth()) 178 rda = cda.refine() 179 for n1, n2 in zip(self.da.getSizes(), rda.getSizes()): 180 self.assertTrue(abs(n1-n2)<=1) 181 182 def testRefineHierarchy(self): 183 levels = self.da.refineHierarchy(2) 184 self.assertTrue(isinstance(levels, list)) 185 self.assertEqual(len(levels), 2) 186 for item in levels: 187 self.assertTrue(isinstance(item, PETSc.DM)) 188 189 def testCoarsenHierarchy(self): 190 levels = self.da.coarsenHierarchy(2) 191 self.assertTrue(isinstance(levels, list)) 192 self.assertEqual(len(levels), 2) 193 for item in levels: 194 self.assertTrue(isinstance(item, PETSc.DM)) 195 196 def testCreateInterpolation(self): 197 da = self.da 198 if da.dim == 1: return 199 rda = da.refine() 200 mat, vec = da.createInterpolation(rda) 201 202 def testCreateInjection(self): 203 da = self.da 204 if da.dim == 1: return 205 rda = da.refine() 206 scatter = da.createInjection(rda) 207 208 def testzeroRowsColumnsStencil(self): 209 da = self.da 210 A = da.createMatrix() 211 x = da.createGlobalVector() 212 x.set(2.0) 213 A.setDiagonal(x) 214 diag1 = x.duplicate() 215 A.getDiagonal(diag1) 216 if self.SIZES != 2: #only coded test for 2D case 217 return 218 istart,iend, jstart, jend = da.getRanges() 219 self.assertTrue(x.equal(diag1)) 220 zeroidx = [] 221 for i in range(istart,iend): 222 for j in range(jstart,jend): 223 row = PETSc.Mat.Stencil() 224 row.index = (i,j) 225 zeroidx = zeroidx + [row] 226 diag2 = x.duplicate() 227 diag2.set(1.0) 228 A.zeroRowsColumnsStencil(zeroidx, 1.0, x, diag2) 229 ans = x.duplicate() 230 ans.set(2.0) 231 self.assertTrue(ans.equal(diag2)) 232 233 234MIRROR = PETSc.DMDA.BoundaryType.MIRROR 235GHOSTED = PETSc.DMDA.BoundaryType.GHOSTED 236PERIODIC = PETSc.DMDA.BoundaryType.PERIODIC 237TWIST = PETSc.DMDA.BoundaryType.TWIST 238 239SCALE = 4 240 241class BaseTestDA_1D(BaseTestDA): 242 SIZES = [100*SCALE] 243 244class BaseTestDA_2D(BaseTestDA): 245 SIZES = [9*SCALE,11*SCALE] 246 247class BaseTestDA_3D(BaseTestDA): 248 SIZES = [6*SCALE,7*SCALE,8*SCALE] 249 250# -------------------------------------------------------------------- 251 252class TestDA_1D(BaseTestDA_1D, unittest.TestCase): 253 pass 254class TestDA_1D_W0(TestDA_1D): 255 SWIDTH = 0 256class TestDA_1D_W2(TestDA_1D): 257 SWIDTH = 2 258 259class TestDA_2D(BaseTestDA_2D, unittest.TestCase): 260 pass 261class TestDA_2D_W0(TestDA_2D): 262 SWIDTH = 0 263class TestDA_2D_W0_N2(TestDA_2D): 264 DOF = 2 265 SWIDTH = 0 266class TestDA_2D_W2(TestDA_2D): 267 SWIDTH = 2 268class TestDA_2D_W2_N2(TestDA_2D): 269 DOF = 2 270 SWIDTH = 2 271class TestDA_2D_PXY(TestDA_2D): 272 SIZES = [13*SCALE,17*SCALE] 273 DOF = 2 274 SWIDTH = 5 275 BOUNDARY = (PERIODIC,)*2 276class TestDA_2D_GXY(TestDA_2D): 277 SIZES = [13*SCALE,17*SCALE] 278 DOF = 2 279 SWIDTH = 5 280 BOUNDARY = (GHOSTED,)*2 281class TestDA_2D_TXY(TestDA_2D): 282 SIZES = [13*SCALE,17*SCALE] 283 DOF = 2 284 SWIDTH = 5 285 BOUNDARY = (TWIST,)*2 286 287class TestDA_3D(BaseTestDA_3D, unittest.TestCase): 288 pass 289class TestDA_3D_W0(TestDA_3D): 290 SWIDTH = 0 291class TestDA_3D_W0_N2(TestDA_3D): 292 DOF = 2 293 SWIDTH = 0 294class TestDA_3D_W2(TestDA_3D): 295 SWIDTH = 2 296class TestDA_3D_W2_N2(TestDA_3D): 297 DOF = 2 298 SWIDTH = 2 299class TestDA_3D_PXYZ(TestDA_3D): 300 SIZES = [11*SCALE,13*SCALE,17*SCALE] 301 DOF = 2 302 SWIDTH = 3 303 BOUNDARY = (PERIODIC,)*3 304class TestDA_3D_GXYZ(TestDA_3D): 305 SIZES = [11*SCALE,13*SCALE,17*SCALE] 306 DOF = 2 307 SWIDTH = 3 308 BOUNDARY = (GHOSTED,)*3 309class TestDA_3D_TXYZ(TestDA_3D): 310 SIZES = [11*SCALE,13*SCALE,17*SCALE] 311 DOF = 2 312 SWIDTH = 3 313 BOUNDARY = (TWIST,)*3 314 315# -------------------------------------------------------------------- 316 317DIM = (1,2,3,) 318DOF = (None,1,2,3,4,5,) 319BOUNDARY_TYPE = ( 320 None, 321 "none", (0,)*3, 0, 322 "ghosted", (GHOSTED,)*3, GHOSTED, 323 "periodic", (PERIODIC,)*3, PERIODIC, 324 "twist", (TWIST,)*3, TWIST, 325 ) 326STENCIL_TYPE = (None,"star","box") 327STENCIL_WIDTH = (None,0,1,2,3) 328 329 330DIM = (1,2,3) 331DOF = (None,2,5) 332BOUNDARY_TYPE = (None,"none","periodic","ghosted","twist") 333STENCIL_TYPE = (None,"box") 334STENCIL_WIDTH = (None,1,2) 335 336class TestDACreate(unittest.TestCase): 337 pass 338counter = 0 339for dim in DIM: 340 for dof in DOF: 341 for boundary in BOUNDARY_TYPE: 342 if isinstance(boundary, tuple): 343 boundary = boundary[:dim] 344 for stencil in STENCIL_TYPE: 345 for width in STENCIL_WIDTH: 346 kargs = dict(sizes=[8*SCALE]*dim, 347 dim=dim, dof=dof, 348 boundary_type=boundary, 349 stencil_type=stencil, 350 stencil_width=width) 351 def testCreate(self, kargs=kargs): 352 kargs = dict(kargs) 353 da = PETSc.DMDA().create(**kargs) 354 da.destroy() 355 setattr(TestDACreate, 356 "testCreate%04d"%counter, 357 testCreate) 358 del testCreate, kargs 359 counter += 1 360del counter, dim, dof, boundary, stencil, width 361 362class TestDADuplicate(unittest.TestCase): 363 pass 364counter = 0 365for dim in DIM: 366 for dof in DOF: 367 for boundary in BOUNDARY_TYPE: 368 if isinstance(boundary, tuple): 369 boundary = boundary[:dim] 370 for stencil in STENCIL_TYPE: 371 for width in STENCIL_WIDTH: 372 kargs = dict(dim=dim, dof=dof, 373 boundary_type=boundary, 374 stencil_type=stencil, 375 stencil_width=width) 376 def testDuplicate(self, kargs=kargs): 377 kargs = dict(kargs) 378 dim = kargs.pop('dim') 379 dof = kargs['dof'] 380 boundary = kargs['boundary_type'] 381 stencil = kargs['stencil_type'] 382 width = kargs['stencil_width'] 383 da = PETSc.DMDA().create([8*SCALE]*dim) 384 newda = da.duplicate(**kargs) 385 self.assertEqual(newda.dim, da.dim) 386 self.assertEqual(newda.sizes, da.sizes) 387 self.assertEqual(newda.proc_sizes, 388 da.proc_sizes) 389 self.assertEqual(newda.ranges, da.ranges) 390 self.assertEqual(newda.corners, da.corners) 391 if (newda.boundary_type == da.boundary_type 392 and 393 newda.stencil_width == da.stencil_width): 394 self.assertEqual(newda.ghost_ranges, 395 da.ghost_ranges) 396 self.assertEqual(newda.ghost_corners, 397 da.ghost_corners) 398 if dof is None: 399 dof = da.dof 400 if boundary is None: 401 boundary = da.boundary_type 402 elif boundary == "none": 403 boundary = (0,) * dim 404 elif boundary == "mirror": 405 boundary = (MIRROR,) * dim 406 elif boundary == "ghosted": 407 boundary = (GHOSTED,) * dim 408 elif boundary == "periodic": 409 boundary = (PERIODIC,) * dim 410 elif boundary == "twist": 411 boundary = (TWIST,) * dim 412 elif isinstance(boundary, int): 413 boundary = (boundary,) * dim 414 if stencil is None: 415 stencil = da.stencil[0] 416 if width is None: 417 width = da.stencil_width 418 self.assertEqual(newda.dof, dof) 419 self.assertEqual(newda.boundary_type, 420 boundary) 421 if dim == 1: 422 self.assertEqual(newda.stencil, 423 (stencil, width)) 424 newda.destroy() 425 da.destroy() 426 setattr(TestDADuplicate, 427 "testDuplicate%04d"%counter, 428 testDuplicate) 429 del testDuplicate, kargs 430 counter += 1 431del counter, dim, dof, boundary, stencil, width 432 433# -------------------------------------------------------------------- 434 435if PETSc.COMM_WORLD.getSize() > 1: 436 del TestDA_1D_W0 437 del TestDA_2D_W0, TestDA_2D_W0_N2 438 del TestDA_3D_W0, TestDA_3D_W0_N2 439 440# -------------------------------------------------------------------- 441 442if __name__ == '__main__': 443 unittest.main() 444 445# -------------------------------------------------------------------- 446