Array indexing after mask: bug or feature?

[Duchstf]

so I've been trying to do this:

truth_x_target = truth_x_target[mask][:,0]

However, this returns an array with the same first dimension as mask, whereas I want to basically "select the first element of each event" after the array is filtered. Is this a bug or feature of the awkward array?

Thank you for your help!

Duc.

[jpivarski]

Could you give more details, like what the type of truth_x_target and mask are, and what you expected?

When I tried an example, I didn't find anything unexpected:

>>> array = ak.Array([[0.0, 1.1, 2.2], [], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9]])
>>> array.type
5 * var * float64
>>> array.tolist()
[[0.0, 1.1, 2.2], [], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9]]

For a singly jagged array of numbers like this, there are two possible masks: flat booleans and (singly) jagged booleans. Here's an example of a flat boolean mask:

>>> mask1 = ak.num(array) > 1
>>> mask1.type
5 * bool
>>> mask1.tolist()
[True, False, True, False, True]

and here's an example of a jagged boolean mask:

>>> mask2 = array * 10 % 2 == 0
>>> mask2.type
5 * var * bool
>>> mask2.tolist()
[[True, False, True], [], [False, True], [False], [True, False, True, False]]

When you slice with a flat boolean mask, it has to have the same length as the array and it removes lists from it:

>>> array[mask1].type
3 * var * float64
>>> array[mask1].tolist()
[[0.0, 1.1, 2.2], [3.3, 4.4], [6.6, 7.7, 8.8, 9.9]]

When you slice with a jagged boolean mask, it has to have all the same lengths for each of its nested lists, and it removes numbers from those nested lists:

>>> array[mask2].type
5 * var * float64
>>> array[mask2].tolist()
[[0.0, 2.2], [], [4.4], [], [6.6, 8.8]]

What happens after that depends on the type and lengths of what the first step returned. For instance, you can have mask1 slice the first dimension and 0 slice the second dimension in one step or in two:

>>> array[mask1, 0]
<Array [0, 3.3, 6.6] type='3 * float64'>
>>> array[mask1][:, 0]
<Array [0, 3.3, 6.6] type='3 * float64'>

In my example, I can't get the first item of each list after mask2 because it leaves some nested lists empty. So I'll start with an equivalent example.

>>> array = ak.Array([[0.0, 1.1, 2.2], [3.3, 4.4], [5.5], [6.6, 7.7, 8.8, 9.9]])
>>> mask2 = array * 10 % 2 == 1
>>> array[mask2].type
4 * var * float64
>>> array[mask2].tolist()
[[1.1], [3.3], [5.5], [7.7, 9.9]]

Now we can pick the second element, but only as a second step:

>>> array[mask2][:, 0].type
4 * float64
>>> array[mask2][:, 0].tolist()
[1.1, 3.3, 5.5, 7.7]

not as a first step:

>>> array[mask2, 0]
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/jpivarski/miniconda3/lib/python3.8/site-packages/awkward/highlevel.py", line 1005, in __getitem__
    return ak._util.wrap(self._layout[where], self._behavior)
ValueError: in NumpyArray, too many dimensions in slice

(https://github.com/scikit-hep/awkward-1.0/blob/1.0.2rc4/src/libawkward/array/NumpyArray.cpp#L3925)

because mask1 goes one level down, letting the thing that comes after it apply to the second dimension, whereas mask2 goes two levels down, so putting a 0 after it tries to slice the third dimension, and this array doesn't have a third dimension.

Does that help?

[Duchstf]

Ah ok, I checked back my codes and I think that's not what went wrong, but it's similar

So initially I have a mask array of type:

mask = 137178 * bool

and then I have two more arrays (a and b), both of them has the same first dimension. And I want to select the first elements in a, so I do this:

a = a[mask][:,0]
b = b[mask]

I would expect them to have the same first dimension but a now has type 464 * float32 (which is what I expected). But now b has type 137178 * var * float32. If I use b = b[mask][:], now it returns an array of type 464 * float32, and this is what I want.

So I based on your response I might be experiencing a bug here? since we would expect the mask to remove lists from b.

[jpivarski]

I still don't understand what the problem is: if I make a and b as having type N * var * int64,

>>> a = ak.Array([[0, 1, 2], [3, 4], [5], [6, 7, 8, 9]])
>>> b = ak.Array([[0, 1, 2], [3, 4], [5], [6, 7, 8, 9]])
>>> mask = ak.Array([True, False, True, True])
>>> a
<Array [[0, 1, 2], [3, 4, ... 5], [6, 7, 8, 9]] type='4 * var * int64'>
>>> b
<Array [[0, 1, 2], [3, 4, ... 5], [6, 7, 8, 9]] type='4 * var * int64'>
>>> mask
<Array [True, False, True, True] type='4 * bool'>

then I do a = a[mask][:, 0] and either b = b[mask] or b = b[mask][:], and either way, the result is the same:

>>> a = a[mask][:, 0]
>>> b = b[mask][:]
>>> a
<Array [0, 5, 6] type='3 * int64'>
>>> b
<Array [[0, 1, 2], [5], [6, 7, 8, 9]] type='3 * var * int64'>

That is, a is replaced by a non-jagged array of shorter length because (1) the mask drops nested lists, (2) the [:, 0] picks the first element from each remaining list, and (3) we only do the mask step to b. Composing with a [:] slice does nothing. (For Python lists, it copies, but in Awkward Array, there's no distinction between copying and viewing: it's a view in this case.)

If b were not jagged to start with, then mask would drop elements from it, making it still be non-jagged. Actually, it's always true that a boolean-array slice doesn't change the data type, other than changing length (which may be considered part of the type or not, depending on your point of view).

I don't know what's happening in your case, and I only partially see what you mean by the possibility that something's wrong: it would be wrong for a jagged array to be turned into a non-jagged array by a mask.

But now b has type 137178 * var * float32. If I use b = b[mask][:], now it returns an array of type 464 * float32

That one sentence sounds odd.

Anyway, whatever problem you're seeing might be related to the fact that you're reusing a and b as variable names: you're replacing the original a and b with derived quantities using the same names. That sort of thing can lead to confusion if the code is run more than once. Because most operations view, rather than copy, there usually isn't much to be gained in terms of memory use by losing access to the original arrays.

[Duchstf]

Anyway, whatever problem you're seeing might be related to the fact that you're reusing a and b as variable names: you're replacing the original a and b with derived quantities using the same names. That sort of thing can lead to confusion if the code is run more than once. Because most operations view, rather than copy, there usually isn't much to be gained in terms of memory use by losing access to the original arrays.

Ah that might be it!! Sorry for the confusion, thank you so much!!!