<jupyter_start><jupyter_text>2019-12-20-coffea-demo 1. IntroductionThis demo of the new Awkward Array was presented on December 20, 2019, before the final 1.0 version was released. Some interfaces may have changed. To run this notebook, make sure you have version 0.1.36 ([GitHub](https://github.com/scikit-hep/awkward-1.0/releases/tag/0.1.36), [pip](https://pypi.org/project/awkward1/0.1.36/)) by installing```bashpip install 'awkward1==0.1.36'```The basic concepts of Awkward arrays are presented on the [old Awkward README](https://github.com/scikit-hep/awkward-0.x/tree/0.12.17readme) and the motivation for a 1.0 rewrite are presented on the [new Awkward README](https://github.com/scikit-hep/awkward-1.0/tree/0.1.32readme).<jupyter_code># The base of the GitHub repo is two levels up from this notebook.
import sys
import os
sys.path.insert(0, os.path.join(os.getcwd(), "..", ".."))<jupyter_output><empty_output><jupyter_text>2. High-level array classThe biggest user-facing change is that, instead of mixing NumPy arrays and `JaggedArray` objects, the new Awkward has a single `Array` class for data analysis.<jupyter_code>import numpy as np
import awkward1 as ak

array1 = ak.Array([[1.1, 2.2, 3.3], [], [4.4, 5.5]])
array1
array2 = ak.Array([{"x": 0, "y": []}, {"x": 1, "y": [1.1]}, {"x": 2, "y": [1.1, 2.2]}])
array2<jupyter_output><empty_output><jupyter_text>The same `Array` class is used for all data structures, such as the array of lists in `array1` and the array of records in `array2`. (Incidentally, the width of that string representation is exactly large enough to fit into GitHub and StackOverflow text boxes without scrolling!)There won't be any user-level functions that apply to some data types and not others. The result of an operation is likely type-dependent, but its accessibility is not. (At this time, the only existing operations are conversions like `ak.tolist` and descriptions like `ak.typeof`.)<jupyter_code>ak.tolist(array1)
ak.tojson(array1)
ak.tolist(array2)
ak.tojson(array2)
ak.typeof(array1)
ak.typeof(array2)<jupyter_output><empty_output><jupyter_text>Data types are described using the [datashape language](https://datashape.readthedocs.io/en/latest/). Some Awkward features are [not expressible](https://github.com/blaze/datashape/issues/237) in the current datashape specification, so they're expressed in an extension of the language using the same _style_ of syntax.  3. Low-level array classesThe old `JaggedArray` and `Table` are still available, but you have to ask for them explicitly with `layout`. They're not "private" or "internal implementations" (there's no underscore in `layout`): they're public for frameworks like Coffea but hidden from data analysts.As such, their string representations have more low-level detail: the contents of indexes, rather than what they mean as high-level types. (The XML formatting is just an elaboration on Python's angle-bracket convention for `repr` and the fact that we need to denote nesting.)<jupyter_code>array1.layout
array2.layout<jupyter_output><empty_output><jupyter_text>These classes are defined in C++ and wrapped by pybind11. The `awkward1.Array` class is pure Python. Many of the same operations work for layout classes, though less attention has been paid to its interface.<jupyter_code>ak.typeof(array1)
ak.typeof(array1.layout)
ak.tojson(array1)
ak.tojson(array1.layout)
array1.layout.tojson()<jupyter_output><empty_output><jupyter_text>4. Behavioral mix-insIn the original Awkward library, we added behaviors to objects and arrays of objects, like computing `pt` or boosting for arrays of Lorentz vectors, by letting structure classes such as `JaggedArray` multiply inherit from classes providing the implementations. That technique was not fully thought-through: it was easy to lose an array's "Lorentzness" when slicing it or performing other operations. It also relies on a Python language feature that can't pass through C++.It has since become clear that behavioral mix-ins aren't an obscure use-case but a primary one, so its implementation requires more thought. Adding behaviors to arrays is now a "first-class feature," built into the array types themselves.<jupyter_code>class PointClass(ak.Record):
    def __repr__(self):
        return "<Point({}, {})>".format(self["x"], self["y"])
    
    def mag(self):
        return abs(np.sqrt(self["x"]**2 + self["y"]**2))

ak.namespace["Point"] = PointClass
array3 = ak.Array([{"x": 1, "y": 1.1}, {"x": 2, "y": 2.2}, {"x": 3, "y": 3.3}])
array3
array3.layout.type<jupyter_output><empty_output><jupyter_text>Types can have arbitrary parameters, which modify their meaning. These types are JSON-encoded and passed through C++ or wherever the arrays get sent.<jupyter_code>pointtype = array3.layout.type
pointtype["__class__"] = "Point"
pointtype
pointtype["__str__"] = "PointType[{}, {}]".format(pointtype.field("x"), pointtype.field("y"))
pointtype
# There will be a better interface for assigning types...
array4 = ak.Array(array3.layout, type=ak.ArrayType(pointtype, len(array3.layout)))
array4
[x.mag() for x in array4]<jupyter_output><empty_output><jupyter_text>The elements of this array are `PointClass` instances because the `__class__` parameter is `"Point"`, a name that is recognized in Awkward's class namespace. The global namespace is in `ak.namespace`, but custom ones can be passed into the `Array` constructor to turn on/off or change behaviors.<jupyter_code>ak.namespace<jupyter_output><empty_output><jupyter_text>As you can see, variable-length strings are also implemented as mix-ins. Apart from this type annotation, a string is just a jagged array of 8-bit integers.<jupyter_code>array5 = ak.Array(["Daisy", "Daisy", "give", "me", "your", "answer", "do."])
array5
array5.layout
ak.tolist(array5.layout)
# Slice it!
ak.tolist(array5.layout[:, 1:])
# Slice it!
array5[:, 1:]<jupyter_output><empty_output><jupyter_text>Like all behavioral mix-ins, the string interpretation is _only_ applied in the high-level `Array` view, not the layout classes. Thus, a C++ function that generates or uses jagged arrays of Lorentz vectors (e.g. a nice FastJet interface?) does not depend on Python. It only has to manipulate a map of strings.The old Awkward also had an `ObjectArray`, which generated Python objects on demand, such as individual Lorentz vectors, and these had to have the same set of methods as arrays of Lorentz vectors. Keeping those coordinated was difficult. Now, however, the individual objects don't disinherit from the Awkward arrays they come from: the strings above are merely a view (which is why the slice worked). Instead of `Methods` and `ObjectArrays`, we now have a unified mechanism.For instance, this `PointClass` object,<jupyter_code>array4[2]<jupyter_output><empty_output><jupyter_text>is still an Awkward `Record`.<jupyter_code>array4[2].layout<jupyter_output><empty_output><jupyter_text>5. Agreement with NumPyAwkward array represents a superset of NumPy's core, so it must return the same results as NumPy. This was tricky in the old Awkward, when we restricted ourselves to vectorized functions, and this led to hidden limitations: slices were limited to depth `2`, concatenation was limited to `axis <= 1`, and `choose(n)` was limited to `n < 5`. But now that we can write compiled for loops, there are no such limitations.<jupyter_code>deepnumpy = np.arange(2*3*5*7).reshape(2, 3, 5, 7)
deepawkward = ak.Array(deepnumpy)
deepawkward
deepnumpy[1:, :2, [4, 1, 1, -2], ::-1]
deepawkward[1:, :2, [4, 1, 1, -2], ::-1]
ak.tolist(deepnumpy[1:, :2, [4, 1, 1, -2], ::-1]) == ak.tolist(deepawkward[1:, :2, [4, 1, 1, -2], ::-1])<jupyter_output><empty_output><jupyter_text>6. Creating arraysA few of the examples above create arrays by passing them to the `Array` constructor. This is like old Awkward's `fromiter` function. In fact, new Awkward has a `fromiter` function, but it's implicitly called by the `Array` constructor.<jupyter_code># Calls ak.fromiter, which converts rowwise → columnar data.
ak.Array([[1.1, 2.2, 3.3], [], [4.4, 5.5]])
# Calls ak.fromjson, which deserializes.
ak.Array("[[1.1, 2.2, 3.3], [], [4.4, 5.5]]")
# Calls ak.fromnumpy, which views.
nparray = np.array([[1.1, 2.2, 3.3], [4.4, 5.5, 6.6]])
akarray = ak.Array(nparray)
akarray
nparray[0, 1] = 999
akarray<jupyter_output><empty_output><jupyter_text>You can also build these manually from the low-level layouts, but it's a lot of work!<jupyter_code>content = ak.layout.NumpyArray(np.array([1.1, 2.2, 3.3, 4.4, 5.5]))
offsets = ak.layout.Index64(np.array([0, 3, 3, 5], dtype=np.int64))   # match 64-bit to 64-bit to avoid copy
listoffsetarray = ak.layout.ListOffsetArray64(offsets, content)
listoffsetarray
ak.Array(listoffsetarray)<jupyter_output><empty_output><jupyter_text>7. FillableArrayThe `fromiter` algorithm has been expanded into a builder interface, so that you can accumulate Awkward arrays gradually.<jupyter_code>builder = ak.FillableArray()

for i in range(10):
    builder.beginrecord()
    builder.field("x")
    builder.real(np.random.normal())
    builder.field("y")
    builder.beginlist()
    for j in range(np.random.poisson(2.5)):
        builder.integer(np.random.randint(0, 10))
    builder.endlist()
    builder.endrecord()

builder<jupyter_output><empty_output><jupyter_text>This is not a regular array, but you can `snapshot` it to get one (and keep filling the `builder`). A `snapshot` does not copy array data: if you take several snapshots while filling, they _might_ share data. (And they _might_ not, if it has allocated new buffers to grow beyond its reserved space!)<jupyter_code>array6 = builder.snapshot()
array6
ak.tolist(array6)
ak.typeof(array6)<jupyter_output><empty_output><jupyter_text>The array that you produce can have nested structure, as shown above. The structure was determined by the order in which `builder` methods were called.You can write algorithms that build arrays as if you were printing out JSON:   * call `beginlist()` instead of printing `"["`,   * call `endlist()` instead of printing `"]"`,   * call `beginrecord()` instead of printing `"{"`,   * call `endrecord()` instead of printing `"}"`,   * call `field(key)` instead of printing `"key":`, etc.<jupyter_code>deepbuilder = ak.FillableArray()

def deepnesting(depth):
    if depth == 0:
        deepbuilder.integer(np.random.randint(0, 10))
    else:
        deepbuilder.beginlist()
        for j in range(np.random.poisson(2.5)):
            deepnesting(depth - 1)
        deepbuilder.endlist()

deepnesting(5)
ak.tolist(deepbuilder.snapshot())
ak.typeof(deepbuilder)
deepbuilder.snapshot().layout<jupyter_output><empty_output><jupyter_text>Both `fromiter` and `fromjson` are implemented using `FillableArray`, the latter using the RapidJSON C++ library for deserialization.<jupyter_code># !wget https://scikit-hep.org/uproot/examples/HZZ.json
hzz = ak.fromjson("HZZ.json")
hzz
for key in hzz.layout.keys():
    print("{:18s} {}".format(key, hzz[key].type))<jupyter_output>jets               2421 * var * {"px": float64, "py": float64, "pz": float64, "E": float64, "id": bool}
muons              2421 * var * {"px": float64, "py": float64, "pz": float64, "E": float64, "q": int64, "iso": float64}
electrons          2421 * var * {"px": float64, "py": float64, "pz": float64, "E": float64, "q": int64, "iso": float64}
photons            2421 * var * {"px": float64, "py": float64, "pz": float64, "E": float64, "iso": float64}
MET                2421 * {"x": float64, "y": float64}
MC_hadronic_b      2421 * {"px": float64, "py": float64, "pz": float64}
MC_leptonic_b      2421 * {"px": float64, "py": float64, "pz": float64}
MC_hadronicW_q     2421 * {"px": float64, "py": float64, "pz": float64}
MC_hadronicW_qbar  2421 * {"px": float64, "py": float64, "pz": float64}
MC_lepton          2421 * {"px": float64, "py": float64, "pz": float64, "pdgid": int64}
MC_neutrino        2421 * {"px": float64, "py": float64, "pz": float64}
num_PV             2421 * int64
trigger_isomu[...]<jupyter_text>The loop over Python objects or JSON nodes was moved from Python into C++, so it's faster. However, the implementation requires vtable-lookups (the type is discovered at runtime), so it's not _a lot_ faster. There's room for specialized methods when the type is known in advance. (See [src/libawkward/io/root.cpp](https://github.com/scikit-hep/awkward-1.0/blob/main/src/libawkward/io/root.cpp) for a ${\tt std::vector}^N{\tt}$ implementation.)In general, turning rowwise data into columnar data is about 10× faster than it used to be.<jupyter_code>import awkward as oldawkward
import json
aslist = json.load(open("HZZ.json")) * 10
asjson = json.dumps(aslist)
%%timeit -r 3

ak.fromiter(aslist)                       # new fromiter
%%timeit -r 3

oldawkward.fromiter(aslist)               # old fromiter
%%timeit -r 3

ak.fromjson(asjson)                       # new fromjson
%%timeit -r 3

oldawkward.fromiter(json.loads(asjson))   # old equivalent of fromjson<jupyter_output>2.29 s ± 97.2 ms per loop (mean ± std. dev. of 3 runs, 1 loop each)<jupyter_text>8. Awkward arrays in NumbaOne of the motivating goals of the Awkward re-write was to incorporate Numba on the same footing.<jupyter_code>import numba

@numba.jit(nopython=True)
def muon_sumpt(events):
    out = np.zeros(len(events), np.float64)
    i = 0
    for event in events:
        for muon in event["muons"]:
            out[i] += np.sqrt(muon["px"]**2 + muon["py"]**2)
        i += 1
    return out
hzz = ak.Array(json.load(open("HZZ.json")) * 100).layout
muon_sumpt(hzz)<jupyter_output><empty_output><jupyter_text>Notice that we can write for loops on event _records_ and muon _records_. We don't have to take apart `JaggedArrays` and write algorithms on offsets and indexes.Incidentally, my first version of the above raised segfaults because the `i += 1` was in the inner loop, rather than the outer loop (indentation error). Since it's Numba, I could debug it by running the pure Python version.<jupyter_code>muon_sumpt.py_func
muon_sumpt.py_func(hzz)
%%timeit -r 3

muon_sumpt(hzz)                           # in Numba
%%timeit -r 1

muon_sumpt.py_func(hzz)                   # pure Python<jupyter_output>11.4 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)<jupyter_text>So, a 50× speedup over Python without changing any code. Debug in Python, accelerate with Numba. Awkward arrays could be a benefit to Numba users in general: Numba can handle complex data types by converting Python objects to and from equivalent structs, but that puts a translation burden at the entry and exit of every Numba function. Awkward leaves the data in the same form (big array buffers), transforming only its handles to the data. (JSON → Awkward arrays → Numba could become a useful workflow in industry.)  Side note...It looks like there's a performance bug in the current implementation: if we remove all particles but muons from what we pass through Numba, we see a 3× speedup relative to leaving them in that scales with the size of the dataset (100× vs 1000×). That shouldn't happen: unused fields are supposed to be ignored in the compiled code. Once everything is operational, we'll investigate these performance issues.<jupyter_code>%%timeit -r 3

muon_sumpt(hzz.astype(None)[["muons"]])   # in Numba, passing only muons through
%%timeit -r 3

muon_sumpt(hzz.astype(None))              # in Numba, passing everything through<jupyter_output>232 ms ± 5.31 ms per loop (mean ± std. dev. of 3 runs, 1 loop each)<jupyter_text>9. FillableArray in NumbaThe example above took an Awkward array into a Numba function and did some processing on it. To write data out, we can use `FillableArrays`.<jupyter_code>@numba.jit(nopython=True)
def make_data(builder):
    for i in range(10):
        builder.beginrecord()

        builder.field("x")
        builder.real(i*1.1)

        builder.field("y")
        builder.beginlist()
        for j in range(i):
            builder.integer(j)
        builder.endlist()

        builder.endrecord()

    return builder
builder = ak.layout.FillableArray()
make_data(builder)
builder.snapshot()
ak.tolist(builder.snapshot())<jupyter_output><empty_output><jupyter_text>Since you can walk over data structures and create data structures (and later, assign fields to datasets like the old `Table`), you have complete freedom to manipulate data   * at compiled-code speeds,   * without having to leave the Python environment,   * without having to rethink your algorithm in terms of array-at-a-time functions.(This _supplements_ the array-at-a-time approach introduced last year.)  10. Awkward arrays in C++Since everything has been implemented in C++, Awkward 1.0 can be used in C++ programs. More importantly, we will (someday) be able to create Awkward arrays in C++ and access them in Python or vice-versa.<jupyter_code>open("test-program.cpp", "w").write("""

#include <iostream>

#include "awkward/fillable/FillableArray.h"
#include "awkward/fillable/FillableOptions.h"

namespace ak = awkward;

int main(int, char**) {
  ak::FillableArray builder(ak::FillableOptions(1024, 2.0));
  for (int i = 0;  i < 10;  i++) {
    builder.beginrecord();

    builder.field_fast("x");    // (field_fast means don't check the whole string, just its pointer)
    builder.real(i*1.1);

    builder.field_fast("y");
    builder.beginlist();
    for (int j = 0;  j < i;  j++) {
      builder.integer(j);
    }
    builder.endlist();

    builder.endrecord();
  }
  
  std::cout << builder.snapshot()->tojson(false, 1) << std::endl;
  return 0;
}
""")
import pygments.formatters
import pygments.lexers.c_cpp
print(pygments.highlight(open("test-program.cpp").read(),
                         pygments.lexers.c_cpp.CppLexer(),
                         pygments.formatters.Terminal256Formatter()))
!g++ -I../../include -L../../awkward1 test-program.cpp -lawkward-static -lawkward-cpu-kernels-static -o test-program
!./test-program<jupyter_output>[{"x":0.0,"y":[]},{"x":1.1,"y":[0]},{"x":2.2,"y":[0,1]},{"x":3.3,"y":[0,1,2]},{"x":4.4,"y":[0,1,2,3]},{"x":5.5,"y":[0,1,2,3,4]},{"x":6.6,"y":[0,1,2,3,4,5]},{"x":7.7,"y":[0,1,2,3,4,5,6]},{"x":8.8,"y":[0,1,2,3,4,5,6,7]},{"x":9.9,"y":[0,1,2,3,4,5,6,7,8]}]<jupyter_text>11. Identities: database-like index for arraysIn the [PartiQL toy language](https://github.com/jpivarski/PartiQLreadme), it became apparent that set operations, in which unique records are identified by reference, rather than by value, are important. They provide such operations as joins and lossless unions.No set operations have been implemented, but implementing them will require an index that tracks particle identities through all other operations. This concept of an index is the primary distinction between an array library like NumPy and a relational library like Pandas. In Awkward, this index is called an `Identity` and can optionally be attached to arrays.**Note:** this interface is the most likely to change. Identities have only been implemented at this early stage so that they don't have to be painfully retrofitted later.<jupyter_code>hzzlayout = ak.fromjson("HZZ.json").layout
hzzlayout.setid()
hzzlayout.id
hzzlayout.field("muons").content.field("px").id
np.asarray(hzzlayout.field("muons").content.field("px").id)<jupyter_output><empty_output><jupyter_text>An `Identity` is a 2-dimensional array with the same structure as a Pandas row `MultiIndex` with a `fieldloc` for the nested columns. They're equivalent to paths from root (wherever you called `setid`) to the element in question.<jupyter_code>hzzlayout[1000, "muons", 1].location<jupyter_output><empty_output><jupyter_text>As a nice side-effect of having indexes, we can give better error messages about where an indexing error occurs. You might use `Identities` just for debugging.<jupyter_code># Indexing error with an Identity:
try:
    hzzlayout[1000, "muons", 2]
except Exception as err:
    print(err)
# Indexing error without an Identity:
try:
    ak.fromjson("HZZ.json").layout[1000, "muons", 2]
except Exception as err:
    print(err)<jupyter_output>in ListArray64 attempting to get 2, index out of range<jupyter_text>When the array goes through any kind of transformation, such as the boolean filter below, the `Identity` is similarly selected.<jupyter_code>mask = np.random.randint(0, 100, len(hzzlayout)) == 0
mask
selected = hzzlayout[mask]
np.asarray(selected.id)<jupyter_output><empty_output>
StarcoderdataJupyterStructuredCleanDedup