Advice: How to slice N-dim array with (N+1)-dim array?

[Superharz]

Package versions:

awkward: 2.6.4

Problem

I would like to use a (N+1)-dim array to select elements form a N-dim array.

Example

I am having an array a with shape N * var * int64 of the following kind:

a = ak.Array([
 [11, -11, 11, -11, 11, 11],
 [11, -11, 11, -11],
 [],
 [11, -11],
 [11],
 [],
 [11],
 [11, -11, 11, -11, 11],
 [],
 [-11, 11]
])

I have an array b of shape N * var * var * int64 with indices that correspond to elements in a:

b = ak.Array([
 [[4], []],
 [[0], [1], [3]],
 [[], []],
 [[], [], []],
 [[0], []],
 [[], []],
 [[0], []],
 [[4], []],
 [[], []],
 [[], []]]
])

I would like to select the elements from a with b. Expected result:

a[b]
 [
 [[11], []],
 [[11], [-11], [-11]],
 [[], []],
 [[], [], []],
 [[11], []],
 [[], []],
 [[11], []],
 [[11], []],
 [[], []],
 [[], []]
]

Error

However, I get the error:

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
Cell In[253], line 1
----> 1 a[b]

File [~\miniforge3\envs\uni\Lib\site-packages\awkward\highlevel.py:1065](http://localhost:8888/lab/tree/ztautau-polarization/src/panda/~/miniforge3/envs/uni/Lib/site-packages/awkward/highlevel.py#line=1064), in Array.__getitem__(self, where)
    635 """
    636 Args:
    637     where (many types supported; see below): Index of positions to
   (...)
   1061 have the same dimension as the array being indexed.
   1062 """
   1063 with ak._errors.SlicingErrorContext(self, where):
   1064     return wrap_layout(
-> 1065         prepare_layout(self._layout[where]),
   1066         self._behavior,
   1067         allow_other=True,
   1068         attrs=self._attrs,
   1069     )

File [~\miniforge3\envs\uni\Lib\site-packages\awkward\contents\content.py:512](http://localhost:8888/lab/tree/ztautau-polarization/src/panda/~/miniforge3/envs/uni/Lib/site-packages/awkward/contents/content.py#line=511), in Content.__getitem__(self, where)
    511 def __getitem__(self, where):
--> 512     return self._getitem(where)

File [~\miniforge3\envs\uni\Lib\site-packages\awkward\contents\content.py:565](http://localhost:8888/lab/tree/ztautau-polarization/src/panda/~/miniforge3/envs/uni/Lib/site-packages/awkward/contents/content.py#line=564), in Content._getitem(self, where)
    562         return out._getitem_at(0)
    564 elif isinstance(where, ak.highlevel.Array):
--> 565     return self._getitem(where.layout)
    567 # Convert between nplikes of different backends
    568 elif (
    569     isinstance(where, ak.contents.Content)
    570     and where.backend is not self._backend
    571 ):

File [~\miniforge3\envs\uni\Lib\site-packages\awkward\contents\content.py:640](http://localhost:8888/lab/tree/ztautau-polarization/src/panda/~/miniforge3/envs/uni/Lib/site-packages/awkward/contents/content.py#line=639), in Content._getitem(self, where)
    637     return where.to_NumpyArray(np.int64)
    639 elif isinstance(where, Content):
--> 640     return self._getitem((where,))
    642 elif is_sized_iterable(where):
    643     # Do we have an array
    644     nplike = nplike_of_obj(where, default=None)

File [~\miniforge3\envs\uni\Lib\site-packages\awkward\contents\content.py:557](http://localhost:8888/lab/tree/ztautau-polarization/src/panda/~/miniforge3/envs/uni/Lib/site-packages/awkward/contents/content.py#line=556), in Content._getitem(self, where)
    548 nextwhere = ak._slicing.prepare_advanced_indexing(items, backend)
    550 next = ak.contents.RegularArray(
    551     this,
    552     this.length,
    553     1,
    554     parameters=None,
    555 )
--> 557 out = next._getitem_next(nextwhere[0], nextwhere[1:], None)
    559 if out.length is not unknown_length and out.length == 0:
    560     return out._getitem_nothing()

File [~\miniforge3\envs\uni\Lib\site-packages\awkward\contents\regulararray.py:698](http://localhost:8888/lab/tree/ztautau-polarization/src/panda/~/miniforge3/envs/uni/Lib/site-packages/awkward/contents/regulararray.py#line=697), in RegularArray._getitem_next(self, head, tail, advanced)
    682     assert head.offsets.nplike is index_nplike
    683     self._maybe_index_error(
    684         self._backend[
    685             "awkward_RegularArray_getitem_jagged_expand",
   (...)
    696         slicer=head,
    697     )
--> 698     down = self._content._getitem_next_jagged(
    699         multistarts, multistops, head._content, tail
    700     )
    702     return RegularArray(
    703         down, headlength, self._length, parameters=self._parameters
    704     )
    706 elif isinstance(head, ak.contents.IndexedOptionArray):

File [~\miniforge3\envs\uni\Lib\site-packages\awkward\contents\listoffsetarray.py:423](http://localhost:8888/lab/tree/ztautau-polarization/src/panda/~/miniforge3/envs/uni/Lib/site-packages/awkward/contents/listoffsetarray.py#line=422), in ListOffsetArray._getitem_next_jagged(self, slicestarts, slicestops, slicecontent, tail)
    417 def _getitem_next_jagged(
    418     self, slicestarts: Index, slicestops: Index, slicecontent: Content, tail
    419 ) -> Content:
    420     out = ak.contents.ListArray(
    421         self.starts, self.stops, self._content, parameters=self._parameters
    422     )
--> 423     return out._getitem_next_jagged(slicestarts, slicestops, slicecontent, tail)

File [~\miniforge3\envs\uni\Lib\site-packages\awkward\contents\listarray.py:472](http://localhost:8888/lab/tree/ztautau-polarization/src/panda/~/miniforge3/envs/uni/Lib/site-packages/awkward/contents/listarray.py#line=471), in ListArray._getitem_next_jagged(self, slicestarts, slicestops, slicecontent, tail)
    461 outoffsets = ak.index.Index64.empty(
    462     slicestarts.length + 1,
    463     self._backend.index_nplike,
    464 )
    465 assert (
    466     outoffsets.nplike is self._backend.index_nplike
    467     and slicestarts.nplike is self._backend.index_nplike
   (...)
    470     and self._stops.nplike is self._backend.index_nplike
    471 )
--> 472 self._maybe_index_error(
    473     self._backend[
    474         "awkward_ListArray_getitem_jagged_descend",
    475         outoffsets.dtype.type,
    476         slicestarts.dtype.type,
    477         slicestops.dtype.type,
    478         self._starts.dtype.type,
    479         self._stops.dtype.type,
    480     ](
    481         outoffsets.data,
    482         slicestarts.data,
    483         slicestops.data,
    484         slicestarts.length,
    485         self._starts.data,
    486         self._stops.data,
    487     ),
    488     slicer=ListArray(slicestarts, slicestops, slicecontent),
    489 )
    491 as_list_offset_array = self.to_ListOffsetArray64(False)
    492 next_content = as_list_offset_array._content[
    493     as_list_offset_array.offsets[0] : as_list_offset_array.offsets[-1]
    494 ]

File [~\miniforge3\envs\uni\Lib\site-packages\awkward\contents\content.py:277](http://localhost:8888/lab/tree/ztautau-polarization/src/panda/~/miniforge3/envs/uni/Lib/site-packages/awkward/contents/content.py#line=276), in Content._maybe_index_error(self, error, slicer)
    275 else:
    276     message = self._backend.format_kernel_error(error)
--> 277     raise ak._errors.index_error(self, slicer, message)

IndexError: cannot slice ListArray (of length 10) with [[[4], []], [[0], [1], [3]], [[], []], ..., [[4], []], [[], []], [[], []]]: jagged slice inner length differs from array inner length (in compiled code: https://github.com/scikit-hep/awkward/blob/awkward-cpp-33/awkward-cpp/src/cpu-kernels/awkward_ListArray_getitem_jagged_descend.cpp#L27)

This error occurred while attempting to slice

    <Array [[11, -11, 11, -11, 11, 11], ..., [...]] type='10 * var * int64'>

with

    <Array [[[4], []], [[0], ...], ..., [[], []]] type='10 * var * var * int64'>

Other approaches

I have tried to use ak.flatten() with ak.unflatten():

counts = ak.flatten(ak.num(b, axis = -1))
ak.unflatten(a[ak.flatten(b, axis = -1)], counts = counts, axis = -1)

[[[11], []],
 [[11], [-11], [-11], [], [], [], [], []],
 [],
 [],
 [[11], [], [], []],
 [],
 [[11], []],
 [[11], [], [], [], [], []],
 [],
 []]
------------------------------------------
type: 10 * var * var * int64

However, this does not work with the empty lists. I found a similar issue in #905.

The other approach is to use ak.cartesian():

car = ak.cartesian({"index":b, "var":ak.singletons(a)})
car["var"][b]

[[[11], []],
 [[11], [-11], [-11]],
 [[], []],
 [[], [], []],
 [[11], []],
 [[], []],
 [[11], []],
 [[11], []],
 [[], []],
 [[], []]]
----------------------------
type: 10 * var * var * int64

This works. But I am wondering if there is some advanced indexing to just use a[b]?

[agoose77]

It looks like you want to use Awkward's ragged indexing to pull items out of a using b, and to capture the structure of b when doing so?

Awkward's ragged indexing requires that the structures of a and b match (apart from the length of the deepest sublists, and support for missing values in the indexing array), so we need to do a bit of massaging:

First we drop the final dimension of our index array so that it's compatible with a:

ix = ak.flatten(b, axis=-1)

Then we build an indexing result:

temp = a[ix]

Then we restore the structure of b, first by determining the number of new values in each sublist we are creating:

n_items = ak.ravel(ak.num(b, axis=-1))

Before finally unflattening our result

result = ak.unflatten(temp, n_items, axis=-1)

Does this behave as you expect?

[Superharz]

Hi @agoose77,

thank you for your comment.

Your way is equivalent to my first approach with flatten and unflatten.

However, it also has problems with the empty sub-lists, which are added to the wrong outer list. Your approach:

[[[11], []],
 [[11], [-11], [-11], [], [], [], [], []],
 [],
 [],
 [[11], [], [], []],
 [],
 [[11], []],
 [[11], [], [], [], [], []],
 [],
 []]
------------------------------------------
type: 10 * var * var * int64

The intended result:

[[[11], []],
 [[11], [-11], [-11]],
 [[], []],
 [[], [], []],
 [[11], []],
 [[], []],
 [[11], []],
 [[11], []],
 [[], []],
 [[], []]]
----------------------------
type: 10 * var * var * int64

[alexandertuna]

Hi @agoose77, I'll just mention that I'm facing the same issue as @Superharz!

Here's my smaller example, following @agoose77's suggestion:

x = ak.Array([
    [0.0, 10.0, 20.0],
    [30.0, 40.0, 50.0],
])
y = ak.Array([
    [ [0], [1], [1, 2]],
    [ [], [1], [2]],
])

def fancy_index(a, b):
    ix = ak.flatten(b, axis=-1)
    temp = a[ix]
    n_items = ak.ravel(ak.num(b, axis=-1))
    return ak.unflatten(temp, n_items, axis=-1)

fancy_index(x, y).show()

gives:

[
  [[0], [10], [10, 20], []],
  [[40], [50]]
]

i.e. the structure of the index array (y) wasn't restored properly.

Same as @Superharz, I believe this occurs when unflatten encounters an empty sub-array at the boundary of an array.