def test_is_compatible_degenerate_block_cyclic(self): size = 19937 gdd_block_cyclic = ({ 'dist_type': 'c', 'proc_grid_size': 1, 'block_size': 7, 'size': size, }, ) gdd_block = ({ 'dist_type': 'b', 'proc_grid_size': 1, 'bounds': [0, size], }, ) gdd_cyclic = ({ 'dist_type': 'c', 'proc_grid_size': 1, 'size': size, }, ) dist_block_cyclic = Distribution.from_global_dim_data( self.context, gdd_block_cyclic) dist_block = Distribution.from_global_dim_data(self.context, gdd_block) dist_cyclic = Distribution.from_global_dim_data( self.context, gdd_cyclic) self.assertTrue(dist_block_cyclic.is_compatible(dist_block)) self.assertTrue(dist_block_cyclic.is_compatible(dist_cyclic)) self.assertTrue(dist_block.is_compatible(dist_block_cyclic)) self.assertTrue(dist_cyclic.is_compatible(dist_block_cyclic))
def test_block_redist_2D_one_to_many(self): source_dist = Distribution(self.context, (9, 9), ('b', 'b'), (1, 1), targets=[2]) dest_dist = Distribution(self.context, (9, 9), ('b', 'b'), (2, 2), targets=range(4)) plan = source_dist.get_redist_plan(dest_dist) expected = [ { 'source_rank': 2, 'dest_rank': 0, 'indices': [(0, 5, 1), (0, 5, 1)] }, { 'source_rank': 2, 'dest_rank': 1, 'indices': [(0, 5, 1), (5, 9, 1)] }, { 'source_rank': 2, 'dest_rank': 2, 'indices': [(5, 9, 1), (0, 5, 1)] }, { 'source_rank': 2, 'dest_rank': 3, 'indices': [(5, 9, 1), (5, 9, 1)] }, ] for p, e in zip(plan, expected): self.assertEqual(p, e)
def test_2D_bn(self): nrows, ncols = 31, 53 cm = Distribution(self.context, (nrows, ncols), {0: 'b'}, (4, 1)) chunksize = (nrows // 4) + 1 for _ in range(100): r, c = randrange(nrows), randrange(ncols) rank = r // chunksize self.assertSequenceEqual(cm.owning_ranks((r, c)), [rank])
def test_no_empty_local_arrays_3_targets(self): for n in range(1, 20): dist = Distribution(self.context, shape=(n, ), dist=('b', ), targets=self.context.targets[:3]) for ls in dist.localshapes(): self.assertNotIn(0, ls)
def test_reduce_0D(self): N = 10 ** 5 dist = Distribution(self.context, (N,)) new_dist = dist.reduce(axes=[0]) self.assertEqual(new_dist.dist, ()) self.assertSequenceEqual(new_dist.shape, ()) self.assertEqual(new_dist.grid_shape, ()) self.assertEqual(set(new_dist.targets), set(dist.targets[:1]))
def test_2D_bn(self): nrows, ncols = 31, 53 cm = Distribution(self.context, (nrows, ncols), {0: "b"}, (4, 1)) chunksize = (nrows // 4) + 1 for _ in range(100): r, c = randrange(nrows), randrange(ncols) rank = r // chunksize self.assertSequenceEqual(cm.owning_ranks((r, c)), [rank])
def test_reduce_0D(self): N = 10**5 dist = Distribution(self.context, (N, )) new_dist = dist.reduce(axes=[0]) self.assertEqual(new_dist.dist, ()) self.assertSequenceEqual(new_dist.shape, ()) self.assertEqual(new_dist.grid_shape, ()) self.assertEqual(set(new_dist.targets), set(dist.targets[:1]))
def test_2D_cc(self): nrows, ncols = 3, 5 nprocs_per_dim = 2 cm = Distribution(self.context, (nrows, ncols), ("c", "c"), (nprocs_per_dim, nprocs_per_dim)) for r in range(nrows): for c in range(ncols): rank = (r % nprocs_per_dim) * nprocs_per_dim + (c % nprocs_per_dim) actual = cm.owning_ranks((r, c)) self.assertSequenceEqual(actual, [rank])
def test_block_redist_many_to_one(self): source_dist = Distribution(self.context, (40,), ("b",), (2,), targets=[1, 2]) dest_dist = Distribution(self.context, (40,), ("b",), (1,), targets=[0]) plan = source_dist.get_redist_plan(dest_dist) expected = [ {"source_rank": 1, "dest_rank": 0, "indices": [(0, 20, 1)]}, {"source_rank": 2, "dest_rank": 0, "indices": [(20, 40, 1)]}, ] self.assertEqual(plan, expected)
def test_full_slice_with_int_2d(self): d0 = Distribution(context=self.context, shape=(15, 20)) s = (slice(None), 4) d1 = d0.slice(s) self.assertEqual(len(d0.maps) - 1, len(d1.maps)) self.assertSequenceEqual(d1.dist, d0.dist[:-1]) self.assertEqual(d1.shape, (15,))
def test_full_slice_with_int_2d(self): d0 = Distribution(context=self.context, shape=(15, 20)) s = (slice(None), 4) d1 = d0.slice(s) self.assertEqual(len(d0.maps) - 1, len(d1.maps)) self.assertSequenceEqual(d1.dist, d0.dist[:-1]) self.assertEqual(d1.shape, (15, ))
def test_from_partial_slice_2d(self): d0 = Distribution(context=self.context, shape=(15, 20)) s = (slice(3, 7), 4) d1 = d0.slice(s) self.assertEqual(len(d0.maps) - 1, len(d1.maps)) self.assertSequenceEqual(d1.dist, d0.dist[:-1]) for m, expected in zip(d1.maps, ([(0, 1), (1, 4)], [(0, 1)])): self.assertSequenceEqual(m.bounds, expected)
def test_create_target_subset(self): shape = (100, 100) subtargets = self.context.targets[::2] distribution = Distribution(self.context, shape=shape, targets=subtargets) darr = self.context.ones(distribution) lss = darr.localshapes() self.assertEqual(len(lss), len(subtargets)) ddpr = distribution.get_dim_data_per_rank() self.assertEqual(len(ddpr), len(subtargets))
def test_from_partial_slice_1d(self): d0 = Distribution(context=self.context, shape=(15, )) s = (slice(0, 3), ) d1 = d0.slice(s) self.assertEqual(len(d0.maps), len(d1.maps)) self.assertSequenceEqual(d1.dist, d0.dist) self.assertSequenceEqual(d1.targets, [0]) self.assertSequenceEqual(d1.shape, (3, ))
def test_from_full_slice_with_step_1d_0(self): d0 = Distribution(context=self.context, shape=(15, )) s = (slice(None, None, 2), ) d1 = d0.slice(s) self.assertEqual(len(d0.maps), len(d1.maps)) self.assertSequenceEqual(d1.dist, d0.dist) self.assertSequenceEqual(d1.targets, d0.targets) self.assertEqual(d1.maps[0].bounds[0][0], d0.maps[0].bounds[0][0])
def test_from_full_slice_with_step_1d_0(self): d0 = Distribution(context=self.context, shape=(15,)) s = (slice(None, None, 2),) d1 = d0.slice(s) self.assertEqual(len(d0.maps), len(d1.maps)) self.assertSequenceEqual(d1.dist, d0.dist) self.assertSequenceEqual(d1.targets, d0.targets) self.assertEqual(d1.maps[0].bounds[0][0], d0.maps[0].bounds[0][0])
def test_from_partial_slice_1d(self): d0 = Distribution(context=self.context, shape=(15,)) s = (slice(0, 3),) d1 = d0.slice(s) self.assertEqual(len(d0.maps), len(d1.maps)) self.assertSequenceEqual(d1.dist, d0.dist) self.assertSequenceEqual(d1.targets, [0]) self.assertSequenceEqual(d1.shape, (3,))
def test_global_flat_indices(self): dist0 = Distribution(self.context, (40, ), ('b', ), (2, ), targets=[1, 3]) self.assertSequenceEqual([[(0, 20)], [(20, 40)]], [ global_flat_indices(ddpr) for ddpr in dist0.get_dim_data_per_rank() ]) dist1 = Distribution(self.context, (5, 8), ('b', 'n'), (2, ), targets=[0, 2]) self.assertSequenceEqual([[(0, 24)], [(24, 40)]], [ global_flat_indices(ddpr) for ddpr in dist1.get_dim_data_per_rank() ]) dist2 = Distribution(self.context, (5, 8), ('b', 'b'), (2, 2), targets=[0, 1, 2, 3]) self.assertSequenceEqual( [[(0, 4), (8, 12), (16, 20)], [(4, 8), (12, 16), (20, 24)], [(24, 28), (32, 36)], [(28, 32), (36, 40)]], [ global_flat_indices(ddpr) for ddpr in dist2.get_dim_data_per_rank() ])
def test_2D_bb(self): nrows, ncols = 3, 5 nprocs_per_dim = 2 cm = Distribution(self.context, (nrows, ncols), ("b", "b"), (nprocs_per_dim, nprocs_per_dim)) row_chunks = nrows // nprocs_per_dim + 1 col_chunks = ncols // nprocs_per_dim + 1 for r in range(nrows): for c in range(ncols): rank = (r // row_chunks) * nprocs_per_dim + (c // col_chunks) actual = cm.owning_ranks((r, c)) self.assertSequenceEqual(actual, [rank])
def test_2D_cc(self): nrows, ncols = 3, 5 nprocs_per_dim = 2 cm = Distribution(self.context, (nrows, ncols), ('c', 'c'), (nprocs_per_dim, nprocs_per_dim)) for r in range(nrows): for c in range(ncols): rank = ((r % nprocs_per_dim) * nprocs_per_dim + (c % nprocs_per_dim)) actual = cm.owning_ranks((r, c)) self.assertSequenceEqual(actual, [rank])
def test_redist_incompatible_sizes(self): source_da = self.context.empty((10, ), dtype=numpy.int32) with self.assertRaises(ValueError): source_da.distribute_as(Distribution(self.context, (9, ))) with self.assertRaises(ValueError): source_da.distribute_as(Distribution(self.context, (3, 4))) with self.assertRaises(ValueError): source_da.distribute_as(Distribution(self.context, (3, 4000)))
def test_3D(self): a = self.context.empty((3, 3, 3)) db = Distribution(self.context, (3, 3, 3), dist=('n', 'b', 'n')) b = self.context.empty(db) dc = Distribution(self.context, (3, 3, 3), dist=('n', 'n', 'b')) c = self.context.empty(dc) self.assertEqual(a.grid_shape, (3, 1, 1)) self.assertEqual(b.grid_shape, (1, 3, 1)) self.assertEqual(c.grid_shape, (1, 1, 3))
def test_redist_unsupported_dist_types(self): source_dist = Distribution(self.context, (10, 20), ('n', 'c')) source_da = self.context.empty(source_dist) with self.assertRaises(NotImplementedError): source_da.distribute_as(source_da.distribution) source_da = self.context.empty((10, 20), ('b', 'b')) dest_dist = Distribution(self.context, (10, 20), ('b', 'c')) with self.assertRaises(NotImplementedError): source_da.distribute_as(dest_dist)
def test_from_full_slice_2d(self): d0 = Distribution(context=self.context, shape=(15, 20)) s = (slice(None), slice(None)) d1 = d0.slice(s) self.assertEqual(len(d0.maps), len(d1.maps)) self.assertSequenceEqual(d1.dist, d0.dist) for m0, m1 in zip(d0.maps, d1.maps): if m0.dist == "b": self.assertSequenceEqual(m0.bounds, m1.bounds) self.assertSequenceEqual(d1.targets, d0.targets)
def test_block_redist_2D_one_to_many(self): source_dist = Distribution(self.context, (9, 9), ("b", "b"), (1, 1), targets=[2]) dest_dist = Distribution(self.context, (9, 9), ("b", "b"), (2, 2), targets=range(4)) plan = source_dist.get_redist_plan(dest_dist) expected = [ {"source_rank": 2, "dest_rank": 0, "indices": [(0, 5, 1), (0, 5, 1)]}, {"source_rank": 2, "dest_rank": 1, "indices": [(0, 5, 1), (5, 9, 1)]}, {"source_rank": 2, "dest_rank": 2, "indices": [(5, 9, 1), (0, 5, 1)]}, {"source_rank": 2, "dest_rank": 3, "indices": [(5, 9, 1), (5, 9, 1)]}, ] for p, e in zip(plan, expected): self.assertEqual(p, e)
def test_from_full_slice_2d(self): d0 = Distribution(context=self.context, shape=(15, 20)) s = (slice(None), slice(None)) d1 = d0.slice(s) self.assertEqual(len(d0.maps), len(d1.maps)) self.assertSequenceEqual(d1.dist, d0.dist) for m0, m1 in zip(d0.maps, d1.maps): if m0.dist == 'b': self.assertSequenceEqual(m0.bounds, m1.bounds) self.assertSequenceEqual(d1.targets, d0.targets)
def test_2D_bb(self): nrows, ncols = 3, 5 nprocs_per_dim = 2 cm = Distribution(self.context, (nrows, ncols), ('b', 'b'), (nprocs_per_dim, nprocs_per_dim)) row_chunks = nrows // nprocs_per_dim + 1 col_chunks = ncols // nprocs_per_dim + 1 for r in range(nrows): for c in range(ncols): rank = (r // row_chunks) * nprocs_per_dim + (c // col_chunks) actual = cm.owning_ranks((r, c)) self.assertSequenceEqual(actual, [rank])
def test_redist_reshape_split_targets(self): dist0 = Distribution(self.context, (40, ), ('b', ), (1, ), targets=[0]) dist1 = Distribution(self.context, (5, 8), ('b', 'n'), (2, ), targets=[0, 1]) da_src = self.context.empty(dist0) da_src.fill(-10) da_dest = da_src.distribute_as(dist1) expected = numpy.empty((5, 8)) expected.fill(-10) assert_array_equal(da_dest.tondarray(), expected)
def test_is_compatible_degenerate_block_cyclic(self): size = 19937 gdd_block_cyclic = ({"dist_type": "c", "proc_grid_size": 1, "block_size": 7, "size": size},) gdd_block = ({"dist_type": "b", "proc_grid_size": 1, "bounds": [0, size]},) gdd_cyclic = ({"dist_type": "c", "proc_grid_size": 1, "size": size},) dist_block_cyclic = Distribution.from_global_dim_data(self.context, gdd_block_cyclic) dist_block = Distribution.from_global_dim_data(self.context, gdd_block) dist_cyclic = Distribution.from_global_dim_data(self.context, gdd_cyclic) self.assertTrue(dist_block_cyclic.is_compatible(dist_block)) self.assertTrue(dist_block_cyclic.is_compatible(dist_cyclic)) self.assertTrue(dist_block.is_compatible(dist_block_cyclic)) self.assertTrue(dist_cyclic.is_compatible(dist_block_cyclic))
def test_set_index_error(self): distribution = Distribution(self.context, (10, ), dist={0: 'c'}) dap = self.context.empty(distribution) with self.assertRaises(IndexError): dap[11] = 55 with self.assertRaises(IndexError): dap[-11] = 55
def test_iteration(self): size = 10 distribution = Distribution(self.context, (size, ), dist={0: 'c'}) dap = self.context.empty(distribution) dap.fill(10) for val in dap: self.assertEqual(val, 10)
def test_gh_435_regression_with_var(self): dist = Distribution(self.context, shape=(14, ), dist=('b'), targets=range(4)) darr = self.context.ones(dist) darr.var()
def test_tondarray(self): distribution = Distribution(self.context, (3, 3)) dap = self.context.empty(distribution) ndarr = numpy.arange(9).reshape(3, 3) for (i, j), val in numpy.ndenumerate(ndarr): dap[i, j] = ndarr[i, j] numpy.testing.assert_array_equal(dap.tondarray(), ndarr)
def test_from_global_dim_data_irregular_block(self): bounds = (0, 2, 3, 4, 10) glb_dim_data = ({"dist_type": "b", "bounds": bounds},) distribution = Distribution.from_global_dim_data(self.context, glb_dim_data) distarr = DistArray(distribution, dtype=int) distarr.toarray()
def test_load_nu(self): distribution = Distribution.from_dim_data_per_rank(self.context, nu_test_data) da = self.context.load_npy(self.output_path, distribution) for i in range(da.shape[0]): for j in range(da.shape[1]): self.assertEqual(da[i, j], self.expected[i, j])
def fromndarray(self, arr, distribution=None): """Create a DistArray from an ndarray. Parameters ---------- distribution : Distribution object, optional If a Distribution object is not provided, one is created with `Distribution(arr.shape)`. Returns ------- DistArray A DistArray distributed as specified, using the values and dtype from `arr`. """ if distribution is None: distribution = Distribution(self, arr.shape) out = self.empty(distribution, dtype=arr.dtype) try: out[...] = arr except AttributeError: # no slicing for a given map type; do it the slow way for index, value in numpy.ndenumerate(arr): out[index] = value return out
def test_global_dim_data_local_dim_data_equivalence(self): rows, cols = 5, 9 glb_dim_data = ( {"dist_type": "c", "block_size": 2, "size": rows, "proc_grid_size": 2}, {"dist_type": "c", "block_size": 2, "proc_grid_size": 2, "size": cols}, ) distribution = Distribution.from_global_dim_data(self.context, glb_dim_data) actual = distribution.get_dim_data_per_rank() expected = [ ( {"block_size": 2, "dist_type": "c", "proc_grid_rank": 0, "proc_grid_size": 2, "size": rows, "start": 0}, {"block_size": 2, "dist_type": "c", "proc_grid_rank": 0, "proc_grid_size": 2, "size": cols, "start": 0}, ), ( {"block_size": 2, "dist_type": "c", "proc_grid_rank": 0, "proc_grid_size": 2, "size": rows, "start": 0}, {"block_size": 2, "dist_type": "c", "proc_grid_rank": 1, "proc_grid_size": 2, "size": cols, "start": 2}, ), ( {"block_size": 2, "dist_type": "c", "proc_grid_rank": 1, "proc_grid_size": 2, "size": rows, "start": 2}, {"block_size": 2, "dist_type": "c", "proc_grid_rank": 0, "proc_grid_size": 2, "size": cols, "start": 0}, ), ( {"block_size": 2, "dist_type": "c", "proc_grid_rank": 1, "proc_grid_size": 2, "size": rows, "start": 2}, {"block_size": 2, "dist_type": "c", "proc_grid_rank": 1, "proc_grid_size": 2, "size": cols, "start": 2}, ), ] self.assertSequenceEqual(actual, expected)
def test_from_global_dim_data_bc(self): """ Test creation of a block-cyclic array. """ rows, cols = 5, 9 global_dim_data = ( # dim 0 { 'dist_type': 'c', 'proc_grid_size': 2, 'size': rows, 'block_size': 2, }, # dim 1 { 'dist_type': 'c', 'proc_grid_size': 2, 'size': cols, 'block_size': 2, }, ) distribution = Distribution.from_global_dim_data( self.context, global_dim_data) distarr = DistArray(distribution, dtype=int) distarr.toarray() las = distarr.get_localarrays() local_shapes = [la.local_shape for la in las] self.assertSequenceEqual(local_shapes, [(3, 5), (3, 4), (2, 5), (2, 4)])
def test_grid_rank(self): # regression test for issue #235 d = Distribution(self.context, (4, 4, 4), dist=('b', 'n', 'b'), grid_shape=(1, 1, 4)) a = self.context.empty(d) self.assertEqual(a.grid_shape, (1, 1, 4))
def test_redist_identity(self): source_dist = dest_dist = Distribution(self.context, (10, 10), ('b', 'b'), (1, 1), targets=[0]) source_da = self.context.empty(source_dist, dtype=numpy.int32) source_da.fill(-42) dest_da = source_da.distribute_as(dest_dist) assert_array_equal(source_da.tondarray(), dest_da.tondarray())
def test___init__(self): shape = (5, 5) distribution = Distribution(self.context, shape, ('b', 'c')) da = DistArray(distribution, dtype=int) da.fill(42) nda = numpy.empty(shape, dtype=int) nda.fill(42) assert_array_equal(da.tondarray(), nda)
def test_from_global_dim_data_irregular_block(self): bounds = (0, 2, 3, 4, 10) glb_dim_data = ({'dist_type': 'b', 'bounds': bounds}, ) distribution = Distribution.from_global_dim_data( self.context, glb_dim_data) distarr = DistArray(distribution, dtype=int) distarr.toarray()
def test_harder_distribution(self): a = numpy.asarray(numpy.arange(30), dtype=numpy.int16) gdd = (dict(dist_type="b", bounds=[0, 3, 12, 21, 30]),) dist = Distribution.from_global_dim_data(context=self.context, global_dim_data=gdd) da = self.context.fromndarray(a, distribution=dist) dtype = numpy.int32 with self.assertRaises(ValueError): da.view(dtype=dtype)
def local_view(larr, ddpr, dtype): from distarray.localapi.maps import Distribution if len(ddpr) == 0: dim_data = () else: dim_data = ddpr[larr.comm_rank] ldist = Distribution(comm=larr.comm, dim_data=dim_data) lview = larr.view(ldist, dtype=dtype) return proxyize(lview)
def test_from_global_dim_data_uu(self): rows = 6 cols = 20 row_ixs = numpy.random.permutation(range(rows)) col_ixs = numpy.random.permutation(range(cols)) row_indices = [row_ixs[: rows // 2], row_ixs[rows // 2 :]] col_indices = [col_ixs[: cols // 4], col_ixs[cols // 4 :]] glb_dim_data = ({"dist_type": "u", "indices": row_indices}, {"dist_type": "u", "indices": col_indices}) distribution = Distribution.from_global_dim_data(self.context, glb_dim_data) distarr = DistArray(distribution, dtype=int) distarr.toarray()
def test_redist_2D(self): nrows, ncols = 7, 13 source_dist = Distribution(self.context, (nrows, ncols), ("b", "b"), (2, 2), targets=range(4)) dest_gdd = ( {"dist_type": "b", "bounds": [0, nrows // 3, nrows]}, {"dist_type": "b", "bounds": [0, ncols // 3, ncols]}, ) dest_dist = Distribution.from_global_dim_data(self.context, dest_gdd, targets=range(4)) source_da = self.context.empty(source_dist, dtype=numpy.int32) source_da.fill(-42) dest_da = source_da.distribute_as(dest_dist) assert_array_equal(source_da.tondarray(), dest_da.tondarray())
def test_reduce(self): nr, nc, nd = 10 ** 5, 10 ** 6, 10 ** 4 dist = Distribution(self.context, (nr, nc, nd), ("b", "c", "n"), grid_shape=(2, 2, 1)) new_dist0 = dist.reduce(axes=[0]) self.assertEqual(new_dist0.dist, ("c", "n")) self.assertSequenceEqual(new_dist0.shape, (nc, nd)) self.assertEqual(new_dist0.grid_shape, dist.grid_shape[1:]) self.assertLess(set(new_dist0.targets), set(dist.targets)) new_dist1 = dist.reduce(axes=[1]) self.assertEqual(new_dist1.dist, ("b", "n")) self.assertSequenceEqual(new_dist1.shape, (nr, nd)) self.assertEqual(new_dist1.grid_shape, dist.grid_shape[:1] + dist.grid_shape[2:]) self.assertLess(set(new_dist1.targets), set(dist.targets)) new_dist2 = dist.reduce(axes=[2]) self.assertEqual(new_dist2.dist, ("b", "c")) self.assertSequenceEqual(new_dist2.shape, (nr, nc)) self.assertEqual(new_dist2.grid_shape, dist.grid_shape[:-1]) self.assertEqual(set(new_dist2.targets), set(dist.targets))
def fromfunction(self, function, shape, **kwargs): """Create a DistArray from a function over global indices. Unlike numpy's `fromfunction`, the result of distarray's `fromfunction` is restricted to the same Distribution as the index array generated from `shape`. See numpy.fromfunction for more details. """ self.push_function(function.__name__, function, targets=self.targets) def _local_fromfunction(func_name, comm, ddpr, kwargs): from distarray.localapi import fromfunction from distarray.localapi.maps import Distribution from importlib import import_module main = import_module('__main__') if len(ddpr): dim_data = ddpr[comm.Get_rank()] else: dim_data = () func = getattr(main, func_name) dist = Distribution(comm, dim_data=dim_data) local_arr = fromfunction(func, dist, **kwargs) return proxyize(local_arr) dist = kwargs.get('dist', None) grid_shape = kwargs.get('grid_shape', None) distribution = Distribution(context=self, shape=shape, dist=dist, grid_shape=grid_shape) ddpr = distribution.get_dim_data_per_rank() da_name = self.apply(_local_fromfunction, (function.__name__, distribution.comm, ddpr, kwargs), targets=distribution.targets) return DistArray.from_localarrays(da_name[0], distribution=distribution)
def test_is_compatible(self): nr, nc, nd = 10 ** 5, 10 ** 6, 10 ** 4 cm0 = Distribution(self.context, (nr, nc, nd), ("b", "c", "n")) self.assertTrue(cm0.is_compatible(cm0)) cm1 = Distribution(self.context, (nr, nc, nd), ("b", "c", "n")) self.assertTrue(cm1.is_compatible(cm1)) self.assertTrue(cm0.is_compatible(cm1)) self.assertTrue(cm1.is_compatible(cm0)) nr -= 1 nc -= 1 nd -= 1 cm2 = Distribution(self.context, (nr, nc, nd), ("b", "c", "n")) self.assertFalse(cm1.is_compatible(cm2)) self.assertFalse(cm2.is_compatible(cm1))
def test_from_global_dim_data_bu(self): rows = 9 row_break_point = rows // 2 cols = 10 col_indices = numpy.random.permutation(range(cols)) col_break_point = len(col_indices) // 3 indices = [col_indices[:col_break_point], col_indices[col_break_point:]] glb_dim_data = ( {"dist_type": "b", "bounds": (0, row_break_point, rows)}, {"dist_type": "u", "indices": indices}, ) distribution = Distribution.from_global_dim_data(self.context, glb_dim_data) distarr = DistArray(distribution, dtype=int) distarr.toarray()
def test_from_global_dim_data_1d(self): total_size = 40 list_of_indices = [ [29, 38, 18, 19, 11, 33, 10, 1, 22, 25], [5, 15, 34, 12, 16, 24, 23, 39, 6, 36], [0, 7, 27, 4, 32, 37, 21, 26, 9, 17], [35, 14, 20, 13, 3, 30, 2, 8, 28, 31], ] glb_dim_data = ({"dist_type": "u", "indices": list_of_indices},) distribution = Distribution.from_global_dim_data(self.context, glb_dim_data) distarr = DistArray(distribution, dtype=int) for i in range(total_size): distarr[i] = i localarrays = distarr.get_localarrays() for i, arr in enumerate(localarrays): assert_allclose(arr, list_of_indices[i])
def test_from_global_dim_data_bc(self): """ Test creation of a block-cyclic array. """ rows, cols = 5, 9 global_dim_data = ( # dim 0 {"dist_type": "c", "proc_grid_size": 2, "size": rows, "block_size": 2}, # dim 1 {"dist_type": "c", "proc_grid_size": 2, "size": cols, "block_size": 2}, ) distribution = Distribution.from_global_dim_data(self.context, global_dim_data) distarr = DistArray(distribution, dtype=int) distarr.toarray() las = distarr.get_localarrays() local_shapes = [la.local_shape for la in las] self.assertSequenceEqual(local_shapes, [(3, 5), (3, 4), (2, 5), (2, 4)])
def test_global_flat_indices(self): dist0 = Distribution(self.context, (40,), ("b",), (2,), targets=[1, 3]) self.assertSequenceEqual( [[(0, 20)], [(20, 40)]], [global_flat_indices(ddpr) for ddpr in dist0.get_dim_data_per_rank()] ) dist1 = Distribution(self.context, (5, 8), ("b", "n"), (2,), targets=[0, 2]) self.assertSequenceEqual( [[(0, 24)], [(24, 40)]], [global_flat_indices(ddpr) for ddpr in dist1.get_dim_data_per_rank()] ) dist2 = Distribution(self.context, (5, 8), ("b", "b"), (2, 2), targets=[0, 1, 2, 3]) self.assertSequenceEqual( [[(0, 4), (8, 12), (16, 20)], [(4, 8), (12, 16), (20, 24)], [(24, 28), (32, 36)], [(28, 32), (36, 40)]], [global_flat_indices(ddpr) for ddpr in dist2.get_dim_data_per_rank()], )
def test_load_nu(self): distribution = Distribution.from_dim_data_per_rank(self.context, nu_test_data) da = self.context.load_hdf5(self.output_path, distribution, key="test") assert_array_equal(self.expected, da)
def test_not_compatible(self): dist_b1 = Distribution(self.context, (10,), ("b",), (1,), targets=[0]) dist_b2 = Distribution(self.context, (9,), ("b",), (1,), targets=[0]) self.assertFalse(dist_b1.is_compatible(dist_b2)) self.assertFalse(dist_b2.is_compatible(dist_b1)) dist_b3 = Distribution(self.context, (10,), ("b",), (2,), targets=[0, 1]) self.assertFalse(dist_b1.is_compatible(dist_b3)) self.assertFalse(dist_b3.is_compatible(dist_b1)) dist_b4 = Distribution(self.context, (10,), ("c",), (2,), targets=[0, 1]) self.assertFalse(dist_b4.is_compatible(dist_b3)) self.assertFalse(dist_b3.is_compatible(dist_b4)) gdd_unstructured = ({"dist_type": "u", "indices": [range(10)]},) dist_u = Distribution.from_global_dim_data(self.context, gdd_unstructured) self.assertFalse(dist_u.is_compatible(dist_b1)) self.assertFalse(dist_b1.is_compatible(dist_u))
def test_is_compatible_degenerate(self): dist_bc = Distribution(self.context, (10, 10), ("b", "c"), (1, 1), targets=[0]) dist_cb = Distribution(self.context, (10, 10), ("c", "b"), (1, 1), targets=[0]) self.assertTrue(dist_bc.is_compatible(dist_cb)) self.assertTrue(dist_cb.is_compatible(dist_bc))
def test_is_compatible_nodist(self): # See GH issue #461. dist_bcn = Distribution(self.context, (10, 10, 10), ("b", "c", "n"), (1, 1, 1), targets=[0]) dist_nnn = Distribution(self.context, (10, 10, 10), ("n", "n", "n"), (1, 1, 1), targets=[0]) self.assertTrue(dist_bcn.is_compatible(dist_nnn)) self.assertTrue(dist_nnn.is_compatible(dist_bcn))