def getHighestPtBoson(pdgs,pts,theId):
idxs = np.arange(parents.content.size)
idxs = awkward.JaggedArray(parents.starts,parents.stops,idxs)
idxs = idxs - parents.starts
Vidxs = idxs[np.abs(pdgs)==theId]
Vpts = pts[np.abs(pdgs)==theId]
return Vidxs[Vpts.argmax()]
def getChildrenOfType(parents,pdgs,parId,dauId):
pars = getParentsOfType(parents,pdgs,parId)
idxs = np.arange(parents.content.size)
idxs = awkward.JaggedArray(parents.starts,parents.stops,idxs)
idxs = idxs - parents.starts
idxs = idxs[(pars != 0) & (np.abs(pdgs[idxs]) == dauId)]
return idxs
def muon_control_region(gghbbcuts, dataset, gencat, presel_weight, eventInfo,
leadingak8jet, ak4jets_Mbtag, looseMuons,
looseElectrons, looseTaus, hasTightVJet, plots):
leadingLooseMuon = looseMuons[looseMuons.pt.argmax()]
AK8jet_muon_matches = leadingak8jet.fastmatch(leadingLooseMuon,
matchfunc=matchByDPhi)
AK8jet_AK4Mbjet_antimatches = leadingak8jet.fastmatch(
ak4jets_Mbtag, matchfunc=antiMatchByDR)
#for name,vari in jet_systs:
# attr = "pt"
# if len(vari) > 0: attr += "_"+vari
systematic = "central"
mucrweight = ( #jet selection
((leadingak8jet.pt > gghbbcuts.PTCUTMUCR).sum() > 0) &
((leadingak8jet.msd_corr_8 > gghbbcuts.MASSCUT).sum() > 0)
& hasTightVJet &
#muon selection
((np.abs(leadingLooseMuon.eta) < 2.1).sum() > 0) &
((leadingLooseMuon.pt > gghbbcuts.MUONPTCUT).sum() > 0) &
#matches and cross cleaning
(AK8jet_muon_matches.sum() == 0) &
(AK8jet_AK4Mbjet_antimatches.sum() > 0) &
#lepton vetos
(looseMuons.counts == 1) & (looseElectrons.counts == 0) &
(looseTaus.counts == 0))
#fill plots
jetweight = leadingak8jet.weight.sum()
#jetweight[jetweight > 0.] = 1.0 #for now
weight = jetweight * presel_weight
weight_mucr = weight * mucrweight
fill_plots_mucr(dataset, gencat, systematic, leadingak8jet, weight_mucr,
plots)
def tagDecayJagged(decayEndpoint,decayEndpointStatus,require,reject,parentage,status,pdgId,absPdg = True):
pdgMatch = pdgId
if absPdg: pdgMatch = np.abs(pdgId)
return ((status == decayEndpointStatus) &
(pdgMatch == decayEndpoint) &
((parentage & require) == require) &
((parentage & reject) == 0))
def getHadronicVIndices(VPdgId,parentage,parents,status,pdgId,statusValue=23):
inChain = activePdgIds
mask = inChain[VPdgId]
events = tagDecayJagged(1,statusValue,mask,0,parentage,status,pdgId).astype('u4')
events = events + 2*tagDecayJagged(2,statusValue,mask,0,parentage,status,pdgId)
events = events + 3*tagDecayJagged(3,statusValue,mask,0,parentage,status,pdgId)
events = events + 4*tagDecayJagged(4,statusValue,mask,0,parentage,status,pdgId)
events = events + 5*tagDecayJagged(5,statusValue,mask,0,parentage,status,pdgId)
if abs(VPdgId) == 24:
mask = inChain[-VPdgId]
events = events + tagDecayJagged(1,statusValue,mask,0,parentage,status,pdgId)
events = events + 2*tagDecayJagged(2,statusValue,mask,0,parentage,status,pdgId)
events = events + 3*tagDecayJagged(3,statusValue,mask,0,parentage,status,pdgId)
events = events + 4*tagDecayJagged(4,statusValue,mask,0,parentage,status,pdgId)
events = events + 5*tagDecayJagged(5,statusValue,mask,0,parentage,status,pdgId)
daughters_flat, = np.where(events.content != 0)
dau_pdgIds = np.abs(pdgId.content[daughters_flat])
#here we assume that V always decays into two things!!!!
good_parents = parents[events != 0]
good_dau_pdg = awkward.JaggedArray.fromcounts(np.full(daughters_flat.size//2,2),dau_pdgIds)
good_dau_pdg = awkward.JaggedArray.fromcounts(good_parents.counts//2,good_dau_pdg)
#good_daughters = awkward.JaggedArray.fromcounts(np.full(daughters_flat.size/2,2),daughters_flat)
#good_daughters = awkward.JaggedArray.fromcounts(good_parents.counts/2,good_daughters)
offset_Vs = (good_parents + parents.starts).content[::2]
offset_Vs = awkward.JaggedArray.fromcounts(good_parents.counts//2,offset_Vs) - parents.starts
return offset_Vs,good_dau_pdg.max() #return the up-type quark in each decay pair
def _fast_mass(p4):
""" quick mass calculation for caching """
px = p4.x
py = p4.y
pz = p4.z
en = p4.t
p3mag2 = (px * px + py * py + pz * pz)
return np.sqrt(np.abs(en * en - p3mag2))
def getParentsOfTypeFlat(parents_last,parents,pdgs,theId,out,lastOnly=False):
#print 'step'
parent_pdgs = pdgs[parents_last]
notfound = ( (np.abs(parent_pdgs) != theId) & (parents_last != 0) )
parents[notfound] = parents[parents[notfound]]
if np.any(parents_last != parents):
parents_last[notfound] = parents[notfound]
getParentsOfTypeFlat(parents_last,parents,pdgs,theId,lastOnly)
else:
if lastOnly:
parents[parent_pdgs == pdgs] = 0
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)
def __init__(self):
self.p4 = thep4
self.px = px
self.py = py
self.pz = pz
self.en = energy
self.pt = np.hypot(px, py)
self.phi = np.arctan2(py, px)
self.eta = np.arctanh(pz / np.sqrt(px * px + py * py + pz * pz))
self.mass = np.sqrt(
np.abs(energy * energy - (px * px + py * py + pz * pz)))
self.blah = energy * px
self.count = counts
def matchByDPhi(first, second, deltaPhiCut=2. / 3. * math.pi):
args = first.phi._argcross(second.phi)
argsnested = awkward.JaggedArray.fromcounts(
first.phi.counts,
awkward.JaggedArray.fromcounts(
first.phi._broadcast(second.phi.counts).flatten(), args._content))
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
dphis = (phi0s - phi1s + math.pi) % (2 * math.pi) - math.pi
passdphi = (np.abs(dphis) < deltaPhiCut)
passdphi = awkward.JaggedArray.fromoffsets(offsets_inner, passdphi)
return awkward.JaggedArray.fromoffsets(offsets_outer, passdphi.any())
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._broadcast(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 MuonIDSF(evaluator,mupt,mueta):
pt = mupt
eta = np.abs(mueta)
if isinstance(mupt,awkward.JaggedArray):
assert (mupt.offsets==mueta.offsets).all()
pt = mupt.flatten()
eta = mueta.flatten()
else:
assert mupt.size == mueta.size
eff = evaluator['NUM_SoftID_DEN_genTracks/abseta_pt_value'](eta,pt)
err = evaluator['NUM_SoftID_DEN_genTracks/abseta_pt_error'](eta,pt)
lo = eff - err
hi = eff + err
if isinstance(mupt,awkward.JaggedArray):
eff = awkward.JaggedArray.fromoffsets(mupt.offsets,eff)
lo = awkward.JaggedArray.fromoffsets(mupt.offsets,lo)
hi = awkward.JaggedArray.fromoffsets(mupt.offsets,hi)
return eff,hi,lo
def MuonTrigSF(evaluator,mupt,mueta):
pt = mupt
eta = np.abs(mueta)
if isinstance(mupt,awkward.JaggedArray):
assert (mupt.offsets==mueta.offsets).all()
pt = mupt.flatten()
eta = mueta.flatten()
else:
assert mupt.size == mueta.size
eff = evaluator['Mu50_PtEtaBins/efficienciesDATA/pt_abseta_DATA'](pt,eta)
err = evaluator['Mu50_PtEtaBins/efficienciesDATA/pt_abseta_DATA_error'](pt,eta)
lo = eff - err
hi = eff + err
if isinstance(mupt,awkward.JaggedArray):
eff = awkward.JaggedArray.fromoffsets(mupt.offsets,eff)
lo = awkward.JaggedArray.fromoffsets(mupt.offsets,lo)
hi = awkward.JaggedArray.fromoffsets(mupt.offsets,hi)
return eff,hi,lo
def test_root_scalefactors():
extractor = lookup_tools.extractor()
extractor.add_weight_sets([
"testSF2d scalefactors_Tight_Electron tests/samples/testSF2d.histo.root"
])
extractor.finalize()
evaluator = extractor.make_evaluator()
counts, test_eta, test_pt = dummy_jagged_eta_pt()
# test flat eval
test_out = evaluator["testSF2d"](test_eta, test_pt)
# 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, 0.9529984, 0.9458763, 0.9529984,
0.80655736, 0.80655736, 0.80655736, 0.98976982, 0.97466666, 0.98199672,
0.86332178, 1.03286386, 0.94072163, 1.03398061, 0.82857144, 0.80655736,
1.00775194, 0.80655736
])
diff = np.abs(test_out - expected_output)
print("Max diff: %.16f" % diff.max())
print("Median diff: %.16f" % np.median(diff))
print("Diff over threshold rate: %.1f %%" %
(100 * (diff >= 1.e-8).sum() / diff.size))
assert (diff < 1.e-8).all()
def PUPPIweight(ak8pt,ak8eta):
pt = ak8pt
eta = ak8eta
if isinstance(ak8pt,awkward.JaggedArray):
assert (ak8pt.offsets==ak8eta.offsets).all()
pt = ak8pt.flatten()
eta = ak8eta.flatten()
else:
assert ak8pt.size == ak8eta.size
genCorr = corrGEN(pt)
recoCorr = np.ones_like(eta)
etaCut = np.abs(eta) < 1.3
recoCorr[etaCut] = corrRECO_cen(pt[etaCut])
recoCorr[~etaCut] = corrRECO_for(pt[~etaCut])
total = genCorr*recoCorr
if isinstance(ak8pt,awkward.JaggedArray):
total = awkward.JaggedArray.fromoffsets(ak8pt.offsets,total)
return total
def parseGeneratorHistory(gp_pdgId_in,gp_parent_in):
inChain = activePdgIds
#index manipulation
offsets = gp_pdgId_in.offsets
parents = gp_pdgId_in.parents
pstarts = offsets[parents].astype('i4')
gp_pdgId = gp_pdgId_in.content
gp_ancestor = gp_parent_in.content + pstarts
gp_ancestor_valid = (gp_parent_in.content >= 0)
#create parentage bitmaps
gp_pdgId_mapped = np.zeros(shape=gp_pdgId.shape, dtype='u4')
for pdgId, bit in inChain.items():
if abs(pdgId) == 24:
gp_pdgId_mapped[gp_pdgId==pdgId] = bit
else:
gp_pdgId_mapped[np.abs(gp_pdgId)==pdgId] = bit
gp_proc = np.zeros(shape=gp_pdgId.shape, dtype='u4')
pdg_tmp = np.empty_like(gp_pdgId_mapped)
parent_tmp = np.empty_like(gp_ancestor)
niter = 0
while np.any(gp_ancestor_valid) and niter < 50:
np.take(gp_pdgId_mapped, gp_ancestor, out=pdg_tmp, mode='clip')
np.take(gp_parent_in.content, gp_ancestor, out=parent_tmp, mode='clip')
np.bitwise_or(gp_proc, pdg_tmp, where=gp_ancestor_valid, out=gp_proc)
np.bitwise_and(gp_ancestor_valid, parent_tmp>=0, where=gp_ancestor_valid, out=gp_ancestor_valid)
np.add(parent_tmp, pstarts, out=gp_ancestor)
niter += 1
#print 'Parsed ancestor tree in %d iterations'%niter
if niter == 50 and np.any(gp_ancestor_valid):
raise Exception('reached 50 iterations, gen particles not trustable')
return awkward.JaggedArray.fromoffsets(offsets, gp_proc)
def _default_match(combs, deltaRCut=10000, deltaPtCut=10000):
""" default matching function for match(), match in deltaR / deltaPt """
passPtCut = ((np.abs(combs.i0.pt - combs.i1.pt) / combs.i0.pt) <
deltaPtCut)
mask = (combs.i0.delta_r(combs.i1) < deltaRCut) & passPtCut
return mask.any()
def plot1d(hist,
ax=None,
clear=True,
overlay=None,
stack=False,
overflow='none',
line_opts=None,
fill_opts=None,
error_opts=None,
overlay_overflow='none',
density=False,
binwnorm=None):
"""
hist: Hist object with maximum of two dimensions
ax: matplotlib Axes object (if None, one is created)
clear: clear Axes before drawing (if passed); if False, this function will skip drawing the legend
overlay: the axis of hist to overlay (remaining one will be x axis)
stack: whether to stack or overlay the other dimension (if one exists)
overflow: overflow behavior of plot axis (see Hist.sum() docs)
The draw options are passed as dicts to the relevant matplotlib function, with some exceptions in case
it is especially common or useful. If none of *_opts is specified, nothing will be plotted!
Pass an empty dict (e.g. line_opts={}) for defaults
line_opts: options to plot a step
Special options interpreted by this function and not passed to matplotlib:
(none)
fill_opts: to plot a filled area
Special options interpreted by this function and not passed to matplotlib:
(none)
error_opts: to plot an errorbar
Special options interpreted by this function and not passed to matplotlib:
'emarker' (default: '') marker to place at cap of errorbar
overlay_overflow: overflow behavior of dense overlay axis, if one exists
density: Convert sum weights to probability density (i.e. integrates to 1 over domain of axis) (NB: conflicts with binwnorm)
binwnorm: Convert sum weights to bin-width-normalized, with units equal to supplied value (usually you want to specify 1.)
"""
if ax is None:
fig, ax = plt.subplots(1, 1)
else:
if not isinstance(ax, plt.Axes):
raise ValueError("ax must be a matplotlib Axes object")
if clear:
ax.clear()
fig = ax.figure
if hist.dim() > 2:
raise ValueError(
"plot1d() can only support up to two dimensions (one for axis, one to stack or overlay)"
)
if overlay is None and hist.dim() > 1:
raise ValueError(
"plot1d() can only support one dimension without an overlay axis chosen"
)
if density and binwnorm is not None:
raise ValueError("Cannot use density and binwnorm at the same time!")
if binwnorm is not None:
if not isinstance(binwnorm, numbers.Number):
raise ValueError("Bin width normalization not a number, but a %r" %
binwnorm.__class__)
if line_opts is None and fill_opts is None and error_opts is None:
if stack:
fill_opts = {}
else:
line_opts = {}
error_opts = {}
axis = hist.axes()[0]
if overlay is not None:
overlay = hist.axis(overlay)
if axis == overlay:
axis = hist.axes()[1]
if isinstance(axis, SparseAxis):
raise NotImplementedError("Plot a sparse axis (e.g. bar chart)")
elif isinstance(axis, DenseAxis):
ax.set_xlabel(axis.label)
ax.set_ylabel(hist.label)
edges = axis.edges(overflow=overflow)
centers = axis.centers(overflow=overflow)
stack_sumw, stack_sumw2 = None, None
primitives = {}
identifiers = hist.identifiers(
overlay,
overflow=overlay_overflow) if overlay is not None else [None]
for i, identifier in enumerate(identifiers):
if identifier is None:
sumw, sumw2 = hist.values(sumw2=True, overflow=overflow)[()]
elif isinstance(overlay, SparseAxis):
sumw, sumw2 = hist.project(overlay, identifier).values(
sumw2=True, overflow=overflow)[()]
else:
sumw, sumw2 = hist.values(sumw2=True, overflow='allnan')[()]
the_slice = (i if
overflow_behavior(overlay_overflow).start is None
else i + 1, overflow_behavior(overflow))
if hist._idense(overlay) == 1:
the_slice = (the_slice[1], the_slice[0])
sumw = sumw[the_slice]
sumw2 = sumw2[the_slice]
if (density or binwnorm is not None) and np.sum(sumw) > 0:
overallnorm = np.sum(
sumw) * binwnorm if binwnorm is not None else 1.
binnorms = overallnorm / (np.diff(edges) * np.sum(sumw))
sumw = sumw * binnorms
sumw2 = sumw2 * binnorms**2
label = str(identifier)
primitives[identifier] = []
first_color = None
if stack:
if stack_sumw is None:
stack_sumw, stack_sumw2 = sumw.copy(), sumw2.copy()
else:
stack_sumw += sumw
stack_sumw2 += sumw2
if line_opts is not None:
opts = {'where': 'post', 'label': label}
opts.update(line_opts)
l, = ax.step(x=edges,
y=np.r_[stack_sumw, stack_sumw[-1]],
**opts)
first_color = l.get_color()
primitives[identifier].append(l)
if fill_opts is not None:
opts = {'step': 'post', 'label': label}
if first_color is not None:
opts['color'] = first_color
opts.update(fill_opts)
f = ax.fill_between(x=edges,
y1=np.r_[stack_sumw - sumw,
stack_sumw[-1] - sumw[-1]],
y2=np.r_[stack_sumw, stack_sumw[-1]],
**opts)
if first_color is None:
first_color = f.get_facecolor()[0]
primitives[identifier].append(f)
# error_opts for stack is interpreted later
else:
if line_opts is not None:
opts = {'where': 'post', 'label': label}
opts.update(line_opts)
l, = ax.step(x=edges, y=np.r_[sumw, sumw[-1]], **opts)
first_color = l.get_color()
primitives[identifier].append(l)
if fill_opts is not None:
opts = {'step': 'post', 'label': label}
if first_color is not None:
opts['color'] = first_color
opts.update(fill_opts)
f = ax.fill_between(x=edges,
y1=np.r_[sumw, sumw[-1]],
**opts)
if first_color is None:
first_color = f.get_facecolor()[0]
primitives[identifier].append(f)
if error_opts is not None:
opts = {'linestyle': 'none', 'label': label}
if first_color is not None:
opts['color'] = first_color
opts.update(error_opts)
emarker = opts.pop('emarker', '')
err = np.abs(poisson_interval(sumw, sumw2) - sumw)
errbar = ax.errorbar(x=centers, y=sumw, yerr=err, **opts)
plt.setp(errbar[1], 'marker', emarker)
primitives[identifier].append(errbar)
if stack_sumw is not None and error_opts is not None:
err = poisson_interval(stack_sumw, stack_sumw2)
opts = {'step': 'post', 'label': 'Sum unc.'}
opts.update(error_opts)
errbar = ax.fill_between(x=edges,
y1=np.r_[err[0, :], err[0, -1]],
y2=np.r_[err[1, :], err[1, -1]],
**opts)
primitives[StringBin('stack_unc', opts['label'])] = [errbar]
if clear:
if overlay is not None:
handles, labels = list(), list()
for identifier, handlelist in primitives.items():
handles.append(
tuple(h for h in handlelist if h.get_label()[0] != '_'))
labels.append(str(identifier))
primitives['legend'] = ax.legend(title=overlay.label,
handles=handles,
labels=labels)
ax.autoscale(axis='x', tight=True)
ax.set_ylim(0, None)
return fig, ax, primitives
def test_analysis_objects():
counts, px, py, pz, energy = dummy_four_momenta()
thep4 = np.stack((px, py, pz, energy)).T
#test JaggedTLorentzVectorArray
tlva1 = uproot_methods.TLorentzVectorArray(px, py, pz, energy)
tlva2 = uproot_methods.TLorentzVectorArray(thep4[:, 0], thep4[:, 1],
thep4[:, 2], thep4[:, 3])
jtlva1 = JaggedTLorentzVectorArray.fromcounts(counts, tlva1)
jtlva2 = JaggedTLorentzVectorArray.fromcounts(counts, tlva2)
jtlva1_selection1 = jtlva1[jtlva1.counts > 0]
jtlva1_selection2 = jtlva1_selection1[jtlva1_selection1.pt > 5]
jtlva2_selection1 = jtlva2[jtlva2.counts > 0]
jtlva2_selection2 = jtlva1_selection1[jtlva2_selection1.pt > 5]
diffx = np.abs(jtlva1.x - jtlva2.x)
diffy = np.abs(jtlva1.y - jtlva2.y)
diffz = np.abs(jtlva1.z - jtlva2.z)
difft = np.abs(jtlva1.t - jtlva2.t)
assert (diffx < 1e-8).flatten().all()
assert (diffy < 1e-8).flatten().all()
assert (diffz < 1e-8).flatten().all()
assert (difft < 1e-8).flatten().all()
#test JaggedCandidateArray
jca1 = JaggedCandidateArray.candidatesfromcounts(counts, p4=thep4)
jca2 = JaggedCandidateArray.candidatesfromcounts(counts, p4=thep4)
assert ((jca1.offsets == jca2.offsets).all())
addon1 = jca1.zeros_like()
addon2 = jca2.ones_like()
jca1['addon'] = addon1
jca2['addon'] = addon2
jca1.add_attributes(addonFlat=addon1.flatten(), addonJagged=addon1)
diffm = np.abs(jca1.p4.mass - jca2.p4.mass)
assert ((jca1.offsets == jca2.offsets).all())
diffpt = np.abs(jca1.p4.pt - jca2.p4.pt)
assert ((jca1.offsets == jca2.offsets).all())
eta2 = jca2.p4.eta
eta1 = jca1.p4.eta
print(np.sum(eta1.counts), np.sum(eta2.counts))
diffeta_temp = np.abs(eta1 - eta2)
diffeta = np.abs(jca1.p4.eta - jca2.p4.eta)
assert ((jca1.offsets == jca2.offsets).all())
assert (diffm < 1e-8).flatten().all()
assert (diffpt < 1e-8).flatten().all()
assert (diffeta < 1e-8).flatten().all()
#test fast functions
fastfs = ['pt', 'eta', 'phi', 'mass']
for func in fastfs:
func1 = getattr(jca1, func)
func2 = getattr(jca1.p4, func)
dfunc = np.abs(func1 - func2)
assert (dfunc < 1e-8).flatten().all()
adistinct = jca1.distincts()
apair = jca1.pairs()
across = jca1.cross(jca2)
assert 'p4' in adistinct.columns
assert 'p4' in apair.columns
assert 'p4' in across.columns
admsum = (adistinct.i0.p4 + adistinct.i1.p4).mass
apmsum = (apair.i0.p4 + apair.i1.p4).mass
acmsum = (across.i0.p4 + across.i1.p4).mass
diffadm = np.abs(adistinct.p4.mass - admsum)
diffapm = np.abs(apair.p4.mass - apmsum)
diffacm = np.abs(across.p4.mass - acmsum)
assert (diffadm < 1e-8).flatten().all()
assert (diffapm < 1e-8).flatten().all()
assert (diffacm < 1e-8).flatten().all()
selection11 = jca1[jca1.counts > 0]
selection12 = selection11[selection11.p4.pt > 5]
selection21 = jca2[jca2.counts > 0]
selection22 = selection21[selection21.p4.pt > 5]
diffcnts = selection12.counts - jtlva1_selection2.counts
diffm = np.abs(selection12.p4.mass - jtlva1_selection2.mass)
diffaddon = selection12.addon - selection22.addon
assert (diffcnts == 0).flatten().all()
assert (diffm < 1e-8).flatten().all()
assert (diffaddon == -1).flatten().all()
#test gen-reco matching
gen, reco = gen_reco_TLV()
flat_gen = gen.flatten()
gen_px, gen_py, gen_pz, gen_e = flat_gen.x, flat_gen.y, flat_gen.z, flat_gen.t
flat_reco = reco.flatten()
reco_px, reco_py, reco_pz, reco_e = flat_reco.x, flat_reco.y, flat_reco.z, flat_reco.t
jca_gen = JaggedCandidateArray.candidatesfromcounts(gen.counts,
px=gen_px,
py=gen_py,
pz=gen_pz,
energy=gen_e)
jca_reco = JaggedCandidateArray.candidatesfromcounts(reco.counts,
px=reco_px,
py=reco_py,
pz=reco_pz,
energy=reco_e)
print('gen eta: ', jca_gen.p4.eta, '\n gen phi:', jca_gen.p4.phi)
print('reco eta: ', jca_reco.p4.eta, '\n reco phi:', jca_reco.p4.phi)
print('match mask: ', jca_reco.match(jca_gen, deltaRCut=0.3))
print('arg matches: ', jca_reco.argmatch(jca_gen, deltaRCut=0.3))
argmatch_nocut = jca_gen.argmatch(jca_reco).flatten()
argmatch_dr03 = jca_gen.argmatch(jca_reco, deltaRCut=0.3).flatten()
argmatch_dr03_dpt01 = jca_gen.argmatch(jca_reco,
deltaRCut=0.3,
deltaPtCut=0.1).flatten()
assert (argmatch_nocut.size == 5)
assert (argmatch_dr03[argmatch_dr03 != -1].size == 3)
assert (argmatch_dr03_dpt01[argmatch_dr03_dpt01 != -1].size == 2)
assert (jca_gen.match(jca_reco,
deltaRCut=0.3).flatten().flatten().sum() == 3)
assert (jca_gen.match(jca_reco, deltaRCut=0.3,
deltaPtCut=0.1).flatten().flatten().sum() == 2)
def plotratio(num,
denom,
ax=None,
clear=True,
overflow='none',
error_opts=None,
denom_fill_opts=None,
guide_opts=None,
unc='clopper-pearson',
label=None):
"""
Create a ratio plot, dividing two compatible histograms
num: Hist object with single axis
denom: Hist object with identical axis to num
ax: matplotlib Axes object (if None, one is created)
clear: clear Axes before drawing (if passed); if False, this function will skip drawing the legend
overflow: overflow behavior of plot axis (see Hist.sum() docs)
The draw options are passed as dicts to the relevant matplotlib function, with some exceptions in case
it is especially common or useful. If none of *_opts is specified, nothing will be plotted!
Pass an empty dict (e.g. error_opts={}) for defaults.
error_opts: to plot an errorbar
Special options interpreted by this function and not passed to matplotlib:
'emarker' (default: '') marker to place at cap of errorbar
denom_fill_opts: to plot a filled area centered at 1, representing denominator uncertainty
Special options interpreted by this function and not passed to matplotlib:
(none)
guide_opts: to plot a horizontal guide line at ratio of 1.
Special options interpreted by this function and not passed to matplotlib:
(none)
unc: Uncertainty calculation option
'clopper-pearson': interval for efficiencies
'poisson-ratio': interval for ratio of poisson distributions
'num': poisson interval of numerator scaled by denominator value
(common for data/mc, for better or worse...)
label: associate a label with this entry (note: y axis label set by num.label)
"""
if ax is None:
fig, ax = plt.subplots(1, 1)
else:
if not isinstance(ax, plt.Axes):
raise ValueError("ax must be a matplotlib Axes object")
if clear:
ax.clear()
fig = ax.figure
if not num.compatible(denom):
raise ValueError(
"numerator and denominator histograms have incompatible axis definitions"
)
if num.dim() > 1:
raise ValueError(
"plotratio() can only support one-dimensional histograms")
if error_opts is None and denom_fill_opts is None and guide_opts is None:
error_opts = {}
denom_fill_opts = {}
axis = num.axes()[0]
if isinstance(axis, SparseAxis):
raise NotImplementedError(
"Ratio for sparse axes (labeled axis with errorbars)")
elif isinstance(axis, DenseAxis):
ax.set_xlabel(axis.label)
ax.set_ylabel(num.label)
edges = axis.edges(overflow=overflow)
centers = axis.centers(overflow=overflow)
sumw_num, sumw2_num = num.values(sumw2=True, overflow=overflow)[()]
sumw_denom, sumw2_denom = denom.values(sumw2=True,
overflow=overflow)[()]
rsumw = sumw_num / sumw_denom
if unc == 'clopper-pearson':
rsumw_err = np.abs(
clopper_pearson_interval(sumw_num, sumw_denom) - rsumw)
elif unc == 'poisson-ratio':
# poisson ratio n/m is equivalent to binomial n/(n+m)
rsumw_err = np.abs(
clopper_pearson_interval(sumw_num, sumw_num + sumw_denom) -
rsumw)
elif unc == 'num':
rsumw_err = np.abs(
poisson_interval(rsumw, sumw2_num / sumw_denom**2) - rsumw)
else:
raise ValueError("Unrecognized uncertainty option: %r" % unc)
primitives = {}
if error_opts is not None:
opts = {'label': label, 'linestyle': 'none'}
opts.update(error_opts)
emarker = opts.pop('emarker', '')
errbar = ax.errorbar(x=centers, y=rsumw, yerr=rsumw_err, **opts)
plt.setp(errbar[1], 'marker', emarker)
primitives['error'] = errbar
if denom_fill_opts is not None:
unity = np.ones_like(sumw_denom)
denom_unc = poisson_interval(unity, sumw2_denom / sumw_denom**2)
opts = {
'step': 'post',
'facecolor': (0, 0, 0, 0.3),
'linewidth': 0
}
opts.update(denom_fill_opts)
fill = ax.fill_between(edges, np.r_[denom_unc[0], denom_unc[0,
-1]],
np.r_[denom_unc[1], denom_unc[1,
-1]], **opts)
primitives['denom_fill'] = fill
if guide_opts is not None:
opts = {'linestyle': '--', 'color': (0, 0, 0, 0.5), 'linewidth': 1}
opts.update(guide_opts)
primitives['guide'] = ax.axhline(1., **opts)
if clear:
ax.autoscale(axis='x', tight=True)
ax.set_ylim(0, None)
return fig, ax, primitives
