def test(): original = ak.Array([[1, 2, 3, 4], [], [5, 6, 7], [8, 9]]) assert ak.unflatten(original, [2, 2, 1, 2, 1, 1], axis=1).tolist() == [ [[1, 2], [3, 4]], [], [[5], [6, 7]], [[8], [9]], ] assert ak.unflatten(original, [1, 3, 1, 2, 1, 1], axis=1).tolist() == [ [[1], [2, 3, 4]], [], [[5], [6, 7]], [[8], [9]], ] with pytest.raises(ValueError): ak.unflatten(original, [2, 1, 2, 2, 1, 1], axis=1) assert ak.unflatten(original, [2, 0, 2, 1, 2, 1, 1], axis=1).tolist() == [ [[1, 2], [], [3, 4]], [], [[5], [6, 7]], [[8], [9]], ]
def test(): array = ak.Array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]) assert ak.unflatten(array, 5).tolist() == [[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]] assert ak.unflatten(array, [3, 0, 2, 1, 4]).tolist() == [ [0, 1, 2], [], [3, 4], [5], [6, 7, 8, 9], ] assert ak.unflatten(array, [3, None, 2, 1, 4]).tolist() == [ [0, 1, 2], None, [3, 4], [5], [6, 7, 8, 9], ] original = ak.Array([[0, 1, 2], [], [3, 4], [5], [6, 7, 8, 9]]) counts = ak.num(original) array = ak.flatten(original) assert counts.tolist() == [3, 0, 2, 1, 4] assert array.tolist() == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] assert ak.unflatten(array, counts).tolist() == [ [0, 1, 2], [], [3, 4], [5], [6, 7, 8, 9], ]
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 test_groupby(): array = ak.Array( [ {"x": 1, "y": 1.1}, {"x": 2, "y": 2.2}, {"x": 1, "y": 1.1}, {"x": 3, "y": 3.3}, {"x": 1, "y": 1.1}, {"x": 2, "y": 2.2}, ] ) sorted = array[ak.argsort(array.x)] assert sorted.x.tolist() == [1, 1, 1, 2, 2, 3] assert ak.run_lengths(sorted.x).tolist() == [3, 2, 1] assert ak.unflatten(sorted, ak.run_lengths(sorted.x)).tolist() == [ [{"x": 1, "y": 1.1}, {"x": 1, "y": 1.1}, {"x": 1, "y": 1.1}], [{"x": 2, "y": 2.2}, {"x": 2, "y": 2.2}], [{"x": 3, "y": 3.3}], ] array = ak.Array( [ [{"x": 1, "y": 1.1}, {"x": 2, "y": 2.2}, {"x": 1, "y": 1.1}], [{"x": 3, "y": 3.3}, {"x": 1, "y": 1.1}, {"x": 2, "y": 2.2}], ] ) sorted = array[ak.argsort(array.x)] assert sorted.x.tolist() == [[1, 1, 2], [1, 2, 3]] assert ak.run_lengths(sorted.x).tolist() == [[2, 1], [1, 1, 1]] counts = ak.flatten(ak.run_lengths(sorted.x), axis=None) assert ak.unflatten(sorted, counts, axis=-1).tolist() == [ [[{"x": 1, "y": 1.1}, {"x": 1, "y": 1.1}], [{"x": 2, "y": 2.2}]], [[{"x": 1, "y": 1.1}], [{"x": 2, "y": 2.2}], [{"x": 3, "y": 3.3}]], ]
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(): content = ak.repartition(range(10), 3) assert ak.unflatten(content, [3, 0, 2, 1, 3, 1]).tolist() == [ [0, 1, 2], [], [3, 4], [5], [6, 7, 8], [9], ] original = ak.repartition(ak.Array([[1, 2, 3, 4], [5, 6, 7], [8, 9]]), 2) assert ak.unflatten(original, [2, 2, 1, 2, 1, 1], axis=1).tolist() == [ [[1, 2], [3, 4]], [[5], [6, 7]], [[8], [9]], ] original = ak.repartition(ak.Array([[1, 2, 3, 4], [], [5, 6, 7], [8, 9]]), 2) assert ak.unflatten(original, [2, 2, 1, 2, 1, 1], axis=1).tolist() == [ [[1, 2], [3, 4]], [], [[5], [6, 7]], [[8], [9]], ]
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 test(): a = ak.Array([[1, 2, 3, 4], [5, 6, 7, 8]]) assert ak.unflatten(a, [2, 2, 2, 2], axis=1).tolist() == [ [[1, 2], [3, 4]], [[5, 6], [7, 8]], ] assert ak.unflatten(a, 2, axis=1).tolist() == [[[1, 2], [3, 4]], [[5, 6], [7, 8]]]
def test(): array = ak.Array([[1, 2, 3]]) assert ak.unflatten(array, [2, 1], axis=1).tolist() == [[[1, 2], [3]]] assert ak.unflatten(array, [0, 2, 1], axis=1).tolist() == [[[], [1, 2], [3]]] assert ak.unflatten(array, [0, 0, 2, 1], axis=1).tolist() == [[[], [], [1, 2], [3]]]
def apply_roccor(df, rochester, is_mc): if is_mc: hasgen = ~np.isnan(ak.fill_none(df.Muon.matched_gen.pt, np.nan)) mc_rand = np.random.rand(*ak.to_numpy(ak.flatten(df.Muon.pt)).shape) mc_rand = ak.unflatten(mc_rand, ak.num(df.Muon.pt, axis=1)) corrections = np.array(ak.flatten(ak.ones_like(df.Muon.pt))) errors = np.array(ak.flatten(ak.ones_like(df.Muon.pt))) mc_kspread = rochester.kSpreadMC( df.Muon.charge[hasgen], df.Muon.pt[hasgen], df.Muon.eta[hasgen], df.Muon.phi[hasgen], df.Muon.matched_gen.pt[hasgen], ) mc_ksmear = rochester.kSmearMC( df.Muon.charge[~hasgen], df.Muon.pt[~hasgen], df.Muon.eta[~hasgen], df.Muon.phi[~hasgen], df.Muon.nTrackerLayers[~hasgen], mc_rand[~hasgen], ) errspread = rochester.kSpreadMCerror( df.Muon.charge[hasgen], df.Muon.pt[hasgen], df.Muon.eta[hasgen], df.Muon.phi[hasgen], df.Muon.matched_gen.pt[hasgen], ) errsmear = rochester.kSmearMCerror( df.Muon.charge[~hasgen], df.Muon.pt[~hasgen], df.Muon.eta[~hasgen], df.Muon.phi[~hasgen], df.Muon.nTrackerLayers[~hasgen], mc_rand[~hasgen], ) hasgen_flat = np.array(ak.flatten(hasgen)) corrections[hasgen_flat] = np.array(ak.flatten(mc_kspread)) corrections[~hasgen_flat] = np.array(ak.flatten(mc_ksmear)) errors[hasgen_flat] = np.array(ak.flatten(errspread)) errors[~hasgen_flat] = np.array(ak.flatten(errsmear)) corrections = ak.unflatten(corrections, ak.num(df.Muon.pt, axis=1)) errors = ak.unflatten(errors, ak.num(df.Muon.pt, axis=1)) else: corrections = rochester.kScaleDT(df.Muon.charge, df.Muon.pt, df.Muon.eta, df.Muon.phi) errors = rochester.kScaleDTerror(df.Muon.charge, df.Muon.pt, df.Muon.eta, df.Muon.phi) df["Muon", "pt_roch"] = df.Muon.pt * corrections df["Muon", "pt_roch_up"] = df.Muon.pt_roch + df.Muon.pt * errors df["Muon", "pt_roch_down"] = df.Muon.pt_roch - df.Muon.pt * errors
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_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_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"]}) resosfs = resosf.getScaleFactor(JetPt=test_pt, JetEta=test_eta) resosfs_jag = resosf.getScaleFactor(JetPt=test_pt_jag, JetEta=test_eta_jag)
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 _kExtra(self, kpt, eta, nl, u, s=0, m=0): # if it is a jagged array, save the offsets then flatten everything # needed for the ternary conditions later abseta = abs(eta) kData = self._kRes[s][m][1](abseta) # type 1 is data kMC = self._kRes[s][m][0](abseta) # type 0 is MC mask = kData > kMC x = awkward.zeros_like(kpt) sigma = self._sigma(kpt, eta, nl, s, m) # Rochester cbA = beta, cbN = m, as well as cbM (always 0?) = loc and cbS = scale to transform y = (x-loc)/scale in the pdf method cbA = self._cbA[s][m](abseta, nl) cbN = self._cbN[s][m](abseta, nl) cbS = self._cbS[s][m](abseta, nl) counts = awkward.num(u) u_flat = awkward.flatten(u) loc = awkward.zeros_like(u_flat) cbA_flat = awkward.flatten(cbA) cbN_flat = awkward.flatten(cbN) cbS_flat = awkward.flatten(cbS) invcdf = awkward.unflatten( doublecrystalball.ppf(u_flat, cbA_flat, cbA_flat, cbN_flat, cbN_flat, loc, cbS_flat), counts, ) x = awkward.where( mask, (numpy.sqrt(kData * kData - kMC * kMC) * sigma * invcdf), x, ) result = awkward.where(x > -1, 1.0 / (1.0 + x), awkward.ones_like(kpt)) if isinstance(kpt, numpy.ndarray): result = numpy.array(result) return result
def load(buffer, entry_start: int = None, entry_stop: int = None) -> ak.Array: """Load MIDAS data from a buffer and return the ragged array. :param buffer: object supporting buffer protocol :param entry_start: the first entry to include. If None, start at 0. :param entry_stop: the first entry to exclude. If None, stop at the end of the file. """ if entry_start is None: entry_start = 0 if entry_stop is None: entry_stop = np.iinfo(np.int32).max if entry_start < 0: raise ValueError("Start entry must be positive") if entry_stop < 0: raise ValueError("Stop entry must be positive") if entry_stop < entry_start: raise ValueError("Stop entry must not be less than the start entry") # Create a view over the memory map data = np.frombuffer(buffer, dtype="<B") # Run the VM vm = ForthMachine32(MIDAS_PROGRAM) vm.run({ "data": data, "range": np.array([entry_start, entry_stop], dtype=np.int32) }) # Create the result outputs = vm.outputs.copy() num = outputs.pop("num") return ak.unflatten(ak.zip(outputs), num)
def constituents(self, min_pt): outputs_to_inputs = self.constituent_index(min_pt) shape = ak.num(outputs_to_inputs) total = np.sum(shape) duplicate = ak.unflatten(np.zeros(total, np.int64), shape) prepared = self.data[:, np.newaxis][duplicate] return prepared[outputs_to_inputs]
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 test_evaluator_exceptions(): extractor = lookup_tools.extractor() extractor.add_weight_sets([ "testSF2d scalefactors_Tight_Electron tests/samples/testSF2d.histo.root" ]) counts, test_eta, test_pt = dummy_jagged_eta_pt() test_eta_jagged = ak.unflatten(test_eta, counts) test_pt_jagged = ak.unflatten(test_pt, counts) extractor.finalize() evaluator = extractor.make_evaluator() try: test_out = evaluator["testSF2d"](test_pt_jagged, test_eta) except Exception as e: assert (isinstance(e, TypeError))
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_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) resosfs = resosf.getScaleFactor(JetEta=test_eta) resosfs_jag = resosf.getScaleFactor(JetEta=test_eta_jag) assert (ak.all(resosfs == ak.flatten(resosfs_jag)))
def test(): layout = ak.layout.IndexedArray64( ak.layout.Index64(np.array([3, 1, 0, 2])), ak.layout.ListOffsetArray64( ak.layout.Index64(np.array([0, 3, 6, 9, 12])), ak.layout.NumpyArray(np.array([0, 0, 0, 1, 1, 1, 2, 2, 3, 3, 3, 3])), ), ) assert ak.unflatten( layout, ak.flatten(ak.run_lengths(layout)), axis=1 ).tolist() == [[[3, 3, 3]], [[1, 1, 1]], [[0, 0, 0]], [[2, 2], [3]]]
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 counts2nestedindex(stack): """Turn jagged local counts into doubly-kagged global index into a target Signature: local_counts,target_offsets,!counts2nestedindex Outputs a jagged array with same axis-0 shape as counts axis-1 """ target_offsets = stack.pop() local_counts = stack.pop() out = awkward.unflatten( numpy.arange(target_offsets[-1], dtype=numpy.int64), awkward.flatten(local_counts), ) stack.append(out)
def runNN(model, varsIn, varSet, normMean, normStd): dataset = RootDataset(varsIn=varsIn, varSet=varSet, normMean=normMean, normStd=normStd) nnOutput = getNNOutput(dataset, model) fjets = varsIn["fjets"] counts = ak.num(fjets.pt) svjJetsAK8 = ak.unflatten(nnOutput, counts) wpt = 0.5 darksvjJetsAK8 = fjets[svjJetsAK8 >= wpt] varsIn['nsvjJetsAK8'] = [ak.num(darksvjJetsAK8), "evtw"] varsIn['nnOutput'] = [svjJetsAK8, "fjw"]
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])))
def constituents(self, min_pt): np_results = self._results.to_numpy_with_constituents(min_pt) off = np.insert(np_results[-1], 0, 0) out = ak.Array( ak.layout.ListOffsetArray64( ak.layout.Index64(np_results[0]), ak.layout.NumpyArray(np_results[1]) ) ) outputs_to_inputs = ak.Array( ak.layout.ListOffsetArray64(ak.layout.Index64(off), out.layout) ) shape = ak.num(outputs_to_inputs) total = np.sum(shape) duplicate = ak.unflatten(np.zeros(total, np.int64), shape) prepared = self.data[:, np.newaxis][duplicate] return prepared[outputs_to_inputs][0]
def get_scale_factor(self, jets, passing_cut): '''Starting from a jet collection and a string pointing to the flag defining if the jet is b-tagged or not computes the per-jet weight to be used. Supports only a single WP for the moment''' # First of all flatten everything to make it easier to handle pt = ak.to_numpy(ak.flatten(jets.pt)) eta = ak.to_numpy(ak.flatten(jets.eta)) flav = ak.to_numpy(ak.flatten(jets.hadronFlavour)) pass_wp = ak.to_numpy(ak.flatten(jets[passing_cut])) # Get the MC efficiency eff = self.efficiency_(pt, eta, flav) # for each systematic/central value compute the proper SF # cache the SF values as sometimes they are repeated, there # might also be systematic combinations that are never accessed # but pruning them at the beginning can be hard # use schema to define combinations, lcb is a tuple with the sf keys # for light, charm, bottom for each systematic flavour_sf_cache = {} scale_factors = {} # our final product for key, lcb in self.schema_.items(): # populate cache if needed for i in range(3): flavour_sf_cache[lcb[i]] = flavour_sf_cache.get( # for some reason there is an additional dimension, pass_wp has no effect lcb[i], self.sf_[lcb[i]](eta, pt, pass_wp)) scale_factors[key] = eff * self.match_flav_( flavour_sf_cache[lcb[0]], flavour_sf_cache[lcb[1]], flavour_sf_cache[lcb[2]], flav) # use SF and eff to compute p(data) and p(MC) p_data = { key: np.where(pass_wp, val, 1 - val) for key, val in scale_factors.items() } p_mc = np.where(pass_wp, eff, 1 - eff) # return the unflattened version of the ratio return { key: ak.unflatten(i / p_mc, ak.num(jets.pt)) for key, i in p_data.items() }
def constituents(self, min_pt): self._out = [] self._input_flag = 0 for i in range(len(self._clusterable_level)): np_results = self._results[i].to_numpy_with_constituents(min_pt) off = np.insert(np_results[-1], 0, 0) out = ak.Array( ak.layout.ListOffsetArray64( ak.layout.Index64(np_results[0]), ak.layout.NumpyArray(np_results[1]), ), behavior=self.data.behavior, ) outputs_to_inputs = ak.Array( ak.layout.ListOffsetArray64(ak.layout.Index64(off), out.layout) ) shape = ak.num(outputs_to_inputs) total = np.sum(shape) duplicate = ak.unflatten(np.zeros(total, np.int64), shape) prepared = self._clusterable_level[i][:, np.newaxis][duplicate] self._out.append(prepared[outputs_to_inputs]) res = ak.Array(self._replace_multi()) return res