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_resolution(): from coffea.jetmet_tools import JetResolution counts, test_eta, test_pt = dummy_jagged_eta_pt() test_Rho = np.full_like(test_eta, 100.) 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 = reso.getResolution(JetEta=test_eta_jag, Rho=test_Rho_jag, JetPt=test_pt_jag)
def test_corrected_jets_factory(): import os from coffea.jetmet_tools import CorrectedJetsFactory, CorrectedMETFactory, JECStack events = None from coffea.nanoevents import NanoEventsFactory factory = NanoEventsFactory.from_root( os.path.abspath("tests/samples/nano_dy.root")) events = factory.events() jec_stack_names = [ "Summer16_23Sep2016V3_MC_L1FastJet_AK4PFPuppi", "Summer16_23Sep2016V3_MC_L2Relative_AK4PFPuppi", "Summer16_23Sep2016V3_MC_L2L3Residual_AK4PFPuppi", "Summer16_23Sep2016V3_MC_L3Absolute_AK4PFPuppi", "Spring16_25nsV10_MC_PtResolution_AK4PFPuppi", "Spring16_25nsV10_MC_SF_AK4PFPuppi", ] for key in evaluator.keys(): if "Summer16_23Sep2016V3_MC_UncertaintySources_AK4PFPuppi" in key: jec_stack_names.append(key) jec_inputs = {name: evaluator[name] for name in jec_stack_names} jec_stack = JECStack(jec_inputs) name_map = jec_stack.blank_name_map name_map["JetPt"] = "pt" name_map["JetMass"] = "mass" name_map["JetEta"] = "eta" name_map["JetA"] = "area" jets = events.Jet jets["pt_raw"] = (1 - jets["rawFactor"]) * jets["pt"] jets["mass_raw"] = (1 - jets["rawFactor"]) * jets["mass"] jets["pt_gen"] = ak.values_astype(ak.fill_none(jets.matched_gen.pt, 0), np.float32) jets["rho"] = ak.broadcast_arrays(events.fixedGridRhoFastjetAll, jets.pt)[0] name_map["ptGenJet"] = "pt_gen" name_map["ptRaw"] = "pt_raw" name_map["massRaw"] = "mass_raw" name_map["Rho"] = "rho" jec_cache = cachetools.Cache(np.inf) print(name_map) tic = time.time() jet_factory = CorrectedJetsFactory(name_map, jec_stack) toc = time.time() print("setup corrected jets time =", toc - tic) tic = time.time() prof = pyinstrument.Profiler() prof.start() corrected_jets = jet_factory.build(jets, lazy_cache=jec_cache) prof.stop() toc = time.time() print("corrected_jets build time =", toc - tic) print(prof.output_text(unicode=True, color=True, show_all=True)) tic = time.time() print("Generated jet pt:", corrected_jets.pt_gen) print("Original jet pt:", corrected_jets.pt_orig) print("Raw jet pt:", jets.pt_raw) print("Corrected jet pt:", corrected_jets.pt) print("Original jet mass:", corrected_jets.mass_orig) print("Raw jet mass:", jets["mass_raw"]) print("Corrected jet mass:", corrected_jets.mass) print("jet eta:", jets.eta) for unc in jet_factory.uncertainties(): print(unc) print(corrected_jets[unc].up.pt) print(corrected_jets[unc].down.pt) toc = time.time() print("build all jet variations =", toc - tic) # Test that the corrections were applied correctly from coffea.jetmet_tools import ( FactorizedJetCorrector, JetResolution, JetResolutionScaleFactor, ) scalar_form = ak.without_parameters(jets["pt_raw"]).layout.form corrector = FactorizedJetCorrector( **{name: evaluator[name] for name in jec_stack_names[0:4]}) corrs = corrector.getCorrection(JetEta=jets["eta"], Rho=jets["rho"], JetPt=jets["pt_raw"], JetA=jets["area"]) reso = JetResolution( **{name: evaluator[name] for name in jec_stack_names[4:5]}) jets["jet_energy_resolution"] = reso.getResolution( JetEta=jets["eta"], Rho=jets["rho"], JetPt=jets["pt_raw"], form=scalar_form, lazy_cache=jec_cache, ) resosf = JetResolutionScaleFactor( **{name: evaluator[name] for name in jec_stack_names[5:6]}) jets["jet_energy_resolution_scale_factor"] = resosf.getScaleFactor( JetEta=jets["eta"], lazy_cache=jec_cache) # Filter out the non-deterministic (no gen pt) jets def smear_factor(jetPt, pt_gen, jersf): return (ak.full_like(jetPt, 1.0) + (jersf[:, 0] - ak.full_like(jetPt, 1.0)) * (jetPt - pt_gen) / jetPt) test_gen_pt = ak.concatenate( [corrected_jets.pt_gen[0, :-2], corrected_jets.pt_gen[-1, :-1]]) test_raw_pt = ak.concatenate([jets.pt_raw[0, :-2], jets.pt_raw[-1, :-1]]) test_pt = ak.concatenate( [corrected_jets.pt[0, :-2], corrected_jets.pt[-1, :-1]]) test_eta = ak.concatenate([jets.eta[0, :-2], jets.eta[-1, :-1]]) test_jer = ak.concatenate([ jets.jet_energy_resolution[0, :-2], jets.jet_energy_resolution[-1, :-1] ]) test_jer_sf = ak.concatenate([ jets.jet_energy_resolution_scale_factor[0, :-2], jets.jet_energy_resolution_scale_factor[-1, :-1], ]) test_jec = ak.concatenate([corrs[0, :-2], corrs[-1, :-1]]) test_corrected_pt = ak.concatenate( [corrected_jets.pt[0, :-2], corrected_jets.pt[-1, :-1]]) test_corr_pt = test_raw_pt * test_jec test_pt_smear_corr = test_corr_pt * smear_factor(test_corr_pt, test_gen_pt, test_jer_sf) # Print the results of the "by-hand" calculations and confirm that the values match the expected values print("\nConfirm the CorrectedJetsFactory values:") print("Jet pt (gen)", test_gen_pt.tolist()) print("Jet pt (raw)", test_raw_pt.tolist()) print("Jet pt (nano):", test_pt.tolist()) print("Jet eta:", test_eta.tolist()) print("Jet energy resolution:", test_jer.tolist()) print("Jet energy resolution sf:", test_jer_sf.tolist()) print("Jet energy correction:", test_jec.tolist()) print("Corrected jet pt (ref)", test_corr_pt.tolist()) print("Corrected & smeared jet pt (ref):", test_pt_smear_corr.tolist()) print("Corrected & smeared jet pt:", test_corrected_pt.tolist(), "\n") assert ak.all(np.abs(test_pt_smear_corr - test_corrected_pt) < 1e-6) name_map["METpt"] = "pt" name_map["METphi"] = "phi" name_map["JetPhi"] = "phi" name_map["UnClusteredEnergyDeltaX"] = "MetUnclustEnUpDeltaX" name_map["UnClusteredEnergyDeltaY"] = "MetUnclustEnUpDeltaY" tic = time.time() met_factory = CorrectedMETFactory(name_map) toc = time.time() print("setup corrected MET time =", toc - tic) met = events.MET tic = time.time() # prof = pyinstrument.Profiler() # prof.start() corrected_met = met_factory.build(met, corrected_jets, lazy_cache=jec_cache) # prof.stop() toc = time.time() # print(prof.output_text(unicode=True, color=True, show_all=True)) print("corrected_met build time =", toc - tic) tic = time.time() print(corrected_met.pt_orig) print(corrected_met.pt) prof = pyinstrument.Profiler() prof.start() for unc in jet_factory.uncertainties() + met_factory.uncertainties(): print(unc) print(corrected_met[unc].up.pt) print(corrected_met[unc].down.pt) prof.stop() toc = time.time() print("build all met variations =", toc - tic) print(prof.output_text(unicode=True, color=True, show_all=True))