Esempio n. 1
0
def test_atop_chunks():
    x = da.ones((5, 5), chunks=((2, 1, 2), (3, 2)))

    def double(a, axis=0):
        return np.concatenate([a, a], axis=axis)

    y = atop(double, 'ij', x, 'ij',
             adjust_chunks={'i': lambda n: 2 * n}, axis=0, dtype=x.dtype)
    assert y.chunks == ((4, 2, 4), (3, 2))
    assert_eq(y, np.ones((10, 5)))

    y = atop(double, 'ij', x, 'ij',
             adjust_chunks={'j': lambda n: 2 * n}, axis=1, dtype=x.dtype)
    assert y.chunks == ((2, 1, 2), (6, 4))
    assert_eq(y, np.ones((5, 10)))

    x = da.ones((10, 10), chunks=(5, 5))
    y = atop(double, 'ij', x, 'ij', axis=0,
             adjust_chunks={'i': 10}, dtype=x.dtype)
    assert y.chunks == ((10, 10), (5, 5))
    assert_eq(y, np.ones((20, 10)))

    y = atop(double, 'ij', x, 'ij', axis=0,
             adjust_chunks={'i': (10, 10)}, dtype=x.dtype)
    assert y.chunks == ((10, 10), (5, 5))
    assert_eq(y, np.ones((20, 10)))
Esempio n. 2
0
def test_atop_stacked_new_axes_front(concatenate):
    def f(x):
        if isinstance(x, list):
            x = np.concatenate(x)
        return x[None, ...] * np.ones(7)[(slice(None), ) + (None, ) * x.ndim]

    x = da.ones(5, chunks=2)
    y = atop(f,
             'qa',
             x,
             'a',
             new_axes={'q': 7},
             concatenate=concatenate,
             dtype=x.dtype)
    z = atop(f,
             'qab',
             y,
             'ab',
             new_axes={'q': 7},
             concatenate=concatenate,
             dtype=x.dtype)
    assert z.chunks == ((7, ), (7, ), (2, 2, 1))
    assert_eq(z, np.ones((7, 7, 5)))

    w = atop(lambda x: x[:, 0, 0],
             'a',
             z,
             'abc',
             dtype=x.dtype,
             concatenate=True)
    assert w.chunks == ((7, ), )
    assert_eq(w, np.ones((7, )))
Esempio n. 3
0
def test_atop_new_axes():
    def f(x):
        return x[:, None] * np.ones((1, 7))
    x = da.ones(5, chunks=2)
    y = atop(f, 'aq', x, 'a', new_axes={'q': 7}, concatenate=True,
             dtype=x.dtype)
    assert y.chunks == ((2, 2, 1), (7,))
    assert_eq(y, np.ones((5, 7)))

    def f(x):
        return x[None, :] * np.ones((7, 1))
    x = da.ones(5, chunks=2)
    y = atop(f, 'qa', x, 'a', new_axes={'q': 7}, concatenate=True,
             dtype=x.dtype)
    assert y.chunks == ((7,), (2, 2, 1))
    assert_eq(y, np.ones((7, 5)))

    def f(x):
        y = x.sum(axis=1)
        return y[:, None] * np.ones((1, 5))

    x = da.ones((4, 6), chunks=(2, 2))
    y = atop(f, 'aq', x, 'ab', new_axes={'q': 5}, concatenate=True,
             dtype=x.dtype)
    assert y.chunks == ((2, 2), (5,))
    assert_eq(y, np.ones((4, 5)) * 6)
Esempio n. 4
0
def test_atop_chunks():
    x = da.ones((5, 5), chunks=((2, 1, 2), (3, 2)))

    def double(a, axis=0):
        return np.concatenate([a, a], axis=axis)

    y = atop(double, 'ij', x, 'ij',
             adjust_chunks={'i': lambda n: 2 * n}, axis=0, dtype=x.dtype)
    assert y.chunks == ((4, 2, 4), (3, 2))
    assert_eq(y, np.ones((10, 5)))

    y = atop(double, 'ij', x, 'ij',
             adjust_chunks={'j': lambda n: 2 * n}, axis=1, dtype=x.dtype)
    assert y.chunks == ((2, 1, 2), (6, 4))
    assert_eq(y, np.ones((5, 10)))

    x = da.ones((10, 10), chunks=(5, 5))
    y = atop(double, 'ij', x, 'ij', axis=0,
             adjust_chunks={'i': 10}, dtype=x.dtype)
    assert y.chunks == ((10, 10), (5, 5))
    assert_eq(y, np.ones((20, 10)))

    y = atop(double, 'ij', x, 'ij', axis=0,
             adjust_chunks={'i': (10, 10)}, dtype=x.dtype)
    assert y.chunks == ((10, 10), (5, 5))
    assert_eq(y, np.ones((20, 10)))
Esempio n. 5
0
def test_atop_new_axes():
    def f(x):
        return x[:, None] * np.ones((1, 7))
    x = da.ones(5, chunks=2)
    y = atop(f, 'aq', x, 'a', new_axes={'q': 7}, concatenate=True,
             dtype=x.dtype)
    assert y.chunks == ((2, 2, 1), (7,))
    assert_eq(y, np.ones((5, 7)))

    def f(x):
        return x[None, :] * np.ones((7, 1))
    x = da.ones(5, chunks=2)
    y = atop(f, 'qa', x, 'a', new_axes={'q': 7}, concatenate=True,
             dtype=x.dtype)
    assert y.chunks == ((7,), (2, 2, 1))
    assert_eq(y, np.ones((7, 5)))

    def f(x):
        y = x.sum(axis=1)
        return y[:, None] * np.ones((1, 5))

    x = da.ones((4, 6), chunks=(2, 2))
    y = atop(f, 'aq', x, 'ab', new_axes={'q': 5}, concatenate=True,
             dtype=x.dtype)
    assert y.chunks == ((2, 2), (5,))
    assert_eq(y, np.ones((4, 5)) * 6)
Esempio n. 6
0
def test_atop_stacked_new_axes(concatenate):
    def f(x):
        return x[..., None] * np.ones((1, 7))

    x = da.ones(5, chunks=2)
    y = atop(f, 'aq', x, 'a', new_axes={'q': 7}, concatenate=concatenate, dtype=x.dtype)
    z = atop(f, 'abq', y, 'ab', new_axes={'q': 7}, concatenate=concatenate, dtype=x.dtype)
    assert z.chunks == ((2, 2, 1), (7,), (7,))
    assert_eq(z, np.ones((5, 7, 7)))
Esempio n. 7
0
def test_args_delayed():
    x = da.arange(10, chunks=(5, ))
    y = dask.delayed(lambda: 100)()

    z = atop(add, 'i', x, 'i', y, None, dtype=x.dtype)
    assert_eq(z, np.arange(10) + 100)

    z = atop(lambda x, y: x + y, 'i', x, 'i', y=y, dtype=x.dtype)
    assert_eq(z, np.arange(10) + 100)
Esempio n. 8
0
def test_atop_stacked_new_axes(concatenate):
    def f(x):
        return x[..., None] * np.ones((1, 7))

    x = da.ones(5, chunks=2)
    y = atop(f, 'aq', x, 'a', new_axes={'q': 7}, concatenate=concatenate, dtype=x.dtype)
    z = atop(f, 'abq', y, 'ab', new_axes={'q': 7}, concatenate=concatenate, dtype=x.dtype)
    assert z.chunks == ((2, 2, 1), (7,), (7,))
    assert_eq(z, np.ones((5, 7, 7)))
Esempio n. 9
0
def test_atop_stacked_new_axes_same_dim(concatenate):
    def f(x):
        return x[..., None] * np.ones((1, 7))

    x = da.ones(5, chunks=2)
    y = da.zeros(5, chunks=2)
    a = atop(f, 'aq', x, 'a', new_axes={'q': 7}, concatenate=concatenate, dtype=x.dtype)
    b = atop(f, 'aq', y, 'a', new_axes={'q': 7}, concatenate=concatenate, dtype=x.dtype)
    c = a + b
    assert c.chunks == ((2, 2, 1), (7,))
    assert_eq(c, np.ones((5, 7)))
Esempio n. 10
0
def test_atop_stacked_new_axes_same_dim(concatenate):
    def f(x):
        return x[..., None] * np.ones((1, 7))

    x = da.ones(5, chunks=2)
    y = da.zeros(5, chunks=2)
    a = atop(f, 'aq', x, 'a', new_axes={'q': 7}, concatenate=concatenate, dtype=x.dtype)
    b = atop(f, 'aq', y, 'a', new_axes={'q': 7}, concatenate=concatenate, dtype=x.dtype)
    c = a + b
    assert c.chunks == ((2, 2, 1), (7,))
    assert_eq(c, np.ones((5, 7)))
Esempio n. 11
0
def test_atop_stacked_new_axes_front(concatenate):
    def f(x):
        if isinstance(x, list):
            x = np.concatenate(x)
        return x[None, ...] * np.ones(7)[(slice(None),) + (None,) * x.ndim]

    x = da.ones(5, chunks=2)
    y = atop(f, 'qa', x, 'a', new_axes={'q': 7}, concatenate=concatenate, dtype=x.dtype)
    z = atop(f, 'qab', y, 'ab', new_axes={'q': 7}, concatenate=concatenate, dtype=x.dtype)
    assert z.chunks == ((7,), (7,), (2, 2, 1))
    assert_eq(z, np.ones((7, 7, 5)))

    w = atop(lambda x: x[:, 0, 0], 'a', z, 'abc', dtype=x.dtype, concatenate=True)
    assert w.chunks == ((7,),)
    assert_eq(w, np.ones((7,)))
Esempio n. 12
0
def test_atop_kwargs():
    def f(a, b=0):
        return a + b

    x = da.ones(5, chunks=(2,))
    y = atop(f, 'i', x, 'i', b=10, dtype=x.dtype)
    assert_eq(y, np.ones(5) + 10)
Esempio n. 13
0
def test_atop_kwargs():
    def f(a, b=0):
        return a + b

    x = da.ones(5, chunks=(2,))
    y = atop(f, 'i', x, 'i', b=10, dtype=x.dtype)
    assert_eq(y, np.ones(5) + 10)
Esempio n. 14
0
def test_atop_new_axes_2():
    x = da.ones((2, 2), chunks=(1, 1))

    def func(x):
        return np.stack([x, -x], axis=-1)

    y = atop(func, ('x', 'y', 'sign'), x, ('x', 'y'), dtype=x.dtype,
             concatenate=True, new_axes={'sign': 2})

    assert_eq(y, y)
Esempio n. 15
0
def test_atop_new_axes_2():
    x = da.ones((2, 2), chunks=(1, 1))

    def func(x):
        return np.stack([x, -x], axis=-1)

    y = atop(func, ('x', 'y', 'sign'), x, ('x', 'y'), dtype=x.dtype,
             concatenate=True, new_axes={'sign': 2})

    assert_eq(y, y)
Esempio n. 16
0
def test_atop_names():
    x = da.ones(5, chunks=(2,))
    y = atop(add, 'i', x, 'i', dtype=x.dtype)
    assert y.name.startswith('add')
Esempio n. 17
0
def test_atop_names():
    x = da.ones(5, chunks=(2,))
    y = atop(add, 'i', x, 'i', dtype=x.dtype)
    assert y.name.startswith('add')