def test_filter_pt_eta():
with TempTestDir("filter_pt_eta") as dir_name:
name = "test.parquet"
ew = Components.EventWise(os.path.join(dir_name, name),
columns=["PT", "Pseudorapidity"],
contents={
"PT": AwkdArrays.event_floats,
"Pseudorapidity": AwkdArrays.event_floats
})
ew.selected_event = 0
input_outputs = [(0.05, 1., [], []), (0.05, 1., [0], [0]),
(0.05, 1., [0, 1], [0, 1]), (0.15, 1., [], []),
(0.15, 1., [0], []), (0.15, 1., [0, 1], [1])]
for pt, eta, inp, expected in input_outputs:
out = FormJetInputs.filter_pt_eta(ew, np.array(inp), pt, eta)
tst.assert_allclose(out, expected)
pt = ak.from_iter([[0.1, 0.1, 0.1, 0.1, 0.1]])
pz = ak.from_iter([[0., 0.1, 100., -0.1, -100.]])
ew = Components.EventWise(os.path.join(dir_name, name),
columns=["PT", "Pz"],
contents={
"PT": pt,
"Pz": pz
})
ew.selected_event = 0
input_outputs = [(0.05, 1., [], []),
(0.05, 1., [0, 1, 2, 3, 4], [0, 1, 3]),
(0.05, .1, [], []), (0.05, .1, [0, 1, 2, 3, 4], [0])]
for pt, eta, inp, expected in input_outputs:
out = FormJetInputs.filter_pt_eta(ew, np.array(inp), pt, eta)
tst.assert_allclose(out, expected)
def test_filter_obs():
with TempTestDir("filter_obs") as dir_name:
name = "test.parquet"
ew = Components.EventWise(os.path.join(dir_name, name),
columns=["Particle_Tower", "Particle_Track"],
contents={
"Particle_Track": AwkdArrays.event_ints,
"Particle_Tower": AwkdArrays.event_ints
})
ew.selected_event = 0
input_outputs = [([], []), ([0], [0]), ([0, 1], [0, 1])]
for inp, expected in input_outputs:
out = FormJetInputs.filter_obs(ew, np.array(inp))
tst.assert_allclose(out, expected)
all_unseen = ak.from_iter([[-1, -1], [-1]])
ew = Components.EventWise(os.path.join(dir_name, name),
columns=["Particle_Tower", "Particle_Track"],
contents={
"Particle_Track": AwkdArrays.event_ints,
"Particle_Tower": all_unseen
})
ew.selected_event = 0
input_outputs = [([], []), ([0], [0]), ([0, 1], [0, 1])]
for inp, expected in input_outputs:
out = FormJetInputs.filter_obs(ew, np.array(inp))
tst.assert_allclose(out, expected)
ew = Components.EventWise(os.path.join(dir_name, name),
columns=["Particle_Tower", "Particle_Track"],
contents={
"Particle_Track": all_unseen,
"Particle_Tower": AwkdArrays.event_ints
})
ew.selected_event = 0
input_outputs = [([], []), ([0], [0]), ([0, 1], [0, 1])]
for inp, expected in input_outputs:
out = FormJetInputs.filter_obs(ew, np.array(inp))
tst.assert_allclose(out, expected)
ew = Components.EventWise(os.path.join(dir_name, name),
columns=["Particle_Tower", "Particle_Track"],
contents={
"Particle_Track": all_unseen,
"Particle_Tower": all_unseen
})
ew.selected_event = 0
input_outputs = [([], []), ([0], []), ([0, 1], [])]
for inp, expected in input_outputs:
out = FormJetInputs.filter_obs(ew, np.array(inp))
tst.assert_allclose(out, expected)
one_unseen = ak.from_iter([[-1, 1], [2]])
ew = Components.EventWise(os.path.join(dir_name, name),
columns=["Particle_Tower", "Particle_Track"],
contents={
"Particle_Track": one_unseen,
"Particle_Tower": all_unseen
})
ew.selected_event = 0
input_outputs = [([], []), ([0], []), ([0, 1], [1])]
for inp, expected in input_outputs:
out = FormJetInputs.filter_obs(ew, np.array(inp))
tst.assert_allclose(out, expected)
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
def test_filter_neutrinos():
with TempTestDir("filter_neutrinos") as dir_name:
name = "test.parquet"
ew = Components.EventWise(os.path.join(dir_name, name),
columns=["MCPID"],
contents={"MCPID": [[], [1, 12]]})
ew.selected_event = 0
out = FormJetInputs.filter_neutrinos(ew, np.array([]))
assert len(out) == 0
ew.selected_event = 1
out = FormJetInputs.filter_neutrinos(ew, np.array([]))
assert len(out) == 0
out = FormJetInputs.filter_neutrinos(ew, np.array([1]))
assert len(out) == 0
out = FormJetInputs.filter_neutrinos(ew, np.array([0, 1]))
assert len(out) == 1
assert out[0] == 0
def test_filter_ends():
with TempTestDir("filter_ends") as dir_name:
name = "test.parquet"
ew = Components.EventWise(os.path.join(dir_name, name),
columns=["Children"],
contents={"Children": AwkdArrays.empty_jets})
ew.selected_event = 0
input_outputs = [([], []), ([0], [0]), ([0, 1], [0, 1])]
for inp, expected in input_outputs:
out = FormJetInputs.filter_ends(ew, np.array(inp))
tst.assert_allclose(out, expected)
ew = Components.EventWise(os.path.join(dir_name, name),
columns=["Children"],
contents={"Children": AwkdArrays.empty_jet})
ew.selected_event = 0
input_outputs = [([], []), ([0], [0]), ([0, 1], [0])]
for inp, expected in input_outputs:
out = FormJetInputs.filter_ends(ew, np.array(inp))
tst.assert_allclose(out, expected)
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]])