How to restructure arrays by adding fields

[fstrug]

Link to page: https://awkward-array.org/how-to-restructure-add-fields.html

I need some further documentation so that I can understand how to add a new "branch" to an awkward array that comes from an opened TTree using uproot. It seems that TTrees opened in uproot follow a similar structure to this example:

array1 = ak.zip({"x": [1, 2, 3, 4, 5], "y": [1.1, 2.2, 3.3, 4.4, 5.5]})

How would I go about adding a new field to this array such that it is properly indexed/formatted with the original array. If we try:

array1 = ak.zip({"x": [1, 2, 3, 4, 5], "y": [1.1, 2.2, 3.3, 4.4, 5.5]})
array2 = ak.zip({"z" : [6,7,8,9,10]})
print(repr(array1))
print(repr(array2))
**Output:**
<Array [{x: 1, y: 1.1}, ... {x: 5, y: 5.5}] type='5 * {"x": int64, "y": float64}'>
<Array [{z: 6}, {z: 7}, ... {z: 9}, {z: 10}] type='5 * {"z": int64}'>

Now if we try combine array1 and array2 together.

new_array = ak.zip([array1, array2])
print(repr(new_array))

We get this output
<Array [({x: 1, y: 1.1}, ... {z: 10})] type='5 * ({"x": int64, "y": float64}, {"...'>
Instead of something like:
<Array [{x: 1, y: 1.1,z:6}, . . ., {x:5, y:5.5, z:10} ] type='5 * ({"x": int64, "y": float6...'>
How can we add the field "z" and get the desired format?

[jpivarski]

What you're getting is nested: "tuple of (record of x and y) and (record of z)". A "tuple" is a record without field names (these two fields are "0" and "1"), so this is records-within-records because you constructed a record array out of two record arrays.

You can flatten the field structure like this:

new_array = ak.zip({"x": array1.x, "y": array1.y, "z": array2.z})

or more automatically like this:

new_array = ak.zip(dict(zip(ak.fields(array1), ak.unzip(array1))) | dict(zip(ak.fields(array2), ak.unzip(array2))))

using the dict1 | dict2 syntax from Python 3.9+, or if you have an older Python, there's {**dict1, **dict2}.

[fstrug]

I am able to get this working, but I'm not having the same luck when applying this to my data. I get an error when trying to reconstruct the array from itself (I haven't been able to add a 'field' yet because of this). With the following:

filename = "root://kodiak-se.baylor.edu//store/user/lpcsusyhad/Stop_production/Summer16_94X_v3/PostProcessed_11Apr2019_fastsimv5_v6p1_v6p5/SMS_T2tt_mStop_150to250_fastsim_\
2016/SMS_T2tt_mStop150_mLSP0_fastsim_2016_Skim_070621_1_082136_0.root:Events"
events = uproot.open(filename)
batch_size = 1000
for batch in events.iterate(step_size=batch_size, library='ak'):
    recreated_batch = ak.zip(dict(zip(ak.fields(batch), ak.unzip(batch))))

Gives me the following.

Traceback (most recent call last):
  File "/uscms_data/d3/fstrug/SusyAnalysis_CMSSW_12_0_4/CMSSW_12_0_4/src/Studies/acc_study/SB_acc_study.py", line 81, in <module>
    recreated_batch = ak.zip(dict(zip(ak.fields(batch), ak.unzip(batch))))
  File "/uscms_data/d3/fstrug/SusyAnalysis_CMSSW_12_0_4/CMSSW_12_0_4/src/python3_9/lib/python3.9/site-packages/awkward/operations/structure.py", line 664, in zip
    out = ak._util.broadcast_and_apply(
  File "/uscms_data/d3/fstrug/SusyAnalysis_CMSSW_12_0_4/CMSSW_12_0_4/src/python3_9/lib/python3.9/site-packages/awkward/_util.py", line 1166, in broadcast_and_apply
    out = apply(broadcast_pack(inputs, isscalar), 0, user)
  File "/uscms_data/d3/fstrug/SusyAnalysis_CMSSW_12_0_4/CMSSW_12_0_4/src/python3_9/lib/python3.9/site-packages/awkward/_util.py", line 919, in apply
    outcontent = apply(nextinputs, depth + 1, user)
  File "/uscms_data/d3/fstrug/SusyAnalysis_CMSSW_12_0_4/CMSSW_12_0_4/src/python3_9/lib/python3.9/site-packages/awkward/_util.py", line 969, in apply
    nextinputs.append(x.broadcast_tooffsets64(offsets).content)
ValueError: in ListOffsetArray64, cannot broadcast nested list

This tree does have branches that contain vectors such as jet_pts for an event, but I am not sure if this is what the error is referring to.

[jpivarski]

What it's referring to is that the lists you're trying to zip together have a different size. For instance, you could zip and array of "muon pt" with "muon eta" because even though every event has a different number of muons, it has the same number of muon pts as it has muon etas. But it wouldn't be able to zip "muon pt" with "jet pt" because there can be a different number of muons as jets.

(The v2 error message is better than this one... But what's in production is still version 1.)

Anyway, you can set a depth_limit for the ak.zip (see documentation), which keeps it from zipping through all dimensions. You can definitely put "muon pt" in the same record array as "jet pt", as long as they're like this:

{"muon_pt": [#, #, #], "jet_pt": [#, #, #, #, #]}

The down-side is that this also prevents the muon fields from being zipped. That is, you get this:

{"muon_pt": [#, #, #], "muon_eta": [#, #, #]}

rather than this:

{"muon": {"pt": #, "eta": #}}

which is what I'd want if I were you. (You get to think of muons as objects, rather than collections of unconnected attributes.)

You can get what you really want if you do it manually:

ak.zip({
    "muon": ak.zip({
        "pt": batch["Muon_pt"],
        "eta": batch["Muon_eta"],
        "phi": batch["Muon_phi"],
        "mass": batch["Muon_mass"],
    }),
    "jet": ak.zip({
        "pt": batch["Jet_pt"],
        "eta": batch["Jet_eta"],
        "phi": batch["Jet_phi"],
        "mass": batch["Jet_mass"],
    }),
})

The reason this has to be done manually is because it needs input from you about which branches have the same number of items in each event and which don't. If there's a naming convention that guarantees this, then you can script that.