Using an index based predefined mask

[williamnash]

Hello, I just started using Awkward array for HEP and had a quick question about the best way to reconfigure my existing code to use the package.

I am starting with a flat tree which contains objects which link to each other via preset indices (e.g. a reco muon will contain an index for the gen muon it is matched to). I haven't yet found a simple way to combine the two jagged arrays of gen particles and muons in such a way that I can select only those gen muons among all the gen particles. For instance if I start with

gen_particles = ak.Array([[a, b, c], [a, b, c, d, e], [a, b, c]])
muon_gen_ids = ak.Array([[0,1],[1,3], [0]])

I want select the subset of gen particles with ids which belong to the list of muon_gen_ids. In this case I want the array

gen_muons = ak.Array([[a,b],[b,d], [a]])

I know how to select for the gen_particle_ids using the local_index() function, and imagine the solution would involve the mask() function, but haven't figure our a solution that doesn't involve a for-loop (which would be slower than my current implementation of just indexing into the array). Any suggestions?

Thanks!
Will

[jpivarski]

As it turns out, it's easier than you think:

>>> gen_particles = ak.Array([["a", "b", "c"], ["a", "b", "c", "d", "e"], ["a", "b", "c"]])
>>> muon_gen_ids = ak.Array([[0, 1], [1, 3], [0]])
>>> gen_particles[muon_gen_ids]
<Array [['a', 'b'], ['b', 'd'], ['a']] type='3 * var * string'>

What you want is to use the ids as a slice. This is like NumPy's advanced indexing, but generalized to the non-rectilinear case.

(Thanks for the very clear example and desired result! It helped me understand your problem quickly.)

[williamnash]

Wow super slick. Thank you!