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

I see now that your question was cross-posted to https://github.com/scikit-hep/awkward/discussions/2017, where Angus found your issue: you said `muons.Muon["pt2"]`, rather than `muons["pt2"]` (and I must have auto-corrected in my head). Also, the GitHub Discussions tend to be more active than StackOverflow for Uproot/Awkward questions; I saw this because I've set up an alert in StackOverflow, but I don't think anyone else who works on these projects has done that. — Jim Pivarski, Dec 19 '22 at 17:57
ahh, thanks for letting me know, then I will keep this in GitHub. — Zhenxuan, Dec 21 '22 at 07:47

score 0 · Answer 1 · answered Dec 19 '22 at 17:53

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.

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

1 Answers1