Masking and flattening (using an argmax as a integer-array slice)

[tamasgal]

I have some problems understanding how to "flatten" nested arrays/records when applying masks.

I think it's best explained with an actual use-case (with reduced complexity). Given a set of events (in this case 3), each event contains a variable lengths of particle tracks where I want to pick exactly one with the highest likelihood (lik).

from skhep_testdata import data_path
import uproot
import awkward as ak

f = uproot.open(data_path("uproot-issue431b.root"))
# grab the tracks of the first three events and only the lik-parameter for simplicity
tracks = f["E/Evt/trks"].arrays(["lik"], aliases={"lik": "trks.lik"})[:3]

tracks now is an array of records, which is basically a list of dictionaries:

>>> tracks.tolist()  # truncated output
[{'lik': [294.64, 294.64, ..., 67.77, 67.77]},
 {'lik': [96.75, 96.75, ..., 39.18, 38.87]},
 {'lik': [560.27, 560.27, ..., 118.72, 117.80]}]

I create an integer-mask for the tracks with the highest likelihood (in this example it's always the first, but it's not generally the case). I have to use keepdims=True to match the nested scheme (maybe that's where I enter the one-way route?):

mask = ak.argmax(tracks.lik, axis=-1, keepdims=True)
# <Array [[0], [0], [0]] type='3 * var * ?int64'>

and use that mask to actually select the tracks:

selected_tracks = tracks[mask]
# <Array [{lik: [295]}, ... 96.8]}, {lik: [560]}] type='3 * {"lik": var * ?float64}'>

As seen above, the masking works fine but it is nested. This is expected since the mask with [[0], [0], [0]] has ndim==2 and since it's an integer mask, it will not change the dimensions. However, I cannot find a way to flatten the resulting array of records since one obviously cannot flatten a record. This means that we need to put [0] after each attribute to access the value, although we know that each element is a single record (with single entries for each field):

selected_tracks.lik
# <Array [[295], [96.8], [560]] type='3 * var * ?float64'>
ak.flatten(selected_tracks.lik)  # this is the desired output but there are dozens of fields
# <Array [295, 96.8, 560] type='3 * ?float64'>

# or when accessing single entries:
selected_tracks[0].lik
# <Array [295] type='1 * ?float64'>
selected_tracks[0].lik[0]
# 295

Is there a better approach to reduce the dimensions of the fields of nested records, so that selected_tracks.ATTRIBUTES returns a flat list without the need to flatten them explicitly?

I am pretty sure it's kind of an corner-case but unfortunately many of our files are structured like this and I am trying to wrap all this into a low-level user-library which needs to be high-performant and user-friendly at the same time.

[jpivarski]

It's not such a corner case—other users might not encounter exactly this case, but probably a similar-enough one.

Since you know each list returned by ak.argmax with keepdims=True has exactly one element (even if that element is None), you can use [:, 0] as a slice to remove the dimension by picking the first (and only) element:

>>> tracks[mask].tolist()
[{'lik': [294.6407542676734]}, {'lik': [96.75133289411137]}, {'lik': [560.2775306614813]}]
>>> tracks[mask][:, 0].tolist()
[{'lik': 294.6407542676734}, {'lik': 96.75133289411137}, {'lik': 560.2775306614813}]

I did a quick check to verify that this also works with empty lists (missing values in the mask).

>>> fewer_tracks = tracks[tracks.lik > 100]
>>> ak.num(fewer_tracks)
<Array [{lik: 20}, {lik: 0}, {lik: 55}] type='3 * {"lik": int64}'>
>>> mask = ak.argmax(fewer_tracks.lik, axis=1, keepdims=True)
>>> mask
<Array [[0], [None], [0]] type='3 * var * ?int64'>
>>> fewer_tracks[mask].tolist()
[{'lik': [294.6407542676734]}, {'lik': [None]}, {'lik': [560.2775306614813]}]
>>> fewer_tracks[mask][:, 0].tolist()
[{'lik': 294.6407542676734}, {'lik': None}, {'lik': 560.2775306614813}]

What the [:, 0] slice is taking advantage of is the fact that "slicing columns" (record fields selected by string names) commutes with "slicing rows" (such as : for everything in the first dimension and 0 for the first in each list of the second dimension).

>>> deep = ak.Array([[{"x": {"a": 1}, "y": [{"b": 2}]}]])
>>> deep.type
1 * var * {"x": {"a": int64}, "y": var * {"b": int64}}
>>> deep[:, 0].tolist()
[{'x': {'a': 1}, 'y': [{'b': 2}]}]
>>> deep[0, 0].tolist()
{'x': {'a': 1}, 'y': [{'b': 2}]}
>>> deep[0, 0, "x"].tolist()
{'a': 1}
>>> deep[0, "x", 0].tolist()
{'a': 1}
>>> deep["x", 0, 0].tolist()
{'a': 1}
>>> deep[0, 0, "y"].tolist()
[{'b': 2}]
>>> deep[0, 0, "y", 0].tolist()
{'b': 2}
>>> deep[0, "y", 0, 0].tolist()
{'b': 2}
>>> deep["y", 0, 0, 0].tolist()
{'b': 2}

To be fully correct, this is Commutivity with a Caveat (I wish I had a better word): column-string selections can move anywhere to the left (shallower) than a deepest position, but not to the right, as illustrated by this failure:

>>> deep[0, 0, 0, "y"].tolist()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/jpivarski/irishep/awkward-1.0/awkward/highlevel.py", line 943, in __getitem__
    return ak._util.wrap(self._layout[where], self._behavior)
ValueError: in NumpyArray, too many dimensions in slice

because the first [0, 0, 0 can't be applied to the "x" field.

[tamasgal]

Ah dear, of course I know the [:, 0] syntax but I have no idea why I have not came up with it in this case. I somehow thought it's special since those are Records 🤦 Maybe too much coding the last few days...

Yes, this all make sense, I can wrap the [:, 0] thing in the high-level accessors where I know that the result contains one element.

Btw. it's a bit funny since I am pretty sure that this behaviour was a bit different in past. At least I have some unit-tests which at some point failed but I cannot put together the exact uproot/awkward versions anymore, so I have no idea where the slicing behaviour changed, but it was at some point returning single elements. I know that it was a combination of uproot3 and an earlier awkward1 version, probably around 0.3.x or early 0.4.x.

[jpivarski]

Uproot 3 can't be combined with Awkward 1, unless you do ak.from_awkward0.

It is the case that Awkward 0's argmax returned a jagged array of indexes so that it could be used as a slice without the extra [:, 0], but this was in conflict with NumPy's behavior for the same inputs, and Awkward 1 is constrained to only generalize NumPy, not change it.

>>> import awkward as ak
>>> import awkward0
>>> ak.argmax(ak.from_iter([[1, 2, 3], [], [4, 5]]), axis=1)
<Array [2, None, 1] type='3 * ?int64'>
>>> awkward0.fromiter([[1, 2, 3], [], [4, 5]]).argmax()
<JaggedArray [[2] [] [1]] at 0x7fe4ba0b4e20>

These are two different ways of expressing missing values: numbers with None on the one hand and lists of zero or one elements on the other hand, and different languages prefer one or the other. In Scala, for instance, missing values in data of type Option[X], which are also called None, are equivalent to empty lists, and the non-missing values, Some(x), are equivalent to non-missing lists.

So Awkward 1 has ways to switch between the two representations:

>>> a = ak.argmax(ak.from_iter([[1, 2, 3], [], [4, 5]]), axis=1)
>>> ak.singletons(a)
<Array [[2], [], [1]] type='3 * var * int64'>
>>> ak.firsts(ak.singletons(a))
<Array [2, None, 1] type='3 * ?int64'>

These were introduced to make use of the output of ak.argmax and #203 is an open issue to make that easier. However, @nsmith- noticed that it's much simpler to pass keepdims=True and slice with [:, 0].

[tamasgal]

Yep, we used ak.from_awkward0.

Yes, makes sense, thanks! I also agree that keepdims=True and [:, 0] is a good choice.

[tamasgal]

This is the first job which failed: https://git.km3net.de/km3py/km3io/-/jobs/104059 with

awkward==1.0.0rc2
awkward0==0.15.0
uproot3==3.14.1
uproot3-methods==0.10.0

And the last which worked https://git.km3net.de/km3py/km3io/-/jobs/101157 with

awkward==0.14.0
awkward1==0.4.3
uproot==3.13.0
uproot-methods==0.8.0

So somewhere between these things have suddenly failed, which were using awkward1 arrays:

>       assert best.rec_stages[0].tolist() == [1, 3, 5, 4]
E       AssertionError: assert [[1, 3, 5, 4]] == [1, 3, 5, 4]

[tamasgal]

...anyways, it's not so important since the [:, 0] is fine ;)