How to delete some events (know the index) in analysis using awkward??

[shaoweisong]

Hi experts,
I am using a preselection framework to do the gen level study. After selecting gen particles, there are 800486 events left. 119 events have 4 genparticles, so the lenth of these event arrays is 4 (like event 538). And the length of other event arrays is 2 (like event 3), which means there are 2 particles in each of these events.

And for now, I need to remove the event arrays with 4 particles. I know in numpy array we can use np.delete(arr, index_list,axis=0) to do this kind of thing. But What should I do to remove the events corresponding to index_list in awkward array?

[agoose77]

Finding where sublists have a particular length

Before we look at removing the rows (sublists) that have lengths different from four, let's first find where those sublists are. Then we can ask Awkward Array to remove them.

Consider an example list particle_x:

>>> particle_x = ak.Array([
    [1.0, 2.0, 3.0, 4.0],
    [5.0, 6.0],
    []
])
>>> particle_x
<Array [[1, 2, 3, 4], [5, 6], []] type='3 * var * float64'>

First of all, let's measure the length of each sublist in particle_x. If I compute ak.num(particle_x, axis=-1), Awkward will return a list of integers that corresponds to the lengths of these sublists:

>>> n_particles = ak.num(particle_x, axis=-1)
>>> n_particles
<Array [4, 2, 0] type='3 * int64'>

Now, if I want to only take the sublists that have length 4, I can build a boolean array from this result by asking whether the result of ak.num() is equal to 4:

>>> has_four_particles = n_particles == 4
>>> has_four_particles
<Array [True, False, False] type='3 * bool'>

Keeping only the required sublists

Now that we know where the lists that we want to keep are, we need to work out what to do with them. Here, you need to decide whether you need the result of "deleting" the wrong-sized lists to have the same length as your original array. This might be if you want to compare some aspect of your transformed array with the original.

If you do care about the length of your array, then you can use array.mask[...] to set the values you do not care about to None, e.g.:

>>> particle_x.mask[has_four_particles]
<Array [[1, 2, 3, 4], None, None] type='3 * option[var * float64]'>

You can see that only the first list is retained, whilst the subsequent lists are set to None. This is why the type of the Array has changed from '3 * var * float64' to '3 * option[var * float64]' - the option refers to potentially missing values.

If you don't care about keeping the resulting list the same length, you can just directly filter your array like so:

>>> particle_x[has_four_particles]
<Array [[1, 2, 3, 4]] type='1 * var * float64'>

You can see that this array has the same number of dimensions (2), but only one sublist ([1, 2, 3, 4]).

Final Solution

Therefore, in your case, you want to compute

gen_lv[
    ak.num(gen_lv, axis=-1) == 4
]

[shaoweisong]

Thank you so much for your answer :)