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

To help you get started, we’ve selected a few coffea examples, based on popular ways it is used in public projects.

Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.

github CoffeaTeam / coffea / tests / test_lookup_tools.py View on Github external
extractor.finalize(reduce_list=['testSF2d'])

    evaluator = extractor.make_evaluator()

    counts, test_eta, test_pt = dummy_jagged_eta_pt()
    
    # test flat eval
    test_out = evaluator["testSF2d"](test_eta, test_pt)

    # print it
    print(evaluator["testSF2d"])
    
    # test structured eval
    test_eta_jagged = awkward.JaggedArray.fromcounts(counts, test_eta)
    test_pt_jagged = awkward.JaggedArray.fromcounts(counts, test_pt)
    test_out_jagged = evaluator["testSF2d"](test_eta_jagged, test_pt_jagged)

    assert (test_out_jagged.counts==counts).all()
    assert (test_out==test_out_jagged.flatten()).all()

    # From make_expected_lookup.py
    expected_output = np.array([
       0.90780139, 0.82748538, 0.86332178, 0.86332178, 0.97981155,
       0.79701495, 0.88245934, 0.82857144, 0.91884059, 0.97466666,
       0.94072163, 1.00775194, 0.82748538, 1.00775194, 0.97203946,
       0.98199672, 0.80655736, 0.90893763, 0.88245934, 0.79701495,
       0.82748538, 0.82857144, 0.91884059, 0.90893763, 0.97520661,
       0.97520661, 0.82748538, 0.91884059, 0.97203946, 0.88245934,
       0.79701495, 0.9458763 , 1.00775194, 0.80655736, 1.00775194,
       1.00775194, 0.98976982, 0.98976982, 0.86332178, 0.94072163,
       0.80655736, 0.98976982, 0.96638656, 0.9458763 , 0.90893763,
github CoffeaTeam / coffea / tests / dummy_distributions.py View on Github external
def gen_reco_TLV():
    gen_pt = awkward.JaggedArray.fromiter([[10.0, 20.0, 30.0], [], [40.0, 50.0]])
    reco_pt = awkward.JaggedArray.fromiter([[20.2, 10.1, 30.3, 50.5], [50.5], [60]])

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

    gen_phi = awkward.JaggedArray.fromiter([[-1.5, 0.0, 1.5], [], [0.78, -0.78]])
    reco_phi = awkward.JaggedArray.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)

    return (gen, reco)
github CoffeaTeam / coffea / coffea / analysis_objects / JaggedCandidateMethods.py View on Github external
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())
github CoffeaTeam / coffea / coffea / analysis_objects / JaggedCandidateMethods.py View on Github external
def cross(self, other, nested=False):
        """
        This method calls the cross method of JaggedArray to get all pairs
        per-event with another JaggedCandidateArray.

        The resulting JaggedArray of that call is dressed with the jagged candidate
        array four-momentum and cached fast access pt/eta/phi/mass.
        """
        outs = super(JaggedCandidateMethods, self).cross(other, nested)
        outs['p4'] = outs.i0['p4'] + outs.i1['p4']
        thep4 = outs['p4']
        outs['__fast_pt'] = awkward.JaggedArray.fromoffsets(outs.offsets, _fast_pt(thep4.content))
        outs['__fast_eta'] = awkward.JaggedArray.fromoffsets(outs.offsets, _fast_eta(thep4.content))
        outs['__fast_phi'] = awkward.JaggedArray.fromoffsets(outs.offsets, _fast_phi(thep4.content))
        outs['__fast_mass'] = awkward.JaggedArray.fromoffsets(outs.offsets, _fast_mass(thep4.content))
        return self.fromjagged(outs)
github CoffeaTeam / coffea / coffea / analysis_objects / JaggedCandidateArray.py View on Github external
from .JaggedCandidateMethods import JaggedCandidateMethods
from ..util import awkward

JaggedCandidateArray = awkward.Methods.mixin(JaggedCandidateMethods, awkward.JaggedArray)
github CoffeaTeam / coffea / coffea / lookup_tools / jme_standard_function.py View on Github external
bin_indices.append(masked_bin_eval(bin_indices[0],
                                               self._bins[binname],
                                               bin_vals[binname]))

        bin_tuple = tuple(bin_indices)

        # get clamp values and clip the inputs
        eval_values = []
        for eval_name in self._eval_vars:
            clamp_mins = None
            if self._eval_clamp_mins[eval_name].content.size == 1:
                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))
github CoffeaTeam / coffea / coffea / lookup_tools / dense_evaluated_lookup.py View on Github external
def _evaluate(self, *args):
        indices = []
        for arg in args:
            if type(arg) == awkward.JaggedArray:
                raise Exception('JaggedArray in inputs')
        if self._dimension == 1:
            indices.append(np.clip(np.searchsorted(self._axes, args[0], side='right') - 1, 0, self._values.shape[0] - 1))
        else:
            for dim in range(self._dimension):
                indices.append(np.clip(np.searchsorted(self._axes[dim], args[dim], side='right') - 1,
                                       0, self._values.shape[len(self._axes) - dim - 1] - 1))
        indices.reverse()
        return numba_apply_1d(self._values[tuple(indices)], args[self._feval_dim])