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_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_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_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_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))
