Slice by dask array of ints #3407

crusaderky · 2018-04-15T14:54:30Z

Closes #3396 (together with #3405)

Strengths:

Supports chunking on the index
Very parsimonious of RAM
distributed-friendly - you'll get an output chunk for every chunk of the index, and the only thing that will travel across the network will be the contribution of every input chunk to every output chunk.

Limitations:

0D and 1D indexes only
no fancy indexing with multiple >=1D arrays (just like there isn't any with numpy indexes)

crusaderky · 2018-04-15T14:57:03Z

docs/source/array-slicing.rst

-*  Slicing one ``dask.array`` with another ``x[x > 0]``
+*  Slicing one `~dask.array.Array` with a `~dask.array.Array` of bools ``x[x > 0]``
+*  Slicing one `~dask.array.Array` with a zero or one-dimensional `~dask.array.Array`
+   of ints ``x[x.map_blocks(np.argsort)]``


This is a problematic example as argsort won't behave as expected if x is chunked on the axis it operates. Should be replaced with argtopk as soon as it is merged (#3405).

crusaderky · 2018-04-15T14:57:39Z

dask/array/tests/test_slicing.py

+    # Unsupported: 0d numpy array slicers (#3406)
+    # np.testing.assert_array_equal(x[idx0, idx0], d[idx0, idx0_d])
+    # np.testing.assert_array_equal(x[idx0, idx], d[idx0, idx_d])
+


Uncomment after fixing #3406

crusaderky · 2018-04-15T15:00:34Z

dask/array/slicing.py

+
+    if np.isnan(x.chunks[axis]).any():
+        raise NotImplementedError("Slicing an array with unknown chunks with a "
+                                  "dask.array of ints is not supported")


This is feasible but would horribly complicate everything because the graph will have to figure out dynamically the offset and the chunks (see below). I could not think of any strong reason to break my head on it as I think it's a very marginal use case.

crusaderky · 2018-04-15T15:01:36Z

dask/array/slicing.py

+    token = tokenize(x, idx, axis)
+    name1 = 'slice_with_int_dask_array_chunk-' + token
+    name2 = 'slice_with_int_dask_array_aggregate-' + token
+


Suggestions are welcome to use some kind of wrapper function instead of manually building the graph.
I tried with atop but didn't have much success.

jakirkham · 2018-04-16T13:38:55Z

This looks pretty similar to PR ( #3210 ). Might be worth taking a look.

jakirkham · 2018-04-24T15:39:51Z

Is it possible to do something like, x[x.nonzero()], with this implementation?

crusaderky · 2018-04-24T19:49:43Z

@jakirkham yes.

>>> import dask.array as da
>>> a = da.arange(-3, 3, chunks=3)
>>> a[a.nonzero()].compute()
array([-3, -2, -1,  1,  2])

[EDIT] to be more precise: the index can have NaN chunks, but the sliced array can't.

jakirkham · 2018-04-24T19:59:04Z

Cool, thanks for checking. Could you please add a test along those lines?

crusaderky · 2018-04-24T20:03:18Z

@jakirkham I added it before you even had time to ask :)

crusaderky · 2018-05-06T20:44:34Z

@jakirkham any update on this?

jakirkham · 2018-05-09T05:45:46Z

Sorry, haven't had time to give it another look. Am pretty busy for the next month. Though this isn't necessarily a bad thing. Would be a great time to get some other thoughts on this. Expect @shoyer's got some good ideas.

crusaderky · 2018-05-16T22:57:01Z

Found a bug - I'll revert once it's ironed out

crusaderky · 2018-05-19T20:19:15Z

@jakirkham fixed bug and enhanced tests - ready for merge...

jakirkham · 2018-06-27T18:01:55Z

Sounds good. If you have any thoughts on take_along_axis, feel free to share them in issue ( #3663 ).

crusaderky · 2018-06-27T18:23:13Z

@djhoese fixed uint index

shoyer

It looks like the only thing major thing to add here is handling negative indices, or at least being sure that we give a graceful error message?

shoyer · 2018-06-27T18:28:21Z

docs/source/changelog.rst

@@ -7,7 +7,7 @@ Changelog
 Array
 +++++

- Allow slicing a Dask Array by another one-dimensional Dask Array of integers
+-


you can probably add this back? :)

I removed to reduce merge collisions, in accordance to the new policy of the dask project

shoyer

OK, this looks good to me. I'll merge this in day or two unless anyone else has comments.

jakirkham · 2018-06-27T19:24:47Z

Does this work with vindex now? If so, can we have a test (assuming it is not already there and I've overlooked it)?

crusaderky · 2018-06-27T19:45:01Z

@jakirkham no test, and I didn't check. I'll look into it ASAP.

crusaderky · 2018-06-27T22:47:04Z

@jakirkham vindex doesn't work, in the sense that the index is silently computed at definition time.
I gave it a brief look (I never saw vindex before) and it looks like something very much NOT trivial. I would humbly suggest avoiding scope creep and opening a new PR after this one is merged.

I wrote initial unit tests for it; feel free to grab them: crusaderky@9cf947a

jakirkham · 2018-06-30T16:59:32Z

So I think you can steal the code out of PR ( #3210 ) to accomplish this. It already has the correct dispatch logic in vindex when Dask Arrays are present. Would just drop the atop call and following code, which can replaced with a simple getitem call.

crusaderky · 2018-07-05T09:48:42Z

@jakirkham @shoyer to me the functionality in vindex is a separate, albeit related, feature. Can we please merge this one and open a new PR if there aren't specific critiques?!?

shoyer · 2018-07-05T17:02:26Z

Agreed, vindex is something separate.

Thanks!

mrocklin · 2018-07-05T17:15:51Z

Thanks for your patience here @crusaderky ! Thanks for taking the time to review @shoyer !

…

On Thu, Jul 5, 2018 at 1:02 PM, Stephan Hoyer ***@***.***> wrote: Agreed, vindex is something separate. Thanks! — You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub <#3407 (comment)>, or mute the thread <https://github.com/notifications/unsubscribe-auth/AASszG4uv3cavNY-pWOhE9GYZxrMUqOfks5uDkakgaJpZM4TVgaw> .

crusaderky · 2018-07-05T20:14:43Z

Thanks everybody for your help!

Slice by dask array of ints

340f1c8

crusaderky commented Apr 15, 2018

View reviewed changes

Reimplement slice_with_int_dask_array using atop

c939abe

jcrist mentioned this pull request Apr 23, 2018

Risk attribution: a[argsort(b)[:k]] #3396

Closed

gimperiale added 2 commits April 23, 2018 22:37

Merge remote-tracking branch 'dask/master' into slice_by_da

5e008b9

Use argtopk in slice-by-int example

f86d89b

crusaderky mentioned this pull request Apr 23, 2018

WIP: Support using vindex with Dask Arrays #3210

Closed

Test a[a.nonzero()]

17fde8b

Merge remote-tracking branch 'dask/master' into slice_by_da

d70eadb

crusaderky mentioned this pull request May 11, 2018

slice by list loses chunking #1430

Closed

crusaderky changed the title ~~Slice by dask array of ints~~ WIP: Slice by dask array of ints May 16, 2018

gimperiale added 3 commits May 19, 2018 16:58

Fixed algorithm

1d9da1a

Merge remote-tracking branch 'dask/master' into slice_by_da

360973f

Minor refactor and lint

f771967

crusaderky closed this May 19, 2018

crusaderky reopened this May 19, 2018

crusaderky changed the title ~~WIP: Slice by dask array of ints~~ Slice by dask array of ints May 19, 2018

TomAugspurger mentioned this pull request May 22, 2018

da.map_blocks chunk alignment #3519

Closed

Merge remote-tracking branch 'dask/master' into slice_by_da

e5cc649

crusaderky added 4 commits June 27, 2018 18:23

Merge remote-tracking branch 'dask/master' into slice_by_da

9890d7b

Prevent changelog conflicts

741b63a

Use tuples instead of lists to slice ndarray

ea6aa74

atop: use tuples of ints instead of letters

f0d5049

crusaderky added 3 commits June 27, 2018 19:02

Remove ambiguous dtype=int (int32 on Windows and int64 on Linux)

7503cbe

docstring/comment improvements

43e7b58

Support for uint index

7f0fc84

shoyer reviewed Jun 27, 2018

View reviewed changes

crusaderky added 2 commits June 27, 2018 19:40

Support negative indices

b432577

Better docstrings for kernel functions

b3f8963

crusaderky mentioned this pull request Jun 27, 2018

Add NumPy's new take_along_axis #3663

Open

crusaderky added 2 commits June 27, 2018 20:00

Note about potential search algorithm

8b93559

Merge remote-tracking branch 'dask/master' into slice_by_da

f928659

shoyer approved these changes Jun 27, 2018

View reviewed changes

Additional unit tests

18d53c5

shoyer merged commit 8bc7003 into dask:master Jul 5, 2018

crusaderky deleted the slice_by_da branch July 5, 2018 20:15

stsievert mentioned this pull request Jul 27, 2018

WIP: ENH: Basic implementation of SGD dask/dask-glm#69

Open

3 tasks

This was referenced Oct 30, 2018

Indexing a dask array with another dask array #3096

Closed

Handling Dask Arrays with vindex #4156

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Slice by dask array of ints #3407

Slice by dask array of ints #3407

crusaderky commented Apr 15, 2018 •

edited

crusaderky Apr 15, 2018

crusaderky Apr 15, 2018

crusaderky Apr 15, 2018

crusaderky Apr 15, 2018

jakirkham commented Apr 16, 2018

jakirkham commented Apr 24, 2018

crusaderky commented Apr 24, 2018 •

edited

jakirkham commented Apr 24, 2018

crusaderky commented Apr 24, 2018

crusaderky commented May 6, 2018

jakirkham commented May 9, 2018

crusaderky commented May 16, 2018

crusaderky commented May 19, 2018

jakirkham commented Jun 27, 2018

crusaderky commented Jun 27, 2018

shoyer left a comment

shoyer Jun 27, 2018

crusaderky Jun 27, 2018

shoyer left a comment

jakirkham commented Jun 27, 2018 •

edited

crusaderky commented Jun 27, 2018

crusaderky commented Jun 27, 2018

jakirkham commented Jun 30, 2018

crusaderky commented Jul 5, 2018

shoyer commented Jul 5, 2018

mrocklin commented Jul 5, 2018 via email

crusaderky commented Jul 5, 2018

Slice by dask array of ints #3407

Slice by dask array of ints #3407

Conversation

crusaderky commented Apr 15, 2018 • edited

crusaderky Apr 15, 2018

Choose a reason for hiding this comment

crusaderky Apr 15, 2018

Choose a reason for hiding this comment

crusaderky Apr 15, 2018

Choose a reason for hiding this comment

crusaderky Apr 15, 2018

Choose a reason for hiding this comment

jakirkham commented Apr 16, 2018

jakirkham commented Apr 24, 2018

crusaderky commented Apr 24, 2018 • edited

jakirkham commented Apr 24, 2018

crusaderky commented Apr 24, 2018

crusaderky commented May 6, 2018

jakirkham commented May 9, 2018

crusaderky commented May 16, 2018

crusaderky commented May 19, 2018

jakirkham commented Jun 27, 2018

crusaderky commented Jun 27, 2018

shoyer left a comment

Choose a reason for hiding this comment

shoyer Jun 27, 2018

Choose a reason for hiding this comment

crusaderky Jun 27, 2018

Choose a reason for hiding this comment

shoyer left a comment

Choose a reason for hiding this comment

jakirkham commented Jun 27, 2018 • edited

crusaderky commented Jun 27, 2018

crusaderky commented Jun 27, 2018

jakirkham commented Jun 30, 2018

crusaderky commented Jul 5, 2018

shoyer commented Jul 5, 2018

mrocklin commented Jul 5, 2018 via email

crusaderky commented Jul 5, 2018

crusaderky commented Apr 15, 2018 •

edited

crusaderky commented Apr 24, 2018 •

edited

jakirkham commented Jun 27, 2018 •

edited