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