def cluster_sequence(file_name, names, classes, params, jet_pt_cut, dijet_mass, pileup_file=None, mean_pileup_per_event=50, z_cuts=[0.2, 0.4]): if pileup_file is not None: ew = NonwritingEventWise.from_file(file_name) else: ew = Components.EventWise.from_file(file_name) if pileup_file is not None: print("Adding pileup") ew = remove_excess_data(ew) pileup = NonwritingEventWise.from_file(pileup_file) ew = AddPileup.add_pileup(ew, pileup, mean_pileup_per_event) filter_functions = "remove pileup" else: filter_functions = "ignore pileup" print("Adding jet inputs") FormJetInputs.create_jetInputs(ew, filter_functions=filter_functions) print("Adding jets") for name, cla, param in zip(names, classes, params): finished = False while not finished: finished = FormJets.cluster_multiapply(ew, cla, param, name, np.inf) n_jets = len(names) * 20 if z_cuts: print("Doing z_cuts") ParallelFormJets.batch_masks(ew, z_cuts, n_jets) ParallelFormJets.batch_filter(ew, n_jets) print("Calculating scores") if 'top' in str(dijet_mass): ParallelFormJets.batch_semileptonic(ew, n_jets) ParallelFormJets.batch_correct_semileptonic_masses(ew, n_jets) else: CompareClusters.append_scores(ew, dijet_mass=dijet_mass, end_time=np.inf) ParallelFormJets.batch_correct_masses(ew, jet_pt_cut, n_jets) return ew
def test_append_scores(): # set up as for filter jets params = {} # event 0 params['Jet_Label'] = [ak.from_iter([])] params['Jet_Parent'] = [ak.from_iter([])] params['Jet_Child1'] = [ak.from_iter([])] params['Jet_PT'] = [ak.from_iter([])] params['Energy'] = [ak.from_iter([])] params['Px'] = [ak.from_iter([])] params['Py'] = [ak.from_iter([])] params['Pz'] = [ak.from_iter([])] params['Children'] = [ak.from_iter([])] params['Parents'] = [ak.from_iter([])] params['MCPID'] = [ak.from_iter([])] params['JetInputs_SourceIdx'] = [ak.from_iter([])] params['TagIndex'] = [ak.from_iter([])] # event 1 params['Jet_Label'] = [ak.from_iter([[0], [1, 2, 3, 4, 5]])] params['Jet_Parent'] += [ak.from_iter([[-1], [-1, 1, 1, 2, 2]])] params['Jet_Child1'] += [ak.from_iter([[-1], [1, 2, -1, -1, -1]])] params['Jet_PT'] += [ak.from_iter([[ 50., ], [0.2, 0.1, 0., 0., .1]])] params['JetInputs_SourceIdx'] += [ak.from_iter(np.arange(6))] params['Energy'] += [ ak.from_iter([30., 10., 20., 70., 20., 10., 45., 56., 40., 25.]) ] params['Px'] += [ak.from_iter([3., 0., 2., 1., 2., -1., 0., 3., -1., 0.])] params['Py'] += [ak.from_iter([3., 0., 2., 2., 2., 1., -1., -3., 0., -1.])] params['Pz'] += [ak.from_iter([3., 0., 2., 0., 2., 2., -5., -2., 1., 0.])] params['Children'] += [ ak.from_iter([[], [3], [], [5], [], [], [2, 7, 8, 9], [], [], [], []]) ] params['Parents'] += [ ak.from_iter([[], [], [6], [1], [], [3], [], [6], [6], [6], []]) ] params['MCPID'] += [ak.from_iter([4, -5, 5, 3, 2, 1, -5, -1, 7, 11, 12])] params['TagIndex'] += [ak.from_iter([1, 6])] # event 2 params['Jet_Label'] = [ak.from_iter([[0], [1, 2, 3]])] params['Jet_Parent'] += [ak.from_iter([[-1], [-1, 1, 1]])] params['Jet_Child1'] += [ak.from_iter([[-1], [1, -1, -1]])] params['Jet_PT'] += [ak.from_iter([[ 50., ], [21., 0.1, 0.]])] params['JetInputs_SourceIdx'] += [ak.from_iter(np.arange(6))] params['Energy'] += [ ak.from_iter([30., 10., 20., 70., 20., 10., 45., 56., 40., 25.]) ] params['Px'] += [ak.from_iter([3., 0., 2., 1., 2., -1., 0., 3., -1., 0.])] params['Py'] += [ak.from_iter([3., 0., 2., 2., 2., 1., -1., -3., 0., -1.])] params['Pz'] += [ak.from_iter([3., 0., 2., 0., 2., 2., -5., -2., 1., 0.])] params['Children'] += [ ak.from_iter([[], [3], [], [5], [], [], [2, 7, 8, 9], [], [], [], []]) ] params['Parents'] += [ ak.from_iter([[], [], [6], [1], [], [3], [], [6], [6], [6], []]) ] params['MCPID'] += [ak.from_iter([4, -5, 5, 3, 2, 1, -5, -1, 7, 11, 12])] params['TagIndex'] += [ak.from_iter([1, 6])] params = {key: ak.from_iter(val) for key, val in params.items()} with TempTestDir("tst") as dir_name: # this will raise a value error if given an empty eventWise file_name = "test.parquet" ew = Components.EventWise(os.path.join(dir_name, file_name)) ew.append(**params) # mock JetQuality.quality_width_fracton and get_detectable_comparisons with unittest.mock.patch('jet_tools.JetQuality.quality_width_fracton', new=fake_quality_width_fraction): with unittest.mock.patch( 'jet_tools.CompareClusters.get_detectable_comparisons', new=fake_detectable_comparisons): with unittest.mock.patch( 'jet_tools.TrueTag.add_detectable_fourvector', new=fake_empty): import jet_tools CompareClusters.append_scores(ew) tst.assert_allclose(ew.Jet_Bork, [np.inf, np.inf, np.inf]) tst.assert_allclose(ew.Jet_Quack[0], [0.5, np.nan]) tst.assert_allclose(ew.Jet_Quack[1], [0.5, -np.inf]) tst.assert_allclose(ew.Jet_Quack[2], []) assert np.isnan(ew.Jet_AveBork) tst.assert_allclose(ew.Jet_AveQuack, 0.5) assert np.isnan(ew.Jet_QualityWidth) tst.assert_allclose(ew.Jet_QualityFraction, 3 / Constants.dijet_mass)