RecordArray masks differently than normal Array

[lnoehte]

Hi,
I just ran into something unexpected when masking a RecordArray. If I have to mask RecordArrays differently compared to normal Arrays, let me know. I didn't find anything about that, yet.
I am using awkward version 1.8.0.

To give an example which behavior I would expect, I will mask an Array first:

import awkward as ak
array1 = ak.Array(
    [
        [53.1, 30.2],
        [],
        [91.7, 88.6, 20.3],
    ]
)
print("array1:")
print(array1)

mask1 = array1 > 40
print("mask1:")
print(mask1)

print()
print("slicing array1 with mask1:")
print(array1[mask1])
print("masking array1 with mask1:")
print(array1.mask[mask1])
print(ak.mask(array1, mask1))

The output is the following:

array1:
[[53.1, 30.2], [], [91.7, 88.6, 20.3]]
mask1:
[[True, False], [], [True, True, False]]

slicing array1 with mask1:
[[53.1], [], [91.7, 88.6]]
masking array1 with mask1:
[[53.1, None], [], [91.7, 88.6, None]]
[[53.1, None], [], [91.7, 88.6, None]]

This is exactly what I was expecting to get.

Now if I do the same thing with a RecordArray:

array2 = ak.Array(
    [
        {
            'Muon_pt': [53.1, 30.2],
        },
        {
            'Muon_pt': [],
        },
        {
            'Muon_pt': [91.7, 88.6, 20.3],
        }
    ]
)
print("array2:")
print(array2)
print(array2.tolist())

mask2 = array2.Muon_pt > 40
print("mask2:")
print(mask2)

print()
print("slicing array2 with mask2:")
print(array2[mask2])
print("masking array2 with mask2:")
print(array2.mask[mask2])
print(ak.mask(array2, mask2))

print()
print("to list:")
print()
print("slicing array2 with mask2:")
print(array2[mask2].tolist())
print("masking array2 with mask2:")
print(array2.mask[mask2].tolist())
print(ak.mask(array2, mask2).tolist())

The output is:

array2:
[{Muon_pt: [53.1, 30.2]}, {Muon_pt: []}, {Muon_pt: [91.7, 88.6, 20.3]}]
[{'Muon_pt': [53.1, 30.2]}, {'Muon_pt': []}, {'Muon_pt': [91.7, 88.6, 20.3]}]
mask2:
[[True, False], [], [True, True, False]]

slicing array2 with mask2:
[{Muon_pt: [53.1]}, {Muon_pt: []}, {Muon_pt: [91.7, 88.6]}]
masking array2 with mask2:
[[{Muon_pt: [53.1, 30.2]}, None], ... 20.3]}, {Muon_pt: [91.7, 88.6, 20.3]}, None]]
[[{Muon_pt: [53.1, 30.2]}, None], ... 20.3]}, {Muon_pt: [91.7, 88.6, 20.3]}, None]]

to list:

slicing array2 with mask2:
[{'Muon_pt': [53.1]}, {'Muon_pt': []}, {'Muon_pt': [91.7, 88.6]}]
masking array2 with mask2:
[[{'Muon_pt': [53.1, 30.2]}, None], [], [{'Muon_pt': [91.7, 88.6, 20.3]}, {'Muon_pt': [91.7, 88.6, 20.3]}, None]]
[[{'Muon_pt': [53.1, 30.2]}, None], [], [{'Muon_pt': [91.7, 88.6, 20.3]}, {'Muon_pt': [91.7, 88.6, 20.3]}, None]]

But what I was expecting was:

[{'Muon_pt': [53.1, None]}, {'Muon_pt': []}, {'Muon_pt': [91.7, 88.6, None]}]

What I see is, that the shape of array2 changed and is now the shape of mask2. I would expect ak.mask to preserve the shape of the initial array.

I don't know if this is the normal behavior of RecordArrays but I would be happy to find a way to get the expected output.

Thanks for all comments and suggestions.

[agoose77]

What's happening here is that the mask and the array are first broadcast together. In this case, this changes the type of the array:

>>> array2.type
3 * {"Muon_pt": var * float64}

>>> mask2.type
3 * var * bool

>>> ak.broadcast_arrays(array2,mask2)
[<Array [[{Muon_pt: [53.1, 30.2, ... 20.3]}]] type='3 * var * {"Muon_pt": var * f...'>,
 <Array [[{Muon_pt: [True, True, ... False]}]] type='3 * var * {"Muon_pt": var * ...'>]

It's not so obvious why array2.mask[mask2] doesn't produce the result that you expect until you consider what would happen if your record had more than a single field. In that case, if e.g. it had a list field called x, how should this be masked if x had a different shape to Muon_pt?

More specifically, you need to provide Awkward with your mask in a RecordArray, e.g.

mask2 = ak.zip({
   "Muon_pt": array2.Muon_pt > 40
}, depth_limit=1)

This will ensure that the broadcast result matches what you expect to happen.

Whether this is the right approach for your problem, though, I can't say! :)

[agoose77]

Pinging @jpivarski because I feel like there's more to add here, but right now my brain is not 100% working!

[jpivarski]

This looks good! Sorry that I was too busy to help out earlier.

[lnoehte]

Thanks a lot for the quick response.

When I saw the shape of mask2 in array2 I already suspected something happening when broadcasting, that I didn't expect.

I also tried it with more than one field, like

array2 = ak.Array(
    [
        {
            'Muon_pt': [53.1, 30.2],
            'Muon_eta': [1.1, 0.2],
        },
        {
            'Muon_pt': [],
            'Muon_eta': [],
        },
        {
            'Muon_pt': [91.7, 88.6, 20.3],
            'Muon_eta': [0.6, -0.5],
        }
    ]
)

I tried to keep the first example simple, because it seems to yield the same result:

array2:
[{Muon_pt: [53.1, 30.2], Muon_eta: [1.1, 0.2], ... 20.3], Muon_eta: [0.6, -0.5]}]
[{'Muon_pt': [53.1, 30.2], 'Muon_eta': [1.1, 0.2]}, {'Muon_pt': [], 'Muon_eta': []}, {'Muon_pt': [91.7, 88.6, 20.3], 'Muon_eta': [0.6, -0.5]}]
mask2:
[[True, False], [], [True, True, False]]

slicing array2 with mask2:
[{Muon_pt: [53.1], Muon_eta: [1.1]}, ... 91.7, 88.6], Muon_eta: [0.6, -0.5]}]
masking array2 with mask2:
[[{Muon_pt: [53.1, 30.2], Muon_eta: [1.1, 0.2], ... Muon_eta: [0.6, -0.5]}, None]]
[[{Muon_pt: [53.1, 30.2], Muon_eta: [1.1, 0.2], ... Muon_eta: [0.6, -0.5]}, None]]

to list:

slicing array2 with mask2:
[{'Muon_pt': [53.1], 'Muon_eta': [1.1]}, {'Muon_pt': [], 'Muon_eta': []}, {'Muon_pt': [91.7, 88.6], 'Muon_eta': [0.6, -0.5]}]
masking array2 with mask2:
[[{'Muon_pt': [53.1, 30.2], 'Muon_eta': [1.1, 0.2]}, None], [], [{'Muon_pt': [91.7, 88.6, 20.3], 'Muon_eta': [0.6, -0.5]}, {'Muon_pt': [91.7, 88.6, 20.3], 'Muon_eta': [0.6, -0.5]}, None]]
[[{'Muon_pt': [53.1, 30.2], 'Muon_eta': [1.1, 0.2]}, None], [], [{'Muon_pt': [91.7, 88.6, 20.3], 'Muon_eta': [0.6, -0.5]}, {'Muon_pt': [91.7, 88.6, 20.3], 'Muon_eta': [0.6, -0.5]}, None]]

But ultimately, I have a bunch of fields in that RecordArray.

Thanks a lot for the example using ak.zip.
If I see it correctly, this only works if I have only the "Muon_pt" field.
I quickly tried to make it work with the new array2 from above by having two identical fields with the keys Muon_pt and Muon_eta in the mask but till now I was unsuccessful with that approach.

I am also wondering if I really need to have duplicates of this array2.Muon_pt > 40 array for each field that exists in array2 in the mask2, considering that I actually have a few more than just two.

[agoose77]

Your "manual zip" example looks as I would expect it - you're correctly masking out the records (['pt', 'eta', 'phi', 'tightId', 'pfRelIso04_all', 'charge']) according to the value of eta. The option type masks-out those records that do not satisfy your mask.

The difference between slicing and masking is subtle - well observed! I realise that you looked into this in your Binder example; I was looking at this quite late last night, so I focussed on the particular challenge of masking.

Ultimately slicing and masking follow different paths. The masking code broadcasts the mask against the array, which stops at the Muon record as discussed above. Meanwhile, the slicing path doesn't use broadcasting; it traverses into the array being sliced, following a set of rules.

At this point, I will ping @jpivarski to clarify our intentions here, because slicing and masking are highly similar operations, and I wonder whether this discrepancy is a good thing. I feel like we've covered this topic before, but I couldn't find anything in our issue tracker. Jim – do we want slicing and masking to be different here? Is it intentional that the mask to stop at the first record, or is it just an artifact of stopping broadcasting at records? i.e. do we need to change this, or is it deliberate?

[lnoehte]

No worries, I've seen that is was quite late yesterday. I am super grateful to get so much support here.

Concerning the "manual zip", it masks out the entries. I just realized that I might have read the structure incorrectly there and though it added a None at the wrong place when it might have done the right thing.

I'll take another look at the "manual zip" method and come back to you later today. I hope with success. If not, probably with more questions ;)

[lnoehte]

OK, I believe it looks fine to me as well.
I must have been confused from all the curly brackets. :D

One small question about the type. What does the ? mean? Is it a sign that it's a mixed type because of the None added?

[agoose77]

?xxx is the same as option[xxx]. It's just easier to read than option[...], so we use it where possible. The type grammar is heavily based upon datashape

[lnoehte]

Thanks a lot.
I cannot mark replies as answers.
Despite it being my "answer" (more a question), I'll mark this one here, since your first reply basically contains the answer to this discussion. I hope people will find it there.