How to understand events, objects, and numbers in an awkward array?

[ZhenxuanZhang-Jensen]

Hi,
I'm trying to understand what awkward array can do and use it in my specific particle physics analysis. I'm quite confused about how many exactly things can save in one awkward array. For instance, I want to zip couples of events and each event has couples of particles and each particle has couples of values, then I would code like this: array=ak.zip( { 'muon_pt':[[1,2,3],[4,5],[6]], 'muon_eta':[[1,2,3],[4,5],[6]] , 'muon_phi':[[1,2,3],[4,5],[6]],'muon_E':[[1,2,3],[4,5],[6]],'electron_pt':[[1,2],[3],[4]], 'electron_eta':[[1,2],[3],[4]] , 'electron_phi':[[1,2],[3],[4]],'electron_E':[[1,2],[3],[4]] } ) I want to zip three events to an awaward array with two objects muon & electron 4D infos with different numbers of particle in each event. But if I zip like this, then it will show the very common error as left boardcasting: ValueError: in ListOffsetArray64, cannot broadcast nested list . As far as I'm concerned, I think awkward array is very compatible to save jagged array, this issue should be sovled with some different math magic that I don't know, right?
Please help me if you have any comments or suggection to solve this problem, many thanks.

[agoose77]

When you're posting on GitHub (which uses GitHub Flavoured Markdown), you can use triple backticks ``` to embed entire blocks of code!

The error that your seeing when executing

array = ak.zip(
    {
        "muon_pt": [[1, 2, 3], [4, 5], [6]],
        "muon_eta": [[1, 2, 3], [4, 5], [6]],
        "muon_phi": [[1, 2, 3], [4, 5], [6]],
        "muon_E": [[1, 2, 3], [4, 5], [6]],
        "electron_pt": [[1, 2], [3], [4]],
        "electron_eta": [[1, 2], [3], [4]],
        "electron_phi": [[1, 2], [3], [4]],
        "electron_E": [[1, 2], [3], [4]],
    }
)

occurs when Awkward tries to form a record of "muon_pt", "muon_eta", ..., "electron_E" at the deepest level of nesting. This is the default behaviour. However, it looks like you want each record to contain lists of values. This corresponds to setting a depth_limit in ak.zip:

array = ak.zip(
    {
        "muon_pt": [[1, 2, 3], [4, 5], [6]],
        "muon_eta": [[1, 2, 3], [4, 5], [6]],
        "muon_phi": [[1, 2, 3], [4, 5], [6]],
        "muon_E": [[1, 2, 3], [4, 5], [6]],
        "electron_pt": [[1, 2], [3], [4]],
        "electron_eta": [[1, 2], [3], [4]],
        "electron_phi": [[1, 2], [3], [4]],
        "electron_E": [[1, 2], [3], [4]],
    },
    depth_limit=1
)

This gives you array.type of

3 * {
    "muon_pt": var * int64,
    "muon_eta": var * int64,
    "muon_phi": var * int64,
    "muon_E": var * int64,
    "electron_pt": var * int64,
    "electron_eta": var * int64,
    "electron_phi": var * int64,
    "electron_E": var * int64,
}

I.e. each event has a list of muon_pt, a list of muon_eta, etc.

Some additional notes:

Awkward Array has support for four vectors using the vector library. This library leverages the Awkward Array's behaviour mechanism for adding methods to an Array.

To be able to use these behaviours, you need to have properly named fields, i.e. only E, pt, eta, etc. instead of muon_eta, muon_pt. You can, of course, create these later by zipping the vector fields into a new object, but if you get to choose how to build this array, I'd do it directly:

muon = ak.zip(
    {
        "pt": [[1, 2, 3], [4, 5], [6]],
        "eta": [[1, 2, 3], [4, 5], [6]],
        "phi": [[1, 2, 3], [4, 5], [6]],
        "E": [[1, 2, 3], [4, 5], [6]],
    },
    depth_limit=1,
    with_name="Momentum4D"
)
electron = ak.zip(
    {
        "pt": [[1, 2], [3], [4]],
        "eta": [[1, 2], [3], [4]],
        "phi": [[1, 2], [3], [4]],
        "E": [[1, 2], [3], [4]],
    },
    depth_limit=1,
    with_name="Momentum4D"
)
array = ak.zip(
    {
        "muon": muon,
        "electron": electron
    },
    depth_limit=1
)

In the above code sample, we create two arrays electron and muon that contain your different particle information, and then zip them together into a single array.

See these talks for some examples:
https://github.com/jpivarski-talks/2021-07-06-pyhep-uproot-awkward-tutorial/blob/main/evaluated/uproot-awkward-tutorial.ipynb
https://github.com/jpivarski-talks/2021-09-13-star-uproot-awkward-tutorial/blob/main/tutorial-EVALUATED.ipynb

However, when you place arrays into fields (muon, electron), you restrict the kinds of operations that you can perform with them. Sometimes it's preferable to have a single array, and add a new field that identifies the particle type e.g. pdgid.

[ZhenxuanZhang-Jensen]

Thank you so much for the answer and also the suggestion.

[agoose77]

I've updated my answer with some more information :)

[ZhenxuanZhang-Jensen]

yes, this should be very helpful, thank you again.