Combinations within an array of p4 vectors

[srosenzweig09]

Hi, I have an array of p4 vectors and I am trying to calculate the DeltaR between all vectors in the array. I build the array like this:

for pt, eta, phi in zip(sixb_pt, sixb_eta, sixb_phi): 
    signal.append({"pt":pt, "eta":eta, "phi":phi, "mass":4e-9})
signal_p4.append(ak.Array(signal, with_name="MomentumObject4D"))

for each event in a ROOT file. For now, there are exactly 6 entries in the array but this may vary later. The output is a list of these arrays. Then, I try to find the combinations of pairs but I get an error when running the following code:

combos = ak.combinations(signal_p4s[0], 2) lefts, rights = ak.unzip(combos)

The error is as follows:

ValueError                                Traceback (most recent call last)
<ipython-input-24-8e6af5ac9919> in <module>
----> 1 combos = ak.combinations(signal_p4s[0], 2)
      2 lefts, rights = ak.unzip(combos)
      3 # lefts.deltaR(rights)

/eos/user/s/srosenzw/miniconda3/envs/work/lib/python3.8/site-packages/awkward1/operations/structure.py in combinations(array, n, replacement, axis, fields, parameters, with_name, highlevel)
   2614         array, allow_record=False, allow_other=False
   2615     )
-> 2616     out = layout.combinations(
   2617         n, replacement=replacement, keys=fields, parameters=parameters, axis=axis
   2618     )

ValueError: 'axis' out of range for combinations

(https://github.com/scikit-hep/awkward-1.0/blob/0.4.5/src/libawkward/array/NumpyArray.cpp#L2855)

Any advice on how to handle this? Thank you!

[jpivarski]

The default axis for ak.combinations is 1 because a typical use-case is lists of particles per event. I don't think that's the data type you're building—since signal consists of records ({"pt": pt, "eta": eta, "phi": phi, "mass": mass}) and not lists of records, axis=1 is beyond the depth of this array.

If I'm interpreting the data type correctly, axis=0 would work for this array, but the idea is not to find all combinations of particles in a single event, it's to find all combinations of particles in all events, with a single function call.

Instead of putting this inside an event loop, collect all your data as lists of vectors in all events (one list per event). Then call ak.combinations on that with the default axis=1.

Here's an example of doing that without having to create a bunch of Python objects (which use a lot of memory). It uses ak.ArrayBuilder:

builder = ak.ArrayBuilder()

for event in events:
    sixb_pt = event["sixb_pt"]    # however you get these from the events
    sixb_eta = event["sixb_eta"]
    sixb_phi = event["sixb_phi"]

    with builder.list():
        for pt, eta, phi in zip(sixb_pt, sixb_eta, sixb_phi):
            with builder.record("Momentum4D"):   # not MomentumObject4D
                builder.field("pt"); builder.real(pt)
                builder.field("eta"); builder.real(eta)
                builder.field("phi"); builder.real(phi)

signal_p4 = builder.snapshot()

combos = ak.combinations(signal_p4, 2)
lefts, rights = ak.unzip(combos)

deltaRs = lefts.deltaR(rights)

(The name of the record type is not "MomentumObject4D" because it is not a Python object.)

[srosenzweig09]

I think I should mention that type(lefts) provides the output awkward1.highlevel.Array, in case that aids in understanding what is happening. Should lefts have this type?

[jpivarski]

If installing the behaviors globally, vector.register_awkward() only has to be called once, and it affects all ak.Arrays (which is ak.highlevel.Array) that are created thereafter, including through snapshot().

I just noticed that your awkward module name has a "1" in it: awkward1. It might be an old version (when 1.x still had to be qualified). If you're using pip, try clearing out all of your old packages and get the new ones:

pip uninstall awkward awkward1 uproot-methods uproot uproot4  # repeatedly, until there are none left
pip install awkward uproot   # if you need uproot

and import them without qualification:

import awkward as ak   # no "1"
import uproot   # no "4"

Edit: To directly answer your question, yes, there's only one Python type for all Awkward Arrays, regardless of the type of data they contain. So combos is an ak.Array of 2-tuples and lefts is an ak.Array of just the left halfs of each 2-tuple.

[srosenzweig09]

Alright, so I did a cleanse and I think it was a step in the right direction! I now see that

However, I am getting this error now:

So for some reason, it thinks I am trying to use azimuthal coordinates despite entering "pt," "eta," and "phi." Any ideas about what's going on here? Thank you!!

[jpivarski]

The azimuthal coordinates that it is (and should be) looking for are "pt" and "phi". It's not finding them because it's calling "from_fields" instead of "from_momentum_fields," which recognizes "pt" as a synonym for "rho". The line numbers in your stack trace don't line up with the latest version of the code—what version are you using?

I'm looking at the changes since the last released version, and there are quite a few, including one about Awkward Array and momentum synonyms (#84). It may be that we've been fixing bugs and just not releasing. I just released 0.8.2; try that.

[srosenzweig09]

That did it! I wiped uproot and awkward again but originally, I hadn't wiped vector. So I wiped all three (just in case) and reinstalled them and now it works! Thanks so much for the suggestion! I really appreciate your time and energy!