how to add new field in a 'zip' jagged array

Question

I want to add a new field in an already zipped jagged array. For example, if I zip 4D info into a muons object, then I can call pt,eta,phi,charge like this: muons.Muon.pt. However, if I want to add a new field such as 2*pt into this muons object, then I can't do this with muons.Muon['pt2'] = 2 * arrays['Muon_pt"]. Is there anything I misunderstand or how can I add a new field in this jagged array? Could you please help me? thanks

muons = ak.zip({
    "pt": arrays["Muon_pt"],
    "eta": arrays["Muon_eta"],
    "phi": arrays["Muon_phi"],
    "charge": arrays["Muon_charge"],
})

I think I can add a new field in the zipped jagged array, like: muons.Muon['pt2'] then I can call this field with muons.Muon.pt2

Answer 1

What you've described is the right way to go about it, and it should work.

As a walkthrough, we first get some arrays,

>>> import uproot
>>> import skhep_testdata
>>> import awkward as ak
>>> arrays = uproot.open(skhep_testdata.data_path("uproot-HZZ.root"))["events"].arrays()

Then zip them into a structure,

>>> muons = ak.zip({
...     "px": arrays["Muon_Px"],
...     "py": arrays["Muon_Py"],
...     "pz": arrays["Muon_Pz"],
... })
>>> muons.show(type=True)
type: 2421 * var * {
    px: float32,
    py: float32,
    pz: float32
}
[[{px: -52.9, py: -11.7, pz: -8.16}, {px: 37.7, py: 0.693, pz: ..., ...}],
 [{px: -0.816, py: -24.4, pz: 20.2}],
 [{px: 49, py: -21.7, pz: 11.2}, {px: 0.828, py: 29.8, pz: 37}],
 [{px: 22.1, py: -85.8, pz: 404}, {px: 76.7, py: -14, pz: 335}],
 [{px: 45.2, py: 67.2, pz: -89.7}, {px: 39.8, py: 25.4, pz: 20.1}],
 [{px: 9.23, py: 40.6, pz: -14.6}, {px: -5.79, py: -30.3, pz: 43}],
 [{px: 12.5, py: -42.5, pz: -124}, {px: 29.5, py: -4.45, pz: -26.4}],
 [{px: 34.9, py: -16, pz: 156}],
 [{px: -53.2, py: 92, pz: 35.6}, {px: 11.5, py: -4.42, pz: -17.5}],
 [{px: -67, py: 53.2, pz: 54.4}, {px: -18.1, py: -35.1, pz: 58}],
 ...,
 [{px: 14.9, py: 32, pz: -156}],
 [{px: -24.2, py: -35, pz: -19.2}],
 [{px: -9.2, py: -42.2, pz: -64.3}],
 [{px: 34.5, py: 28.8, pz: -151}, {px: -31.6, py: -10.4, pz: -111}],
 [{px: -39.3, py: -14.6, pz: 61.7}],
 [{px: 35.1, py: -14.2, pz: 161}],
 [{px: -29.8, py: -15.3, pz: -52.7}],
 [{px: 1.14, py: 63.6, pz: 162}],
 [{px: 23.9, py: -35.7, pz: 54.7}]]

And now add the new field:

>>> muons["px2"] = 2 * arrays["Muon_Px"]

By construction, this px2 has the same number of elements (deeply, for all sublists) as px. That should also be true in your pt example. So the above line should successfully add the new field; if it doesn't, then I think you want to submit a bug report.

Here's what the new muons looks like for me:

>>> muons.show(type=True)
type: 2421 * var * {
    px: float32,
    py: float32,
    pz: float32,
    px2: float32
}
[[{px: -52.9, py: -11.7, pz: -8.16, px2: -106}, {px: 37.7, py: ..., ...}],
 [{px: -0.816, py: -24.4, pz: 20.2, px2: -1.63}],
 [{px: 49, py: -21.7, pz: 11.2, px2: 98}, {px: 0.828, py: 29.8, ...}],
 [{px: 22.1, py: -85.8, pz: 404, px2: 44.2}, {px: 76.7, py: -14, ...}],
 [{px: 45.2, py: 67.2, pz: -89.7, px2: 90.3}, {px: 39.8, py: 25.4, ...}],
 [{px: 9.23, py: 40.6, pz: -14.6, px2: 18.5}, {px: -5.79, py: -30.3, ...}],
 [{px: 12.5, py: -42.5, pz: -124, px2: 25.1}, {px: 29.5, py: -4.45, ...}],
 [{px: 34.9, py: -16, pz: 156, px2: 69.8}],
 [{px: -53.2, py: 92, pz: 35.6, px2: -106}, {px: 11.5, py: -4.42, ...}],
 [{px: -67, py: 53.2, pz: 54.4, px2: -134}, {px: -18.1, py: -35.1, ...}],
 ...,
 [{px: 14.9, py: 32, pz: -156, px2: 29.8}],
 [{px: -24.2, py: -35, pz: -19.2, px2: -48.3}],
 [{px: -9.2, py: -42.2, pz: -64.3, px2: -18.4}],
 [{px: 34.5, py: 28.8, pz: -151, px2: 69}, {px: -31.6, py: -10.4, ...}],
 [{px: -39.3, py: -14.6, pz: 61.7, px2: -78.6}],
 [{px: 35.1, py: -14.2, pz: 161, px2: 70.1}],
 [{px: -29.8, py: -15.3, pz: -52.7, px2: -59.5}],
 [{px: 1.14, py: 63.6, pz: 162, px2: 2.28}],
 [{px: 23.9, py: -35.7, pz: 54.7, px2: 47.8}]]

px2 has been added to the record type and you can see it in the values. This is Uproot 5 and Awkward Array 2, by the way, but this should also have worked in the previous major version.