Repository: dask/dask

Issue:
ValueError: Can't specify drop_axis and new_axis together
I am using `map_blocks` with a function that analyzes a 3-dimensional array (a big time series of images)
and returns an array where the first dimension is dropped and two new dimensions are added, e.g. `tyx -> yxab`.

The following example works with dask 0.11.1 but fails with the current version 0.25.2, raising a ValueError:

```Python
import numpy as np
import dask.array as da


def function(image_timeseries, radius, angles, bins):
    """Return new array: drop first dimension, append 2 new dimensions."""
    times, length, width = image_timeseries.shape
    return np.zeros(shape=(length-2*radius, width-2*radius, angles, bins),
                    dtype='float32')


def map_blocks(function, image_timeseries, radius=6, angles=32, bins=32,
               chunk=32):
    """Run function on small overlapping blocks of image timeseries."""
    times, length, width = image_timeseries.shape

    # create dask ghost array of overlapping blocks
    data = da.from_array(image_timeseries, chunks=(times, chunk, chunk))
    data = da.ghost.ghost(data, depth=(0, radius, radius), boundary=(0, 1, 1))

    result = da.core.map_blocks(
        function,
        data,
        dtype='float32',
        chunks=(chunk, chunk, angles, bins),
        drop_axis=0,
        new_axis=(2, 3),
        radius=radius,
        angles=angles,
        bins=bins
        ).compute()

    # trim off the borders
    return result[radius:-radius, radius:-radius]


image_timeseries = np.zeros((2**10, 256, 256), dtype='uint16')
result = map_blocks(function, image_timeseries)
print(result.shape)  # (244, 244, 32, 32)
```

```Python
File "dask_issue.py", line 38, in <module>
  result = map_blocks(function, image_timeseries)
File "dask_issue.py", line 30, in map_blocks
  bins=bins
File "X:\Python36\Lib\site-packages\dask\array\core.py", line 614, in map_blocks
  raise ValueError("Can't specify drop_axis and new_axis together")

builtins.ValueError: Can't specify drop_axis and new_axis together
```


Discussion:
What is your intention with both dropping axes and adding new ones at the same time?
> What is your intention with both dropping axes and adding new ones at the same time?

The intention is to make the `map_blocks` function aware that the callable function removes the first dimension and creates two new dimensions at the end wrt the input array.
That is what I understand from the [documentation](http://dask.pydata.org/en/latest/array-api.html#dask.array.core.map_blocks) and it worked as expected with dask 0.11.1. Am I missing something?

Resolving patch:
diff --git a/dask/array/core.py b/dask/array/core.py
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -513,7 +513,8 @@ def map_blocks(func, *args, **kwargs):
     drop_axis : number or iterable, optional
         Dimensions lost by the function.
     new_axis : number or iterable, optional
-        New dimensions created by the function.
+        New dimensions created by the function. Note that these are applied
+        after ``drop_axis`` (if present).
     name : string, optional
         The key name to use for the array. If not provided, will be determined
         by a hash of the arguments.
@@ -610,9 +611,6 @@ 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)]
     other = [(i, a) for i, a in enumerate(args) if not isinstance(a, Array)]
 
@@ -665,15 +663,22 @@ def map_blocks(func, *args, **kwargs):
                           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]
-    elif new_axis:
+    if new_axis:
+        new_axis = sorted(new_axis)
+        for i in new_axis:
+            if not 0 <= i <= len(numblocks):
+                ndim = len(numblocks)
+                raise ValueError("Can't add axis %d when current "
+                                 "axis are %r. Missing axis: "
+                                 "%r" % (i, list(range(ndim)),
+                                         list(range(ndim, i))))
+            numblocks.insert(i, 1)
         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):
-            numblocks.insert(i, 1)
 
     if chunks:
         if len(chunks) != len(numblocks):
@@ -701,7 +706,7 @@ def map_blocks(func, *args, **kwargs):
                                  "`chunks` kwarg.")
         if drop_axis:
             chunks2 = [c for (i, c) in enumerate(chunks2) if i not in drop_axis]
-        elif new_axis:
+        if new_axis:
             for i in sorted(new_axis):
                 chunks2.insert(i, (1,))
 


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
@@ -2106,9 +2106,9 @@ def test_map_blocks_with_changed_dimension():
     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
+    # Adding axis with a gap
     with pytest.raises(ValueError):
-        d.map_blocks(lambda b: b, drop_axis=1, new_axis=1)
+        d.map_blocks(lambda b: b, new_axis=(3, 4))
 
     d = da.from_array(x, chunks=(4, 8))
     e = d.map_blocks(lambda b: b.sum(axis=1), drop_axis=1, dtype=d.dtype)
@@ -2127,6 +2127,19 @@ def test_map_blocks_with_changed_dimension():
     assert e.chunks == ((1,), (4, 4), (4, 4), (1,))
     assert_eq(e, x[None, :, :, None])
 
+    # Both new_axis and drop_axis
+    d = da.from_array(x, chunks=(8, 4))
+    e = d.map_blocks(lambda b: b.sum(axis=0)[:, None, None],
+                     drop_axis=0, new_axis=(1, 2), dtype=d.dtype)
+    assert e.chunks == ((4, 4), (1,), (1,))
+    assert_eq(e, x.sum(axis=0)[:, None, None])
+
+    d = da.from_array(x, chunks=(4, 8))
+    e = d.map_blocks(lambda b: b.sum(axis=1)[:, None, None],
+                     drop_axis=1, new_axis=(1, 2), dtype=d.dtype)
+    assert e.chunks == ((4, 4), (1,), (1,))
+    assert_eq(e, x.sum(axis=1)[:, None, None])
+
 
 def test_broadcast_chunks():
     assert broadcast_chunks(((5, 5),), ((5, 5),)) == ((5, 5),)