How to use the awkward.JaggedArray function in awkward

mc_rand = np.load('tests/samples/nano_dy_rochester_rand.npy')

    # test against nanoaod
    events = NanoEvents.from_file(os.path.abspath('tests/samples/nano_dimuon.root'))

    data_k = rochester.kScaleDT(events.Muon.charge, events.Muon.pt, events.Muon.eta, events.Muon.phi)
    assert(all(np.isclose(data_k.flatten(), official_data_k)))
    data_err = rochester.kScaleDTerror(events.Muon.charge, events.Muon.pt, events.Muon.eta, events.Muon.phi)
    data_err = np.array(data_err.flatten(), dtype=float)
    assert(all(np.isclose(data_err, official_data_err, atol=1e-8)))

    # test against mc
    events = NanoEvents.from_file(os.path.abspath('tests/samples/nano_dy.root'))

    hasgen = ~np.isnan(events.Muon.matched_gen.pt.fillna(np.nan))
    mc_rand = JaggedArray.fromoffsets(hasgen.offsets, mc_rand)
    mc_kspread = rochester.kSpreadMC(events.Muon.charge[hasgen], events.Muon.pt[hasgen], events.Muon.eta[hasgen], events.Muon.phi[hasgen],
                                     events.Muon.matched_gen.pt[hasgen])
    mc_ksmear = rochester.kSmearMC(events.Muon.charge[~hasgen], events.Muon.pt[~hasgen], events.Muon.eta[~hasgen], events.Muon.phi[~hasgen],
                                   events.Muon.nTrackerLayers[~hasgen], mc_rand[~hasgen])
    mc_k = np.ones_like(events.Muon.pt.flatten())
    mc_k[hasgen.flatten()] = mc_kspread.flatten()
    mc_k[~hasgen.flatten()] = mc_ksmear.flatten()
    assert(all(np.isclose(mc_k, official_mc_k)))

    mc_errspread = rochester.kSpreadMCerror(events.Muon.charge[hasgen], events.Muon.pt[hasgen], events.Muon.eta[hasgen], events.Muon.phi[hasgen],
                                            events.Muon.matched_gen.pt[hasgen])
    mc_errsmear = rochester.kSmearMCerror(events.Muon.charge[~hasgen], events.Muon.pt[~hasgen], events.Muon.eta[~hasgen], events.Muon.phi[~hasgen],
                                          events.Muon.nTrackerLayers[~hasgen], mc_rand[~hasgen])
    mc_err = np.ones_like(events.Muon.pt.flatten())
    mc_err[hasgen.flatten()] = mc_errspread.flatten()
    mc_err[~hasgen.flatten()] = mc_errsmear.flatten()

def test_jagged_zip(self):
        a = awkward.fromiter([[1.1, 2.2, 3.3], [], [4.4, 5.5]])
        b = awkward.JaggedArray([1, 5, 5], [4, 5, 7], [999, 10, 20, 30, 999, 40, 50, 999])
        c = numpy.array([100, 200, 300])
        d = 1000
        assert awkward.JaggedArray.zip(one=a, two=b).tolist() == [[{"one": 1.1, "two": 10}, {"one": 2.2, "two": 20}, {"one": 3.3, "two": 30}], [], [{"one": 4.4, "two": 40}, {"one": 5.5, "two": 50}]]
        assert awkward.JaggedArray.zip(one=b, two=a).tolist() == [[{"one": 10, "two": 1.1}, {"one": 20, "two": 2.2}, {"one": 30, "two": 3.3}], [], [{"one": 40, "two": 4.4}, {"one": 50, "two": 5.5}]]
        assert awkward.JaggedArray.zip(one=b, two=c).tolist() == [[{"one": 10, "two": 100}, {"one": 20, "two": 100}, {"one": 30, "two": 100}], [], [{"one": 40, "two": 300}, {"one": 50, "two": 300}]]
        assert awkward.JaggedArray.zip(one=b, two=d).tolist() == [[{"one": 10, "two": 1000}, {"one": 20, "two": 1000}, {"one": 30, "two": 1000}], [], [{"one": 40, "two": 1000}, {"one": 50, "two": 1000}]]
        assert a.zip(b).tolist() == [[(1.1, 10), (2.2, 20), (3.3, 30)], [], [(4.4, 40), (5.5, 50)]]
        assert b.zip(a).tolist() == [[(10, 1.1), (20, 2.2), (30, 3.3)], [], [(40, 4.4), (50, 5.5)]]
        assert b.zip(c).tolist() == [[(10, 100), (20, 100), (30, 100)], [], [(40, 300), (50, 300)]]
        assert b.zip(d).tolist() == [[(10, 1000), (20, 1000), (30, 1000)], [], [(40, 1000), (50, 1000)]]

def test_lorentzvector_jagged(self):
        TLorentzVectorJagged = type("TLorentzVectorJagged", (awkward.JaggedArray, uproot_methods.classes.TLorentzVector.ArrayMethods), {})
        a = TLorentzVectorJagged.fromoffsets([0, 3, 3, 5, 10], TLorentzVectorArray(numpy.zeros(10), numpy.arange(10), numpy.zeros(10), numpy.zeros(10)))
        a._generator = uproot_methods.classes.TLorentzVector.TLorentzVector
        a._args = ()
        a._kwargs = {}
        assert a.tolist() == [[TLorentzVector(0, 0, 0, 0), TLorentzVector(0, 1, 0, 0), TLorentzVector(0, 2, 0, 0)], [], [TLorentzVector(0, 3, 0, 0), TLorentzVector(0, 4, 0, 0)], [TLorentzVector(0, 5, 0, 0), TLorentzVector(0, 6, 0, 0), TLorentzVector(0, 7, 0, 0), TLorentzVector(0, 8, 0, 0), TLorentzVector(0, 9, 0, 0)]]
        assert a.x.tolist() == [[0.0, 0.0, 0.0], [], [0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0]]
        assert a.y.tolist() == [[0, 1, 2], [], [3, 4], [5, 6, 7, 8, 9]]
        assert (a + TLorentzVector(1000, 2000, 0, 0)).tolist() == [[TLorentzVector(1000, 2000, 0, 0), TLorentzVector(1000, 2001, 0, 0), TLorentzVector(1000, 2002, 0, 0)], [], [TLorentzVector(1000, 2003, 0, 0), TLorentzVector(1000, 2004, 0, 0)], [TLorentzVector(1000, 2005, 0, 0), TLorentzVector(1000, 2006, 0, 0), TLorentzVector(1000, 2007, 0, 0), TLorentzVector(1000, 2008, 0, 0), TLorentzVector(1000, 2009, 0, 0)]]
        assert (a + TLorentzVectorArray(numpy.full(4, 1000), numpy.arange(1000, 5000, 1000), numpy.zeros(4), numpy.zeros(4))).tolist() == [[TLorentzVector(1000, 1000, 0, 0), TLorentzVector(1000, 1001, 0, 0), TLorentzVector(1000, 1002, 0, 0)], [], [TLorentzVector(1000, 3003, 0, 0), TLorentzVector(1000, 3004, 0, 0)], [TLorentzVector(1000, 4005, 0, 0), TLorentzVector(1000, 4006, 0, 0), TLorentzVector(1000, 4007, 0, 0), TLorentzVector(1000, 4008, 0, 0), TLorentzVector(1000, 4009, 0, 0)]]

if b <= a:
            a, b = 0, 0
        if a < 0:
            a = 0
        elif a > length:
            a = length
        if b < 0:
            b = 0
        elif b > length:
            b = length

        starts[i] = array.starts[i] + a
        stops[i] = array.starts[i] + b

    next = getitem_next(array.content, tail, advanced)
    return awkward.JaggedArray(starts, stops, next)

def test_jagged_zip(self):
        a = awkward.fromiter([[1.1, 2.2, 3.3], [], [4.4, 5.5]])
        b = awkward.JaggedArray([1, 5, 5], [4, 5, 7], [999, 10, 20, 30, 999, 40, 50, 999])
        c = numpy.array([100, 200, 300])
        d = 1000
        assert awkward.JaggedArray.zip(one=a, two=b).tolist() == [[{"one": 1.1, "two": 10}, {"one": 2.2, "two": 20}, {"one": 3.3, "two": 30}], [], [{"one": 4.4, "two": 40}, {"one": 5.5, "two": 50}]]
        assert awkward.JaggedArray.zip(one=b, two=a).tolist() == [[{"one": 10, "two": 1.1}, {"one": 20, "two": 2.2}, {"one": 30, "two": 3.3}], [], [{"one": 40, "two": 4.4}, {"one": 50, "two": 5.5}]]
        assert awkward.JaggedArray.zip(one=b, two=c).tolist() == [[{"one": 10, "two": 100}, {"one": 20, "two": 100}, {"one": 30, "two": 100}], [], [{"one": 40, "two": 300}, {"one": 50, "two": 300}]]
        assert awkward.JaggedArray.zip(one=b, two=d).tolist() == [[{"one": 10, "two": 1000}, {"one": 20, "two": 1000}, {"one": 30, "two": 1000}], [], [{"one": 40, "two": 1000}, {"one": 50, "two": 1000}]]
        assert a.zip(b).tolist() == [[(1.1, 10), (2.2, 20), (3.3, 30)], [], [(4.4, 40), (5.5, 50)]]
        assert b.zip(a).tolist() == [[(10, 1.1), (20, 2.2), (30, 3.3)], [], [(40, 4.4), (50, 5.5)]]
        assert b.zip(c).tolist() == [[(10, 100), (20, 100), (30, 100)], [], [(40, 300), (50, 300)]]
        assert b.zip(d).tolist() == [[(10, 1000), (20, 1000), (30, 1000)], [], [(40, 1000), (50, 1000)]]

subres = []
            suberr = []
            for im in range(len(pt[ie])):
                if isData:
                    subres += [roccor.kScaleDT(int(charge[ie][im]), float(pt[ie][im]), float(eta[ie][im]), float(phi[ie][im]))]
                    suberr += [roccor.kScaleDTerror(int(charge[ie][im]), float(pt[ie][im]), float(eta[ie][im]), float(phi[ie][im]))]
                else:
                    if gid[ie][im]>=0:
                        subres += [roccor.kSpreadMC(int(charge[ie][im]), float(pt[ie][im]), float(eta[ie][im]), float(phi[ie][im]), float(gpt[ie][gid[ie][im]]))]
                        suberr += [roccor.kSpreadMCerror(int(charge[ie][im]), float(pt[ie][im]), float(eta[ie][im]), float(phi[ie][im]), float(gpt[ie][gid[ie][im]]))]
                    else:
                        subres += [roccor.kSmearMC(int(charge[ie][im]), float(pt[ie][im]), float(eta[ie][im]), float(phi[ie][im]), int(nl[ie][im]), float(u[ie][im]))]
                        suberr += [roccor.kSmearMCerror(int(charge[ie][im]), float(pt[ie][im]), float(eta[ie][im]), float(phi[ie][im]), int(nl[ie][im]), float(u[ie][im]))]
            res += [subres]
            err += [suberr]
    res = JaggedArray.fromiter(res)
    err = JaggedArray.fromiter(err)
    outres = res.flatten()
    outerr = err.flatten()
    np.save(fname.replace('.root','_rochester.npy'), outres)
    np.save(fname.replace('.root','_rochester_err.npy'), outerr)
    if not isData:
        outrand = np.concatenate([ui.flatten() for ui in fullu])
        np.save(fname.replace('.root','_rochester_rand.npy'), outrand)

import uproot_methods
    import awkward

    # histogram creation and manipulation
    from coffea import hist

    fin = uproot.open("HZZ.root")
    tree = fin["events"]

    arrays = {k.replace('Electron_', ''): v for k, v in tree.arrays("Electron_*", namedecode='ascii').items()}
    p4 = uproot_methods.TLorentzVectorArray.from_cartesian(arrays.pop('Px'),
                                                           arrays.pop('Py'),
                                                           arrays.pop('Pz'),
                                                           arrays.pop('E'),
                                                           )
    electrons = awkward.JaggedArray.zip(p4=p4, **arrays)

    arrays = {k.replace('Muon_', ''): v for k, v in tree.arrays("Muon_*", namedecode='ascii').items()}
    p4 = uproot_methods.TLorentzVectorArray.from_cartesian(
        arrays.pop('Px'),
        arrays.pop('Py'),
        arrays.pop('Pz'),
        arrays.pop('E'),
    )
    muons = awkward.JaggedArray.zip(p4=p4, **arrays)

    # Two types of axes exist presently: bins and categories
    lepton_kinematics = hist.Hist("Events",
                                  hist.Cat("flavor", "Lepton flavor"),
                                  hist.Bin("pt", "$p_{T}$", 19, 10, 100),
                                  hist.Bin("eta", r"$\eta$", [-2.5, -1.4, 0, 1.4, 2.5]),
                                  )

for j in range(len(head)):
            norm = head[j]
            if norm < 0:
                norm += length
            if norm < 0 or norm >= length:
                raise IndexError("advanced index is out of bounds in JaggedArray")
            index[k] = array.starts[i] + norm
            nextadvanced[k] = j
            k += 1
        offsets[i + 1] = k

    starts = offsets[:-1]
    stops = offsets[1:]
    next = getitem_next(array.content[index], tail, nextadvanced)
    return awkward.JaggedArray(starts, stops, next)

def test_3D_jagged(wrapped_tree):
    fake_3d = [[np.arange(i + 1) + j
                for i in range(j % 3)]
               for j in range(len(wrapped_tree))]
    fake_3d = JaggedArray.fromiter(fake_3d)
    wrapped_tree.new_variable("Fake3D", fake_3d)
    assert isinstance(fake_3d.count(), JaggedArray)
    assert all((fake_3d.copy().count() == fake_3d.count()).all())

    aliased = expressions.evaluate(wrapped_tree, "Fake3D")
    assert (aliased == fake_3d).all().all().all()

    doubled = expressions.evaluate(wrapped_tree, "Fake3D * 2")
    assert (doubled == fake_3d * 2).all().all().all()
    assert len(doubled[0, :, :]) == 0
    assert doubled[1, 0, :] == [2]
    assert doubled[2, 0, :] == [4]
    assert all(doubled[2, 1, :] == [4, 6])

    doubled = expressions.evaluate(wrapped_tree, "Fake3D + Fake3D")
    assert (doubled == fake_3d * 2).all().all().all()
    assert len(doubled[0, :, :]) == 0
    assert doubled[1, 0, :] == [2]

def _lazy_nested_crossref(self, indices, destination):
        # to be used for stitching a set of indices into a doubly-jagged mapping
        # e.g. Jet_electronIdx1, Jet_electronIdx2
        if not isinstance(destination, NanoCollection):
            raise ValueError("Destination must be a NanoCollection")
        if not isinstance(destination.array, awkward.JaggedArray):
            raise ValueError("Cross-references imply jagged destination")
        if not isinstance(self.array, awkward.JaggedArray):
            raise NotImplementedError
        JaggedArray = destination._get_mixin(destination._get_methods(), awkward.JaggedArray)
        # repair type now that we've materialized
        for idx in indices:
            idx.type.takes = self.array.offsets[-1]
        content = numpy.zeros(len(self.array.content) * len(indices), dtype=awkward.JaggedArray.INDEXTYPE)
        for i, index in enumerate(indices):
            content[i::len(indices)] = numpy.array(index)
        globalindices = awkward.JaggedArray.fromoffsets(
            self.array.offsets,
            awkward.JaggedArray.fromoffsets(
                numpy.arange((len(self.array.content) + 1) * len(indices), step=len(indices)),
                content,
            )
        )
        globalindices = globalindices[globalindices >= 0] + destination.array.starts