Repository: dask/dask

Issue:
`drop_axis` keyword of `map_blocks()` doesn't work
```
In [5]: dask.__version__
Out[5]: '0.10.0'
```

Trying to use `map_blocks()` while reducing dimension:

```
import numpy as np
import dask.array as da
x = da.random.normal(10, 0.1, size=(20000, 3), chunks=(1000, 3))
def func(x):
    return x[:, 0]
x.map_blocks(func, drop_axis=1).compute()
```

gives `KeyError`:

```
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-1-54ce56daec8f> in <module>()
      4 def func(x):
      5     return np.dot(x, np.arange(3))
----> 6 x.map_blocks(func, drop_axis=1).compute()

~/.virtualenvs/stats/lib/python2.7/site-packages/dask/base.pyc in compute(self, **kwargs)
     84             Extra keywords to forward to the scheduler ``get`` function.
     85         """
---> 86         return compute(self, **kwargs)[0]
     87
     88     @classmethod

~/.virtualenvs/stats/lib/python2.7/site-packages/dask/base.pyc in compute(*args, **kwargs)
    173         dsk = merge([opt(merge([v.dask for v in val]),
    174                          [v._keys() for v in val], **kwargs)
--> 175                     for opt, val in groups.items()])
    176     else:
    177         dsk = merge(var.dask for var in variables)

~/.virtualenvs/stats/lib/python2.7/site-packages/dask/array/optimization.pyc in optimize(dsk, keys, **kwargs)
     20     fast_functions = kwargs.get('fast_functions',
     21                                 set([getarray, np.transpose]))
---> 22     dsk2, dependencies = cull(dsk, keys)
     23     dsk4, dependencies = fuse(dsk2, keys, dependencies)
     24     dsk5 = optimize_slices(dsk4)

~/.virtualenvs/stats/lib/python2.7/site-packages/dask/optimize.pyc in cull(dsk, keys)
     43     while stack:
     44         key = stack.pop()
---> 45         out[key] = dsk[key]
     46         deps = get_dependencies(dsk, key, as_list=True)  # fuse needs lists
     47         dependencies[key] = deps

KeyError: ('func-3b07b50eab1c1e07dce5009fbcf46bd9', 6, 0)
```



Discussion:
For advanced uses of `map_blocks` you'll be happier if you supply the chunkshape of your output blocks explicitly.

``` python
In [4]: import numpy as np
   ...: import dask.array as da
   ...: x = da.random.normal(10, 0.1, size=(20000, 3), chunks=(1000, 3))
   ...: def func(x):
   ...:     return x[:, 0]
   ...: x.map_blocks(func, drop_axis=1, chunks=(1000,)).compute()
   ...: 
   ...: 
Out[4]: 
array([ 10.3135116 ,   9.87219944,   9.87745014, ...,   9.96419933,
        10.15943915,   9.95456686])
```

Of course, it would be nice to receive a nicer error message here or perhaps have it try to figure it out for you.  Solutions to this would be most welcome.

This works! Thank you @mrocklin. I think it should be easy for dask would figure it out the default chunkshape for me (at least for the `drop_axis` case)? 

Glad to hear that things are working out well.

I agree that in this case it's pretty clear what to do.  Sometimes it's tricky to make rules that are generally applicable in all cases but this is probably doable.  

The code for `map_blocks` and `drop_axis` might be approachable if you're interested in contributing: https://github.com/dask/dask/blob/master/dask/array/core.py#L507-L581


Resolving patch:
diff --git a/dask/array/core.py b/dask/array/core.py
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -17,7 +17,7 @@
 
 from toolz.curried import (pipe, partition, concat, pluck, join, first,
                            memoize, map, groupby, valmap, accumulate, merge,
-                           reduce, interleave, sliding_window)
+                           reduce, interleave, sliding_window, assoc)
 import numpy as np
 
 from . import chunk
@@ -447,7 +447,7 @@ def map_blocks(func, *args, **kwargs):
     dtype: np.dtype
         Datatype of resulting array
     chunks: tuple (optional)
-        chunk shape of resulting blocks if the function does not preserve shape
+        Chunk shape of resulting blocks if the function does not preserve shape
     drop_axis: number or iterable (optional)
         Dimensions lost by the function
     new_axis: number or iterable (optional)
@@ -510,7 +510,7 @@ def map_blocks(func, *args, **kwargs):
     >>> y.numblocks
     (10,)
 
-    If these must match (up to broadcasting rules) then we can map arbitrary
+    If these match (up to broadcasting rules) then we can map arbitrary
     functions across blocks
 
     >>> def func(a, b):
@@ -551,63 +551,95 @@ def map_blocks(func, *args, **kwargs):
     if isinstance(new_axis, Number):
         new_axis = [new_axis]
 
+    if drop_axis and new_axis:
+        raise ValueError("Can't specify drop_axis and new_axis together")
+
     arrs = [a for a in args if isinstance(a, Array)]
     args = [(i, a) for i, a in enumerate(args) if not isinstance(a, Array)]
 
-    arginds = [(a, tuple(range(a.ndim))[::-1]) for a in arrs]
-
-    numblocks = dict([(a.name, a.numblocks) for a, _ in arginds])
-    argindsstr = list(concat([(a.name, ind) for a, ind in arginds]))
+    argpairs = [(a.name, tuple(range(a.ndim))[::-1]) for a in arrs]
+    numblocks = {a.name: a.numblocks for a in arrs}
+    arginds = list(concat(argpairs))
     out_ind = tuple(range(max(a.ndim for a in arrs)))[::-1]
 
+    try:
+        spec = getargspec(func)
+        block_id = ('block_id' in spec.args or
+                    'block_id' in getattr(spec, 'kwonly_args', ()))
+    except:
+        block_id = False
+
+    if block_id:
+        kwargs['block_id'] = '__dummy__'
+
     if args:
-        dsk = top(partial_by_order, name, out_ind, *argindsstr,
+        dsk = top(partial_by_order, name, out_ind, *arginds,
                   numblocks=numblocks, function=func, other=args,
                   **kwargs)
     else:
-        dsk = top(func, name, out_ind, *argindsstr, numblocks=numblocks,
+        dsk = top(func, name, out_ind, *arginds, numblocks=numblocks,
                   **kwargs)
 
-    # If func has block_id as an argument then swap out func
-    # for func with block_id partialed in
-    try:
-        spec = getargspec(func)
-    except:
-        spec = None
-    if spec:
-        args = spec.args
-        try:
-            args += spec.kwonlyargs
-        except AttributeError:
-            pass
-        if 'block_id' in args:
-            for k in dsk.keys():
-                dsk[k] = (partial(func, block_id=k[1:]),) + dsk[k][1:]
+    # If func has block_id as an argument, add it to the kwargs for each call
+    if block_id:
+        for k in dsk.keys():
+            dsk[k] = dsk[k][:-1] + (assoc(dsk[k][-1], 'block_id', k[1:]),)
 
-    numblocks = list(arrs[0].numblocks)
+    if len(arrs) == 1:
+        numblocks = list(arrs[0].numblocks)
+    else:
+        dims = broadcast_dimensions(argpairs, numblocks)
+        numblocks = [b for (_, b) in reversed(list(dims.items()))]
 
     if drop_axis:
+        if any(numblocks[i] > 1 for i in drop_axis):
+            raise ValueError("Can't drop an axis with more than 1 block. "
+                             "Please use `atop` instead.")
         dsk = dict((tuple(k for i, k in enumerate(k)
                           if i - 1 not in drop_axis), v)
                    for k, v in dsk.items())
         numblocks = [n for i, n in enumerate(numblocks) if i not in drop_axis]
-
-    if new_axis:
+    elif new_axis:
         dsk, old_dsk = dict(), dsk
         for key in old_dsk:
             new_key = list(key)
             for i in new_axis:
                 new_key.insert(i + 1, 0)
             dsk[tuple(new_key)] = old_dsk[key]
-        for i in sorted(new_axis, reverse=False):
+        for i in sorted(new_axis):
             numblocks.insert(i, 1)
 
-    if chunks is not None and chunks and not isinstance(chunks[0], tuple):
-        chunks = [nb * (bs,) for nb, bs in zip(numblocks, chunks)]
-    if chunks is not None:
-        chunks = tuple(chunks)
+    if chunks:
+        if len(chunks) != len(numblocks):
+            raise ValueError("Provided chunks have {0} dims, expected {1} "
+                             "dims.".format(len(chunks), len(numblocks)))
+        chunks2 = []
+        for i, (c, nb) in enumerate(zip(chunks, numblocks)):
+            if isinstance(c, tuple):
+                if not len(c) == nb:
+                    raise ValueError("Dimension {0} has {1} blocks, "
+                                     "chunks specified with "
+                                     "{2} blocks".format(i, nb, len(c)))
+                chunks2.append(c)
+            else:
+                chunks2.append(nb * (c,))
     else:
-        chunks = broadcast_chunks(*[a.chunks for a in arrs])
+        if len(arrs) == 1:
+            chunks2 = list(arrs[0].chunks)
+        else:
+            try:
+                chunks2 = list(broadcast_chunks(*[a.chunks for a in arrs]))
+            except:
+                raise ValueError("Arrays in `map_blocks` don't align, can't "
+                                 "infer output chunks. Please provide "
+                                 "`chunks` kwarg.")
+        if drop_axis:
+            chunks2 = [c for (i, c) in enumerate(chunks2) if i not in drop_axis]
+        elif new_axis:
+            for i in sorted(new_axis):
+                chunks2.insert(i, (1,))
+
+    chunks = tuple(chunks2)
 
     return Array(merge(dsk, *[a.dask for a in arrs]), name, chunks, dtype)
 


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
@@ -780,8 +780,8 @@ def test_map_blocks2():
     x = np.arange(10, dtype='i8')
     d = from_array(x, chunks=(2,))
 
-    def func(block, block_id=None):
-        return np.ones_like(block) * sum(block_id)
+    def func(block, block_id=None, c=0):
+        return np.ones_like(block) * sum(block_id) + c
 
     out = d.map_blocks(func, dtype='i8')
     expected = np.array([0, 0, 1, 1, 2, 2, 3, 3, 4, 4], dtype='i8')
@@ -789,6 +789,12 @@ def func(block, block_id=None):
     assert_eq(out, expected)
     assert same_keys(d.map_blocks(func, dtype='i8'), out)
 
+    out = d.map_blocks(func, dtype='i8', c=1)
+    expected = expected + 1
+
+    assert_eq(out, expected)
+    assert same_keys(d.map_blocks(func, dtype='i8', c=1), out)
+
 
 def test_map_blocks_with_constants():
     d = da.arange(10, chunks=3)
@@ -1852,15 +1858,38 @@ def test_map_blocks_with_changed_dimension():
 
     e = d.map_blocks(lambda b: b.sum(axis=0), chunks=(4,), drop_axis=0,
                      dtype=d.dtype)
-    assert e.ndim == 1
     assert e.chunks == ((4, 4),)
     assert_eq(e, x.sum(axis=0))
 
+    # Provided chunks have wrong shape
+    with pytest.raises(ValueError):
+        d.map_blocks(lambda b: b.sum(axis=0), chunks=(7, 4), drop_axis=0)
+
+    with pytest.raises(ValueError):
+        d.map_blocks(lambda b: b.sum(axis=0), chunks=((4, 4, 4),), drop_axis=0)
+
+    # Can't drop axis with more than 1 block
+    with pytest.raises(ValueError):
+        d.map_blocks(lambda b: b.sum(axis=1), drop_axis=1, dtype=d.dtype)
+
+    # Can't use both drop_axis and new_axis
+    with pytest.raises(ValueError):
+        d.map_blocks(lambda b: b, drop_axis=1, new_axis=1)
+
+    d = da.from_array(x, chunks=(4, 8))
+    e = d.map_blocks(lambda b: b.sum(axis=1), drop_axis=1, dtype=d.dtype)
+    assert e.chunks == ((4, 3),)
+    assert_eq(e, x.sum(axis=1))
+
     x = np.arange(64).reshape((8, 8))
     d = da.from_array(x, chunks=(4, 4))
     e = d.map_blocks(lambda b: b[None, :, :, None],
                      chunks=(1, 4, 4, 1), new_axis=[0, 3], dtype=d.dtype)
-    assert e.ndim == 4
+    assert e.chunks == ((1,), (4, 4), (4, 4), (1,))
+    assert_eq(e, x[None, :, :, None])
+
+    e = d.map_blocks(lambda b: b[None, :, :, None],
+                     new_axis=[0, 3], dtype=d.dtype)
     assert e.chunks == ((1,), (4, 4), (4, 4), (1,))
     assert_eq(e, x[None, :, :, None])