<jupyter_start><jupyter_text>3. Awkward Array  What about an array of lists?<jupyter_code>import awkward as ak
import numpy as np
import uproot
events = uproot.open("data/HZZ.root:events")
events.show()
events["Muon_Px"].array()
events["Muon_Px"].array(entry_stop=20).tolist()<jupyter_output><empty_output><jupyter_text>This is what Awkward Array was made for. NumPy's equivalent is cumbersome and inefficient.<jupyter_code>jagged_numpy = events["Muon_Px"].array(entry_stop=20, library="np")
jagged_numpy<jupyter_output><empty_output><jupyter_text>What if I want the first item in each list as an array?<jupyter_code>np.array([x[0] for x in jagged_numpy])<jupyter_output><empty_output><jupyter_text>This violates the rule from [1-python-performance.ipynb](1-python-performance.ipynb): don't iterate in Python.<jupyter_code>jagged_numpy[:, 0]
jagged_awkward = events["Muon_Px"].array(entry_stop=20, library="ak")
jagged_awkward
jagged_awkward[:, 0]
jagged_awkward[:5, :1]<jupyter_output><empty_output><jupyter_text>Jaggedness in PandasIt can be done by putting the distinction between events in the DataFrame index.<jupyter_code>events.arrays(filter_name="Muon_*", library="pd")<jupyter_output><empty_output><jupyter_text>But if you want multiple particles, they can't be in the same DataFrame.(A DataFrame has only one index; how would you relate jet subentry 1 with muon subentry 1?)<jupyter_code>dataframes = events.arrays(filter_name="/(Muon_|Jet_).*/", library="pd")
len(dataframes)
dataframes[0]
dataframes[1]<jupyter_output><empty_output><jupyter_text>Again, that's why we have Awkward Array.<jupyter_code>array = events.arrays(filter_name="/(Muon_|Jet_).*/", library="ak", how="zip")
array
ak.firsts(array.Jet)
array.Jet.Px
array.Muon
array.Muon.Px
ak.num(array.Jet), ak.num(array.Muon)<jupyter_output><empty_output><jupyter_text>Awkward Array is a general-purpose library: NumPy-like idioms on JSON-like data   Main idea: slicing through structure is computationally inexpensiveSlicing by field name doesn't modify any large buffers and [ak.zip](https://awkward-array.readthedocs.io/en/latest/_auto/ak.zip.html) only scans them to ensure they're compatible (not even that if `depth_limit=1`).<jupyter_code>array = events.arrays()
array<jupyter_output><empty_output><jupyter_text>Think of this as zero-cost:<jupyter_code>array.Muon_Px, array.Muon_Py, array.Muon_Pz<jupyter_output><empty_output><jupyter_text>Think of this as zero-cost:<jupyter_code>ak.zip({"px": array.Muon_Px, "py": array.Muon_Py, "pz": array.Muon_Pz})<jupyter_output><empty_output><jupyter_text>(The above is a manual version of `how="zip"`.) NumPy ufuncs work on these arrays (if they're "[broadcastable](https://awkward-array.readthedocs.io/en/latest/_auto/ak.broadcast_arrays.html)").<jupyter_code>np.sqrt(array.Muon_Px**2 + array.Muon_Py**2)
np.add
array.Muon_Px + array.Muon_Py
import numba as nb
@nb.vectorize([nb.float32(nb.float32, nb.float32)])
def add(x, y):
    return x + y
add
add(array.Muon_Px, array.Muon_Py)
@nb.jit
def anything_else(asdf):
    asdf<jupyter_output><empty_output><jupyter_text>And there are specialized operations that only make sense in a variable-length context.    ak.cartesianak.combinations<jupyter_code>ak.cartesian((array.Muon_Px, array.Jet_Px))
ak.combinations(array.Muon_Px, 3)<jupyter_output><empty_output><jupyter_text>Arrays can have custom [behavior](https://awkward-array.readthedocs.io/en/latest/ak.behavior.html)The following come from the new [Vector](https://github.com/scikit-hep/vectorreadme) library.<jupyter_code>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<jupyter_output><empty_output><jupyter_text>This is an array of lists of vectors, and methods like `pt`, `eta`, `phi` apply through the whole array.<jupyter_code>muons.pt
muons.eta
muons.phi
muons.rapidity
muons[0, 0].pt<jupyter_output><empty_output><jupyter_text>Let's try an example: ΔR(muons, jets)<jupyter_code>jets = ak.zip({"px": array.Jet_Px, "py": array.Jet_Py, "pz": array.Jet_Pz, "E": array.Jet_E}, with_name="Momentum4D")
jets
ak.num(muons), ak.num(jets)
ak.cartesian((muons, jets)).type
ms, js = ak.unzip(ak.cartesian((muons, jets)))
ms, js
ak.num(ms), ak.num(js)
ms.deltaR(js)<jupyter_output><empty_output><jupyter_text>And another: muon pairs (all combinations, not just the first two per event).<jupyter_code>ak.num(muons)
m1, m2 = ak.unzip(ak.combinations(muons, 2))
m1, m2
ak.num(m1), ak.num(m2)
m1 + m2
(m1 + m2).mass
import hist

hist.Hist.new.Reg(120, 0, 120, name="mass").Double().fill(
    ak.flatten((m1 + m2).mass)
).plot()<jupyter_output><empty_output><jupyter_text>It doesn't matter which coordinates were used to construct it<jupyter_code>array2 = uproot.open(
    "root://eospublic.cern.ch//eos/opendata/cms/derived-data/AOD2NanoAODOutreachTool/Run2012BC_DoubleMuParked_Muons.root:Events"
).arrays(["Muon_pt", "Muon_eta", "Muon_phi", "Muon_charge"], entry_stop=100000)
array2.type
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.from_string("mu").mass / 1000.0
muons2
muons2.type<jupyter_output><empty_output><jupyter_text>As long as you use properties (dots, not strings in brackets), you don't need to care what coordinates it's based on.<jupyter_code>muons2.px
muons2.py
muons2.pz
muons2.E
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()<jupyter_output><empty_output><jupyter_text>Awkward Arrays and Vector in NumbaRemember Numba, the JIT-compiler from [1-python-performance.ipynb](1-python-performance.ipynb)? Awkward Array and Vector have been implemented in Numba's compiler.<jupyter_code>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)
array
array.type
muons
muons.type
ak.to_pandas(muons)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1382860763503881, 'fraction_numerical': 0.01694915254237288, 'mean_word_length': 6.316338354577057, 'pattern_counts': {'":': 15, '<': 77, '<?xml version=': 0, '>': 78, 'https://': 4, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '19', 'dir': 'jupyter-structured-clean-dedup', 'id': '426757', 'license': 'BSD-3-Clause', 'n_tokens_mistral': 2525, 'n_tokens_neox': 2397, 'n_words': 484, 'path': 'evaluated/3-awkward-array.ipynb', 'repo_name': 'jpivarski/2021-06-14-uproot-awkward-columnar-hats', 'size': '6313'}