Repository: dask/dask

Issue:
dask.array flatten and reshape
The numpy method `flatten` can be used to reshape an n-dimensional array into a one-dimensional array.

``` python
In [1]: import numpy as np

In [2]: x = np.array([[1, 2, 3], [4, 5, 6]])

In [3]: x
Out[3]: 
array([[1, 2, 3],
       [4, 5, 6]])

In [4]: x.flatten()
Out[4]: array([1, 2, 3, 4, 5, 6])
```

We could add such a method onto `dask.array.Array` if we are willing to give up caring about byte order entirely.  

This might be a good on-ramp issue.  It requires applying the `flatten` method onto each block, creating a new flattened dictionary, manipulating `chunks`, etc..  Furthermore we could use this as a learning opportunity to improve documentation on how to [extend dask array](http://dask.pydata.org/en/latest/array-extend.html).

cc @kwmsmith



Discussion:
The natural generalization of `flatten` is `reshape`, which could also be useful.

At the very least, it would be nice to have support for reshaping _within blocks_. For example, if I have an array with shape `(100, 100, ...)` and blocks `(100, 100, 10)`, I should be able to reshape it to `(10000, ...)`. For example, this would let us generalize point-indexing (https://github.com/ContinuumIO/dask/pull/439) to fully general vectorized indexing.

@cowlicks notes that my use case (reshaping within blocks) could be solved by `map_blocks`. This is almost, but not quite true -- `map_blocks` can't handle changed array dimensions (see #263, #445).


Resolving patch:
diff --git a/dask/array/__init__.py b/dask/array/__init__.py
--- a/dask/array/__init__.py
+++ b/dask/array/__init__.py
@@ -2,9 +2,10 @@
 
 from ..utils import ignoring
 from .core import (Array, stack, concatenate, take, tensordot, transpose,
-        from_array, choose, where, coarsen, insert, broadcast_to, fromfunction,
-        unique, store, squeeze, topk, bincount, histogram, map_blocks, atop,
-        to_hdf5, dot, cov, array, to_npy_stack, from_npy_stack)
+        from_array, choose, where, coarsen, insert, broadcast_to, ravel,
+        reshape, fromfunction, unique, store, squeeze, topk, bincount,
+        histogram, map_blocks, atop, to_hdf5, dot, cov, array, to_npy_stack,
+        from_npy_stack)
 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,
diff --git a/dask/array/core.py b/dask/array/core.py
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -877,6 +877,18 @@ def T(self):
     def transpose(self, axes=None):
         return transpose(self, axes)
 
+    @wraps(np.ravel)
+    def ravel(self):
+        return ravel(self)
+
+    flatten = ravel
+
+    @wraps(np.reshape)
+    def reshape(self, *shape):
+        if len(shape) == 1 and not isinstance(shape[0], Number):
+            shape = shape[0]
+        return reshape(self, shape)
+
     @wraps(topk)
     def topk(self, k):
         return topk(k, self)
@@ -2095,6 +2107,65 @@ def broadcast_to(x, shape):
     return Array(merge(dsk, x.dask), name, chunks, dtype=x.dtype)
 
 
+@wraps(np.ravel)
+def ravel(array):
+    if array.ndim == 0:
+        return array[None]
+    elif array.ndim == 1:
+        return array
+    elif all(len(c) == 1 for c in array.chunks[1:]):
+        # we can simply map np.ravel over the chunks
+        name = 'ravel-' + tokenize(array)
+        trailing_size = int(np.prod([c[0] for c in array.chunks[1:]]))
+        chunks = (tuple(c * trailing_size for c in array.chunks[0]),)
+        dsk = dict(((name, key[1]), (np.ravel, key))
+                   for key in core.flatten(array._keys()))
+        return Array(merge(dsk, array.dask), name, chunks, dtype=array.dtype)
+    else:
+        # we need to do an expensive shuffling of the data
+        return concatenate([ravel(a) for a in array])
+
+
+def unravel(array, shape):
+    """ Given a 1D array, reshape it to the given shape
+    """
+    assert array.ndim == 1
+    if len(shape) == 1:
+        return array
+    else:
+        trailing_size = int(np.prod(shape[1:]))
+        if all(c % trailing_size == 0 for c in array.chunks[0]):
+            # we can call np.reshape on each chunk
+            name = 'unravel-' + tokenize(array)
+            chunks = ((tuple(c // trailing_size for c in array.chunks[0]),)
+                      + tuple((c,) for c in shape[1:]))
+            dsk = dict(((name, key[1]) + (0,) * (len(shape) - 1),
+                        (np.reshape, key, (c,) + shape[1:]))
+                       for key, c in zip(array._keys(), chunks[0]))
+            return Array(merge(dsk, array.dask), name, chunks, dtype=array.dtype)
+        else:
+            # we need to shuffle
+            # nb. this doesn't always work, stack requires aligned chunks
+            return stack([unravel(array[n * trailing_size
+                                        : (n + 1) * trailing_size], shape[1:])
+                          for n in range(shape[0])])
+
+
+@wraps(np.reshape)
+def reshape(array, shape):
+    shape = tuple(shape)
+    known_sizes = [s for s in shape if s != -1]
+    if len(known_sizes) < len(shape):
+        if len(known_sizes) - len(shape) > 1:
+            raise ValueError('can only specify one unknown dimension')
+        missing_size = int(array.size / np.prod(known_sizes))
+        shape = tuple(missing_size if s == -1 else s for s in shape)
+
+    if np.prod(shape) != array.size:
+        raise ValueError('total size of new array must be unchanged')
+    return unravel(ravel(array), shape)
+
+
 def offset_func(func, offset, *args):
     """  Offsets inputs by offset
 


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -579,6 +579,97 @@ def test_broadcast_to():
     assert raises(ValueError, lambda: broadcast_to(a, (3,)))
 
 
+def test_ravel():
+    x = np.random.randint(10, size=(4, 6))
+
+    # 2d
+    # these should use the shortcut
+    for chunks in [(4, 6), (2, 6)]:
+        a = from_array(x, chunks=chunks)
+        assert eq(x.ravel(), a.ravel())
+        assert len(a.ravel().dask) == len(a.dask) + len(a.chunks[0])
+    # these cannot
+    for chunks in [(4, 2), (2, 2)]:
+        a = from_array(x, chunks=chunks)
+        assert eq(x.ravel(), a.ravel())
+        assert len(a.ravel().dask) > len(a.dask) + len(a.chunks[0])
+
+    # 0d
+    assert eq(x[0, 0].ravel(), a[0, 0].ravel())
+
+    # 1d
+    a_flat = a.ravel()
+    assert a_flat.ravel() is a_flat
+
+    # 3d
+    x = np.random.randint(10, size=(2, 3, 4))
+    for chunks in [2, 4, (2, 3, 2), (1, 3, 4)]:
+        a = from_array(x, chunks=chunks)
+        assert eq(x.ravel(), a.ravel())
+
+    assert eq(x.flatten(), a.flatten())
+    assert eq(np.ravel(x), da.ravel(a))
+
+
+def test_unravel():
+    x = np.random.randint(10, size=24)
+
+    # these should use the shortcut
+    for chunks, shape in [(24, (3, 8)),
+                          (24, (12, 2)),
+                          (6, (4, 6)),
+                          (6, (4, 3, 2)),
+                          (6, (4, 6, 1)),
+                          (((6, 12, 6),), (4, 6))]:
+        a = from_array(x, chunks=chunks)
+        unraveled = unravel(a, shape)
+        assert eq(x.reshape(*shape), unraveled)
+        assert len(unraveled.dask) == len(a.dask) + len(a.chunks[0])
+
+    # these cannot
+    for chunks, shape in [(6, (2, 12)),
+                          (6, (1, 4, 6)),
+                          (6, (2, 1, 12))]:
+        a = from_array(x, chunks=chunks)
+        unraveled = unravel(a, shape)
+        assert eq(x.reshape(*shape), unraveled)
+        assert len(unraveled.dask) > len(a.dask) + len(a.chunks[0])
+
+    assert raises(AssertionError, lambda: unravel(unraveled, (3, 8)))
+    assert unravel(a, a.shape) is a
+
+
+def test_reshape():
+    shapes = [(24,), (2, 12), (2, 3, 4)]
+    for original_shape in shapes:
+        for new_shape in shapes:
+            for chunks in [2, 4, 12]:
+                x = np.random.randint(10, size=original_shape)
+                a = from_array(x, chunks)
+                assert eq(x.reshape(new_shape), a.reshape(new_shape))
+
+    assert raises(ValueError, lambda: reshape(a, (100,)))
+    assert eq(x.reshape(*new_shape), a.reshape(*new_shape))
+    assert eq(np.reshape(x, new_shape), reshape(a, new_shape))
+
+    # verify we can reshape a single chunk array without too many tasks
+    x = np.random.randint(10, size=(10, 20))
+    a = from_array(x, 20)  # all one chunk
+    reshaped = a.reshape((20, 10))
+    assert eq(x.reshape((20, 10)), reshaped)
+    assert len(reshaped.dask) == len(a.dask) + 2
+
+
+def test_reshape_unknown_dimensions():
+    for original_shape in [(24,), (2, 12), (2, 3, 4)]:
+        for new_shape in [(-1,), (2, -1), (-1, 3, 4)]:
+            x = np.random.randint(10, size=original_shape)
+            a = from_array(x, 4)
+            assert eq(x.reshape(new_shape), a.reshape(new_shape))
+
+    assert raises(ValueError, lambda: reshape(a, (-1, -1)))
+
+
 def test_full():
     d = da.full((3, 4), 2, chunks=((2, 1), (2, 2)))
     assert d.chunks == ((2, 1), (2, 2))