def test_3D_jagged(wrapped_tree):
fake_3d = [[np.arange(i + 1) + j for i in range(j % 3)]
for j in range(len(wrapped_tree))]
fake_3d = JaggedArray.fromiter(fake_3d)
wrapped_tree.new_variable("Fake3D", fake_3d)
assert isinstance(fake_3d.count(), JaggedArray)
assert all((fake_3d.copy().count() == fake_3d.count()).all())
aliased = expressions.evaluate(wrapped_tree, "Fake3D")
assert (aliased == fake_3d).all().all().all()
doubled = expressions.evaluate(wrapped_tree, "Fake3D * 2")
assert (doubled == fake_3d * 2).all().all().all()
assert len(doubled[0, :, :]) == 0
assert doubled[1, 0, :] == [2]
assert doubled[2, 0, :] == [4]
assert all(doubled[2, 1, :] == [4, 6])
doubled = expressions.evaluate(wrapped_tree, "Fake3D + Fake3D")
assert (doubled == fake_3d * 2).all().all().all()
assert len(doubled[0, :, :]) == 0
assert doubled[1, 0, :] == [2]
assert doubled[2, 0, :] == [4]
assert all(doubled[2, 1, :] == [4, 6])
fake_3d_2 = [[np.arange(i + 3) + j for i in range(j % 2)]
for j in range(len(wrapped_tree))]
fake_3d_2 = JaggedArray.fromiter(fake_3d_2)
wrapped_tree.new_variable("SecondFake3D", fake_3d_2)
with pytest.raises(RuntimeError) as e:
expressions.evaluate(wrapped_tree, "SecondFake3D + Fake3D")
assert "different jaggedness" in str(e)
def NestNestObjArrayToJagged(objarr):
"""uproot read vector<vector<number>> TBranch
as objectArray, this function convert it
to JaggedJaggedArray
"""
# jaggedArray of lists
jaggedList = JaggedArray.fromiter(objarr)
# flat to 1 level
_jagged = JaggedArray.fromiter(jaggedList.content)
return JaggedArray.fromoffsets(jaggedList.offsets, _jagged)
def test_jagged_nth_3D(jagged_1):
fake_3d = [[np.arange(i + 1) + j for i in range(j % 3)] for j in range(5)]
fake_3d = JaggedArray.fromiter(fake_3d)
get_second = reductions.JaggedNth(1, np.nan)
reduced = get_second(fake_3d)
assert len(reduced[0]) == 0
assert len(reduced[1]) == 1
assert np.isnan(reduced[1])
assert len(reduced[2]) == 2
assert np.isnan(reduced[2][0])
assert reduced[2][1] == 3
assert len(reduced[3]) == 0
assert len(reduced[4]) == 1
assert np.isnan(reduced[4])
get_first = reductions.JaggedNth(0, np.nan)
reduced = get_first(fake_3d)
assert len(reduced[0]) == 0
assert len(reduced[1]) == 1
assert reduced[1][0] == 1
assert len(reduced[2]) == 2
assert reduced[2][0] == 2
assert reduced[2][1] == 2
assert len(reduced[3]) == 0
assert len(reduced[4]) == 1
assert reduced[4] == 4
def jagged_1():
return JaggedArray.fromiter([[0.0, 1.1, 2.2], [3.3, 4.4], [5.5],
[6.6, 7.7, 8.8], [9.9, 10.0, 11.0], []])
float(phi[ie][im]),
float(gpt[ie][gid[ie][im]]))
]
else:
subres += [
roccor.kSmearMC(int(charge[ie][im]),
float(pt[ie][im]),
float(eta[ie][im]),
float(phi[ie][im]),
int(nl[ie][im]), float(u[ie][im]))
]
suberr += [
roccor.kSmearMCerror(int(charge[ie][im]),
float(pt[ie][im]),
float(eta[ie][im]),
float(phi[ie][im]),
int(nl[ie][im]),
float(u[ie][im]))
]
res += [subres]
err += [suberr]
res = JaggedArray.fromiter(res)
err = JaggedArray.fromiter(err)
outres = res.flatten()
outerr = err.flatten()
np.save(fname.replace('.root', '_rochester.npy'), outres)
np.save(fname.replace('.root', '_rochester_err.npy'), outerr)
if not isData:
outrand = np.concatenate([ui.flatten() for ui in fullu])
np.save(fname.replace('.root', '_rochester_rand.npy'), outrand)
loose_muons = muons[loose_muon_selection]
loose_photons = photons[loose_photon_selection]
selected_taus = taus[tau_selection]
selected_jets = jets[jet_selection]
# end seletion
# clean leptons
e_combinations = loose_electrons.p4.cross(selected_jets.p4, nested=True)
mask = (e_combinations.i0.delta_r(e_combinations.i1) < 0.3).any()
clean_electrons = loose_electrons[~mask]
m_combinations = loose_muons.p4.cross(selected_jets.p4, nested=True)
mask = (m_combinations.i0.delta_r(m_combinations.i1) < 0.3).any()
clean_muons = loose_muons[mask]
clean_leptons = JaggedArray.fromiter([clean_electrons, clean_muons])
# once merge is done
# mask = loose_electrons.p4.match(selected_jets.p4, 0.3)
# clean electrons = loose_electrons[~mask]
# end cleaning
# weights evaluation
e_counts = clean_electrons.counts
e_sfTrigg = np.ones(clean_electrons.size)
e_sfTrigg[e_counts > 0] = 1 - evaluator["hEffEtaPt"](
clean_electrons.eta[e_counts > 0, 0], clean_electrons.pt[e_counts > 0,
0])
e_sfTrigg[e_counts > 1] = 1 - (1 - evaluator["hEffEtaPt"](
clean_electrons.eta[e_counts > 1, 0],
clean_electrons.pt[e_counts > 1, 0])) * (
