How to use the coffea.util.awkward.JaggedArray function in coffea

def test_evaluator_exceptions():
    extractor = lookup_tools.extractor()
    extractor.add_weight_sets(["testSF2d scalefactors_Tight_Electron tests/samples/testSF2d.histo.root"])

    counts, test_eta, test_pt = dummy_jagged_eta_pt()
    test_eta_jagged = awkward.JaggedArray.fromcounts(counts, test_eta)
    test_pt_jagged = awkward.JaggedArray.fromcounts(counts, test_pt)

    extractor.finalize()
    evaluator = extractor.make_evaluator()
    
    try:
        test_out = evaluator["testSF2d"](test_pt_jagged, test_eta)
    except Exception as e:
        assert(e.args[0] == 'do not mix JaggedArrays and numpy arrays when calling a derived class of lookup_base')

def test_evaluator_exceptions():
    extractor = lookup_tools.extractor()
    extractor.add_weight_sets(["testSF2d scalefactors_Tight_Electron tests/samples/testSF2d.histo.root"])

    counts, test_eta, test_pt = dummy_jagged_eta_pt()
    test_eta_jagged = awkward.JaggedArray.fromcounts(counts, test_eta)
    test_pt_jagged = awkward.JaggedArray.fromcounts(counts, test_pt)

    extractor.finalize()
    evaluator = extractor.make_evaluator()
    
    try:
        test_out = evaluator["testSF2d"](test_pt_jagged, test_eta)
    except Exception as e:
        assert(e.args[0] == 'do not mix JaggedArrays and numpy arrays when calling a derived class of lookup_base')

def process(self, df):
        output = self.accumulator.identity()

        dataset = df['dataset']

        print(os.environ)

        muon = None
        if isinstance(df['Muon_pt'], akd.JaggedArray):
            muon = CandArray.candidatesfromcounts(counts=df['Muon_pt'].counts,
                                                  pt=df['Muon_pt'].content,
                                                  eta=df['Muon_eta'].content,
                                                  phi=df['Muon_phi'].content,
                                                  mass=df['Muon_mass'].content)
        else:
            muon = CandArray.candidatesfromcounts(counts=df['nMuon'],
                                                  pt=df['Muon_pt'],
                                                  eta=df['Muon_eta'],
                                                  phi=df['Muon_phi'],
                                                  mass=df['Muon_mass'])

        dimuon = muon.distincts()

        output['pt'].fill(dataset=dataset, pt=muon.pt.flatten())
        output['mass'].fill(dataset=dataset, mass=dimuon.mass.flatten())

jagged_counts = np.maximum(bins[bin_order[1]].counts - 1, 0)  # need counts-1 since we only care about Nbins
    for i in range(nEvalVars):
        var_order.append(layout[i + offset_name])
        if not interpolatedFunc:
            clamp_mins[layout[i + offset_name]] = awkward.JaggedArray.fromcounts(jagged_counts, np.atleast_1d(pars[columns[i + offset_col]]))
            clamp_maxs[layout[i + offset_name]] = awkward.JaggedArray.fromcounts(jagged_counts, np.atleast_1d(pars[columns[i + offset_col + 1]]))
            assert(clamp_mins[layout[i + offset_name]].valid())
            assert(clamp_maxs[layout[i + offset_name]].valid())
            offset_col += 1

    # now get the parameters, which we will look up with the clamped values
    parms = []
    parm_order = []
    offset_col = 2 * nBinnedVars + 1 + int(not interpolatedFunc) * 2 * nEvalVars
    for i in range(nParms):
        jag = awkward.JaggedArray.fromcounts(jagged_counts, pars[columns[i + offset_col]])
        assert(jag.valid())
        parms.append(jag)
        parm_order.append('p%i' % (i))

    wrapped_up = {}
    wrapped_up[(name, 'jme_standard_function')] = (formula,
                                                   (bins, bin_order),
                                                   (clamp_mins, clamp_maxs, var_order),
                                                   (parms, parm_order))
    return wrapped_up

def _default_argmatch(combs, deltaRCut=10000, deltaPtCut=10000):
    """ default matching function for argmatch(), match in deltaR / deltaPt """
    deltaPts = (np.abs(combs.i0.pt - combs.i1.pt) / combs.i0.pt)
    deltaRs = combs.i0.delta_r(combs.i1)
    indexOfMin = deltaRs.argmin()
    indexOfMinOutShape = indexOfMin.flatten(axis=1)
    passesCut = (deltaRs[indexOfMin] < deltaRCut) & (deltaPts[indexOfMin] < deltaPtCut)
    passesCutOutShape = passesCut.flatten(axis=1)
    flatPass = passesCutOutShape.flatten()
    flatIdxMin = indexOfMinOutShape.flatten()
    flatIdxMin[~flatPass] = -1
    return awkward.JaggedArray.fromoffsets(passesCutOutShape.offsets, flatIdxMin)

import awkward

from ..util import awkward as util_awkward
import uproot_methods
from uproot_methods.classes.TLorentzVector import ArrayMethods, TLorentzVectorArray, TLorentzVector

JaggedArrayMethods = ArrayMethods.mixin(ArrayMethods, util_awkward.JaggedArray)


class Dangerousness(object):
    # NB: This is dangerous
    def __getattr__(self, what):
        if what in self.columns:
            return self[what]
        thewhat = getattr(super(Dangerousness, self), what)
        return self.fromjagged(thewhat)


class SaiyanCommonMethods(Dangerousness, ArrayMethods):
    def match(self, cands, value):
        if isinstance(self, awkward.JaggedArray):
            combinations = self.cross(cands, nested=True)
            mask = (combinations.i0.delta_r(combinations.i1) < value)

clamp_mins = self._eval_clamp_mins[eval_name].content[0]
            else:
                idxs = flatten_idxs(bin_tuple, self._eval_clamp_mins[eval_name])
                clamp_mins = self._eval_clamp_mins[eval_name].content[idxs]
                if isinstance(clamp_mins, awkward.JaggedArray):
                    if clamp_mins.content.size == 1:
                        clamp_mins = clamp_mins.content[0]
                    else:
                        clamp_mins = clamp_mins.flatten()
            clamp_maxs = None
            if self._eval_clamp_maxs[eval_name].content.size == 1:
                clamp_maxs = self._eval_clamp_maxs[eval_name].content[0]
            else:
                idxs = flatten_idxs(bin_tuple, self._eval_clamp_maxs[eval_name])
                clamp_maxs = self._eval_clamp_maxs[eval_name].content[idxs]
                if isinstance(clamp_maxs, awkward.JaggedArray):
                    if clamp_maxs.content.size == 1:
                        clamp_maxs = clamp_maxs.content[0]
                    else:
                        clamp_maxs = clamp_maxs.flatten()
            eval_values.append(np.clip(eval_vals[eval_name], clamp_mins, clamp_maxs))

        # get parameter values
        parm_values = []
        if len(self._parms) > 0:
            idxs = flatten_idxs(bin_tuple, self._parms[0])
            parm_values = [parm.content[idxs] for parm in self._parms]

        return self._formula(*tuple(parm_values + eval_values))

- px, py, pz, mass
            - px, py, pz, energy

        .. note:: Constructing by (pt, eta, phi, mass) will generate a Lorentz vector in that basis.

        Example::

            cands = JaggedCandidateArray.candidatesfromcounts(counts=counts,
                                                              pt=column1,
                                                              eta=column2,
                                                              phi=column3,
                                                              mass=column4,
                                                              ...)

        """
        offsets = awkward.JaggedArray.counts2offsets(counts)
        return cls.candidatesfromoffsets(offsets, **kwargs)

def _default_fastmatch(first, second, deltaRCut=10000):
    drCut2 = deltaRCut**2
    args = first.eta._argcross(second.eta)
    argsnested = awkward.JaggedArray.fromcounts(first.eta.counts,
                                                awkward.JaggedArray.fromcounts(first.eta.tojagged(second.eta.counts).flatten(),
                                                                               args._content))
    eta0s = first.eta.content[argsnested.content.content.i0]
    eta1s = second.eta.content[argsnested.content.content.i1]
    phi0s = first.phi.content[argsnested.content.content.i0]
    phi1s = second.phi.content[argsnested.content.content.i1]
    offsets_outer = argsnested.offsets
    offsets_inner = argsnested.content.offsets
    detas = np.abs(eta0s - eta1s)
    dphis = (phi0s - phi1s + math.pi) % (2 * math.pi) - math.pi
    passdr = ((detas**2 + dphis**2) < drCut2)
    passdr = awkward.JaggedArray.fromoffsets(offsets_inner, passdr)
    return awkward.JaggedArray.fromoffsets(offsets_outer, passdr.any())

"""
        Add new attributes to this JaggedCandidateArray.
        New values may be given in jagged or flat-column form.
        They will be interpreted using the offsets of the
        JaggedCandidateArray they are being inserted into.

        Example::

            cands.add_attributes( name1 = column1,
                                  name2 = column2,
                                  ... )

        """
        for key, item in kwargs.items():
            if isinstance(item, awkward.JaggedArray):
                self[key] = awkward.JaggedArray.fromoffsets(self.offsets, item.content)
            elif isinstance(item, np.ndarray):
                self[key] = awkward.JaggedArray.fromoffsets(self.offsets, item)