Repository: dask/dask

Issue:
Slicing with lists in multiple axes
From the From the [dask docs](http://dask.pydata.org/en/latest/slicing.html):

> Dask.array supports most of the NumPy slicing syntax.
> ...
> It does not currently support the following:
> 
>    Slicing one dask.array with another x[x > 0]
>    **Slicing with lists in multiple axes x[[1, 2, 3], [3, 2, 1]]**
> 
> Both of these are straightforward to add though. If you have a use case then raise an issue.

Here's that issue.

My use case is for point-wise indexing in xray: https://github.com/xray/xray/issues/475

A simple use case using dask arrays:

``` python
x = da.ones((10, 100), chunks=(10, 10))
points = x[[1, 2, 3], [3, 2, 1]]
```

currently raises this error:

``` python
NotImplementedError: Don't yet support nd fancy indexing
```

cc @shoyer 



Discussion:
Am I correct in assuming that in this case you'd like the shape of the output to be `(3,)` and not `(3, 3)`?

And if so is it important to you that this be the default array slicing syntax or are you ok with some other custom method?

Correct.  I'm looking for the default numpy slicing _behavior_, how we get there is up for discussion.  My preference would be to use the numpy slicing syntax, but could use a `take` method of some kind if that is substantially easier.

They're both the same level of difficulty to accomplish.  I mostly want to avoid locking dask into one behavior or the other for as long as possible.

NumPy has been discussing adding `vindex` and `oindex` attributes for explicit vectorized vs. outer/orthogonal indexing (see @seberg's [NEP](https://gist.github.com/seberg/976373b6a2b7c4188591) and [PR](https://github.com/numpy/numpy/pull/6075)).

I suggest that dask should follow NumPy's lead here, and consider implementing both attributes -- even if it will only support a limited subset of vectorized indexing. Outer indexing is easier to reason about and optimize (dask already has all the necessary functionality), and it will be useful to have explicit syntax both to support better fusing of getitem calls and to make things simpler for downstream libraries like xray (which uses outer indexing in `__getitem__`).

Dask _certainly_ should not try to exactly replicate NumPy's current indexing behavior, which sometimes but not always reorders axes with array indices, e.g.,

```
In [14]: x = np.zeros((5, 6, 7, 8))

In [15]: x[:, [0, 1], [0, 1]].shape
Out[15]: (5, 2, 8)

In [16]: x[:, [0, 1], :, [0, 1]].shape
Out[16]: (2, 5, 7)
```


Resolving patch:
diff --git a/dask/array/core.py b/dask/array/core.py
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -19,7 +19,8 @@
 from . import chunk
 from .slicing import slice_array
 from . import numpy_compat
-from ..utils import deepmap, ignoring, repr_long_list, concrete, is_integer
+from ..utils import (deepmap, ignoring, repr_long_list, concrete, is_integer,
+        IndexCallable)
 from ..compatibility import unicode, long
 from .. import threaded, core
 from ..context import _globals
@@ -835,6 +836,33 @@ def __getitem__(self, index):
 
         return Array(merge(self.dask, dsk), out, chunks, dtype=self._dtype)
 
+    def _vindex(self, key):
+        if (not isinstance(key, tuple) or
+           not len([k for k in key if isinstance(k, (np.ndarray, list))]) >= 2 or
+           not all(isinstance(k, (np.ndarray, list)) or k == slice(None, None)
+                   for k in key)):
+            raise IndexError("vindex expects only lists and full slices\n"
+            "At least two entries must be a list\n"
+            "For other combinations try doing normal slicing first, followed\n"
+            "by vindex slicing.  Got: \n\t%s" % str(key))
+        if any((isinstance(k, np.ndarray) and k.ndim != 1) or
+               (isinstance(k, list) and k and isinstance(k[0], list))
+               for k in key):
+            raise IndexError("vindex does not support multi-dimensional keys\n"
+                    "Got: %s" % str(key))
+        if len(set(len(k) for k in key if isinstance(k, (list, np.ndarray)))) != 1:
+            raise IndexError("All indexers must have the same length, got\n"
+                    "\t%s" % str(key))
+        key = key + (slice(None, None),) * (self.ndim - len(key))
+        key = [i if isinstance(i, list) else
+               i.tolist() if isinstance(i, np.ndarray) else
+               None for i in key]
+        return _vindex(self, *key)
+
+    @property
+    def vindex(self):
+        return IndexCallable(self._vindex)
+
     @wraps(np.dot)
     def dot(self, other):
         return tensordot(self, other, axes=((self.ndim-1,), (other.ndim-2,)))
@@ -2105,6 +2133,8 @@ def chunks_from_arrays(arrays):
     >>> chunks_from_arrays([1, 1])
     ((1, 1),)
     """
+    if not arrays:
+        return ()
     result = []
     dim = 0
 
@@ -2136,6 +2166,8 @@ def deepfirst(seq):
 def ndimlist(seq):
     if not isinstance(seq, (list, tuple)):
         return 0
+    elif not seq:
+        return 1
     else:
         return 1 + ndimlist(seq[0])
 
@@ -2161,6 +2193,8 @@ def concatenate3(arrays):
     ndim = ndimlist(arrays)
     if not ndim:
         return arrays
+    if not arrays:
+        return np.empty(())
     chunks = chunks_from_arrays(arrays)
     shape = tuple(map(sum, chunks))
 
@@ -2222,3 +2256,163 @@ def to_hdf5(filename, *args, **kwargs):
                                         **kwargs)
                     for dp, x in data.items()]
         store(list(data.values()), dsets)
+
+
+def interleave_none(a, b):
+    """
+
+    >>> interleave_none([0, None, 2, None], [1, 3])
+    (0, 1, 2, 3)
+    """
+    result = []
+    i = j = 0
+    n = len(a) + len(b)
+    while i + j < n:
+        if a[i] is not None:
+            result.append(a[i])
+            i += 1
+        else:
+            result.append(b[j])
+            i += 1
+            j += 1
+    return tuple(result)
+
+
+def keyname(name, i, okey):
+    """
+
+    >>> keyname('x', 3, [None, None, 0, 2])
+    ('x', 3, 0, 2)
+    """
+    return (name, i) + tuple(k for k in okey if k is not None)
+
+
+def _vindex(x, *indexes):
+    """ Point wise slicing
+
+    This is equivalent to numpy slicing with multiple input lists
+
+    >>> x = np.arange(56).reshape((7, 8))
+    >>> x
+    array([[ 0,  1,  2,  3,  4,  5,  6,  7],
+           [ 8,  9, 10, 11, 12, 13, 14, 15],
+           [16, 17, 18, 19, 20, 21, 22, 23],
+           [24, 25, 26, 27, 28, 29, 30, 31],
+           [32, 33, 34, 35, 36, 37, 38, 39],
+           [40, 41, 42, 43, 44, 45, 46, 47],
+           [48, 49, 50, 51, 52, 53, 54, 55]])
+
+    >>> d = from_array(x, chunks=(3, 4))
+    >>> result = _vindex(d, [0, 1, 6, 0], [0, 1, 0, 7])
+    >>> result.compute()
+    array([ 0,  9, 48,  7])
+    """
+    indexes = [list(index) if index is not None else index for index in indexes]
+    bounds = [list(accumulate(add, (0,) + c)) for c in x.chunks]
+    bounds2 = [b for i, b in zip(indexes, bounds) if i is not None]
+    axis = _get_axis(indexes)
+
+    points = list()
+    for i, idx in enumerate(zip(*[i for i in indexes if i is not None])):
+        block_idx = [np.searchsorted(b, ind, 'right') - 1
+                     for b, ind in zip(bounds2, idx)]
+        inblock_idx = [ind - bounds2[k][j]
+                    for k, (ind, j) in enumerate(zip(idx, block_idx))]
+        points.append((i, tuple(block_idx), tuple(inblock_idx)))
+
+    per_block = groupby(1, points)
+    per_block = dict((k, v) for k, v in per_block.items() if v)
+
+    other_blocks = list(product(*[list(range(len(c))) if i is None else [None]
+                                for i, c in zip(indexes, x.chunks)]))
+
+    token = next(tokens)
+    name = 'vindex-slice' + token
+
+    full_slices = [slice(None, None) if i is None else None for i in indexes]
+
+    dsk = dict((keyname(name, i, okey),
+                (_vindex_transpose,
+                  (_vindex_slice, (x.name,) + interleave_none(okey, key),
+                     interleave_none(full_slices, list(zip(*pluck(2, per_block[key]))))),
+                  axis))
+                for i, key in enumerate(per_block)
+                for okey in other_blocks)
+
+    if per_block:
+        dsk2 = dict((keyname('vindex-merge' + token, 0, okey),
+                     (_vindex_merge,
+                       [list(pluck(0, per_block[key])) for key in per_block],
+                       [keyname(name, i, okey) for i in range(len(per_block))]))
+                     for okey in other_blocks)
+    else:
+        dsk2 = dict()
+
+    chunks = [c for i, c in zip(indexes, x.chunks) if i is None]
+    chunks.insert(0, (len(points),) if points else ())
+    chunks = tuple(chunks)
+
+    return Array(merge(x.dask, dsk, dsk2), 'vindex-merge' + token, chunks, x.dtype)
+
+
+def _get_axis(indexes):
+    """ Get axis along which point-wise slicing results lie
+
+    This is mostly a hack because I can't figure out NumPy's rule on this and
+    can't be bothered to go reading.
+
+    >>> _get_axis([[1, 2], None, [1, 2], None])
+    0
+    >>> _get_axis([None, [1, 2], [1, 2], None])
+    1
+    >>> _get_axis([None, None, [1, 2], [1, 2]])
+    2
+    """
+    ndim = len(indexes)
+    indexes = [slice(None, None) if i is None else [0] for i in indexes]
+    x = np.empty((2,) * ndim)
+    x2 = x[tuple(indexes)]
+    return x2.shape.index(1)
+
+
+def _vindex_slice(block, points):
+    """ Pull out point-wise slices from block """
+    points = [p if isinstance(p, slice) else list(p) for p in points]
+    return block[tuple(points)]
+
+def _vindex_transpose(block, axis):
+    """ Rotate block so that points are on the first dimension """
+    axes = [axis] + list(range(axis)) + list(range(axis + 1, block.ndim))
+    return block.transpose(axes)
+
+def _vindex_merge(locations, values):
+    """
+
+    >>> locations = [0], [2, 1]
+    >>> values = [np.array([[1, 2, 3]]),
+    ...           np.array([[10, 20, 30], [40, 50, 60]])]
+
+    >>> _vindex_merge(locations, values)
+    array([[ 1,  2,  3],
+           [40, 50, 60],
+           [10, 20, 30]])
+    """
+    locations = list(map(list, locations))
+    values = list(values)
+
+    n = sum(map(len, locations))
+
+    shape = list(values[0].shape)
+    shape[0] = n
+    shape = tuple(shape)
+
+    dtype = values[0].dtype
+
+    x = np.empty(shape, dtype=dtype)
+
+    ind = [slice(None, None) for i in range(x.ndim)]
+    for loc, val in zip(locations, values):
+        ind[0] = loc
+        x[tuple(ind)] = val
+
+    return x
diff --git a/dask/array/slicing.py b/dask/array/slicing.py
--- a/dask/array/slicing.py
+++ b/dask/array/slicing.py
@@ -1,5 +1,5 @@
 from itertools import count, product
-from toolz import merge, first, accumulate
+from toolz import merge, first, accumulate, groupby, pluck
 from operator import getitem, add
 from math import ceil
 from ..compatibility import long


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
@@ -1220,3 +1220,69 @@ def test_h5py_newaxis():
 
 def test_ellipsis_slicing():
     assert eq(da.ones(4, chunks=2)[...], np.ones(4))
+
+
+def test_point_slicing():
+    x = np.arange(56).reshape((7, 8))
+    d = da.from_array(x, chunks=(3, 4))
+
+    result = d.vindex[[1, 2, 5, 5], [3, 1, 6, 1]]
+    assert eq(result, x[[1, 2, 5, 5], [3, 1, 6, 1]])
+
+    result = d.vindex[[0, 1, 6, 0], [0, 1, 0, 7]]
+    assert eq(result, x[[0, 1, 6, 0], [0, 1, 0, 7]])
+
+
+def test_point_slicing_with_full_slice():
+    from dask.array.core import _vindex_transpose, _get_axis
+    x = np.arange(4*5*6*7).reshape((4, 5, 6, 7))
+    d = da.from_array(x, chunks=(2, 3, 3, 4))
+
+    inds = [
+            [[1, 2, 3], None, [3, 2, 1], [5, 3, 4]],
+            [[1, 2, 3], None, [4, 3, 2], None],
+            [[1, 2, 3], [3, 2, 1]],
+            [[1, 2, 3], [3, 2, 1], [3, 2, 1], [5, 3, 4]],
+            [[], [], [], None],
+            [np.array([1, 2, 3]), None, np.array([4, 3, 2]), None],
+            [None, None, [1, 2, 3], [4, 3, 2]],
+            [None, [0, 2, 3], None, [0, 3, 2]],
+            ]
+
+    for ind in inds:
+        slc = [i if isinstance(i, (np.ndarray, list)) else slice(None, None)
+                for i in ind]
+        result = d.vindex[tuple(slc)]
+
+        # Rotate the expected result accordingly
+        axis = _get_axis(ind)
+        expected = _vindex_transpose(x[tuple(slc)], axis)
+
+        assert eq(result, expected)
+
+        # Always have the first axis be the length of the points
+        k = len(next(i for i in ind if isinstance(i, (np.ndarray, list))))
+        assert result.shape[0] == k
+
+
+def test_vindex_errors():
+    d = da.ones((5, 5, 5), chunks=(3, 3, 3))
+    assert raises(IndexError, lambda: d.vindex[0])
+    assert raises(IndexError, lambda: d.vindex[[1, 2, 3]])
+    assert raises(IndexError, lambda: d.vindex[[1, 2, 3], [1, 2, 3], 0])
+    assert raises(IndexError, lambda: d.vindex[[1], [1, 2, 3]])
+    assert raises(IndexError, lambda: d.vindex[[1, 2, 3], [[1], [2], [3]]])
+
+def test_vindex_merge():
+    from dask.array.core import _vindex_merge
+    locations = [1], [2, 0]
+    values = [np.array([[1, 2, 3]]),
+              np.array([[10, 20, 30], [40, 50, 60]])]
+
+    assert (_vindex_merge(locations, values) == np.array([[40, 50, 60],
+                                                          [1, 2, 3],
+                                                          [10, 20, 30]])).all()
+
+
+def test_empty_array():
+    assert eq(np.arange(0), da.arange(0, chunks=5))