How to use the coffea.util.numpy.abs 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_analysis_objects.py View on Github external
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)
    acrossn = jca1.cross(jca2, nested=True)
    achoose2 = jca1.choose(2)
    achoose3 = jca1.choose(3)
    
    assert 'p4' in adistinct.columns
    assert 'p4' in apair.columns
    assert 'p4' in across.columns
    assert 'p4' in acrossn.columns
    assert 'p4' in achoose2.columns
    assert 'p4' in achoose3.columns
github CoffeaTeam / coffea / tests / dummy_distributions.py View on Github external
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
github CoffeaTeam / coffea / tests / test_analysis_objects.py View on Github external
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)
github CoffeaTeam / coffea / tests / test_processor.py View on Github external
def test_weights_partial():
    from coffea.processor import Weights
    counts, _, _ = dummy_jagged_eta_pt()
    w1 = np.random.normal(loc=1.0, scale=0.01, size=counts.size)
    w2 = np.random.normal(loc=1.3, scale=0.05, size=counts.size)

    weights = Weights(counts.size, storeIndividual=True)
    weights.add('w1', w1)
    weights.add('w2', w2)

    test_exclude_none = weights.weight()
    assert(np.all(np.abs(test_exclude_none - w1 * w2)<1e-6))

    test_exclude1 = weights.partial_weight(exclude=['w1'])
    assert(np.all(np.abs(test_exclude1 - w2)<1e-6))

    test_include1 = weights.partial_weight(include=['w1'])
    assert(np.all(np.abs(test_include1 - w1)<1e-6))

    test_exclude2 = weights.partial_weight(exclude=['w2'])
    assert(np.all(np.abs(test_exclude2 - w1)<1e-6))

    test_include2 = weights.partial_weight(include=['w2'])
    assert(np.all(np.abs(test_include2 - w2)<1e-6))

    test_include_both = weights.partial_weight(include=['w1','w2'])
    assert(np.all(np.abs(test_include_both - w1 * w2)<1e-6))

    # Check that exception is thrown if arguments are incompatible
    error_raised = False
    try:
github CoffeaTeam / coffea / tests / test_processor.py View on Github external
assert(np.all(np.abs(test_central - (exp_weight)) < 1e-6))

    test_up = weight.weight('testUp')
    exp_up = scale_central * scale_central * 1.10

    assert(np.all(np.abs(test_up - (exp_up)) < 1e-6))

    test_down = weight.weight('testDown')
    exp_down = scale_central * scale_central * 0.95

    assert(np.all(np.abs(test_down - (exp_down)) < 1e-6))

    test_shift_up = weight.weight('testUp')

    assert(np.all(np.abs(test_shift_up - (exp_up)) < 1e-6))
    
    test_shift_down = weight.weight('testDown')

    assert(np.all(np.abs(test_shift_down - (exp_down)) < 1e-6))
github CoffeaTeam / coffea / tests / test_analysis_objects.py View on Github external
#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
github CoffeaTeam / coffea / tests / test_analysis_objects.py View on Github external
achoose3 = jca1.choose(3)
    
    assert 'p4' in adistinct.columns
    assert 'p4' in apair.columns
    assert 'p4' in across.columns
    assert 'p4' in acrossn.columns
    assert 'p4' in achoose2.columns
    assert 'p4' in achoose3.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
    ach3msum = (achoose3.i0.p4 + achoose3.i1.p4 + achoose3.i2.p4).mass
    diffadm = np.abs(adistinct.p4.mass - admsum)
    diffapm = np.abs(apair.p4.mass - apmsum)
    diffacm = np.abs(across.p4.mass - acmsum)
    diffachm = np.abs(achoose2.p4.mass - admsum)
    diffach3m = np.abs(achoose3.p4.mass - ach3msum)
    
    assert (diffadm < 1e-8).flatten().all()
    assert (diffapm < 1e-8).flatten().all()
    assert (diffacm < 1e-8).flatten().all()
    assert (diffachm < 1e-8).flatten().all()
    assert (diffach3m < 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
github CoffeaTeam / coffea / tests / test_lookup_tools.py View on Github external
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()
github CoffeaTeam / coffea / coffea / analysis_objects / JaggedCandidateMethods.py View on Github external
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)
github CoffeaTeam / coffea / coffea / analysis_objects / JaggedCandidateMethods.py View on Github external
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))