Uproot Awkward Columnar HATS#

Originally presented as part of CMS HATS training on June 14, 2021.






What about an array of lists?#

import skhep_testdata
import awkward as ak
import numpy as np
import uproot
events = uproot.open(skhep_testdata.data_path("uproot-HZZ.root"))["events"]
events.show()
name                 | typename                 | interpretation                
---------------------+--------------------------+-------------------------------
NJet                 | int32_t                  | AsDtype('>i4')
Jet_Px               | float[]                  | AsJagged(AsDtype('>f4'))
Jet_Py               | float[]                  | AsJagged(AsDtype('>f4'))
Jet_Pz               | float[]                  | AsJagged(AsDtype('>f4'))
Jet_E                | float[]                  | AsJagged(AsDtype('>f4'))
Jet_btag             | float[]                  | AsJagged(AsDtype('>f4'))
Jet_ID               | bool[]                   | AsJagged(AsDtype('bool'))
NMuon                | int32_t                  | AsDtype('>i4')
Muon_Px              | float[]                  | AsJagged(AsDtype('>f4'))
Muon_Py              | float[]                  | AsJagged(AsDtype('>f4'))
Muon_Pz              | float[]                  | AsJagged(AsDtype('>f4'))
Muon_E               | float[]                  | AsJagged(AsDtype('>f4'))
Muon_Charge          | int32_t[]                | AsJagged(AsDtype('>i4'))
Muon_Iso             | float[]                  | AsJagged(AsDtype('>f4'))
NElectron            | int32_t                  | AsDtype('>i4')
Electron_Px          | float[]                  | AsJagged(AsDtype('>f4'))
Electron_Py          | float[]                  | AsJagged(AsDtype('>f4'))
Electron_Pz          | float[]                  | AsJagged(AsDtype('>f4'))
Electron_E           | float[]                  | AsJagged(AsDtype('>f4'))
Electron_Charge      | int32_t[]                | AsJagged(AsDtype('>i4'))
Electron_Iso         | float[]                  | AsJagged(AsDtype('>f4'))
NPhoton              | int32_t                  | AsDtype('>i4')
Photon_Px            | float[]                  | AsJagged(AsDtype('>f4'))
Photon_Py            | float[]                  | AsJagged(AsDtype('>f4'))
Photon_Pz            | float[]                  | AsJagged(AsDtype('>f4'))
Photon_E             | float[]                  | AsJagged(AsDtype('>f4'))
Photon_Iso           | float[]                  | AsJagged(AsDtype('>f4'))
MET_px               | float                    | AsDtype('>f4')
MET_py               | float                    | AsDtype('>f4')
MChadronicBottom_px  | float                    | AsDtype('>f4')
MChadronicBottom_py  | float                    | AsDtype('>f4')
MChadronicBottom_pz  | float                    | AsDtype('>f4')
MCleptonicBottom_px  | float                    | AsDtype('>f4')
MCleptonicBottom_py  | float                    | AsDtype('>f4')
MCleptonicBottom_pz  | float                    | AsDtype('>f4')
MChadronicWDecayQ... | float                    | AsDtype('>f4')
MChadronicWDecayQ... | float                    | AsDtype('>f4')
MChadronicWDecayQ... | float                    | AsDtype('>f4')
MChadronicWDecayQ... | float                    | AsDtype('>f4')
MChadronicWDecayQ... | float                    | AsDtype('>f4')
MChadronicWDecayQ... | float                    | AsDtype('>f4')
MClepton_px          | float                    | AsDtype('>f4')
MClepton_py          | float                    | AsDtype('>f4')
MClepton_pz          | float                    | AsDtype('>f4')
MCleptonPDGid        | int32_t                  | AsDtype('>i4')
MCneutrino_px        | float                    | AsDtype('>f4')
MCneutrino_py        | float                    | AsDtype('>f4')
MCneutrino_pz        | float                    | AsDtype('>f4')
NPrimaryVertices     | int32_t                  | AsDtype('>i4')
triggerIsoMu24       | bool                     | AsDtype('bool')
EventWeight          | float                    | AsDtype('>f4')
events["Muon_Px"].array()
[[-52.9, 37.7],
 [-0.816],
 [49, 0.828],
 [22.1, 76.7],
 [45.2, 39.8],
 [9.23, -5.79],
 [12.5, 29.5],
 [34.9],
 [-53.2, 11.5],
 [-67, -18.1],
 ...,
 [14.9],
 [-24.2],
 [-9.2],
 [34.5, -31.6],
 [-39.3],
 [35.1],
 [-29.8],
 [1.14],
 [23.9]]
--------------------------
type: 2421 * var * float32
events["Muon_Px"].array(entry_stop=20).tolist()
[[-52.89945602416992, 37.7377815246582],
 [-0.8164593577384949],
 [48.987831115722656, 0.8275666832923889],
 [22.08833122253418, 76.6919174194336],
 [45.171321868896484, 39.75095748901367],
 [9.228110313415527, -5.793715000152588],
 [12.538717269897461, 29.541839599609375],
 [34.883758544921875],
 [-53.16697311401367, 11.491869926452637],
 [-67.01485443115234, -18.118755340576172],
 [15.983028411865234, 34.68440628051758],
 [-70.51190948486328, -38.028743743896484],
 [58.94381332397461],
 [-15.587870597839355],
 [-122.33011627197266, -1.0597527027130127],
 [-46.70415496826172, 39.020023345947266],
 [51.29465866088867, 17.45092010498047],
 [43.28120040893555],
 [-45.92393493652344, 22.549766540527344],
 [43.29360580444336, -33.28158187866211, -4.376191139221191]]

This is what Awkward Array was made for. NumPy’s equivalent is cumbersome and inefficient.

jagged_numpy = events["Muon_Px"].array(entry_stop=20, library="np")
jagged_numpy
array([array([-52.899456,  37.73778 ], dtype=float32),
       array([-0.81645936], dtype=float32),
       array([48.98783  ,  0.8275667], dtype=float32),
       array([22.088331, 76.69192 ], dtype=float32),
       array([45.17132 , 39.750957], dtype=float32),
       array([ 9.22811 , -5.793715], dtype=float32),
       array([12.538717, 29.54184 ], dtype=float32),
       array([34.88376], dtype=float32),
       array([-53.166973,  11.49187 ], dtype=float32),
       array([-67.014854, -18.118755], dtype=float32),
       array([15.983028, 34.684406], dtype=float32),
       array([-70.51191 , -38.028744], dtype=float32),
       array([58.943813], dtype=float32),
       array([-15.587871], dtype=float32),
       array([-122.33012  ,   -1.0597527], dtype=float32),
       array([-46.704155,  39.020023], dtype=float32),
       array([51.29466, 17.45092], dtype=float32),
       array([43.2812], dtype=float32),
       array([-45.923935,  22.549767], dtype=float32),
       array([ 43.293606, -33.28158 ,  -4.376191], dtype=float32)],
      dtype=object)

What if I want the first item in each list as an array?

np.array([x[0] for x in jagged_numpy])
array([ -52.899456  ,   -0.81645936,   48.98783   ,   22.088331  ,
         45.17132   ,    9.22811   ,   12.538717  ,   34.88376   ,
        -53.166973  ,  -67.014854  ,   15.983028  ,  -70.51191   ,
         58.943813  ,  -15.587871  , -122.33012   ,  -46.704155  ,
         51.29466   ,   43.2812    ,  -45.923935  ,   43.293606  ],
      dtype=float32)

This violates the rule from 1-python-performance.ipynb: don’t iterate in Python.

jagged_awkward = events["Muon_Px"].array(entry_stop=20, library="ak")
jagged_awkward
[[-52.9, 37.7],
 [-0.816],
 [49, 0.828],
 [22.1, 76.7],
 [45.2, 39.8],
 [9.23, -5.79],
 [12.5, 29.5],
 [34.9],
 [-53.2, 11.5],
 [-67, -18.1],
 [16, 34.7],
 [-70.5, -38],
 [58.9],
 [-15.6],
 [-122, -1.06],
 [-46.7, 39],
 [51.3, 17.5],
 [43.3],
 [-45.9, 22.5],
 [43.3, -33.3, -4.38]]
------------------------
type: 20 * var * float32
jagged_awkward[:, 0]
[-52.9,
 -0.816,
 49,
 22.1,
 45.2,
 9.23,
 12.5,
 34.9,
 -53.2,
 -67,
 16,
 -70.5,
 58.9,
 -15.6,
 -122,
 -46.7,
 51.3,
 43.3,
 -45.9,
 43.3]
------------------
type: 20 * float32






Awkward Array is a general-purpose library: NumPy-like idioms on JSON-like data#






Main idea: slicing through structure is computationally inexpensive#

Slicing by field name doesn’t modify any large buffers and ak.zip only scans them to ensure they’re compatible (not even that if depth_limit=1).

array = events.arrays()
array
[{NJet: 0, Jet_Px: [], Jet_Py: [], Jet_Pz: [], Jet_E: [], Jet_btag: [], ...},
 {NJet: 1, Jet_Px: [-38.9], Jet_Py: [19.9], Jet_Pz: [-0.895], Jet_E: ..., ...},
 {NJet: 0, Jet_Px: [], Jet_Py: [], Jet_Pz: [], Jet_E: [], Jet_btag: [], ...},
 {NJet: 3, Jet_Px: [-71.7, 36.6, -28.9], Jet_Py: [93.6, ...], Jet_Pz: ..., ...},
 {NJet: 2, Jet_Px: [3.88, 4.98], Jet_Py: [-75.2, ...], Jet_Pz: [...], ...},
 {NJet: 2, Jet_Px: [-46.3, 27.7], Jet_Py: [27.3, ...], Jet_Pz: [...], ...},
 {NJet: 1, Jet_Px: [-82.1], Jet_Py: [47.7], Jet_Pz: [-262], Jet_E: [279], ...},
 {NJet: 1, Jet_Px: [27.9], Jet_Py: [-32.9], Jet_Pz: [231], Jet_E: [235], ...},
 {NJet: 0, Jet_Px: [], Jet_Py: [], Jet_Pz: [], Jet_E: [], Jet_btag: [], ...},
 {NJet: 1, Jet_Px: [37.9], Jet_Py: [-8.89], Jet_Pz: [-69.2], Jet_E: ..., ...},
 ...,
 {NJet: 1, Jet_Px: [-25.6], Jet_Py: [-20.6], Jet_Pz: [-113], Jet_E: [118], ...},
 {NJet: 2, Jet_Px: [3.43, 38.3], Jet_Py: [45.6, ...], Jet_Pz: [...], ...},
 {NJet: 1, Jet_Px: [34], Jet_Py: [58.9], Jet_Pz: [-17], Jet_E: [70.5], ...},
 {NJet: 0, Jet_Px: [], Jet_Py: [], Jet_Pz: [], Jet_E: [], Jet_btag: [], ...},
 {NJet: 1, Jet_Px: [37.1], Jet_Py: [20.1], Jet_Pz: [226], Jet_E: [230], ...},
 {NJet: 2, Jet_Px: [-33.2, -26.1], Jet_Py: [-59.7, ...], Jet_Pz: [...], ...},
 {NJet: 1, Jet_Px: [-3.71], Jet_Py: [-37.2], Jet_Pz: [41], Jet_E: [56], ...},
 {NJet: 2, Jet_Px: [-36.4, -15.3], Jet_Py: [10.2, ...], Jet_Pz: [...], ...},
 {NJet: 0, Jet_Px: [], Jet_Py: [], Jet_Pz: [], Jet_E: [], Jet_btag: [], ...}]
--------------------------------------------------------------------------------
type: 2421 * {
    NJet: int32,
    Jet_Px: var * float32,
    Jet_Py: var * float32,
    Jet_Pz: var * float32,
    Jet_E: var * float32,
    Jet_btag: var * float32,
    Jet_ID: var * bool,
    NMuon: int32,
    Muon_Px: var * float32,
    Muon_Py: var * float32,
    Muon_Pz: var * float32,
    Muon_E: var * float32,
    Muon_Charge: var * int32,
    Muon_Iso: var * float32,
    NElectron: int32,
    Electron_Px: var * float32,
    Electron_Py: var * float32,
    Electron_Pz: var * float32,
    Electron_E: var * float32,
    Electron_Charge: var * int32,
    Electron_Iso: var * float32,
    NPhoton: int32,
    Photon_Px: var * float32,
    Photon_Py: var * float32,
    Photon_Pz: var * float32,
    Photon_E: var * float32,
    Photon_Iso: var * float32,
    MET_px: float32,
    MET_py: float32,
    MChadronicBottom_px: float32,
    MChadronicBottom_py: float32,
    MChadronicBottom_pz: float32,
    MCleptonicBottom_px: float32,
    MCleptonicBottom_py: float32,
    MCleptonicBottom_pz: float32,
    MChadronicWDecayQuark_px: float32,
    MChadronicWDecayQuark_py: float32,
    MChadronicWDecayQuark_pz: float32,
    MChadronicWDecayQuarkBar_px: float32,
    MChadronicWDecayQuarkBar_py: float32,
    MChadronicWDecayQuarkBar_pz: float32,
    MClepton_px: float32,
    MClepton_py: float32,
    MClepton_pz: float32,
    MCleptonPDGid: int32,
    MCneutrino_px: float32,
    MCneutrino_py: float32,
    MCneutrino_pz: float32,
    NPrimaryVertices: int32,
    triggerIsoMu24: bool,
    EventWeight: float32
}

Think of this as zero-cost:

array.Muon_Px, array.Muon_Py, array.Muon_Pz
(<Array [[-52.9, 37.7], [-0.816], ..., [23.9]] type='2421 * var * float32'>,
 <Array [[-11.7, 0.693], [-24.4], ..., [-35.7]] type='2421 * var * float32'>,
 <Array [[-8.16, -11.3], [20.2], ..., [162], [54.7]] type='2421 * var * float32'>)

Think of this as zero-cost:

ak.zip({"px": array.Muon_Px, "py": array.Muon_Py, "pz": array.Muon_Pz})
[[{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}]]
--------------------------------------------------------------------------
type: 2421 * var * {
    px: float32,
    py: float32,
    pz: float32
}

(The above is a manual version of how="zip".)




NumPy ufuncs work on these arrays (if they’re “broadcastable”).

np.sqrt(array.Muon_Px**2 + array.Muon_Py**2)
[[54.2, 37.7],
 [24.4],
 [53.6, 29.8],
 [88.6, 78],
 [81, 47.2],
 [41.6, 30.8],
 [44.4, 29.9],
 [38.4],
 [106, 12.3],
 [85.5, 39.5],
 ...,
 [35.3],
 [42.6],
 [43.2],
 [45, 33.2],
 [41.9],
 [37.8],
 [33.5],
 [63.6],
 [42.9]]
--------------------------
type: 2421 * var * float32




And there are specialized operations that only make sense in a variable-length context.

ak.cartesian()

ak.combinations()

ak.cartesian((array.Muon_Px, array.Jet_Px))
[[],
 [(-0.816, -38.9)],
 [],
 [(22.1, -71.7), (22.1, 36.6), (22.1, ...), ..., (76.7, 36.6), (76.7, -28.9)],
 [(45.2, 3.88), (45.2, 4.98), (39.8, 3.88), (39.8, 4.98)],
 [(9.23, -46.3), (9.23, 27.7), (-5.79, -46.3), (-5.79, 27.7)],
 [(12.5, -82.1), (29.5, -82.1)],
 [(34.9, 27.9)],
 [],
 [(-67, 37.9), (-18.1, 37.9)],
 ...,
 [(14.9, -25.6)],
 [(-24.2, 3.43), (-24.2, 38.3)],
 [(-9.2, 34)],
 [],
 [(-39.3, 37.1)],
 [(35.1, -33.2), (35.1, -26.1)],
 [(-29.8, -3.71)],
 [(1.14, -36.4), (1.14, -15.3)],
 []]
------------------------------------------------------------------------------
type: 2421 * var * (
    float32,
    float32
)
ak.combinations(array.Muon_Px, 2)
[[(-52.9, 37.7)],
 [],
 [(49, 0.828)],
 [(22.1, 76.7)],
 [(45.2, 39.8)],
 [(9.23, -5.79)],
 [(12.5, 29.5)],
 [],
 [(-53.2, 11.5)],
 [(-67, -18.1)],
 ...,
 [],
 [],
 [],
 [(34.5, -31.6)],
 [],
 [],
 [],
 [],
 []]
--------------------
type: 2421 * var * (
    float32,
    float32
)






Arrays can have custom behavior#

The following come from the new Vector library.

import vector
vector.register_awkward()
muons = ak.zip({"px": array.Muon_Px, "py": array.Muon_Py, "pz": array.Muon_Pz, "E": array.Muon_E}, with_name="Momentum4D")
muons
[[{px: -52.9, py: -11.7, pz: -8.16, E: 54.8}, {px: 37.7, py: 0.693, ...}],
 [{px: -0.816, py: -24.4, pz: 20.2, E: 31.7}],
 [{px: 49, py: -21.7, pz: 11.2, E: 54.7}, {px: 0.828, py: 29.8, ...}],
 [{px: 22.1, py: -85.8, pz: 404, E: 413}, {px: 76.7, py: -14, ...}],
 [{px: 45.2, py: 67.2, pz: -89.7, E: 121}, {px: 39.8, py: 25.4, ...}],
 [{px: 9.23, py: 40.6, pz: -14.6, E: 44.1}, {px: -5.79, py: -30.3, ...}],
 [{px: 12.5, py: -42.5, pz: -124, E: 132}, {px: 29.5, py: -4.45, ...}],
 [{px: 34.9, py: -16, pz: 156, E: 160}],
 [{px: -53.2, py: 92, pz: 35.6, E: 112}, {px: 11.5, py: -4.42, ...}],
 [{px: -67, py: 53.2, pz: 54.4, E: 101}, {px: -18.1, py: -35.1, ...}],
 ...,
 [{px: 14.9, py: 32, pz: -156, E: 160}],
 [{px: -24.2, py: -35, pz: -19.2, E: 46.7}],
 [{px: -9.2, py: -42.2, pz: -64.3, E: 77.4}],
 [{px: 34.5, py: 28.8, pz: -151, E: 157}, {px: -31.6, py: -10.4, ...}],
 [{px: -39.3, py: -14.6, pz: 61.7, E: 74.6}],
 [{px: 35.1, py: -14.2, pz: 161, E: 165}],
 [{px: -29.8, py: -15.3, pz: -52.7, E: 62.4}],
 [{px: 1.14, py: 63.6, pz: 162, E: 174}],
 [{px: 23.9, py: -35.7, pz: 54.7, E: 69.6}]]
--------------------------------------------------------------------------
type: 2421 * var * Momentum4D[
    px: float32,
    py: float32,
    pz: float32,
    E: float32
]

This is an array of lists of vectors, and methods like pt, eta, phi apply through the whole array.

muons.pt
[[54.2, 37.7],
 [24.4],
 [53.6, 29.8],
 [88.6, 78],
 [81, 47.2],
 [41.6, 30.8],
 [44.4, 29.9],
 [38.4],
 [106, 12.3],
 [85.5, 39.5],
 ...,
 [35.3],
 [42.6],
 [43.2],
 [45, 33.2],
 [41.9],
 [37.8],
 [33.5],
 [63.6],
 [42.9]]
--------------------------
type: 2421 * var * float32
muons.eta
[[-0.15, -0.295],
 [0.754],
 [0.207, 1.04],
 [2.22, 2.16],
 [-0.955, 0.414],
 [-0.345, 1.13],
 [-1.76, -0.796],
 [2.11],
 [0.329, -1.15],
 [0.6, 1.18],
 ...,
 [-2.19],
 [-0.437],
 [-1.19],
 [-1.92, -1.92],
 [1.18],
 [2.15],
 [-1.24],
 [1.67],
 [1.06]]
--------------------------
type: 2421 * var * float32
muons.phi
[[-2.92, 0.0184],
 [-1.6],
 [-0.417, 1.54],
 [-1.32, -0.18],
 [0.979, 0.569],
 [1.35, -1.76],
 [-1.28, -0.149],
 [-0.43],
 [2.09, -0.367],
 [2.47, -2.05],
 ...,
 [1.13],
 [-2.17],
 [-1.79],
 [0.696, -2.82],
 [-2.79],
 [-0.384],
 [-2.67],
 [1.55],
 [-0.98]]
--------------------------
type: 2421 * var * float32




Let’s try an example: ΔR(muons, jets)

jets = ak.zip({"px": array.Jet_Px, "py": array.Jet_Py, "pz": array.Jet_Pz, "E": array.Jet_E}, with_name="Momentum4D")
jets
[[],
 [{px: -38.9, py: 19.9, pz: -0.895, E: 44.1}],
 [],
 [{px: -71.7, py: 93.6, pz: 196, E: 230}, ..., {px: -28.9, py: 9.32, ...}],
 [{px: 3.88, py: -75.2, pz: -360, E: 368}, {px: 4.98, py: -39.2, ...}],
 [{px: -46.3, py: 27.3, pz: 51.4, E: 74.4}, {px: 27.7, py: -14.5, ...}],
 [{px: -82.1, py: 47.7, pz: -262, E: 279}],
 [{px: 27.9, py: -32.9, pz: 231, E: 235}],
 [],
 [{px: 37.9, py: -8.89, pz: -69.2, E: 79.9}],
 ...,
 [{px: -25.6, py: -20.6, pz: -113, E: 118}],
 [{px: 3.43, py: 45.6, pz: -130, E: 138}, {px: 38.3, py: -3.39, ...}],
 [{px: 34, py: 58.9, pz: -17, E: 70.5}],
 [],
 [{px: 37.1, py: 20.1, pz: 226, E: 230}],
 [{px: -33.2, py: -59.7, pz: -29, E: 74.9}, {px: -26.1, py: -19.1, ...}],
 [{px: -3.71, py: -37.2, pz: 41, E: 56}],
 [{px: -36.4, py: 10.2, pz: 226, E: 230}, {px: -15.3, py: -27.2, ...}],
 []]
---------------------------------------------------------------------------
type: 2421 * var * Momentum4D[
    px: float32,
    py: float32,
    pz: float32,
    E: float32
]
ak.num(muons), ak.num(jets)
(<Array [2, 1, 2, 2, 2, 2, 2, 1, ..., 1, 2, 1, 1, 1, 1, 1] type='2421 * int64'>,
 <Array [0, 1, 0, 3, 2, 2, 1, 1, ..., 1, 0, 1, 2, 1, 2, 0] type='2421 * int64'>)
ms, js = ak.unzip(ak.cartesian((muons, jets)))
ms, js
(<MomentumArray4D [[], [{px: -0.816, ...}], ..., []] type='2421 * var * Mome...'>,
 <MomentumArray4D [[], [{px: -38.9, ...}], ..., [...], []] type='2421 * var ...'>)
ak.num(ms), ak.num(js)
(<Array [0, 1, 0, 6, 4, 4, 2, 1, ..., 1, 0, 1, 2, 1, 2, 0] type='2421 * int64'>,
 <Array [0, 1, 0, 6, 4, 4, 2, 1, ..., 1, 0, 1, 2, 1, 2, 0] type='2421 * int64'>)
ms.deltaR(js)
[[],
 [2.15],
 [],
 [2.9, 1.99, 2.33, 2.56, 0.972, 3.13],
 [2.82, 3.32, 3.4, 2.21],
 [1.74, 1.9, 1.93, 1.6],
 [2.38, 2.92],
 [0.515],
 [],
 [3.33, 3.1],
 ...,
 [2.7],
 [2.93, 2.38],
 [2.99],
 [],
 [3.23],
 [3.08, 2.55],
 [2.4],
 [1.55, 2.94],
 []]
--------------------------------------
type: 2421 * var * float32




And another: muon pairs (all combinations, not just the first two per event).

ak.num(muons)
[2,
 1,
 2,
 2,
 2,
 2,
 2,
 1,
 2,
 2,
 ...,
 1,
 1,
 1,
 2,
 1,
 1,
 1,
 1,
 1]
------------------
type: 2421 * int64
m1, m2 = ak.unzip(ak.combinations(muons, 2))
m1, m2
(<MomentumArray4D [[{px: -52.9, py: -11.7, ...}], ..., []] type='2421 * var ...'>,
 <MomentumArray4D [[{px: 37.7, py: 0.693, ...}], ..., []] type='2421 * var *...'>)
ak.num(m1), ak.num(m2)
(<Array [1, 0, 1, 1, 1, 1, 1, 0, ..., 0, 1, 0, 0, 0, 0, 0] type='2421 * int64'>,
 <Array [1, 0, 1, 1, 1, 1, 1, 0, ..., 0, 1, 0, 0, 0, 0, 0] type='2421 * int64'>)
m1 + m2
[[{x: -15.2, y: -11, z: -19.5, t: 94.2}],
 [],
 [{x: 49.8, y: 8.08, z: 48.1, t: 102}],
 [{x: 98.8, y: -99.8, z: 739, t: 758}],
 [{x: 84.9, y: 92.7, z: -69.6, t: 172}],
 [{x: 3.43, y: 10.3, z: 28.3, t: 97}],
 [{x: 42.1, y: -47, z: -151, t: 172}],
 [],
 [{x: -41.7, y: 87.6, z: 18.2, t: 133}],
 [{x: -85.1, y: 18.1, z: 112, t: 172}],
 ...,
 [],
 [],
 [],
 [{x: 2.94, y: 18.4, z: -262, t: 273}],
 [],
 [],
 [],
 [],
 []]
-----------------------------------------
type: 2421 * var * Momentum4D[
    x: float32,
    y: float32,
    z: float32,
    t: float32
]
(m1 + m2).mass
[[90.2],
 [],
 [74.7],
 [89.8],
 [94.9],
 [92.1],
 [53.4],
 [],
 [89.8],
 [96],
 ...,
 [],
 [],
 [],
 [76],
 [],
 [],
 [],
 [],
 []]
--------------------------
type: 2421 * var * float32
import hist

hist.Hist.new.Reg(120, 0, 120, name="mass").Double().fill(
    ak.flatten((m1 + m2).mass)
).plot()

None
../_images/72cfed029bfd82c8ce310aa6bde6472b09f7d2108f6f276abe19ecdba0c96b5f.png




It doesn’t matter which coordinates were used to construct it#

array2 = uproot.open(
    "https://github.com/jpivarski-talks/2023-12-18-hsf-india-tutorial-bhubaneswar/raw/main/data/SMHiggsToZZTo4L.root:Events"
).arrays(["Muon_pt", "Muon_eta", "Muon_phi", "Muon_charge"], entry_stop=100000)
import particle

muons2 = ak.zip({"pt": array2.Muon_pt, "eta": array2.Muon_eta, "phi": array2.Muon_phi, "q": array2.Muon_charge}, with_name="Momentum4D")
muons2["mass"] = particle.Particle.findall("mu-")[0].mass / 1000.0
muons2
[[{pt: 63, eta: -0.719, phi: 2.97, q: 1, mass: 0.106}, ..., {pt: 4.05, ...}],
 [],
 [],
 [{pt: 54.3, eta: -1.06, phi: -0.362, q: 1, mass: 0.106}, ..., {pt: 3.49, ...}],
 [],
 [{pt: 38.5, eta: 0.315, phi: 2.05, q: -1, mass: 0.106}, {pt: 47, ...}],
 [{pt: 4.45, eta: -0.986, phi: 1.12, q: 1, mass: 0.106}],
 [],
 [],
 [],
 ...,
 [],
 [{pt: 6.47, eta: -1.08, phi: -2.1, q: -1, mass: 0.106}, {pt: 16.2, ...}],
 [],
 [{pt: 33.8, eta: 1.15, phi: -2.67, q: 1, mass: 0.106}, {pt: 16.4, ...}],
 [{pt: 6.39, eta: -1.47, phi: 0.895, q: 1, mass: 0.106}, ..., {pt: 8.33, ...}],
 [],
 [{pt: 35.9, eta: 1.4, phi: -2.86, q: 1, mass: 0.106}],
 [{pt: 15.1, eta: -0.84, phi: 0.0732, q: 1, mass: 0.106}],
 []]
--------------------------------------------------------------------------------
type: 100000 * var * Momentum4D[
    pt: float32,
    eta: float32,
    phi: float32,
    q: int32,
    mass: float64
]

As long as you use properties (dots, not strings in brackets), you don’t need to care what coordinates it’s based on.

muons2.px
[[-62.1, 19.5, 2.05],
 [],
 [],
 [50.8, 22.4, 13, -1.88, 3.76, 6.11, -1.54],
 [],
 [-17.7, 19.1],
 [1.93],
 [],
 [],
 [],
 ...,
 [],
 [-3.29, -15.6],
 [],
 [-30, -16.2],
 [4, 18.7, -12.8, -8.32],
 [],
 [-34.5],
 [15],
 []]
--------------------------------------------
type: 100000 * var * float32
muons2.py
[[10.9, -32.7, 3.49],
 [],
 [],
 [-19.3, -7.1, 51.2, 3.9, -3.8, -5.76, 3.13],
 [],
 [34.2, -42.9],
 [4.01],
 [],
 [],
 [],
 ...,
 [],
 [-5.57, 4.39],
 [],
 [-15.5, -2.18],
 [4.99, -2.66, 15.1, -0.266],
 [],
 [-10.1],
 [1.1],
 []]
---------------------------------------------
type: 100000 * var * float32
muons2.pz
[[-49.3, -38, -1.32],
 [],
 [],
 [-68.8, -5.88, -1.4, 3.65, -6.45, -8.88, 2.96],
 [],
 [12.3, -5.59],
 [-5.14],
 [],
 [],
 [],
 ...,
 [],
 [-8.39, -8.74],
 [],
 [48, 3.9],
 [-13.1, -42.2, 5.85, 40.8],
 [],
 [68.1],
 [-14.2],
 []]
------------------------------------------------
type: 100000 * var * float32
muons2.E
[[80, 53.8, 4.26],
 [],
 [],
 [87.7, 24.2, 52.9, 5.67, 8.38, 12.2, 4.58],
 [],
 [40.4, 47.3],
 [6.8],
 [],
 [],
 [],
 ...,
 [],
 [10.6, 18.4],
 [],
 [58.7, 16.8],
 [14.6, 46.2, 20.7, 41.7],
 [],
 [76.9],
 [20.7],
 []]
--------------------------------------------
type: 100000 * var * float64
m1, m2 = ak.unzip(ak.combinations(muons2, 2))
hist.Hist.new.Log(200, 0.1, 120, name="mass").Double().fill(
    ak.flatten((m1 + m2).mass)
).plot()

None
../_images/002ace33fda07b8d5ba6bc82aebb97564b9f8056a5f2abe7bbf8cf8756ba5dfe.png




Awkward Arrays and Vector in Numba#

Remember Numba, the JIT-compiler from 1-python-performance.ipynb? Awkward Array and Vector have been implemented in Numba’s compiler.

import numba as nb

@nb.njit
def first_big_dimuon(events):
    for event in events:
        for i in range(len(event)):
            mu1 = event[i]
            for j in range(i + 1, len(event)):
                mu2 = event[j]
                dimuon = mu1 + mu2
                if dimuon.mass > 10:
                    return dimuon
first_big_dimuon(muons2)
MomentumObject4D(pt=47.82986068725586, phi=-2.6673949400531214, eta=-1.3629037141799927, mass=89.49191465603693)