The problem I am facing is to create a new array from a set of indices. Namely, I have a set of particles and jets, for each jet there is a list of indices of which particles belong to the given jet.
import awkward as ak
import vector
particles = ak.Array({
"mass": ak.Array([1, 2, 3, 4, 5]),
"px": ak.Array([1, 2, 3, 4, 5]),
"py": ak.Array([1, 2, 3, 4, 5]),
"pz": ak.Array([1, 2, 3, 4, 5]),
})
particles_p4 = vector.awk(ak.zip({
"mass": particles["mass"],
"x": particles["px"],
"y": particles["py"],
"z": particles["pz"]
}))
jet = ak.Array({
"constituents": ak.Array([[1, 2, 4, 5], [3]]),
"energy": ak.Array([1.2, 3.4])
})
What I would like to get is for each jet the particle_p4 values in a new array like so:
<Array [[{x: 1, y: 1, z: 1, ... z: 4, tau: 4}]] type='2 * var * {"x": int64, "y"...'>
where the first element of that would be:
<Array [{x: 1, y: 1, z: 1, ... z: 5, tau: 5}] type='4 * {"x": int64, "y": int64,...'>
Doing this with a for loop is trivial, however I have the notion that this can be done in a more efficient way with the tools available in Awkward array.
Bonus: What about even more nested arrays, for example where each event has multiple jets?
source https://stackoverflow.com/questions/75226252/creating-a-new-awkward-array-from-indices
Comments
Post a Comment