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 define_inputs(eventWise): """ Parameters ---------- eventWise : Returns ------- """ if 'JetInputs_Energy' in eventWise.columns: use_existing = InputTools.yesNo_question("There are already JetInputs, use these? ") if use_existing: return eventWise else: # make a copy of the eventWise with no clusters or JetInputs path = eventWise.path_name print("Will make a copy of the eventWise without existing clusters to change JetInputs") print(f"Current file name is {eventWise.file_name}.") new_name = InputTools.list_complete("Name the copy; ", ['']) if not new_name.endswith('.awkd'): new_name += '.awkd' new_path = os.path.join(eventWise.dir_name, new_name) shutil.copy(path, new_path) del eventWise eventWise = Components.EventWise.from_file(new_path) # remove any clusters clusters = {name.split('_', 1)[0] for name in eventWise.columns if name.endswith('Parent')} for name in clusters: eventWise.remove_prefix(name) # remove the jet inputs eventWise.remove_prefix('JetInputs') # if we get here there are no current JetInputs if InputTools.yesNo_question("Filter the tracks on pT or eta? "): pt_cut = InputTools.get_literal("What is the minimum pT of the tracks? ", float) eta_cut = InputTools.get_literal("What is the absolute maximum of the tracks? ", float) pt_eta_cut = lambda *args: FormJetInputs.filter_pt_eta(*args, min_pt=pt_cut, max_eta=eta_cut) filter_functions = [FormJetInputs.filter_ends, pt_eta_cut] else: filter_functions = [FormJetInputs.filter_ends] FormJetInputs.create_jetInputs(eventWise, filter_functions=filter_functions, batch_length=BATCH_LENGTH) return eventWise
from jet_tools import ReadHepmc, Components, FormJets, TrueTag, MassPeaks, FormJetInputs import numpy as np hepmc = ReadHepmc.Hepmc(hepmc_dir, 0, np.inf) hepmc.file_name = "ttsemileptonic_pt7.awkd" #hepmc.dir_name = "/home/hadh1g17/jets/megaIgnore" hepmc.dir_name = "../megaIgnore" hepmc.write() print() print("Adding components") print() Components.add_all(hepmc) print() print("Creating jet inputs") print() FormJetInputs.create_jetInputs(hepmc) jet_names = [] jet_name = "AntiKTp8Jet" jet_names.append(jet_name) jet_class = FormJets.Traditional jet_params = { 'DeltaR': 0.8, 'ExpofPTFormat': 'min', 'ExpofPTMultiplier': 0, 'PhyDistance': 'angular' } print() print(f"Clustering {jet_name}") print() FormJets.cluster_multiapply(hepmc, jet_class, jet_params, jet_name, np.inf)
) print("(high rapidities will miss the silicon tracker)") print("and PT < 0.5 (low PT is not reconstructed well enough to use)") filter2 = FormJets.filter_pt_eta filters = [filter1, filter2] else: print("For pure HepMC data, the default settings are fine") filters = None input("press enter to continue\n...") print( "After applying these filters isolate the particles that can be used to form jets;" ) print( " > FormJetInputs.create_jetInputs(eventWise, filters, batch_length=np.inf)" ) FormJetInputs.create_jetInputs(eventWise, filters, batch_length=np.inf) print( "the eventwise now has various attributes that start with the word 'JetInputs_'" ) print("the clustering algorithm will look for these.") print(" > eventWise.JetInputs_Px") print(eventWise.JetInputs_Px) input("press enter to continue\n...") print("~~ Example of clustering a single event ~~") print( "Anti-KT, Cambridge-Aachen and KT jet clustering can all be created with" ) print("FormJet.GeneralisedKT which is a generalise method.") print("To cluster a single event with Anti-KT first select the event ") print(" > eventWise.selected_event = 0")
def test_create_JetInputs(): with TempTestDir("create_JetInputs") as dir_name: name = "test.parquet" columns = [ name for name in FormJets.Clustering.float_columns if "Distance" not in name and "Size" not in name ] columns_unchanged = [c for c in columns ] # because the ew will change the columns list floats = SimpleClusterSamples.two_oposite['floats'] contents = { name: ak.from_iter([x]) for name, x in zip(columns, floats.T) } def return_all(_, current_idx): return current_idx def return_second(_, current_idx): return current_idx[1:2] ew = Components.EventWise(os.path.join(dir_name, name), columns=columns, contents=contents) FormJetInputs.create_jetInputs(ew, filter_functions=[return_all]) for name in columns_unchanged: ji_name = "JetInputs_" + name idx = columns.index(name) assert hasattr(ew, ji_name) tst.assert_allclose(getattr(ew, ji_name).tolist(), [floats.T[idx]], err_msg=f"In two_oposite {name} not matching") ew.remove_prefix("JetInputs") FormJetInputs.create_jetInputs(ew, filter_functions=[return_second]) for name in columns_unchanged: ji_name = "JetInputs_" + name idx = columns.index(name) assert hasattr(ew, ji_name) tst.assert_allclose( getattr(ew, ji_name).tolist(), [floats.T[idx][1:2]]) # test batching columns = [ name for name in FormJets.Clustering.float_columns if "Distance" not in name ] contents = { name: ak.from_iter([x, x, x]) for name, x in zip(columns, floats.T) } ew = Components.EventWise(os.path.join(dir_name, name), columns=columns, contents=contents) FormJetInputs.create_jetInputs(ew, filter_functions=[return_all], batch_length=0) for name in columns_unchanged: ji_name = "JetInputs_" + name assert len(getattr(ew, ji_name)) == 0 FormJetInputs.create_jetInputs(ew, filter_functions=[return_all], batch_length=1) for name in columns_unchanged: ji_name = "JetInputs_" + name assert hasattr(ew, ji_name) idx = columns.index(name) tst.assert_allclose(getattr(ew, ji_name).tolist(), [floats.T[idx]]) FormJetInputs.create_jetInputs(ew, filter_functions=[return_all], batch_length=1) for name in columns_unchanged: ji_name = "JetInputs_" + name assert hasattr(ew, ji_name) idx = columns.index(name) tst.assert_allclose( getattr(ew, ji_name).tolist(), [floats.T[idx], floats.T[idx]])