def test_apply_gufunc_elemwise_01b():
def add(x, y):
return x + y
a = da.from_array(np.array([1, 2, 3]), chunks=2, name='a')
b = da.from_array(np.array([1, 2, 3]), chunks=1, name='b')
with pytest.raises(ValueError):
apply_gufunc(add, "(),()->()", a, b, output_dtypes=a.dtype)
def test_apply_gufunc_01c():
def stats(x):
return np.mean(x, axis=-1), np.std(x, axis=-1)
a = da.random.normal(size=(10, 20, 30), chunks=5)
with assert_raises(ValueError):
apply_gufunc(stats, "(i)->(),()", a,
output_dtypes=2 * (a.dtype,))
def test_apply_gufunc_check_same_dimsizes():
def foo(x, y):
return x + y
a = da.random.normal(size=(3,), chunks=(2,))
b = da.random.normal(size=(4,), chunks=(2,))
with pytest.raises(ValueError) as excinfo:
apply_gufunc(foo, "(),()->()", a, b, output_dtypes=float, allow_rechunk=True)
assert "different lengths in arrays" in str(excinfo.value)
def test_apply_gufunc_output_dtypes_string_many_outputs(vectorize):
def stats(x):
return np.mean(x, axis=-1), np.std(x, axis=-1)
a = da.random.normal(size=(10, 20, 30), chunks=(5, 5, 30))
mean, std = apply_gufunc(stats, "(i)->(),()", a, output_dtypes=("f", "f"), vectorize=vectorize)
assert mean.compute().shape == (10, 20)
assert std.compute().shape == (10, 20)
def test_apply_gufunc_elemwise_01():
def add(x, y):
return x + y
a = da.from_array(np.array([1, 2, 3]), chunks=2, name='a')
b = da.from_array(np.array([1, 2, 3]), chunks=2, name='b')
z = apply_gufunc(add, "(),()->()", a, b, output_dtypes=a.dtype)
assert_eq(z, np.array([2, 4, 6]))
def test_gufunc_two_inputs():
def foo(x, y):
return np.einsum('...ij,...jk->ik', x, y)
a = da.ones((2, 3), chunks=100, dtype=int)
b = da.ones((3, 4), chunks=100, dtype=int)
x = apply_gufunc(foo, "(i,j),(j,k)->(i,k)", a, b, output_dtypes=int)
assert_eq(x, 3 * np.ones((2, 4), dtype=int))
def test_gufunc_mixed_inputs():
def foo(x, y):
return x + y
a = np.ones((2, 1), dtype=int)
b = da.ones((1, 8), chunks=(2, 3), dtype=int)
x = apply_gufunc(foo, "(),()->()", a, b, output_dtypes=int)
assert_eq(x, 2 * np.ones((2, 8), dtype=int))
def test_apply_gufunc_02():
def outer_product(x, y):
return np.einsum("...i,...j->...ij", x, y)
a = da.random.normal(size=( 20, 30), chunks=(5, 30))
b = da.random.normal(size=(10, 1, 40), chunks=(10, 1, 40))
c = apply_gufunc(outer_product, "(i),(j)->(i,j)", a, b, output_dtypes=a.dtype)
assert c.compute().shape == (10, 20, 30, 40)
def test_apply_gufunc_01():
def stats(x):
return np.mean(x, axis=-1), np.std(x, axis=-1)
a = da.random.normal(size=(10, 20, 30), chunks=(5, 5, 30))
mean, std = apply_gufunc(stats, "(i)->(),()", a,
output_dtypes=2 * (a.dtype,))
assert mean.compute().shape == (10, 20)
assert std.compute().shape == (10, 20)
def test_apply_gufunc_elemwise_core():
def foo(x):
assert x.shape == (3,)
return 2 * x
a = da.from_array(np.array([1, 2, 3]), chunks=3, name='a')
z = apply_gufunc(foo, "(i)->(i)", a, output_dtypes=int)
assert z.chunks == ((3,),)
assert_eq(z, np.array([2, 4, 6]))
def test_apply_gufunc_output_dtypes(output_dtypes):
def foo(x):
return y
x = np.random.randn(10)
y = x.astype(int)
dy = apply_gufunc(foo, "()->()", x, output_dtypes=output_dtypes)
#print(x, x.compute())
assert_eq(y, dy)
def test_apply_gufunc_elemwise_loop():
def foo(x):
assert x.shape in ((2,), (1,))
return 2 * x
a = da.from_array(np.array([1, 2, 3]), chunks=2, name='a')
z = apply_gufunc(foo, "()->()", a, output_dtypes=int)
assert z.chunks == ((2, 1),)
assert_eq(z, np.array([2, 4, 6]))
def test_apply_gufunc_elemwise_02():
def addmul(x, y):
assert x.shape in ((2,), (1,))
return x + y, x * y
a = da.from_array(np.array([1, 2, 3]), chunks=2, name='a')
b = da.from_array(np.array([1, 2, 3]), chunks=2, name='b')
z1, z2 = apply_gufunc(addmul, "(),()->(),()", a, b, output_dtypes=2 * (a.dtype,))
assert_eq(z1, np.array([2, 4, 6]))
assert_eq(z2, np.array([1, 4, 9]))
def test_apply_gufunc_two_mixed_outputs():
def foo():
return 1, np.ones((2, 3), dtype=float)
x, y = apply_gufunc(foo, "->(),(i,j)",
output_dtypes=(int, float),
output_sizes={'i': 2, 'j': 3})
assert x.compute() == 1
assert y.chunks == ((2,), (3,))
assert_eq(y, np.ones((2, 3), dtype=float))
def test_apply_gufunc_axes_two_kept_coredims():
a = da.random.normal(size=( 20, 30), chunks=(10, 30))
b = da.random.normal(size=(10, 1, 40), chunks=(5, 1, 40))
def outer_product(x, y):
return np.einsum("i,j->ij", x, y)
c = apply_gufunc(outer_product, "(i),(j)->(i,j)", a, b, vectorize=True)
assert c.compute().shape == (10, 20, 30, 40)
def test_apply_gufunc_axis_02():
def myfft(x):
return np.fft.fft(x, axis=-1)
a = np.random.randn(10, 5)
da_ = da.from_array(a, chunks=2)
m = np.fft.fft(a, axis=0)
dm = apply_gufunc(myfft, "(i)->(i)", da_, axis=0, allow_rechunk=True)
assert_eq(m, dm)
def test_apply_gufunc_axis_03():
def mydiff(x):
return np.diff(x, axis=-1)
a = np.random.randn(3, 6, 4)
da_ = da.from_array(a, chunks=2)
m = np.diff(a, axis=1)
dm = apply_gufunc(mydiff, "(i)->(i)", da_, axis=1, output_sizes={'i': 5}, allow_rechunk=True)
assert_eq(m, dm)
def test_apply_gufunc_axis_keepdims(axis):
def mymedian(x):
return np.median(x, axis=-1)
a = np.random.randn(10, 5)
da_ = da.from_array(a, chunks=2)
m = np.median(a, axis=-1 if not axis else axis, keepdims=True)
dm = apply_gufunc(mymedian, "(i)->()", da_, axis=axis, keepdims=True, allow_rechunk=True)
assert_eq(m, dm)
def test_apply_gufunc_axes_01(axes):
def mystats(x, y):
return np.std(x, axis=-1) * np.mean(y, axis=-1)
a = np.random.randn(10, 5)
b = np.random.randn(5, 6)
da_ = da.from_array(a, chunks=2)
db_ = da.from_array(b, chunks=2)
m = np.std(a, axis=0) * np.mean(b, axis=1)
dm = apply_gufunc(mystats, "(i),(j)->()", da_, db_, axes=axes, allow_rechunk=True)
assert_eq(m, dm)
def test_apply_gufunc_axis_02b():
def myfilter(x, cn=10, axis=-1):
y = np.fft.fft(x, axis=axis)
y[cn:-cn] = 0
nx = np.fft.ifft(y, axis=axis)
return np.real(nx)
a = np.random.randn(3, 6, 4)
da_ = da.from_array(a, chunks=2)
m = myfilter(a, axis=1)
dm = apply_gufunc(myfilter, "(i)->(i)", da_, axis=1, allow_rechunk=True)
assert_eq(m, dm)
def test_apply_gufunc_axes_02():
def matmul(x, y):
return np.einsum("...ij,...jk->...ik", x, y)
a = np.random.randn(3, 2, 1)
b = np.random.randn(3, 7, 5)
da_ = da.from_array(a, chunks=2)
db = da.from_array(b, chunks=3)
m = np.einsum("jiu,juk->uik", a, b)
dm = apply_gufunc(matmul, "(i,j),(j,k)->(i,k)", da_, db, axes=[(1, 0), (0, -1), (-2, -1)], allow_rechunk=True)
assert_eq(m, dm)
def test_apply_gufunc_broadcasting_loopdims():
def foo(x, y):
assert len(x.shape) == 2
assert len(y.shape) == 3
x, y = np.broadcast_arrays(x, y)
return x, y, x * y
a = da.random.normal(size=( 10, 30), chunks=(8, 30))
b = da.random.normal(size=(20, 1, 30), chunks=(3, 1, 30))
x, y, z = apply_gufunc(foo, "(i),(i)->(i),(i),(i)", a, b, output_dtypes=3 * (float,), vectorize=False)
assert x.compute().shape == (20, 10, 30)
assert y.compute().shape == (20, 10, 30)
assert z.compute().shape == (20, 10, 30)
def test_apply_gufunc_infer_dtype():
x = np.arange(50).reshape((5, 10))
y = np.arange(10)
dx = da.from_array(x, chunks=5)
dy = da.from_array(y, chunks=5)
def foo(x, *args, **kwargs):
cast = kwargs.pop('cast', 'i8')
return (x + sum(args)).astype(cast)
dz = apply_gufunc(foo, "(),(),()->()", dx, dy, 1)
z = foo(dx, dy, 1)
assert_eq(dz, z)
dz = apply_gufunc(foo, "(),(),()->()", dx, dy, 1, cast='f8')
z = foo(dx, dy, 1, cast='f8')
assert_eq(dz, z)
dz = apply_gufunc(foo, "(),(),()->()", dx, dy, 1, cast='f8', output_dtypes='f8')
z = foo(dx, dy, 1, cast='f8')
assert_eq(dz, z)
def foo(x):
raise RuntimeError("Woops")
with pytest.raises(ValueError) as e:
apply_gufunc(foo, "()->()", dx)
msg = str(e.value)
assert msg.startswith("`dtype` inference failed")
assert "Please specify the dtype explicitly" in msg
assert 'RuntimeError' in msg
# Multiple outputs
def foo(x, y):
return x + y, x - y
z0, z1 = apply_gufunc(foo, "(),()->(),()", dx, dy)
assert_eq(z0, dx + dy)
assert_eq(z1, dx - dy)
def test_apply_gufunc_two_scalar_output():
def foo():
return 1, 2
x, y = apply_gufunc(foo, "->(),()", output_dtypes=(int, int))
assert x.compute() == 1
assert y.compute() == 2
def test_apply_gufunc_axes_input_validation_01():
def foo(x):
return np.mean(x, axis=-1)
a = da.random.normal(size=(20, 30), chunks=30)
with pytest.raises(ValueError):
apply_gufunc(foo, "(i)->()", a, axes=0)
apply_gufunc(foo, "(i)->()", a, axes=[0])
apply_gufunc(foo, "(i)->()", a, axes=[(0, )])
apply_gufunc(foo, "(i)->()", a, axes=[0, tuple()])
apply_gufunc(foo, "(i)->()", a, axes=[(0, ), tuple()])
with pytest.raises(ValueError):
apply_gufunc(foo, "(i)->()", a, axes=[(0, 1)])
with pytest.raises(ValueError):
apply_gufunc(foo, "(i)->()", a, axes=[0, 0])
def test_apply_gufunc_scalar_output():
def foo():
return 1
x = apply_gufunc(foo, "->()", output_dtypes=int)
assert x.compute() == 1
def test_gufunc_vector_output():
def foo():
return np.array([1, 2, 3], dtype=int)
x = apply_gufunc(foo, "->(i_0)", output_dtypes=int, output_sizes={"i_0": 3})
assert x.chunks == ((3,),)
assert_eq(x, np.array([1, 2, 3]))
def test_apply_gufunc_axes_input_validation_01():
def foo(x):
return np.mean(x, axis=-1)
a = da.random.normal(size=(20, 30), chunks=30)
with pytest.raises(ValueError):
apply_gufunc(foo, "(i)->()", a, axes=0)
apply_gufunc(foo, "(i)->()", a, axes=[0])
apply_gufunc(foo, "(i)->()", a, axes=[(0,)])
apply_gufunc(foo, "(i)->()", a, axes=[0, tuple()])
apply_gufunc(foo, "(i)->()", a, axes=[(0,), tuple()])
with pytest.raises(ValueError):
apply_gufunc(foo, "(i)->()", a, axes=[(0, 1)])
with pytest.raises(ValueError):
apply_gufunc(foo, "(i)->()", a, axes=[0, 0])
def test_apply_gufunc_pass_additional_kwargs():
def foo(x, bar):
assert bar == 2
return x
ret = apply_gufunc(foo, "()->()", 1., output_dtypes="f", bar=2)
assert_eq(ret, np.array(1., dtype="f"))
def test_apply_gufunc_scalar_output():
def foo():
return 1
x = apply_gufunc(foo, "->()", output_dtypes=int)
assert x.compute() == 1
def test_apply_gufunc_output_dtypes_string(vectorize):
def stats(x):
return np.mean(x, axis=-1)
a = da.random.normal(size=(10, 20, 30), chunks=(5, 5, 30))
mean = apply_gufunc(stats, "(i)->()", a, output_dtypes="f", vectorize=vectorize)
assert mean.compute().shape == (10, 20)
def test_apply_gufunc_two_scalar_output():
def foo():
return 1, 2
x, y = apply_gufunc(foo, "->(),()", output_dtypes=(int, int))
assert x.compute() == 1
assert y.compute() == 2
def test_gufunc_vector_output():
def foo():
return np.array([1, 2, 3], dtype=int)
x = apply_gufunc(foo, "->(i_0)", output_dtypes=int, output_sizes={"i_0": 3})
assert x.chunks == ((3,),)
assert_eq(x, np.array([1, 2, 3]))
def test_apply_gufunc_pass_additional_kwargs():
def foo(x, bar):
assert bar == 2
return x
ret = apply_gufunc(foo, "()->()", 1., output_dtypes="f", bar=2)
assert_eq(ret, np.array(1., dtype="f"))
