How to use the coffea.analysis_objects.JaggedCandidateArray.candidatesfromcounts 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.
CoffeaTeam / coffea / tests / test_jetmet_tools.py View on Github 
counts, test_px, test_py, test_pz, test_e = dummy_four_momenta()
test_Rho = np.full(shape=(np.sum(counts),), fill_value=100.)
test_A = np.full(shape=(np.sum(counts),), fill_value=5.)
jets = CandArray.candidatesfromcounts(counts, px=test_px, py=test_py, pz=test_pz, energy=test_e)
jets.add_attributes(ptRaw=jets.pt,
massRaw=jets.mass,
rho=test_Rho,
area=test_A)
fakemet = np.random.exponential(scale=1.0,size=counts.size)
metphi = np.random.uniform(low=-math.pi, high=math.pi, size=counts.size)
syst_up = 0.001*fakemet
syst_down = -0.001*fakemet
met = CandArray.candidatesfromcounts(np.ones_like(counts),
pt=fakemet,
eta=np.zeros_like(counts),
phi=metphi,
mass=np.zeros_like(counts),
MetUnclustEnUpDeltaX=syst_up*np.cos(metphi),
MetUnclustEnUpDeltaY=syst_down*np.sin(metphi))
jec_names = ['Summer16_23Sep2016V3_MC_L1FastJet_AK4PFPuppi',
'Summer16_23Sep2016V3_MC_L2Relative_AK4PFPuppi',
'Summer16_23Sep2016V3_MC_L2L3Residual_AK4PFPuppi',
'Summer16_23Sep2016V3_MC_L3Absolute_AK4PFPuppi']
corrector = FactorizedJetCorrector(**{name: evaluator[name] for name in jec_names})
junc_names = []
for name in dir(evaluator):
if 'Summer16_23Sep2016V3_MC_UncertaintySources_AK4PFPuppi' in name:import numpy as np
import awkward as ak
import math
from coffea.analysis_objects import JaggedCandidateArray as CandArray
from coffea.jetmet_tools import (FactorizedJetCorrector,
JetResolution,
JetResolutionScaleFactor,
JetCorrectionUncertainty,
JetTransformer)
counts, test_px, test_py, test_pz, test_e = dummy_four_momenta()
test_Rho = np.full(shape=(np.sum(counts),), fill_value=100.)
test_A = np.full(shape=(np.sum(counts),), fill_value=5.)
jets = CandArray.candidatesfromcounts(counts, px=test_px, py=test_py, pz=test_pz, energy=test_e)
jets.add_attributes(ptRaw=jets.pt,
massRaw=jets.mass,
rho=test_Rho,
area=test_A)
fakemet = np.random.exponential(scale=1.0,size=counts.size)
metphi = np.random.uniform(low=-math.pi, high=math.pi, size=counts.size)
syst_up = 0.001*fakemet
syst_down = -0.001*fakemet
met = CandArray.candidatesfromcounts(np.ones_like(counts),
pt=fakemet,
eta=np.zeros_like(counts),
phi=metphi,
mass=np.zeros_like(counts),
MetUnclustEnUpDeltaX=syst_up*np.cos(metphi),
MetUnclustEnUpDeltaY=syst_down*np.sin(metphi))phi=jca_reco.phi,
energy=jca_reco.p4.energy)
pxpypzM_test = JaggedCandidateArray.candidatesfromcounts(jca_reco.counts,
px=jca_reco.p4.x,
py=jca_reco.p4.y,
pz=jca_reco.p4.z,
mass=jca_reco.mass)
ptphipzE_test = JaggedCandidateArray.candidatesfromcounts(jca_reco.counts,
pt=jca_reco.pt,
phi=jca_reco.phi,
pz=jca_reco.p4.z,
energy=jca_reco.p4.energy)
pthetaphiE_test = JaggedCandidateArray.candidatesfromcounts(jca_reco.counts,
p=jca_reco.p4.p,
theta=jca_reco.p4.theta,
phi=jca_reco.phi,
energy=jca_reco.p4.energy)
p4cart_test = JaggedCandidateArray.candidatesfromcounts(jca_reco.counts,
p4=jca_reco.p4)
p4ptetaphiM_test = JaggedCandidateArray.candidatesfromcounts(jca_reco.counts,
p4=jca_reco_pt.p4)assert (jca_gen.match(jca_reco,deltaRCut=0.3,deltaPtCut=0.1).flatten().flatten().sum()==2)
# test various four-momentum constructors
ptetaphiE_test = JaggedCandidateArray.candidatesfromcounts(jca_reco.counts,
pt=jca_reco.pt,
eta=jca_reco.eta,
phi=jca_reco.phi,
energy=jca_reco.p4.energy)
pxpypzM_test = JaggedCandidateArray.candidatesfromcounts(jca_reco.counts,
px=jca_reco.p4.x,
py=jca_reco.p4.y,
pz=jca_reco.p4.z,
mass=jca_reco.mass)
ptphipzE_test = JaggedCandidateArray.candidatesfromcounts(jca_reco.counts,
pt=jca_reco.pt,
phi=jca_reco.phi,
pz=jca_reco.p4.z,
energy=jca_reco.p4.energy)
pthetaphiE_test = JaggedCandidateArray.candidatesfromcounts(jca_reco.counts,
p=jca_reco.p4.p,
theta=jca_reco.p4.theta,
phi=jca_reco.phi,
energy=jca_reco.p4.energy)
p4cart_test = JaggedCandidateArray.candidatesfromcounts(jca_reco.counts,
p4=jca_reco.p4)
p4ptetaphiM_test = JaggedCandidateArray.candidatesfromcounts(jca_reco.counts,
p4=jca_reco_pt.p4)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.etajtlva2_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))print('match mask: ', match_mask)
fast_match_mask = jca_reco.fastmatch(jca_gen, deltaRCut=0.3)
print('fastmatch mask: ', fast_match_mask)
assert((match_mask == fast_match_mask).all().all())
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)
# test various four-momentum constructors
ptetaphiE_test = JaggedCandidateArray.candidatesfromcounts(jca_reco.counts,
pt=jca_reco.pt,
eta=jca_reco.eta,
phi=jca_reco.phi,
energy=jca_reco.p4.energy)
pxpypzM_test = JaggedCandidateArray.candidatesfromcounts(jca_reco.counts,
px=jca_reco.p4.x,
py=jca_reco.p4.y,
pz=jca_reco.p4.z,
mass=jca_reco.mass)
ptphipzE_test = JaggedCandidateArray.candidatesfromcounts(jca_reco.counts,
pt=jca_reco.pt,
phi=jca_reco.phi,
pz=jca_reco.p4.z,
energy=jca_reco.p4.energy)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)
jca_reco_pt = JaggedCandidateArray.candidatesfromcounts(reco.counts,
pt=jca_reco.pt.content,eta=jca_reco.eta.content,
phi=jca_reco.phi.content,mass=jca_reco.mass.content)
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)
match_mask = jca_reco.match(jca_gen, deltaRCut=0.3)
print('match mask: ', match_mask)
fast_match_mask = jca_reco.fastmatch(jca_gen, deltaRCut=0.3)
print('fastmatch mask: ', fast_match_mask)
assert((match_mask == fast_match_mask).all().all())
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)coffea
Basic tools and wrappers for enabling not-too-alien syntax when running columnar Collider HEP analysis.
Package Health Score
87 / 100Popular coffea functions
- coffea.analysis_objects.JaggedCandidateArray.candidatesfromcounts
- coffea.hist.Bin
- coffea.hist.Cat
- coffea.hist.Hist
- coffea.processor
- coffea.util.awkward
- coffea.util.awkward.JaggedArray
- coffea.util.awkward.JaggedArray.fromcounts
- coffea.util.awkward.JaggedArray.fromoffsets
- coffea.util.numpy
- coffea.util.numpy.abs
- coffea.util.numpy.array
- coffea.util.numpy.clip
- coffea.util.numpy.random.normal
- coffea.util.numpy.searchsorted
- coffea.util.numpy.sum
- coffea.util.numpy.union1d
- coffea.util.numpy.unique
- coffea.util.numpy.where
- coffea.util.numpy.zeros