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