Merge pull request dask#400 from shoyer/from_func

Add dask.array.from_func

dask/array/__init__.py

@@ -2,9 +2,9 @@
 
 from ..utils import ignoring
 from .core import (Array, stack, concatenate, take, tensordot, transpose,
-        from_array, choose, where, coarsen, insert, broadcast_to,
-        fromfunction, compute, unique, store, squeeze, topk, bincount,
-        map_blocks, atop, to_hdf5)
+        from_array, choose, where, coarsen, insert, broadcast_to, fromfunction,
+        compute, unique, store, squeeze, topk, bincount, map_blocks, atop,
+        to_hdf5)
 from .core import (logaddexp, logaddexp2, conj, exp, log, log2, log10, log1p,
         expm1, sqrt, square, sin, cos, tan, arcsin, arccos, arctan, arctan2,
         hypot, sinh, cosh, tanh, arcsinh, arccosh, arctanh, deg2rad, rad2deg,

dask/array/core.py

@@ -1142,6 +1142,34 @@ def from_array(x, chunks, name=None, lock=False, **kwargs):
     return Array(merge({name: x}, dsk), name, chunks, dtype=x.dtype)
 
 
+def from_func(func, shape, dtype=None, name=None, args=(), kwargs={}):
+    """ Create dask array in a single block by calling a function
+
+    Calling the provided function with func(*args, **kwargs) should return a
+    NumPy array of the indicated shape and dtype.
+
+    Example
+    -------
+
+    >>> a = from_func(np.arange, (3,), np.int64, args=(3,))
+    >>> a.compute()
+    array([0, 1, 2])
+
+    This works particularly well when coupled with dask.array functions like
+    concatenate and stack:
+
+    >>> arrays = [from_func(np.array, (), args=(n,)) for n in range(5)]
+    >>> stack(arrays).compute()
+    array([0, 1, 2, 3, 4])
+    """
+    if args or kwargs:
+        func = partial(func, *args, **kwargs)
+    name = name or next(names)
+    dsk = {(name,) + (0,) * len(shape): (func,)}
+    chunks = tuple((i,) for i in shape)
+    return Array(dsk, name, chunks, dtype)
+
+
 def atop(func, out_ind, *args, **kwargs):
     """ Tensor operation: Generalized inner and outer products
 

dask/array/tests/test_array_core.py

@@ -960,6 +960,15 @@ def test_from_array_with_lock():
     assert eq(e + f, x + x)
 
 
+def test_from_func():
+    x = np.arange(10)
+    d = from_func(lambda n: n * x, (10,), np.int64, kwargs={'n': 2})
+
+    assert d.shape == x.shape
+    assert d.dtype == x.dtype
+    assert eq(d.compute(), 2 * x)
+
+
 def test_topk():
     x = np.array([5, 2, 1, 6])
     d = da.from_array(x, chunks=2)