Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
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
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
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)
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:
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))
#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
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
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 _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 _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))