Reading a branch that has subbranches (the AsGroup Interpretation)

[tamasgal]

Help, I struggle again with the transitioning ;)

With uproot3, there was a weird interpretation (glitch) in one of our file formats, which I/we never fully understood, but the outcome was very useful. The data in "E/Evt/hits" below is split up into subbranches which one can read nicely with array/lazyarray, however, if you access the array of the main branch directly (E/Evt/hits), you get the number of entries per event instead of a large jagged array of structs, since as you can see below, uproot3 recognises it as a vector<int>.
We used this array to quickly identify e.g. empty or large events for further cuts since it loads extremely fast. The same behaviour is btw. observed in E/Evt/mc_hits, E/Evt/trks, E/Evt/mc_trks` and in other branches of different other formats we use as well:

>>> import uproot3  # 3.14.0
>>> from skhep_testdata import data_path
>>> f = uproot3.open(data_path("uproot-issue431b.root"))
>>> f["E/Evt/hits"].show()
hits                       TStreamerSTL               asdtype('>i4')
hits.id                    TStreamerBasicType         asjagged(asdtype('>i4'))
hits.dom_id                TStreamerBasicType         asjagged(asdtype('>i4'))
hits.channel_id            TStreamerBasicType         asjagged(asdtype('>u4'))
hits.tdc                   TStreamerBasicType         asjagged(asdtype('>u4'))
hits.tot                   TStreamerBasicType         asjagged(asdtype('>u4'))
hits.trig                  TStreamerBasicType         asjagged(asdtype('>i4'))
hits.pmt_id                TStreamerBasicType         asjagged(asdtype('>i4'))
hits.t                     TStreamerObjectAny         asjagged(asdtype('>f8'))
hits.a                     TStreamerBasicType         asjagged(asdtype('>f8'))
hits.pos.x                 TStreamerBasicType         asjagged(asdtype('>f8'))
hits.pos.y                 TStreamerBasicType         asjagged(asdtype('>f8'))
hits.pos.z                 TStreamerBasicType         asjagged(asdtype('>f8'))
hits.dir.x                 TStreamerBasicType         asjagged(asdtype('>f8'))
hits.dir.y                 TStreamerBasicType         asjagged(asdtype('>f8'))
hits.dir.z                 TStreamerBasicType         asjagged(asdtype('>f8'))
hits.pure_t                TStreamerBasicType         asjagged(asdtype('>f8'))
hits.pure_a                TStreamerBasicType         asjagged(asdtype('>f8'))
hits.type                  TStreamerBasicType         asjagged(asdtype('>i4'))
hits.origin                TStreamerBasicType         asjagged(asdtype('>i4'))
hits.pattern_flags         TStreamerBasicType         asjagged(asdtype('>u4'))
>>> f["E/Evt/hits/hits.tot"].array()
<JaggedArray [[24 30 22 ... 38 26 23] [29 26 22 ... 26 28 24] [27 19 13 ... 27 24 16] ... [22 22 9 ... 27 32 27] [30 32 17 ... 30 24 29] [27 41 36 ... 29 24 28]] at 0x00011421d640>
>>> f["E/Evt/hits"].array()
array([176, 125, 318, 157,  83,  60,  71,  84, 255, 105], dtype=int32)

So far so good... The "problem" is that uproot4 is able to correctly figure out how to parse E/Evt/hits as an awkward.Record array (vector<Hit>) and the only way I have figured out so far to get the lengths of each sub-array is to use ak.num or ak.count, both of them however require to load a bunch of data into memory which takes time and... well, memory ;) other than that, they count all sub-arrays. I can reduce the footprint by only retrieving any of the subbranches, e.g. ak.count(f["E/Evt/hits/hits.channel_id"].array(), axis=1) but that does not feel good (and still takes much longer compared to the uproot3 glitch.

It's btw. interesting that the interpretation is still showing AsJagged(AsDtype('>i4')) and int32_t[], so I am a bit confused why it works.

>>> import uproot4

>>> from skhep_testdata import data_path

>>> import awkward1 as ak

>>> f = uproot4.open(data_path("uproot-issue431b.root"))

>>> f["E/Evt/hits"].show()
name                 | typename                 | interpretation
---------------------+--------------------------+-------------------------------
hits                 | vector<Hit>              | AsGroup(<TBranchElement 'hits'hits.id              | int32_t[]                | AsJagged(AsDtype('>i4'))
hits.dom_id          | int32_t[]                | AsJagged(AsDtype('>i4'))
hits.channel_id      | uint32_t[]               | AsJagged(AsDtype('>u4'))
hits.tdc             | uint32_t[]               | AsJagged(AsDtype('>u4'))
hits.tot             | uint32_t[]               | AsJagged(AsDtype('>u4'))
hits.trig            | int32_t[]                | AsJagged(AsDtype('>i4'))
hits.pmt_id          | int32_t[]                | AsJagged(AsDtype('>i4'))
hits.t               | double[]                 | AsJagged(AsDtype('>f8'))
hits.a               | double[]                 | AsJagged(AsDtype('>f8'))
hits.pos.x           | double[]                 | AsJagged(AsDtype('>f8'))
hits.pos.y           | double[]                 | AsJagged(AsDtype('>f8'))
hits.pos.z           | double[]                 | AsJagged(AsDtype('>f8'))
hits.dir.x           | double[]                 | AsJagged(AsDtype('>f8'))
hits.dir.y           | double[]                 | AsJagged(AsDtype('>f8'))
hits.dir.z           | double[]                 | AsJagged(AsDtype('>f8'))
hits.pure_t          | double[]                 | AsJagged(AsDtype('>f8'))
hits.pure_a          | double[]                 | AsJagged(AsDtype('>f8'))
hits.type            | int32_t[]                | AsJagged(AsDtype('>i4'))
hits.origin          | int32_t[]                | AsJagged(AsDtype('>i4'))
hits.pattern_flags   | uint32_t[]               | AsJagged(AsDtype('>u4'))

>>> f["E/Evt/hits"].array()  # correctly parses the whole dataset as a record array
<Array [{'hits.id': [0, 0, 0, 0, ... 0, 0, 0]}] type='10 * {"hits.id": var * int...'>

>>> ak.count(f["E/Evt/hits"].array(), axis=1)  # this fails for a typical sized file due to excessive memory usage
<Record ... 157, 83, 60, 71, 84, 255, 105]} type='{"hits.id": var * int64, "hits...'>

>>> ak.count(f["E/Evt/hits/hits.channel_id"].array(), axis=1)  # kind of works, but still too slow
<Record ... 157, 83, 60, 71, 84, 255, 105]} type='{"hits.id": var * int64, "hits...'>

The f["E/Evt/hits"].num_entries gives the overall number of events, but I have not found any other shortcut to the subbranch lengths.

Do you have a hint how to get to this information?

[jpivarski]

This is being read as a group (uproot4.AsGroup interpretation). Sometimes, a branch with subbranches has no data, other times, it has integer counts. I think that TClonesArrays do the latter.

You can check for this with

>>> import uproot4
>>> from skhep_testdata import data_path
>>> f = uproot4.open(data_path("uproot-issue431b.root"))
>>> isinstance(f["E/Evt/hits"].interpretation, uproot4.AsGrouped)
True

and you can get the counts by supplying that interpretation:

>>> f["E/Evt/hits"].array(uproot4.AsDtype(">i4"))
<Array [176, 125, 318, 157, ... 84, 255, 105] type='10 * int32'>

That interpretation won't work in every AsGrouped case, but the only other case I've seen is when there's no data at all associated with the branch-with-subbranches. You can check for that with

>>> f["E/Evt/hits"].num_baskets
1

or catch the failure of the uproot4.AsDtype(">i4") interpretation.

As for counting after the fact, prefer ak.num over ak.count, since the latter is a full-fledged reducer. (The difference becomes observable when axis < array.ndim - 1, but it already costs performance in the cases where they have the same output.)

[tamasgal]

Thanks Jim, the interpretation to integer is working perfectly! The error handling hints are also useful.
I have the feeling I will never fully understand the ROOT intrinsics, even after spending so much time on them 😄