def test_read_nanomc(suffix):
path = os.path.abspath(f'tests/samples/nano_dy.{suffix}')
factory = getattr(NanoEventsFactory, f'from_{suffix}')(path)
events = factory.events()
# test after views first
genroundtrips(events.GenPart.mask[events.GenPart.eta > 0])
genroundtrips(events.mask[ak.any(events.Electron.pt > 50, axis=1)].GenPart)
genroundtrips(events.GenPart)
genroundtrips(events.GenPart[events.GenPart.eta > 0])
genroundtrips(events[ak.any(events.Electron.pt > 50, axis=1)].GenPart)
# sane gen matching (note for electrons gen match may be photon(22))
assert ak.all((abs(events.Electron.matched_gen.pdgId) == 11) | (events.Electron.matched_gen.pdgId == 22))
assert ak.all(abs(events.Muon.matched_gen.pdgId) == 13)
genroundtrips(events.Electron.matched_gen)
crossref(events[ak.num(events.Jet) > 2])
crossref(events)
if suffix == 'root':
assert ak.any(events.Photon.isTight, axis=1).tolist()[:9] == [False, True, True, True, False, False, False, False, False]
if suffix == 'parquet':
assert ak.any(events.Photon.isTight, axis=1).tolist()[:9] == [False, True, False, True, False, False, False, False, True]
def test_count_nonzero():
array = ak.Array(
[
[
[np.datetime64("2022"), np.datetime64("2023"), np.datetime64("2025")],
[],
[np.datetime64("2027"), np.datetime64("2011")],
[np.datetime64("2013")],
],
[],
[[np.datetime64("2017"), np.datetime64("2019")], [np.datetime64("2023")]],
],
check_valid=True,
)
assert ak.count_nonzero(array) == 9
assert ak.to_list(ak.count_nonzero(array, axis=-1)) == [[3, 0, 2, 1], [], [2, 1]]
assert ak.to_list(ak.count_nonzero(array, axis=-2)) == [[3, 2, 1], [], [2, 1]]
assert ak.to_list(ak.all(array, axis=-1)) == [
[True, True, True, True],
[],
[True, True],
]
assert ak.to_list(ak.all(array, axis=-2)) == [[True, True, True], [], [True, True]]
def test_polar_two_vector():
a = ak.zip(
{
"r": [[1, 2], [], [3], [4]],
"phi": [[0.3, 0.4], [], [0.5], [0.6]],
},
with_name="PolarTwoVector",
highlevel=False,
)
a = ak.Array(a, behavior=vector.behavior)
assert record_arrays_equal(
a * 2,
ak.zip({
"r": [[2, 4], [], [6], [8]],
"phi": [[0.3, 0.4], [], [0.5], [0.6]]
}),
)
assert ak.all((a * (-2)).r == [[2, 4], [], [6], [8]])
assert ak.all(
(a * (-2)).phi - ak.Array([[-2.8415926535, -2.7415926535], [],
[-2.6415926535], [-2.5415926535]]) < ATOL)
assert record_arrays_equal(
a / 2,
ak.zip({
"r": [[0.5, 1], [], [1.5], [2]],
"phi": [[0.3, 0.4], [], [0.5], [0.6]]
}),
)
assert ak.all(abs((-a).x + a.x) < ATOL)
assert ak.all(abs((-a).y + a.y) < ATOL)
assert record_arrays_equal(a * (-1), -a)
assert ak.all(a.unit.phi == a.phi)
def test_root_scalefactors():
extractor = lookup_tools.extractor()
extractor.add_weight_sets([
"testSF2d scalefactors_Tight_Electron tests/samples/testSF2d.histo.root"
])
extractor.finalize(reduce_list=["testSF2d"])
evaluator = extractor.make_evaluator()
counts, test_eta, test_pt = dummy_jagged_eta_pt()
# test flat eval
test_out = evaluator["testSF2d"](test_eta, test_pt)
# print it
print(evaluator["testSF2d"])
# test structured eval
test_eta_jagged = ak.unflatten(test_eta, counts)
test_pt_jagged = ak.unflatten(test_pt, counts)
test_out_jagged = evaluator["testSF2d"](test_eta_jagged, test_pt_jagged)
assert ak.all(ak.num(test_out_jagged) == counts)
assert ak.all(ak.flatten(test_out_jagged) == test_out)
print(test_out)
diff = np.abs(test_out - _testSF2d_expected_output)
print("Max diff: %.16f" % diff.max())
print("Median diff: %.16f" % np.median(diff))
print("Diff over threshold rate: %.1f %%" %
(100 * (diff >= 1.0e-8).sum() / diff.size))
assert (diff < 1.0e-8).all()
def crossref(events):
# check some cross-ref roundtrips (some may not be true always but they are for the test file)
assert ak.all(events.Jet.matched_muons.matched_jet.pt == events.Jet.pt)
assert ak.all(events.Electron.matched_photon.matched_electron.r9 ==
events.Electron.r9)
# exercise LorentzVector.nearest
assert ak.all(
events.Muon.matched_jet.delta_r(events.Muon.nearest(events.Jet)) == 0.)
def test_read_nanomc(suffix):
path = os.path.abspath(f"tests/samples/nano_dy.{suffix}")
# parquet files were converted from even older nanoaod
nanoversion = NanoAODSchema.v6 if suffix == "root" else NanoAODSchema.v5
factory = getattr(NanoEventsFactory, f"from_{suffix}")(
path, schemaclass=nanoversion
)
events = factory.events()
# test after views first
genroundtrips(events.GenPart.mask[events.GenPart.eta > 0])
genroundtrips(events.mask[ak.any(events.Electron.pt > 50, axis=1)].GenPart)
genroundtrips(events.GenPart)
genroundtrips(events.GenPart[events.GenPart.eta > 0])
genroundtrips(events[ak.any(events.Electron.pt > 50, axis=1)].GenPart)
# sane gen matching (note for electrons gen match may be photon(22))
assert ak.all(
(abs(events.Electron.matched_gen.pdgId) == 11)
| (events.Electron.matched_gen.pdgId == 22)
)
assert ak.all(abs(events.Muon.matched_gen.pdgId) == 13)
genroundtrips(events.Electron.matched_gen)
crossref(events[ak.num(events.Jet) > 2])
crossref(events)
# test issue 409
assert ak.to_list(events[[]].Photon.mass) == []
if suffix == "root":
assert ak.any(events.Photon.isTight, axis=1).tolist()[:9] == [
False,
True,
True,
True,
False,
False,
False,
False,
False,
]
if suffix == "parquet":
assert ak.any(events.Photon.isTight, axis=1).tolist()[:9] == [
False,
True,
False,
True,
False,
False,
False,
False,
True,
]
def genroundtrips(genpart):
# check genpart roundtrip
assert ak.all(genpart.children.parent.pdgId == genpart.pdgId)
assert ak.all(ak.any(genpart.parent.children.pdgId == genpart.pdgId, axis=-1, mask_identity=True))
# distinctParent should be distinct and it should have a relevant child
assert ak.all(genpart.distinctParent.pdgId != genpart.pdgId)
assert ak.all(ak.any(genpart.distinctParent.children.pdgId == genpart.pdgId, axis=-1, mask_identity=True))
# exercise hasFlags
genpart.hasFlags(['isHardProcess'])
genpart.hasFlags(['isHardProcess', 'isDecayedLeptonHadron'])
def test_jet_resolution():
from coffea.jetmet_tools import JetResolution
counts, test_eta, test_pt = dummy_jagged_eta_pt()
test_Rho = np.full_like(test_eta, 10.0)
test_pt_jag = ak.unflatten(test_pt, counts)
test_eta_jag = ak.unflatten(test_eta, counts)
test_Rho_jag = ak.unflatten(test_Rho, counts)
jer_names = ["Spring16_25nsV10_MC_PtResolution_AK4PFPuppi"]
reso = JetResolution(**{name: evaluator[name] for name in jer_names})
print(reso)
resos = reso.getResolution(JetEta=test_eta, Rho=test_Rho, JetPt=test_pt)
resos_jag = reso.getResolution(JetEta=test_eta_jag,
Rho=test_Rho_jag,
JetPt=test_pt_jag)
assert ak.all(np.abs(resos - ak.flatten(resos_jag)) < 1e-6)
test_pt_jag = test_pt_jag[0:3]
test_eta_jag = test_eta_jag[0:3]
test_Rho_jag = test_Rho_jag[0:3]
test_Rho_jag = ak.concatenate(
[test_Rho_jag[:-1], [ak.concatenate([test_Rho_jag[-1, :-1], 100.0])]])
counts = counts[0:3]
print("Raw jet values:")
print("pT:", test_pt_jag)
print("eta:", test_eta_jag)
print("rho:", test_Rho_jag, "\n")
resos_jag_ref = ak.unflatten(
np.array([
0.21974642,
0.32421591,
0.33702479,
0.27420327,
0.13940689,
0.48134521,
0.26564994,
1.0,
]),
counts,
)
resos_jag = reso.getResolution(JetEta=test_eta_jag,
Rho=test_Rho_jag,
JetPt=test_pt_jag)
print("Reference Resolution (jagged):", resos_jag_ref)
print("Resolution (jagged):", resos_jag)
# NB: 5e-4 tolerance was agreed upon by lgray and aperloff, if the differences get bigger over time
# we need to agree upon how these numbers are evaluated (double/float conversion is kinda random)
assert ak.all(
np.abs(ak.flatten(resos_jag_ref) - ak.flatten(resos_jag)) < 5e-4)
def test_jet_correction_uncertainty():
from coffea.jetmet_tools import JetCorrectionUncertainty
counts, test_eta, test_pt = dummy_jagged_eta_pt()
test_pt_jag = ak.unflatten(test_pt, counts)
test_eta_jag = ak.unflatten(test_eta, counts)
junc_names = ["Summer16_23Sep2016V3_MC_Uncertainty_AK4PFPuppi"]
junc = JetCorrectionUncertainty(
**{name: evaluator[name]
for name in junc_names})
print(junc)
juncs = junc.getUncertainty(JetEta=test_eta, JetPt=test_pt)
juncs_jag = list(
junc.getUncertainty(JetEta=test_eta_jag, JetPt=test_pt_jag))
for i, (level, corrs) in enumerate(juncs):
assert corrs.shape[0] == test_eta.shape[0]
assert ak.all(corrs == ak.flatten(juncs_jag[i][1]))
test_pt_jag = test_pt_jag[0:3]
test_eta_jag = test_eta_jag[0:3]
counts = counts[0:3]
print("Raw jet values:")
print("pT:", test_pt_jag.tolist())
print("eta:", test_eta_jag.tolist(), "\n")
juncs_jag_ref = ak.unflatten(
np.array([
[1.053504214, 0.946495786],
[1.033343349, 0.966656651],
[1.065159157, 0.934840843],
[1.033140127, 0.966859873],
[1.016858652, 0.983141348],
[1.130199999, 0.869800001],
[1.039968468, 0.960031532],
[1.033100002, 0.966899998],
]),
counts,
)
juncs_jag = list(
junc.getUncertainty(JetEta=test_eta_jag, JetPt=test_pt_jag))
for i, (level, corrs) in enumerate(juncs_jag):
print("Index:", i)
print("Correction level:", level)
print("Reference Uncertainties (jagged):", juncs_jag_ref)
print("Uncertainties (jagged):", corrs)
assert ak.all(
np.abs(ak.flatten(juncs_jag_ref) - ak.flatten(corrs)) < 1e-6)
def test_nested_lorentz_vectorization(collection, events):
mask = ak.num(events[collection]) > 0
assert ak.all(ak.num(events[collection].PrunedP4_5, axis=2) == 5)
assert (ak.type(events[collection][mask].PrunedP4_5[
0, 0, 0]).parameters["__record__"] == "LorentzVector")
assert ak.all(ak.num(events[collection].SoftDroppedP4_5, axis=2) == 5)
assert (ak.type(events[collection][mask].SoftDroppedP4_5[
0, 0, 0]).parameters["__record__"] == "LorentzVector")
assert ak.all(ak.num(events[collection].TrimmedP4_5, axis=2) == 5)
assert (ak.type(events[collection][mask].TrimmedP4_5[
0, 0, 0]).parameters["__record__"] == "LorentzVector")
def process(self, events):
cleanjets = events.Jet[ak.all(
events.Jet.metric_table(events.Muon[events.Muon.pt > 10]) >= 0.4,
axis=2)
& ak.all(
events.Jet.metric_table(events.Electron[
events.Electron.pt > 10]) >= 0.4,
axis=2,
)
& (events.Jet.pt > 30)]
return (hist.Hist.new.Reg(
100, 0, 200, name="sumjetpt",
label=r"Jet $\sum p_{T}$ [GeV]").Double().fill(
ak.sum(cleanjets.pt, axis=1)))
def test_root_scalefactors():
extractor = lookup_tools.extractor()
extractor.add_weight_sets([
"testSF2d scalefactors_Tight_Electron tests/samples/testSF2d.histo.root"
])
extractor.finalize(reduce_list=['testSF2d'])
evaluator = extractor.make_evaluator()
counts, test_eta, test_pt = dummy_jagged_eta_pt()
# test flat eval
test_out = evaluator["testSF2d"](test_eta, test_pt)
# print it
print(evaluator["testSF2d"])
# test structured eval
test_eta_jagged = ak.unflatten(test_eta, counts)
test_pt_jagged = ak.unflatten(test_pt, counts)
test_out_jagged = evaluator["testSF2d"](test_eta_jagged, test_pt_jagged)
assert ak.all(ak.num(test_out_jagged) == counts)
assert ak.all(ak.flatten(test_out_jagged) == test_out)
# From make_expected_lookup.py
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
])
print(test_out)
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 _string_equal(one, two):
nplike = ak.nplike.of(one, two)
behavior = ak._util.behaviorof(one, two)
one, two = ak.without_parameters(one).layout, ak.without_parameters(
two).layout
# first condition: string lengths must be the same
counts1 = nplike.asarray(one.count(axis=-1))
counts2 = nplike.asarray(two.count(axis=-1))
out = counts1 == counts2
# only compare characters in strings that are possibly equal (same length)
possible = nplike.logical_and(out, counts1)
possible_counts = counts1[possible]
if len(possible_counts) > 0:
onepossible = one[possible]
twopossible = two[possible]
reduced = ak.all(ak.Array(onepossible) == ak.Array(twopossible),
axis=-1).layout
# update same-length strings with a verdict about their characters
out[possible] = reduced
return ak._util.wrap(ak.layout.NumpyArray(out), behavior)
def gen_jet_pair_mass(df):
gjmass = None
gjets = df.GenJet
gleptons = df.GenPart[(abs(df.GenPart.pdgId) == 13)
| (abs(df.GenPart.pdgId) == 11)
| (abs(df.GenPart.pdgId) == 15)]
gl_pair = ak.cartesian({
"jet": gjets,
"lepton": gleptons
},
axis=1,
nested=True)
_, _, dr_gl = delta_r(
gl_pair["jet"].eta,
gl_pair["lepton"].eta,
gl_pair["jet"].phi,
gl_pair["lepton"].phi,
)
isolated = ak.all((dr_gl > 0.3), axis=-1)
if ak.count(gjets[isolated], axis=None) > 0:
# TODO: convert only relevant fields!
gjet1 = ak.to_pandas(gjets[isolated]).loc[pd.IndexSlice[:, 0],
["pt", "eta", "phi", "mass"]]
gjet2 = ak.to_pandas(gjets[isolated]).loc[pd.IndexSlice[:, 1],
["pt", "eta", "phi", "mass"]]
gjet1.index = gjet1.index.droplevel("subentry")
gjet2.index = gjet2.index.droplevel("subentry")
gjsum = p4_sum(gjet1, gjet2)
gjmass = gjsum.mass
return gjmass
def test_jet_correction_uncertainty_sources():
from coffea.jetmet_tools import JetCorrectionUncertainty
counts, test_eta, test_pt = dummy_jagged_eta_pt()
test_pt_jag = ak.unflatten(test_pt, counts)
test_eta_jag = ak.unflatten(test_eta, counts)
junc_names = []
levels = []
for name in dir(evaluator):
if 'Summer16_23Sep2016V3_MC_UncertaintySources_AK4PFPuppi' in name:
junc_names.append(name)
levels.append(name.split('_')[-1])
#test for underscore in dataera
if 'Fall17_17Nov2017_V6_MC_UncertaintySources_AK4PFchs_AbsoluteFlavMap' in name:
junc_names.append(name)
levels.append(name.split('_')[-1])
junc = JetCorrectionUncertainty(
**{name: evaluator[name]
for name in junc_names})
print(junc)
juncs = junc.getUncertainty(JetEta=test_eta, JetPt=test_pt)
juncs_jag = list(
junc.getUncertainty(JetEta=test_eta_jag, JetPt=test_pt_jag))
for i, (level, corrs) in enumerate(juncs):
assert (level in levels)
assert (corrs.shape[0] == test_eta.shape[0])
tic = time.time()
assert (ak.all(corrs == ak.flatten(juncs_jag[i][1])))
toc = time.time()
def fill_gen_jets(df, output):
gjets = df.GenJet
gleptons = df.GenPart[(abs(df.GenPart.pdgId) == 13)
| (abs(df.GenPart.pdgId) == 11)
| (abs(df.GenPart.pdgId) == 15)]
gl_pair = ak.cartesian({
"jet": gjets,
"lepton": gleptons
},
axis=1,
nested=True)
_, _, dr_gl = delta_r(
gl_pair["jet"].eta,
gl_pair["lepton"].eta,
gl_pair["jet"].phi,
gl_pair["lepton"].phi,
)
isolated = ak.all((dr_gl > 0.3), axis=-1)
gjet1 = ak.to_pandas(gjets[isolated]).loc[pd.IndexSlice[:, 0],
["pt", "eta", "phi", "mass"]]
gjet2 = ak.to_pandas(gjets[isolated]).loc[pd.IndexSlice[:, 1],
["pt", "eta", "phi", "mass"]]
gjet1.index = gjet1.index.droplevel("subentry")
gjet2.index = gjet2.index.droplevel("subentry")
gjsum = p4_sum(gjet1, gjet2)
for var in ["pt", "eta", "phi", "mass"]:
output[f"gjet1_{var}"] = gjet1[var]
output[f"gjet2_{var}"] = gjet2[var]
output[f"gjj_{var}"] = gjsum[var]
output["gjj_dEta"], output["gjj_dPhi"], output["gjj_dR"] = delta_r(
output.gjet1_eta, output.gjet2_eta, output.gjet1_phi, output.gjet2_phi)
return output
def test_parsing(header_version: int, events_per_chunk: int) -> None:
here = Path(__file__).parent
input_filename = here / "parsing" / f"final_state_hadrons_header_v{header_version}.dat"
for i, arrays in enumerate(
parse_ascii.read(filename=input_filename,
events_per_chunk=events_per_chunk,
parser="pandas")):
# Get the reference array
# Create the reference arrays by checking out the parser v1 (e477e0277fa560f9aba82310c02da8177e61c9e4), setting
# the chunk size in skim_ascii, and then calling:
# $ python jetscape_analysis/analysis/reader/skim_ascii.py -i tests/parsing/final_state_hadrons_header_v1.dat -o tests/parsing/events_per_chunk_50/parser_v1_header_v1/test.parquet
# NOTE: The final state hadron files won't exist when you check out that branch, so
# it's best to copy them for your existing branch.
reference_arrays = ak.from_parquet(
Path(
f"{here}/parsing/events_per_chunk_{events_per_chunk}/parser_v1_header_v1/test_{i:02}.parquet"
))
# There are more fields in v2 than in the reference arrays (v1), so only take those
# that are present in reference for comparison.
# NOTE: We have to compare the fields one-by-one because the shapes of the fields
# are different, and apparently don't broadcast nicely with `__eq__`
for field in ak.fields(reference_arrays):
new_field = _rename_columns.get(field, field)
assert ak.all(reference_arrays[field] == arrays[new_field])
# Check for cross section if header v2
if header_version == 2:
assert "cross_section" in ak.fields(arrays)
assert "cross_section_error" in ak.fields(arrays)
def test_spherical_three_vector():
a = ak.zip(
{
"rho": [[1.0, 2.0], [], [3.0], [4.0]],
"theta": [[1.2, 0.7], [], [1.8], [1.9]],
"phi": [[0.3, 0.4], [], [0.5], [0.6]],
},
with_name="SphericalThreeVector",
highlevel=False,
)
a = ak.Array(a, behavior=vector.behavior)
assert ak.all(abs((-a).x + a.x) < ATOL)
assert ak.all(abs((-a).y + a.y) < ATOL)
assert ak.all(abs((-a).z + a.z) < ATOL)
assert record_arrays_equal(a * (-1), -a)
def test_jet_correction_regrouped_uncertainty_sources():
from coffea.jetmet_tools import JetCorrectionUncertainty
counts, test_eta, test_pt = dummy_jagged_eta_pt()
test_pt_jag = ak.unflatten(test_pt, counts)
test_eta_jag = ak.unflatten(test_eta, counts)
junc_names = []
levels = []
for name in dir(evaluator):
if 'Regrouped_Fall17_17Nov2017_V32_MC_UncertaintySources_AK4PFchs' in name:
junc_names.append(name)
if len(name.split('_')) == 9:
levels.append("_".join(name.split('_')[-2:]))
else:
levels.append(name.split('_')[-1])
junc = JetCorrectionUncertainty(
**{name: evaluator[name]
for name in junc_names})
print(junc)
juncs_jag = list(
junc.getUncertainty(JetEta=test_eta_jag, JetPt=test_pt_jag))
for i, tpl in enumerate(
list(junc.getUncertainty(JetEta=test_eta, JetPt=test_pt))):
assert (tpl[0] in levels)
assert (tpl[1].shape[0] == test_eta.shape[0])
assert (ak.all(tpl[1] == ak.flatten(juncs_jag[i][1])))
def test_jet_resolution_sf():
from coffea.jetmet_tools import JetResolutionScaleFactor
counts, test_eta, test_pt = dummy_jagged_eta_pt()
test_pt_jag = ak.unflatten(test_pt, counts)
test_eta_jag = ak.unflatten(test_eta, counts)
jersf_names = ["Spring16_25nsV10_MC_SF_AK4PFPuppi"]
resosf = JetResolutionScaleFactor(
**{name: evaluator[name]
for name in jersf_names})
print(resosf)
# 0-jet compatibility
assert resosf.getScaleFactor(JetEta=test_eta[:0]).shape == (0, 3)
resosfs = resosf.getScaleFactor(JetEta=test_eta)
resosfs_jag = resosf.getScaleFactor(JetEta=test_eta_jag)
assert ak.all(resosfs == ak.flatten(resosfs_jag))
test_pt_jag = test_pt_jag[0:3]
test_eta_jag = test_eta_jag[0:3]
counts = counts[0:3]
print("Raw jet values:")
print("pT:", test_pt_jag)
print("eta:", test_eta_jag, "\n")
resosfs_jag_ref = ak.unflatten(
np.array([
[1.857, 1.928, 1.786],
[1.084, 1.095, 1.073],
[1.364, 1.403, 1.325],
[1.177, 1.218, 1.136],
[1.138, 1.151, 1.125],
[1.364, 1.403, 1.325],
[1.177, 1.218, 1.136],
[1.082, 1.117, 1.047],
]),
counts,
)
resosfs_jag = resosf.getScaleFactor(JetEta=test_eta_jag)
print("Reference Resolution SF (jagged):", resosfs_jag_ref)
print("Resolution SF (jagged):", resosfs_jag)
assert ak.all(
np.abs(ak.flatten(resosfs_jag_ref) - ak.flatten(resosfs_jag)) < 1e-6)
def test_jet_resolution_sf_2d():
from coffea.jetmet_tools import JetResolutionScaleFactor
counts, test_eta, test_pt = dummy_jagged_eta_pt()
test_pt_jag = ak.unflatten(test_pt, counts)
test_eta_jag = ak.unflatten(test_eta, counts)
resosf = JetResolutionScaleFactor(
**{name: evaluator[name]
for name in ["Autumn18_V7_MC_SF_AK4PFchs"]})
print(resosf)
# 0-jet compatibility
assert resosf.getScaleFactor(JetPt=test_pt[:0],
JetEta=test_eta[:0]).shape == (0, 3)
resosfs = resosf.getScaleFactor(JetPt=test_pt, JetEta=test_eta)
resosfs_jag = resosf.getScaleFactor(JetPt=test_pt_jag, JetEta=test_eta_jag)
assert ak.all(resosfs == ak.flatten(resosfs_jag))
test_pt_jag = test_pt_jag[0:3]
test_eta_jag = test_eta_jag[0:3]
counts = counts[0:3]
print("Raw jet values:")
print("pT:", test_pt_jag)
print("eta:", test_eta_jag, "\n")
resosfs_jag_ref = ak.unflatten(
np.array([
[1.11904, 1.31904, 1.0],
[1.1432, 1.2093, 1.0771],
[1.16633, 1.36633, 1.0],
[1.17642, 1.37642, 1.0],
[1.1808, 1.1977, 1.1640],
[1.15965, 1.35965, 1.0],
[1.17661, 1.37661, 1.0],
[1.1175, 1.1571, 1.0778],
]),
counts,
)
resosfs_jag = resosf.getScaleFactor(JetPt=test_pt_jag, JetEta=test_eta_jag)
print("Reference Resolution SF (jagged):", resosfs_jag_ref)
print("Resolution SF (jagged):", resosfs_jag)
assert ak.all(
np.abs(ak.flatten(resosfs_jag_ref) - ak.flatten(resosfs_jag)) < 1e-6)
def allclose(
self,
other: VectorProtocol,
rtol: ScalarCollection = 1e-05,
atol: ScalarCollection = 1e-08,
equal_nan: BoolCollection = False,
) -> BoolCollection:
return ak.all(
self.isclose(other, rtol=rtol, atol=atol, equal_nan=equal_nan))
def test_nested_collection(collection, subcollection, arr_type, element, events):
assert ak.type(events[collection][subcollection])
assert ak.type(events[collection][subcollection + "Counts"])
assert (
ak.type(events[collection][subcollection])
.type.type.type.__str__()
.startswith(arr_type)
)
if element is None:
assert ak.all(
events[collection][subcollection + "Counts"]
== ak.count(events[collection][subcollection], axis=-1)
)
else:
assert ak.all(
events[collection][subcollection + "Counts"]
== ak.count(events[collection][subcollection][element], axis=-1)
)
def test_reducers():
# axis=None reducers are implemented in NumPy.
assert ak.sum(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]])) == 6 + 6j
assert ak.prod(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]])) == -12 + 12j
# axis != None reducers are implemented in libawkward; this should be ReducerSum.
assert ak.sum(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]]),
axis=1).tolist() == [
3 + 3j,
0 + 0j,
3 + 3j,
]
# And this is in ReducerProd.
assert ak.prod(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]]),
axis=1).tolist() == [
0 + 4j,
1 + 0j,
3 + 3j,
]
# ReducerCount, ReducerCountNonzero, ReducerAny, and ReducerAll work.
assert ak.count(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]]),
axis=1).tolist() == [2, 0, 1]
assert ak.count_nonzero(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]]),
axis=1).tolist() == [2, 0, 1]
assert ak.any(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]]),
axis=1).tolist() == [
True,
False,
True,
]
assert ak.all(ak.from_iter([[1 + 1j, 2 + 2j], [], [3 + 3j]]),
axis=1).tolist() == [
True,
True,
True,
]
assert ak.any(ak.from_iter([[1 + 1j, 2 + 2j, 0 + 0j], [], [3 + 3j]]),
axis=1).tolist() == [True, False, True]
assert ak.all(ak.from_iter([[1 + 1j, 2 + 2j, 0 + 0j], [], [3 + 3j]]),
axis=1).tolist() == [False, True, True]
def allclose(
self,
other: VectorProtocol,
rtol: ScalarCollection = 1e-05,
atol: ScalarCollection = 1e-08,
equal_nan: BoolCollection = False,
) -> BoolCollection:
"""
Like ``np.ndarray.allclose``, but for VectorArray4D.
"""
return ak.all(
self.isclose(other, rtol=rtol, atol=atol, equal_nan=equal_nan))
def test_pt_eta_phi_e_lorentz_vector():
a = ak.zip(
{
"pt": [[1, 2], [], [3], [4]],
"eta": [[1.2, 1.4], [], [1.6], [3.4]],
"phi": [[0.3, 0.4], [], [0.5], [0.6]],
"energy": [[50, 51], [], [52], [60]],
},
with_name="PtEtaPhiELorentzVector",
highlevel=False,
)
a = ak.Array(a, behavior=vector.behavior)
assert ak.all((a * (-2)).pt == ak.Array([[2, 4], [], [6], [8]]))
assert ak.all(
(a * (-2)).theta - ak.Array([[2.556488570968, 2.65804615357], [],
[2.74315571762], [3.07487087733]]) < ATOL)
assert ak.all(
(a * (-2)).phi - ak.Array([[-2.8415926535, -2.7415926535], [],
[-2.6415926535], [-2.5415926535]]) < ATOL)
assert record_arrays_equal(
a / 2,
ak.zip({
"pt": [[0.5, 1], [], [1.5], [2]],
"eta": [[1.2, 1.4], [], [1.6], [3.4]],
"phi": [[0.3, 0.4], [], [0.5], [0.6]],
"energy": [[25, 25.5], [], [26], [30]],
}),
)
assert record_arrays_equal(a * (-1), -a)
boosted = a.boost(-a.boostvec)
assert ak.all(abs(boosted.x) < ATOL)
assert ak.all(abs(boosted.y) < ATOL)
assert ak.all(abs(boosted.z) < ATOL)
def test_factorized_jet_corrector():
from coffea.jetmet_tools import FactorizedJetCorrector
counts, test_eta, test_pt = dummy_jagged_eta_pt()
test_Rho = np.full_like(test_eta, 100.)
test_A = np.full_like(test_eta, 5.)
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})
print(corrector)
pt_copy = np.copy(test_pt)
corrs = corrector.getCorrection(JetEta=test_eta,
Rho=test_Rho,
JetPt=test_pt,
JetA=test_A)
assert ((np.abs(pt_copy - test_pt) < 1e-6).all())
test_pt_jag = ak.unflatten(test_pt, counts)
test_eta_jag = ak.unflatten(test_eta, counts)
test_Rho_jag = ak.unflatten(test_Rho, counts)
test_A_jag = ak.unflatten(test_A, counts)
corrs_jag = corrector.getCorrection(JetEta=test_eta_jag,
Rho=test_Rho_jag,
JetPt=test_pt_jag,
JetA=test_A_jag)
assert (ak.all(np.abs(pt_copy - ak.flatten(test_pt_jag)) < 1e-6))
assert (ak.all(np.abs(corrs - ak.flatten(corrs_jag)) < 1e-6))
def studyAcc(ifile, title, acceptances_dict, mass):
dq_events = getData(ifile, "Truth")
# choosing the nevts for denominator (because all files have 10000 evts)
den = len(dq_events)
# choosing 500 evts for denominator
den500 = False
if den500:
dq_events = dq_events[:500]
den = 500
dq_showers = dq_events["Showers"]
dq_electrons = dq_events["Electrons"]
dq_sh_z = ak.fill_none(ak.pad_none(dq_showers.sz, 2, axis=1), 0)
dq_sh_e = ak.fill_none(ak.pad_none(dq_showers.sedep_ecal, 2, axis=1), -1)
dq_e_ge = ak.fill_none(ak.pad_none(dq_electrons.ge, 2, axis=1), -1)
shower_zero_mask = ak.any(dq_sh_e == 0, axis=1)
shower_zero_index = np.where(shower_zero_mask)
eshower_mask = ak.all(dq_sh_e > 0.1 * dq_e_ge, axis=1)
eshower_index = np.where(eshower_mask)
masks = {'shower_zero': shower_zero_index, 'eshower': eshower_index}
dq_evts = dq_events[ak.all(dq_e_ge > -1, axis=1)]
# print(acceptances_dict)
# we can print the mass, the number of events with gen e-, the number of events in the sample, and the denominator
print(mass, len(dq_evts), len(dq_events), den)
for maskstr, m in masks.items():
if maskstr in acceptances_dict:
acceptances_dict[maskstr].append(len(m[0]) / den)
else:
acceptances_dict[maskstr] = [len(m[0]) / den]
#print(maskstr+':', len(m[0])/den)
return acceptances_dict
def get_filters(Flags, year, isMC, accumulator=None):
met_filters_to_use = met_filters[year].copy()
if 'BadPFMuonDzFilter' not in Flags.fields: met_filters_to_use.remove('BadPFMuonDzFilter')
if (year == '2017') or (year == '2018'):
if isMC: met_filters_to_use.remove('eeBadScFilter')
pass_filters = ak.all((Flags[i] for i in met_filters_to_use), axis=0)
if accumulator:
accumulator['cutflow']['pass filters'] += ak.sum(pass_filters)
return pass_filters, accumulator
else:
return pass_filters
def select_events(self, events, photons, metadata):
options = copy.deepcopy(self.selection_options)
for key, value in metadata.items(): # add sample-specific options to selection options
options[key] = value
# Diphoton preselection: NOTE we assume diphoton preselection is common to every analysis
diphoton_events = events # most of dipho preselection already applied, still need to enforce pt/mgg and mgg cuts
selected_photons = photon_selections.create_selected_photons(photons, self.branches, self.debug) # create record manually since it doesn't seem to work for selectedPhoton
diphoton_events, selected_photons = diphoton_selections.diphoton_preselection(diphoton_events, selected_photons, options, self.debug)
events_and_objects = {}
if "HHggTauTau_InclusivePresel" in self.selections:
if "genZStudy" in self.selections and not options["data"]:
gen_events = diphoton_events.GenPart
else:
gen_events = None
selected_events = analysis_selections.ggTauTau_inclusive_preselection(diphoton_events, selected_photons, diphoton_events.Electron, diphoton_events.Muon, diphoton_events.Tau, diphoton_events.Jet, diphoton_events.dR_tautauSVFitLoose, gen_events, diphoton_events.Category_pairsLoose,options, self.debug)
elif self.selections == "ttH_LeptonicPresel":
selected_events = analysis_selections.tth_leptonic_preselection(diphoton_events, selected_photons, diphoton_events.Electron, diphoton_events.Muon, diphoton_events.Jet, options, self.debug)
elif self.selections == "ttH_HadronicPresel":
selected_events = analysis_selections.tth_hadronic_preselection(diphoton_events, selected_photons, diphoton_events.Electron, diphoton_events.Muon, diphoton_events.Jet, options, self.debug)
elif self.selections== "ttH_InclusivePresel":
selected_events = analysis_selections.tth_inclusive_preselection(diphoton_events, selected_photons, diphoton_events.Electron, diphoton_events.Muon, diphoton_events.Jet, options, self.debug)
elif self.selections == "HHggbb_Presel":
options["boosted"] = False
selected_events = analysis_selections.ggbb_preselection(diphoton_events, selected_photons, diphoton_events.Electron, diphoton_events.Muon, diphoton_events.Jet, diphoton_events.FatJet, options, self.debug)
elif "HHggbb_boosted_Presel" in self.selections:
options["boosted"] = True
options["signal"] = True if awkward.all(events["process_id"] < 0) else False # HH_ggbb pID = -1 and VBF_HHggbb pID = -2
if "GenInfo" in self.selections and not options["data"]:
gen_events = diphoton_events.GenPart
else:
gen_events = None
selected_events = analysis_selections.ggbb_preselection(diphoton_events, selected_photons, diphoton_events.Electron, diphoton_events.Muon, diphoton_events.Jet, diphoton_events.FatJet, gen_events, options, self.debug)
else:
print("[LoopHelper] Selection: %s is not currently implemented, please check." % self.selections)
return
return selected_events
