def test_chunked_axis(comm):
grid = Grid(comm, (2,))
dist = Distribution(grid, [
ChunkedAxis([0, 200, 400])
], shape=(400,), strides=(1,))
local = np.arange(20)[None, :].copy()
assert dist.local_shape() == (200,)
# local_to_global
global_ = dist.local_to_global(local)
start = 200 * comm.Get_rank()
assert np.all(global_[0, :] == np.arange(start, start + 20))
# global_to_rank
global_ = np.arange(195, 205)[None, :].copy()
assert np.all(dist.global_to_rank_coords(global_) ==
[0, 0, 0, 0, 0, 1, 1, 1, 1, 1])
def test_chunked_execute(comm):
rank = comm.Get_rank()
grid = Grid(comm, (2, 3))
dist_a = Distribution(grid, [
ChunkedAxis([0, 10, 40]),
ChunkedAxis([0, 5, 10, 15])
], shape=(40, 15), strides=(15, 1))
dist_b = Distribution(grid, [
ChunkedAxis([0, 30, 40]),
ChunkedAxis([0, 10, 13, 15])
], shape=(40, 15), strides=(15, 1))
# Local shape
meq_(comm, [(10, 5), (10, 5), (10, 5),
(30, 5), (30, 5), (30, 5)], dist_a.local_shape())
meq_(comm, [(30, 10), (30, 3), (30, 2),
(10, 10), (10, 3), (10, 2)], dist_b.local_shape(), )
# sendcnts; all results computed on all nodes
sendcounts = np.zeros((6, 6), int)
for i in range(6):
sendcounts[i,:] = distarray.compute_send_counts(i, dist_a, dist_b)
assert np.all(sendcounts == [
[ 50, 0, 0, 0, 0, 0],
[ 50, 0, 0, 0, 0, 0],
[ 0, 30, 20, 0, 0, 0],
[100, 0, 0, 50, 0, 0],
[100, 0, 0, 50, 0, 0],
[ 0, 60, 40, 0, 30, 20]])
# recvcnts; all results computed on all nodes
recvcounts = np.zeros((6, 6), int)
for i in range(6):
recvcounts[i,:] = distarray.compute_send_counts(i, dist_b, dist_a)
assert np.all(recvcounts == sendcounts.T)
plan = RedistributionPlan(dist_a, dist_b, np.double)
arr_a = np.ones(dist_a.local_shape()) * rank
arr_b = np.empty(dist_b.local_shape()) * np.nan
plan.execute(arr_a, arr_b)
#mprint(comm, arr_b)
if rank == 0:
assert np.all(arr_b[:10, :5] == 0)
assert np.all(arr_b[:10, 5:] == 1)
assert np.all(arr_b[10:, :5] == 3)
assert np.all(arr_b[10:, 5:] == 4)
elif rank == 1 or rank == 2:
assert np.all(arr_b[:10, :] == 2)
assert np.all(arr_b[10:, :] == 5)
elif rank == 3:
assert np.all(arr_b[:, :5] == 3)
assert np.all(arr_b[:, 5:] == 4)
elif rank == 4 or rank == 5:
assert np.all(arr_b == 5)