add point slicing to dask.array.core by mrocklin · Pull Request #439 · dask/dask

mrocklin · 2015-07-17T16:20:03Z

This is equivalent to numpy slicing with multiple input lists.

We could use a better name. cc @shoyer @jhamman

Example

>>> 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 = isel(d, [0, 1, 6, 0], [0, 1, 0, 7])
>>> result.compute()
array([ 0,  9, 48,  7])

Fixes #433

shoyer · 2015-07-17T16:24:23Z

I think vindex is definitely a better name than isel (which in xray usually actually means outer indexing, not vectorized indexing).

What happens if I index into a higher dimensional array? e.g., along two axes of a 3D array -- where does the new dimension go? This is assuredly the first thing that we'll try to do in xray :).

mrocklin · 2015-07-17T16:27:46Z

That case is definitely not handled with the PR as is. Looks a bit tricker but doable.

jhamman · 2015-07-17T17:54:59Z

vindex is fine by me. Two use cases that I'm interested if we can add here:

First, as @shoyer mentions - indexing a 3+d array with 2 indexers.

The second takes that a bit further - how will we handle the case where we're indexing on the 1st and 3rd dims:

# numpy syntax / behavior
>>> x = np.arange(60).reshape((5, 4, 3))
>>> x[[1, 2, 3], :, [1, 2, 2]]
array([[13, 16, 19, 22],
       [26, 29, 32, 35],
       [38, 41, 44, 47]])

will vindex(x, [1, 2, 3], slice(None), [1, 2, 2]) work?

mrocklin · 2015-07-21T22:10:04Z

Sorry to delay on this. This might require some thought and time which are in unfortunately short supply until next week.

shoyer · 2015-07-28T07:05:42Z

I just realized that because vectorized indexing (if slice objects are omitted) is equivalent to broadcasting, we can probably write the fully general form of vectorized indexing using atop/elemwise:

import dask.array as da
import numpy as np

a = np.arange(25).reshape(5, 5)
x = da.from_array(a, chunks=3)

idx = da.from_array(np.arange(5), chunks=3)

da.core.elemwise(lambda *args: np.asarray(x)[args], idx, idx).compute()
# array([ 0,  6, 12, 18, 24])

shoyer · 2015-07-28T07:26:24Z

On further consideration, vectorized indexing with top seems much harder than made it out to be in my last comment. Time for me to go to bed, I think...

@shoyer

This is equivalent to numpy slicing with multiple input lists. We could use a better name. cc @shoyer @jhamman Example ------- >>> 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 = isel(d, [0, 1, 6, 0], [0, 1, 0, 7]) >>> result.compute() array([ 0, 9, 48, 7]) Fixes dask#433

mrocklin · 2015-07-28T22:08:12Z

OK. This seems to work. It was hard.

I'll implement a vindex getter.

mrocklin · 2015-07-28T22:19:14Z

In [1]: import dask.array as da

In [2]: import numpy as np

In [3]: x = np.arange(10*10*10).reshape((10, 10, 10))

In [4]: d = da.from_array(x, chunks=(4, 5, 6))

In [5]: d.vindex[[0, 5, 9], :, [0, 5, 9]]
Out[5]: dask.array<vindex-merge-2, shape=(3, 10), chunks=((3,), (5, 5)), dtype=int64>

In [6]: d.vindex[[0, 5, 9], :, [0, 5, 9]].compute()
Out[6]: 
array([[  0,  10,  20,  30,  40,  50,  60,  70,  80,  90],
       [505, 515, 525, 535, 545, 555, 565, 575, 585, 595],
       [909, 919, 929, 939, 949, 959, 969, 979, 989, 999]])

In [7]: d.vindex[:, [0, 5, 9], [0, 5, 9]].compute()
Out[7]: 
array([[  0,  55,  99],
       [100, 155, 199],
       [200, 255, 299],
       [300, 355, 399],
       [400, 455, 499],
       [500, 555, 599],
       [600, 655, 699],
       [700, 755, 799],
       [800, 855, 899],
       [900, 955, 999]])

mrocklin · 2015-07-28T22:21:56Z

The only qualm that I have with this PR is that the code is not particularly clear. Doing this properly was non-trivial and I do not envy the next person who has to come along (quite possibly also me).

To this end further test cases are very welcome. Lets extend the amount of time until the next time someone needs to touch this.

mrocklin · 2015-07-28T22:33:17Z

@jcrist points out that this doesn't yet support numpy arrays

mrocklin · 2015-07-28T22:42:33Z

Fixed

shoyer · 2015-07-28T23:19:20Z

to be clear: this will fail for multi-dimensional indexers, right? if so, we should catch that.

We now raise on multi-dimensional indexers

shoyer · 2015-07-29T05:33:29Z

Is this being used?

mrocklin · 2015-07-29T18:23:21Z

In [1]: import dask.array as da

In [2]: import numpy as np

In [3]: x = np.arange(10*10*10).reshape((10, 10, 10))

In [4]: d = da.from_array(x, chunks=(4, 5, 6))

In [5]: d.vindex[[0, 5, 9], :, [0, 5, 9]]
Out[5]: dask.array<vindex-merge-2, shape=(3, 10), chunks=((3,), (5, 5)), dtype=int64>

In [6]: d.vindex[[0, 5, 9], :, [0, 5, 9]].compute()
Out[6]: 
array([[  0,  10,  20,  30,  40,  50,  60,  70,  80,  90],
       [505, 515, 525, 535, 545, 555, 565, 575, 585, 595],
       [909, 919, 929, 939, 949, 959, 969, 979, 989, 999]])

In [7]: d.vindex[:, [0, 5, 9], [0, 5, 9]].compute()
Out[7]: 
array([[  0, 100, 200, 300, 400, 500, 600, 700, 800, 900],
       [ 55, 155, 255, 355, 455, 555, 655, 755, 855, 955],
       [ 99, 199, 299, 399, 499, 599, 699, 799, 899, 999]])

shoyer · 2015-07-31T00:46:56Z

This example is broken:

import dask.array as da
import numpy as np

a = np.arange(125).reshape(5, 5, 5)
x = da.from_array(a, chunks=3)
x.vindex[[0, 1, 2, 3], [3, 2, 1, 4]] 
# has the wrong shape (4,) and an error is raised when you compute it

shoyer · 2015-07-31T00:49:17Z

This should raise an exception, but doesn't:

x.vindex[:, [0, 1, 2, 3], [0, 1, 2]]

mrocklin · 2015-07-31T01:21:48Z

Fixed the two recent issues in two recent commits.

mrocklin · 2015-07-31T22:02:26Z

Any further comments on this?

jhamman · 2015-07-31T22:10:41Z

@mrocklin - Thanks for putting this together. I just gave it a once over and the original use case seems to be well taken care of.

shoyer · 2015-07-31T23:22:44Z

Looks good to me

mrocklin · 2015-07-31T23:29:42Z

Woot. Thanks for the detailed feedback all.

add point slicing to dask.array.core

mrocklin mentioned this pull request Jul 17, 2015

Slicing with lists in multiple axes #433

Closed

mrocklin added 4 commits July 28, 2015 15:07

work on full slice point-wise indexing

14a35bb

support empty arrays in dask.array

b6fc879

point-wise indexing with None works

360caf7

mrocklin force-pushed the point-slice branch from 24e91c1 to 360caf7 Compare July 28, 2015 22:07

add vindex getter

e299c96

unite vindex error

a016461

mrocklin added 2 commits July 28, 2015 15:35

fix doctests

2916439

support np.ndarrays in point slicing

3d19ae6

shoyer reviewed Jul 28, 2015
View reviewed changes

raise error on multi-dimensional key in vindex

2507600

shoyer reviewed Jul 29, 2015
View reviewed changes

Comment thread dask/array/core.py Outdated

shoyer Jul 29, 2015

Copy link
Copy Markdown

Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this being used?

mrocklin added 3 commits July 29, 2015 09:55

bring points to first dimension in point slicing

de0b2ee

remove dead code

e0fee0f

add docstring

f71fa53

mrocklin added 2 commits July 30, 2015 18:18

point slicing supports incomplete index

fc149c8

vindex: raise error on mismatched lengths

fd62919

mrocklin added a commit that referenced this pull request Jul 31, 2015

Merge pull request #439 from mrocklin/point-slice

a598670

add point slicing to dask.array.core

mrocklin merged commit a598670 into dask:master Jul 31, 2015

mrocklin deleted the point-slice branch July 31, 2015 23:29

shoyer mentioned this pull request Aug 12, 2015

dask.array flatten and reshape #570

Closed

shoyer mentioned this pull request Dec 12, 2016

Generated Dask graph is huge - performance issue? pydata/xarray#1161

Closed

alimanfoo mentioned this pull request Oct 31, 2017

Advanced indexing zarr-developers/zarr-python#172

Merged

6 tasks done

jakirkham mentioned this pull request Feb 22, 2018

Indexing a dask array with another dask array #3096

Closed

phofl pushed a commit to phofl/dask that referenced this pull request Dec 23, 2024

Add Rolling.apply (dask#439)

0302eeb

Conversation

mrocklin commented Jul 17, 2015

Example

Uh oh!

shoyer commented Jul 17, 2015

Uh oh!

mrocklin commented Jul 17, 2015

Uh oh!

jhamman commented Jul 17, 2015

Uh oh!

mrocklin commented Jul 21, 2015

Uh oh!

shoyer commented Jul 28, 2015

Uh oh!

shoyer commented Jul 28, 2015

Uh oh!

mrocklin commented Jul 28, 2015

Uh oh!

mrocklin commented Jul 28, 2015

Uh oh!

mrocklin commented Jul 28, 2015

Uh oh!

mrocklin commented Jul 28, 2015

Uh oh!

mrocklin commented Jul 28, 2015

Uh oh!

shoyer Jul 28, 2015

Choose a reason for hiding this comment

Uh oh!

mrocklin Jul 28, 2015

Choose a reason for hiding this comment

Uh oh!

shoyer Jul 29, 2015

Choose a reason for hiding this comment

Uh oh!

mrocklin commented Jul 29, 2015

Uh oh!

shoyer commented Jul 31, 2015

Uh oh!

shoyer commented Jul 31, 2015

Uh oh!

mrocklin commented Jul 31, 2015

Uh oh!

mrocklin commented Jul 31, 2015

Uh oh!

jhamman commented Jul 31, 2015

Uh oh!

shoyer commented Jul 31, 2015

Uh oh!

mrocklin commented Jul 31, 2015

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants