Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
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,
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)
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())
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)
from .JaggedCandidateMethods import JaggedCandidateMethods
from ..util import awkward
JaggedCandidateArray = awkward.Methods.mixin(JaggedCandidateMethods, awkward.JaggedArray)
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))
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])