How to use awkward - 10 common examples

def fromnumpy(cls, array):
        if array.dtype.kind == "S":
            encoding = None
        elif array.dtype.kind == "U":
            encoding = "utf-32le"
        else:
            raise TypeError("not a string array")

        starts = awkward.util.numpy.arange(                   0,  len(array)      * array.dtype.itemsize, array.dtype.itemsize)
        stops  = awkward.util.numpy.arange(array.dtype.itemsize, (len(array) + 1) * array.dtype.itemsize, array.dtype.itemsize)
        content = array.view(awkward.util.CHARTYPE)

        shorter = awkward.util.numpy.ones(len(array), dtype=awkward.util.BOOLTYPE)
        if array.dtype.kind == "S":
            for checkat in range(array.dtype.itemsize - 1, -1, -1):
                shorter &= (content[checkat::array.dtype.itemsize] == 0)
                stops[shorter] -= 1
                if not shorter.any():
                    break

        elif array.dtype.kind == "U":
            content2 = content.view(awkward.util.numpy.uint32)
            itemsize2 = array.dtype.itemsize >> 2                 # itemsize // 4
            for checkat in range(itemsize2 - 1, -1, -1):
                shorter &= (content2[checkat::itemsize2] == 0)    # all four bytes are zero
                stops[shorter] -= 4
                if not shorter.any():
                    break

def fromnumpy(cls, array):
        if array.dtype.kind == "S":
            encoding = None
        elif array.dtype.kind == "U":
            encoding = "utf-32le"
        else:
            raise TypeError("not a string array")

        starts = awkward.util.numpy.arange(                   0,  len(array)      * array.dtype.itemsize, array.dtype.itemsize)
        stops  = awkward.util.numpy.arange(array.dtype.itemsize, (len(array) + 1) * array.dtype.itemsize, array.dtype.itemsize)
        content = array.view(awkward.util.CHARTYPE)

        shorter = awkward.util.numpy.ones(len(array), dtype=awkward.util.BOOLTYPE)
        if array.dtype.kind == "S":
            for checkat in range(array.dtype.itemsize - 1, -1, -1):
                shorter &= (content[checkat::array.dtype.itemsize] == 0)
                stops[shorter] -= 1
                if not shorter.any():
                    break

        elif array.dtype.kind == "U":
            content2 = content.view(awkward.util.numpy.uint32)
            itemsize2 = array.dtype.itemsize >> 2                 # itemsize // 4
            for checkat in range(itemsize2 - 1, -1, -1):
                shorter &= (content2[checkat::itemsize2] == 0)    # all four bytes are zero

def test_jagged_fillna(self):
        a = awkward.fromiter([[1.1, 2.2, 3.3], [], [4.4, 5.5]])
        assert a.pad(4).fillna(999).tolist() == [[1.1, 2.2, 3.3, 999], [999, 999, 999, 999], [4.4, 5.5, 999, 999]]
        assert a.pad(4, numpy.ma.masked).fillna(999).regular().tolist() == [[1.1, 2.2, 3.3, 999], [999, 999, 999, 999], [4.4, 5.5, 999, 999]]

def test_jagged_cross_argnested(self):
        a = awkward.fromiter([[1.1, 2.2, 3.3], [], [4.4, 5.5]])
        b = awkward.fromiter([[100, 200], [300], [400]])
        c = awkward.fromiter([[999], [999], [999, 888]])

        assert a.cross(b).tolist() == [[(1.1, 100), (1.1, 200), (2.2, 100), (2.2, 200), (3.3, 100), (3.3, 200)], [], [(4.4, 400), (5.5, 400)]]
        assert a.argcross(b).tolist() == [[(0, 0), (0, 1), (1, 0), (1, 1), (2, 0), (2, 1)], [], [(0, 0), (1, 0)]]
        assert a.cross(b, nested=True).tolist() == [[[(1.1, 100), (1.1, 200)], [(2.2, 100), (2.2, 200)], [(3.3, 100), (3.3, 200)]], [], [[(4.4, 400)], [(5.5, 400)]]]
        assert a.argcross(b, nested=True).tolist() == [[[(0, 0), (0, 1)], [(1, 0), (1, 1)], [(2, 0), (2, 1)]], [], [[(0, 0)], [(1, 0)]]]

        assert a.cross(b, nested=True).cross(c, nested=True).tolist()[0] == [[[(ai, bi, ci) for ci in c[0]] for bi in b[0]] for ai in a[0]]
        assert a.cross(b, nested=True).cross(c, nested=True).tolist()[1] == [[[(ai, bi, ci) for ci in c[1]] for bi in b[1]] for ai in a[1]]
        assert a.cross(b, nested=True).cross(c, nested=True).tolist()[2] == [[[(ai, bi, ci) for ci in c[2]] for bi in b[2]] for ai in a[2]]

        assert a.cross(b).cross(c).tolist() == [[(1.1, 100, 999), (1.1, 200, 999), (2.2, 100, 999), (2.2, 200, 999), (3.3, 100, 999), (3.3, 200, 999)], [], [(4.4, 400, 999), (4.4, 400, 888), (5.5, 400, 999), (5.5, 400, 888)]]
        assert a.cross(b, nested=True).cross(c).tolist() == [[[(1.1, 100, 999), (1.1, 200, 999)], [(2.2, 100, 999), (2.2, 200, 999)], [(3.3, 100, 999), (3.3, 200, 999)]], [], [[(4.4, 400, 999), (4.4, 400, 888)], [(5.5, 400, 999), (5.5, 400, 888)]]]
        assert a.cross(b).cross(c, nested=True).tolist() == [[[(1.1, 100, 999)], [(1.1, 200, 999)], [(2.2, 100, 999)], [(2.2, 200, 999)], [(3.3, 100, 999)], [(3.3, 200, 999)]], [], [[(4.4, 400, 999), (4.4, 400, 888)], [(5.5, 400, 999), (5.5, 400, 888)]]]

def test_numba_getitem_jagged_intarray(self):
        a = JaggedArray.fromiter([[1.1, 2.2, 3.3], [], [4.4, 5.5]])
        a2 = JaggedArray.fromcounts([2, 0, 1], a)
        @numba.njit
        def test1(x, i):
            return x[i]
        assert test1(a, awkward.fromiter([[2, 0, 0], [], [1]])).tolist() == [[3.3, 1.1, 1.1], [], [5.5]]
        assert test1(a2, awkward.fromiter([[1, 0], [], [0]])).tolist() == [[[], [1.1, 2.2, 3.3]], [], [[4.4, 5.5]]]
        assert test1(a2, awkward.fromiter([[[2, 0, 0], []], [], [[1]]])).tolist() == [[[3.3, 1.1, 1.1], []], [], [[5.5]]]

def test_physics_matching(self):
        gen_pt   = awkward.fromiter([[10.0, 20.0, 30.0],       [],     [40.0, 50.0]])
        reco_pt  = awkward.fromiter([[20.2, 10.1, 30.3, 50.5], [50.5], [50.5]])

        gen_eta  = awkward.fromiter([[-3.0, -2.0, 2.0],        [],     [-1.0, 1.0]])
        reco_eta = awkward.fromiter([[-2.2, -3.3, 2.2, 0.0],   [0.0],  [1.1]])

        gen_phi  = awkward.fromiter([[-1.5,  0.0, 1.5],        [],     [0.78, -0.78]])
        reco_phi = awkward.fromiter([[ 0.1, -1.4, 1.4, 0.78],  [0.78], [-0.77]])

        gen  = uproot_methods.TLorentzVectorArray.from_ptetaphim(gen_pt, gen_eta, gen_phi, 0.2)
        reco = uproot_methods.TLorentzVectorArray.from_ptetaphim(reco_pt, reco_eta, reco_phi, 0.2)

        ("gen", gen)
        ("reco", reco)

        ("gen.cross(reco)", gen.cross(reco))

        pairing = gen.cross(reco, nested=True)
        ("pairing = gen.cross(reco, nested=True)", gen.cross(reco, nested=True))

        metric = pairing.i0.delta_r(pairing.i1)
        ("metric = pairing.i0.delta_r(pairing.i1)", metric)

assert isinstance(awkward.fromiter(x).tolist()[0]["a"], float)
        x.insert(1, None)
        assert awkward.fromiter(x).tolist() == x
        x.insert(0, None)
        assert awkward.fromiter(x).tolist() == x

        x = [{"a": 1, "b": {"x": 1.1}}, {"a": 2, "b": {"x": 2.2}}, {"a": 3, "b": {"x": 3.3}}]
        assert awkward.fromiter(x).tolist() == x
        x.insert(1, None)
        assert awkward.fromiter(x).tolist() == x
        x.insert(0, None)
        assert awkward.fromiter(x).tolist() == x

        assert awkward.fromiter([{}]).tolist() == [None]
        assert awkward.fromiter([{}, {}, {}]).tolist() == [None, None, None]
        assert awkward.fromiter([{"a": 1, "b": 1.1}, {"a": 2, "b": 2.2}, {}]).tolist() == [{"a": 1, "b": 1.1}, {"a": 2, "b": 2.2}, None]
        assert awkward.fromiter([None, {"a": 1, "b": 1.1}, {"a": 2, "b": 2.2}, {}]).tolist() == [None, {"a": 1, "b": 1.1}, {"a": 2, "b": 2.2}, None]
        assert awkward.fromiter([{"a": 1, "b": 1.1}, None, {"a": 2, "b": 2.2}, {}]).tolist() == [{"a": 1, "b": 1.1}, None, {"a": 2, "b": 2.2}, None]
        assert awkward.fromiter([{"a": 1, "b": 1.1}, {}, {"a": 2, "b": 2.2}]).tolist() == [{"a": 1, "b": 1.1}, None, {"a": 2, "b": 2.2}]
        assert awkward.fromiter([{}, {"a": 1, "b": 1.1}, {"a": 2, "b": 2.2}]).tolist() == [None, {"a": 1, "b": 1.1}, {"a": 2, "b": 2.2}]

x = [1, [], 2]
        assert awkward.fromiter(x).tolist() == x
        x.insert(1, None)
        assert awkward.fromiter(x).tolist() == x
        x.insert(0, None)
        assert awkward.fromiter(x).tolist() == x

        x = [1, [3.14], 2]
        assert awkward.fromiter(x).tolist() == x
        x.insert(1, None)
        assert awkward.fromiter(x).tolist() == x
        x.insert(0, None)
        assert awkward.fromiter(x).tolist() == x

        x = [[], 1, 2]
        assert awkward.fromiter(x).tolist() == x
        x.insert(1, None)
        assert awkward.fromiter(x).tolist() == x
        x.insert(0, None)
        assert awkward.fromiter(x).tolist() == x

        x = [[3.14], 1, 2]
        assert awkward.fromiter(x).tolist() == x
        x.insert(1, None)
        assert awkward.fromiter(x).tolist() == x
        x.insert(0, None)
        assert awkward.fromiter(x).tolist() == x

        x = [1, [], []]
        assert awkward.fromiter(x).tolist() == x
        x.insert(1, None)
        assert awkward.fromiter(x).tolist() == x

def test_jagged_cross_argnested(self):
        a = awkward.fromiter([[1.1, 2.2, 3.3], [], [4.4, 5.5]])
        b = awkward.fromiter([[100, 200], [300], [400]])
        c = awkward.fromiter([[999], [999], [999, 888]])

        assert a.cross(b).tolist() == [[(1.1, 100), (1.1, 200), (2.2, 100), (2.2, 200), (3.3, 100), (3.3, 200)], [], [(4.4, 400), (5.5, 400)]]
        assert a.argcross(b).tolist() == [[(0, 0), (0, 1), (1, 0), (1, 1), (2, 0), (2, 1)], [], [(0, 0), (1, 0)]]
        assert a.cross(b, nested=True).tolist() == [[[(1.1, 100), (1.1, 200)], [(2.2, 100), (2.2, 200)], [(3.3, 100), (3.3, 200)]], [], [[(4.4, 400)], [(5.5, 400)]]]
        assert a.argcross(b, nested=True).tolist() == [[[(0, 0), (0, 1)], [(1, 0), (1, 1)], [(2, 0), (2, 1)]], [], [[(0, 0)], [(1, 0)]]]

        assert a.cross(b, nested=True).cross(c, nested=True).tolist()[0] == [[[(ai, bi, ci) for ci in c[0]] for bi in b[0]] for ai in a[0]]
        assert a.cross(b, nested=True).cross(c, nested=True).tolist()[1] == [[[(ai, bi, ci) for ci in c[1]] for bi in b[1]] for ai in a[1]]
        assert a.cross(b, nested=True).cross(c, nested=True).tolist()[2] == [[[(ai, bi, ci) for ci in c[2]] for bi in b[2]] for ai in a[2]]

        assert a.cross(b).cross(c).tolist() == [[(1.1, 100, 999), (1.1, 200, 999), (2.2, 100, 999), (2.2, 200, 999), (3.3, 100, 999), (3.3, 200, 999)], [], [(4.4, 400, 999), (4.4, 400, 888), (5.5, 400, 999), (5.5, 400, 888)]]
        assert a.cross(b, nested=True).cross(c).tolist() == [[[(1.1, 100, 999), (1.1, 200, 999)], [(2.2, 100, 999), (2.2, 200, 999)], [(3.3, 100, 999), (3.3, 200, 999)]], [], [[(4.4, 400, 999), (4.4, 400, 888)], [(5.5, 400, 999), (5.5, 400, 888)]]]
        assert a.cross(b).cross(c, nested=True).tolist() == [[[(1.1, 100, 999)], [(1.1, 200, 999)], [(2.2, 100, 999)], [(2.2, 200, 999)], [(3.3, 100, 999)], [(3.3, 200, 999)]], [], [[(4.4, 400, 999), (4.4, 400, 888)], [(5.5, 400, 999), (5.5, 400, 888)]]]

        a = awkward.fromiter([[1.1, 2.2, 3.3], [], [4.4, 5.5]])

def test_jagged_parents(self):
        a = awkward.fromiter([[1.1, 2.2, 3.3, 4.4, 5.5], [], [6.6, 7.7, 8.8], [9.9]])
        assert a.parents.tolist() == [0, 0, 0, 0, 0, 2, 2, 2, 3]
        b = a[[False, True, False, True]]
        assert b.parents.tolist() == [-1, -1, -1, -1, -1, -1, -1, -1, 1]