Exemplo n.º 1
0
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])
Exemplo n.º 2
0
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)