Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
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)