[Python] How to Slice An Awkward Array of Records by the Length of the Record Value

I'm new to posting questions on here and to using awkward arrays. Let me know if I should re-format my question or include more information. I am using awkward/2.6.7, awkward_cpp/37, jupyterlab==3.0.17, python 3.11.9. I have an awkward array called "lilTree" of the form (returned from running lilTree.type.show()):

4000 * {
variable1: var * int32,
variable2: var * float32,
variable3: var * float32,
variable4: float32

}


My variables 1-3 values are awkward arrays of floats with lengths varying from 2-4. I think about the outer 4000 as indexing events, and each entry in these awkward arrays corresponds to a jet. I would like to throw away all but the first two jet values for each of these variables. I know that awkward arrays are immutable in that you can't easily slice them like you can with lists (at least not destructively). You can filter them, but I don't want to throw away these longer entries all together or my data will be too sparse. I also have variable4 that is just a float (event-level variable) that complicates things.

My first idea is to "pick apart then put back together": loop through each event and make a dictionary for every variable in that event only two jets saved. Then I could convert these all to awkward arrays again and then concatenate them into a new awkward array with the correct number of entries in each sub-array. This seems really inefficient.

Is there a way to do this that is more efficient than two loops?

myList = []
i=0
while i<lilTree.type.length:
myArray = ak.Array([{"variable4": lilTree["variable4"],
"variable1":lilTree["variable1"][:2] ,
"variable2": lilTree["variable2"][:2],
"variable3": lilTree["variable3"][:2])
myList.append(myArray)
i+=1

newTree = ak.concatenate(myList)


I think this works... but it is very inefficient. Any tips/tricks for me?

Continue reading...