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.

github CoffeaTeam / coffea / tests / test_jetmet_tools.py View on Github external
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:
github CoffeaTeam / coffea / tests / test_jetmet_tools.py View on Github external
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))
github CoffeaTeam / coffea / tests / test_analysis_objects.py View on Github external
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)
github CoffeaTeam / coffea / tests / test_analysis_objects.py View on Github external
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)
github CoffeaTeam / coffea / tests / test_analysis_objects.py View on Github external
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
github CoffeaTeam / coffea / tests / test_analysis_objects.py View on Github external
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))
github CoffeaTeam / coffea / tests / test_analysis_objects.py View on Github external
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)
github CoffeaTeam / coffea / tests / test_analysis_objects.py View on Github external
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)