Repository: dask/dask

Issue:
Risk attribution: a[argsort(b)[:k]]
# Abstract problem
I have ``y = f(x1, x2, ...)``, where ``x1, x2...`` are 1-dimensional arrays with the same shape, and ``f`` is a generic elementwise, embarassingly parallel function (e.g. ``sum()``). As a consequence, ``y`` has the same shape as the inputs.
I need to find:
1. the top k elements of ``y``
2. the elements of ``x1, x2,...`` which produced the top k elements of ``y``

# Real life use case
``x1, x2...`` are the values of financial assets under a Monte Carlo risk simulation, one point per scenario, and y is the value of the portfolio containing them. I need to find
1. the 99% worst value of y, e.g. where the portfolio loses the most money (_Value at Risk_, or _VaR_)
2. the value of each individual asset building the portfolio where the portfolio is the worst 99% (_risk attribution_) 

In qualitative terms, (2) answers the question: _When you're on the brink of bankrupting, what is causing you to do so? Which of the products that you own are a time bomb that you should sell NOW if you fancy surviving a disaster, e.g. a new global market crisis?_

# Solution
The fist point is readily solvable with topk: 
```
k = int(round(y.size * .01)
-(-y).topk(k)[-1]
```
The second point is not possible in dask to my understanding; in numpy it would be: ``x[argsort(y)[k]]``

There's two things that are missing to make it happen in dask:
1. an ``argtopk`` function, which returns the indexes of the top k elements of an array
2. the ability to slice ``a[b]``, where b is a dask array of integers.

With these, the dask solution to the problem would be ``x[(-y).argtopk(k)[-1]]``

Working on a POC - stay tuned...


Discussion:
Attaching POC.
https://gist.github.com/crusaderky/1947d71be82c9ff9979ab7a5bb9938e2

The code for argtopk looks good to me - should be very quick to integrate in dask.
The slice function is ugly, full of caveats and requires numba in order to be performant - help would be appreciated.
See https://github.com/dask/dask/issues/3395 on the need for a recursive flag
Thanks for the effort here!  I recommend submitting this as a PR.  This
will make it easier for people to investigate, try, and comment on the
code.

On Thu, Apr 12, 2018 at 8:55 AM, crusaderky <notifications@github.com>
wrote:

> See #3395 <https://github.com/dask/dask/issues/3395> on the need for a
> recursive flag
>
> —
> You are receiving this because you are subscribed to this thread.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/dask/issues/3396#issuecomment-380813744>, or mute
> the thread
> <https://github.com/notifications/unsubscribe-auth/AASszCxfoK8GLKA4YmN2on4cP6itLr4pks5tn1zrgaJpZM4TQ-h3>
> .
>

Reimplemented slice function:
- removed numba requirement (without sacrificing speed)
- added support for chunked index
- added support for multiple axes on the sliced array
- still has the limitation that the index must be 1-dimensional

https://gist.github.com/crusaderky/5d25a19159c3e0299c94dd1defff69ae

argtopk can be now found at https://github.com/dask/dask/pull/3405

I found ``dask.array.slicing.slice_with_dask_array`` (which is conveniently invoked from ``Array.__getitem__``) but I can't understand what it's supposed to do?!?
[EDIT] I think it only works for dask arrays of bools, am I correct?

Resolving patch:
diff --git a/dask/array/chunk.py b/dask/array/chunk.py
--- a/dask/array/chunk.py
+++ b/dask/array/chunk.py
@@ -294,3 +294,103 @@ def einsum(*operands, **kwargs):
     # Avoid concatenate=True in atop by adding 1's
     # for the contracted dimensions
     return chunk.reshape(chunk.shape + (1,) * ncontract_inds)
+
+
+def slice_with_int_dask_array(x, idx, offset, x_size, axis):
+    """ Chunk function of `slice_with_int_dask_array_on_axis`.
+    Slice one chunk of x by one chunk of idx.
+
+    Parameters
+    ----------
+    x: ndarray, any dtype, any shape
+        i-th chunk of x
+    idx: ndarray, ndim=1, dtype=any integer
+        j-th chunk of idx (cartesian product with the chunks of x)
+    offset: ndarray, shape=(1, ), dtype=int64
+        Index of the first element along axis of the current chunk of x
+    x_size: int
+        Total size of the x da.Array along axis
+    axis: int
+        normalized axis to take elements from (0 <= axis < x.ndim)
+
+    Returns
+    -------
+    x sliced along axis, using only the elements of idx that fall inside the
+    current chunk.
+    """
+    # Needed when idx is unsigned
+    idx = idx.astype(np.int64)
+
+    # Normalize negative indices
+    idx = np.where(idx < 0, idx + x_size, idx)
+
+    # A chunk of the offset dask Array is a numpy array with shape (1, ).
+    # It indicates the index of the first element along axis of the current
+    # chunk of x.
+    idx = idx - offset
+
+    # Drop elements of idx that do not fall inside the current chunk of x
+    idx_filter = (idx >= 0) & (idx < x.shape[axis])
+    idx = idx[idx_filter]
+
+    # np.take does not support slice indices
+    # return np.take(x, idx, axis)
+    return x[tuple(
+        idx if i == axis else slice(None)
+        for i in range(x.ndim)
+    )]
+
+
+def slice_with_int_dask_array_aggregate(idx, chunk_outputs, x_chunks, axis):
+    """ Final aggregation function of `slice_with_int_dask_array_on_axis`.
+    Aggregate all chunks of x by one chunk of idx, reordering the output of
+    `slice_with_int_dask_array`.
+
+    Note that there is no combine function, as a recursive aggregation (e.g.
+    with split_every) would not give any benefit.
+
+    Parameters
+    ----------
+    idx: ndarray, ndim=1, dtype=any integer
+        j-th chunk of idx
+    chunk_outputs: ndarray
+        concatenation along axis of the outputs of `slice_with_int_dask_array`
+        for all chunks of x and the j-th chunk of idx
+    x_chunks: tuple
+        dask chunks of the x da.Array along axis, e.g. ``(3, 3, 2)``
+    axis: int
+        normalized axis to take elements from (0 <= axis < x.ndim)
+
+    Returns
+    -------
+    Selection from all chunks of x for the j-th chunk of idx, in the correct
+    order
+    """
+    # Needed when idx is unsigned
+    idx = idx.astype(np.int64)
+
+    # Normalize negative indices
+    idx = np.where(idx < 0, idx + sum(x_chunks), idx)
+
+    x_chunk_offset = 0
+    chunk_output_offset = 0
+
+    # Assemble the final index that picks from the output of the previous
+    # kernel by adding together one layer per chunk of x
+    # FIXME: this could probably be reimplemented with a faster search-based
+    # algorithm
+    idx_final = np.zeros_like(idx)
+    for x_chunk in x_chunks:
+        idx_filter = (idx >= x_chunk_offset) & (idx < x_chunk_offset + x_chunk)
+        idx_cum = np.cumsum(idx_filter)
+        idx_final += np.where(idx_filter, idx_cum - 1 + chunk_output_offset, 0)
+        x_chunk_offset += x_chunk
+        if idx_cum.size > 0:
+            chunk_output_offset += idx_cum[-1]
+
+    # np.take does not support slice indices
+    # return np.take(chunk_outputs, idx_final, axis)
+    return chunk_outputs[tuple(
+        idx_final if i == axis else slice(None)
+        for i in range(chunk_outputs.ndim)
+    )]
diff --git a/dask/array/core.py b/dask/array/core.py
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -1324,11 +1324,13 @@ def __getitem__(self, index):
         if not isinstance(index, tuple):
             index = (index,)
 
-        from .slicing import normalize_index, slice_with_dask_array
+        from .slicing import normalize_index, slice_with_int_dask_array, slice_with_bool_dask_array
         index2 = normalize_index(index, self.shape)
 
-        if any(isinstance(i, Array) for i in index2):
-            self, index2 = slice_with_dask_array(self, index2)
+        if any(isinstance(i, Array) and i.dtype.kind in 'iu' for i in index2):
+            self, index2 = slice_with_int_dask_array(self, index2)
+        if any(isinstance(i, Array) and i.dtype == bool for i in index2):
+            self, index2 = slice_with_bool_dask_array(self, index2)
 
         if all(isinstance(i, slice) and i == slice(None) for i in index2):
             return self
diff --git a/dask/array/slicing.py b/dask/array/slicing.py
--- a/dask/array/slicing.py
+++ b/dask/array/slicing.py
@@ -241,7 +241,7 @@ def slice_wrap_lists(out_name, in_name, blockdims, index):
 
     # lists and full slices.  Just use take
     if all(isinstance(i, np.ndarray) or i == slice(None, None, None)
-            for i in index):
+           for i in index):
         axis = where_list[0]
         blockdims2, dsk3 = take(out_name, in_name, blockdims,
                                 index[where_list[0]], axis=axis)
@@ -840,7 +840,125 @@ def check_index(ind, dimension):
         raise IndexError(msg % (ind, dimension))
 
 
-def slice_with_dask_array(x, index):
+def slice_with_int_dask_array(x, index):
+    """ Slice x with at most one 1D dask arrays of ints.
+
+    This is a helper function of :meth:`Array.__getitem__`.
+
+    Parameters
+    ----------
+    x: Array
+    index: tuple with as many elements as x.ndim, among which there are
+           one or more Array's with dtype=int
+
+    Returns
+    -------
+    tuple of (sliced x, new index)
+
+    where the new index is the same as the input, but with slice(None)
+    replaced to the original slicer where a 1D filter has been applied and
+    one less element where a zero-dimensional filter has been applied.
+    """
+    from .core import Array
+
+    assert len(index) == x.ndim
+    fancy_indexes = [
+        isinstance(idx, (tuple, list)) or
+        (isinstance(idx, (np.ndarray, Array)) and idx.ndim > 0)
+        for idx in index
+    ]
+    if sum(fancy_indexes) > 1:
+        raise NotImplementedError("Don't yet support nd fancy indexing)")
+
+    out_index = []
+    dropped_axis_cnt = 0
+    for in_axis, idx in enumerate(index):
+        out_axis = in_axis - dropped_axis_cnt
+        if isinstance(idx, Array) and idx.dtype.kind in 'iu':
+            if idx.ndim == 0:
+                idx = idx[np.newaxis]
+                x = slice_with_int_dask_array_on_axis(x, idx, out_axis)
+                x = x[tuple(
+                    0 if i == out_axis else slice(None)
+                    for i in range(x.ndim)
+                )]
+                dropped_axis_cnt += 1
+            elif idx.ndim == 1:
+                x = slice_with_int_dask_array_on_axis(x, idx, out_axis)
+                out_index.append(slice(None))
+            else:
+                raise NotImplementedError(
+                    "Slicing with dask.array of ints only permitted when "
+                    "the indexer has zero or one dimensions")
+        else:
+            out_index.append(idx)
+    return x, tuple(out_index)
+
+
+def slice_with_int_dask_array_on_axis(x, idx, axis):
+    """ Slice a ND dask array with a 1D dask arrays of ints along the given
+    axis.
+
+    This is a helper function of :func:`slice_with_int_dask_array`.
+    """
+    from .core import Array, atop, from_array
+    from . import chunk
+
+    assert 0 <= axis < x.ndim
+
+    if np.isnan(x.chunks[axis]).any():
+        raise NotImplementedError("Slicing an array with unknown chunks with "
+                                  "a dask.array of ints is not supported")
+
+    # Calculate the offset at which each chunk starts along axis
+    # e.g. chunks=(..., (5, 3, 4), ...) -> offset=[0, 5, 8]
+    offset = np.roll(np.cumsum(x.chunks[axis]), 1)
+    offset[0] = 0
+    offset = from_array(offset, chunks=1)
+    # Tamper with the declared chunks of offset to make atop align it with
+    # x[axis]
+    offset = Array(offset.dask, offset.name, (x.chunks[axis], ), offset.dtype)
+
+    # Define axis labels for atop
+    x_axes = tuple(range(x.ndim))
+    idx_axes = (x.ndim, )  # arbitrary index not already in x_axes
+    offset_axes = (axis, )
+    p_axes = x_axes[:axis + 1] + idx_axes + x_axes[axis + 1:]
+    y_axes = x_axes[:axis] + idx_axes + x_axes[axis + 1:]
+
+    # Calculate the cartesian product of every chunk of x vs every chunk of idx
+    p = atop(chunk.slice_with_int_dask_array,
+             p_axes, x, x_axes, idx, idx_axes, offset, offset_axes,
+             x_size=x.shape[axis], axis=axis, dtype=x.dtype)
+
+    # Aggregate on the chunks of x along axis
+    y = atop(chunk.slice_with_int_dask_array_aggregate,
+             y_axes, idx, idx_axes, p, p_axes,
+             concatenate=True, x_chunks=x.chunks[axis], axis=axis,
+             dtype=x.dtype)
+    return y
+
+
+def slice_with_bool_dask_array(x, index):
+    """ Slice x with one or more dask arrays of bools
+
+    This is a helper function of `Array.__getitem__`.
+
+    Parameters
+    ----------
+    x: Array
+    index: tuple with as many elements as x.ndim, among which there are
+           one or more Array's with dtype=bool
+
+    Returns
+    -------
+    tuple of (sliced x, new index)
+
+    where the new index is the same as the input, but with slice(None)
+    replaced to the original slicer when a filter has been applied.
+
+    Note: The sliced x will have nan chunks on the sliced axes.
+    """
     from .core import Array, atop, elemwise
 
     out_index = [slice(None)
@@ -857,8 +975,8 @@ def slice_with_dask_array(x, index):
                 out_index)
 
     if any(isinstance(ind, Array) and ind.dtype == bool and ind.ndim != 1
-            for ind in index):
-        raise NotImplementedError("Slicing with dask.array only permitted when "
+           for ind in index):
+        raise NotImplementedError("Slicing with dask.array of bools only permitted when "
                                   "the indexer has only one dimension or when "
                                   "it has the same dimension as the sliced "
                                   "array")


Test patch:
diff --git a/dask/array/tests/test_slicing.py b/dask/array/tests/test_slicing.py
--- a/dask/array/tests/test_slicing.py
+++ b/dask/array/tests/test_slicing.py
@@ -502,7 +502,92 @@ def test_oob_check():
         x[0, 0]
 
 
-def test_index_with_dask_array():
+@pytest.mark.parametrize('idx_chunks', [None, 3, 2, 1])
+@pytest.mark.parametrize('x_chunks', [None, (3, 5), (2, 3), (1, 2), (1, 1)])
+def test_index_with_int_dask_array(x_chunks, idx_chunks):
+    # test data is crafted to stress use cases:
+    # - pick from different chunks of x out of order
+    # - a chunk of x contains no matches
+    # - only one chunk of x
+    x = np.array([[10, 20, 30, 40, 50],
+                  [60, 70, 80, 90, 100],
+                  [110, 120, 130, 140, 150]])
+    idx = np.array([3, 0, 1])
+    expect = np.array([[40, 10, 20],
+                       [90, 60, 70],
+                       [140, 110, 120]])
+
+    if x_chunks is not None:
+        x = da.from_array(x, chunks=x_chunks)
+    if idx_chunks is not None:
+        idx = da.from_array(idx, chunks=idx_chunks)
+
+    assert_eq(x[:, idx], expect)
+    assert_eq(x.T[idx, :], expect.T)
+
+
+@pytest.mark.parametrize('chunks', [1, 2, 3])
+def test_index_with_int_dask_array_0d(chunks):
+    # Slice by 0-dimensional array
+    x = da.from_array([[10, 20, 30],
+                       [40, 50, 60]], chunks=chunks)
+    idx0 = da.from_array(1, chunks=1)
+    assert_eq(x[idx0, :], x[1, :])
+    assert_eq(x[:, idx0], x[:, 1])
+
+
+@pytest.mark.parametrize('chunks', [1, 2, 3, 4, 5])
+def test_index_with_int_dask_array_nanchunks(chunks):
+    # Slice by array with nan-sized chunks
+    a = da.arange(-2, 3, chunks=chunks)
+    assert_eq(a[a.nonzero()], np.array([-2, -1,  1,  2]))
+    # Edge case: the nan-sized chunks resolve to size 0
+    a = da.zeros(5, chunks=chunks)
+    assert_eq(a[a.nonzero()], np.array([]))
+
+
+@pytest.mark.parametrize('chunks', [2, 4])
+def test_index_with_int_dask_array_negindex(chunks):
+    a = da.arange(4, chunks=chunks)
+    idx = da.from_array([-1, -4], chunks=1)
+    assert_eq(a[idx], np.array([3, 0]))
+
+
+@pytest.mark.parametrize('chunks', [2, 4])
+def test_index_with_int_dask_array_indexerror(chunks):
+    a = da.arange(4, chunks=chunks)
+    idx = da.from_array([4], chunks=1)
+    with pytest.raises(IndexError):
+        a[idx].compute()
+    idx = da.from_array([-5], chunks=1)
+    with pytest.raises(IndexError):
+        a[idx].compute()
+
+
+@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64',
+                                   'uint8', 'uint16', 'uint32', 'uint64'])
+def test_index_with_int_dask_array_dtypes(dtype):
+    a = da.from_array([10, 20, 30, 40], chunks=-1)
+    idx = da.from_array(np.array([1, 2]).astype(dtype), chunks=1)
+    assert_eq(a[idx], np.array([20, 30]))
+
+
+def test_index_with_int_dask_array_nocompute():
+    """ Test that when the indices are a dask array
+    they are not accidentally computed
+    """
+    def crash():
+        raise NotImplementedError()
+
+    x = da.arange(5, chunks=-1)
+    idx = da.Array({('x', 0): (crash,)}, name='x',
+                   chunks=((2,), ), dtype=np.int64)
+    result = x[idx]
+    with pytest.raises(NotImplementedError):
+        result.compute()
+
+
+def test_index_with_bool_dask_array():
     x = np.arange(36).reshape((6, 6))
     d = da.from_array(x, chunks=(3, 3))
     ind = np.asarray([True, True, False, True, False, False], dtype=bool)
@@ -512,7 +597,7 @@ def test_index_with_dask_array():
         assert_eq(x[x_index], d[index])
 
 
-def test_index_with_dask_array_2():
+def test_index_with_bool_dask_array_2():
     x = np.random.random((10, 10, 10))
     ind = np.random.random(10) > 0.5