def get_make_tags(eventWise, jet_name):
"""
Parameters
----------
eventWise :
jet_name :
Returns
-------
"""
tag_name = jet_name + '_'
if InputTools.yesNo_question("Apply pt cut to jets before tagging? "):
jet_pt_cut = InputTools.get_literal("Jet pt cut; ", float)
jet_name += str(int(jet_pt_cut))
else:
jet_pt_cut = None
tag_name += "Tags"
if hasattr(eventWise, tag_name):
if InputTools.yesNo_question(f"{tag_name} already exists, use this? "):
# just return here
return jet_pt_cut
dr = InputTools.get_literal("What is the maximum angle between the tag and the jet (dr)? ", float)
min_tracks = InputTools.get_literal("Min tracks to tag jet; ", int)
TrueTag.add_tags(eventWise, jet_name, dr, batch_length=BATCH_LENGTH, jet_pt_cut=jet_pt_cut, min_tracks=min_tracks, overwrite=True)
return jet_pt_cut
def test_yesNo_question():
truthies = ['y', 'yes', '1', 'true', 'Y', 'YES']
for inp in truthies:
with replace_stdin(io.StringIO(inp)):
found = InputTools.yesNo_question("Msg ")
assert found
falsies = ['n', 'no', '0', 'false', 'N', 'No']
for inp in falsies:
with replace_stdin(io.StringIO(inp)):
found = InputTools.yesNo_question("Msg ")
assert not found
def plot_results(eventWise, jet_name, pretag_jet_pt_cut, img_base):
"""
Parameters
----------
eventWise :
jet_name :
pretag_jet_pt_cut :
img_base :
Returns
-------
"""
tag_before_pt_cut = pretag_jet_pt_cut is None
jet_pt_cut = pretag_jet_pt_cut
if not tag_before_pt_cut:
if InputTools.yesNo_question("Apply pt cut to jets after tagging? "):
jet_pt_cut = InputTools.get_literal("Jet pt cut; ", float)
if InputTools.yesNo_question("Plot shape variables? "):
ShapeVariables.append_tagshapes(eventWise, batch_length=BATCH_LENGTH, jet_pt_cut=jet_pt_cut, tag_before_pt_cut=tag_before_pt_cut)
ShapeVariables.append_jetshapes(eventWise, jet_name, batch_length=BATCH_LENGTH, jet_pt_cut=jet_pt_cut, tag_before_pt_cut=tag_before_pt_cut)
if tag_before_pt_cut:
ShapeVariables.plot_shapevars(eventWise, jet_name)
else:
ShapeVariables.plot_shapevars(eventWise, jet_name, jet_pt_cut)
if img_base:
plt.tight_layout()
fig = plt.gcf()
fig.set_size_inches(10, 10)
plt.savefig(img_base + '_shape.png')
else:
plt.show()
# mass peaks
if not hasattr(eventWise, jet_name + "_Tags"):
raise Exception
#st()
MassPeaks.plot_PT_pairs(eventWise, jet_name, jet_pt_cut=jet_pt_cut, show=not img_base)
if img_base:
plt.tight_layout()
fig = plt.gcf()
fig.set_size_inches(10, 10)
plt.savefig(img_base + '_PTpairs.png')
else:
plt.show()
def pick_class_params():
""" """
cluster_options = FormJets.multiapply_input
cluster_name = InputTools.list_complete("Which form of clustering? ", cluster_options.keys()).strip()
cluster_function = cluster_options[cluster_name]
if cluster_name not in FormJets.cluster_classes:
if cluster_name in ["Fast", "Home"]:
cluster_class = getattr(FormJets, "GeneralisedKT")
else:
raise NotImplementedError
else:
cluster_class = getattr(FormJets, cluster_name)
default_parameters = cluster_class.default_params
chosen_parameters = {}
print(f"Select the parameters for {cluster_name}, blank for default.")
for name, default in default_parameters.items():
selection = InputTools.list_complete(f"{name} (default {default_parameters[name]}); ", [''])
if selection == '':
chosen_parameters[name] = default
continue
try:
selection = ast.literal_eval(selection)
except ValueError:
pass
if not InputTools.yesNo_question(f"Understood {selection}, is this correct? "):
print("fix it manually")
raise Exception
#st()
pass
chosen_parameters[name] = selection
return cluster_name, cluster_function, chosen_parameters
def main():
""" """
from jet_tools import Components, InputTools, FormJets
eventWise_path = InputTools.get_file_name("Name the eventWise: ", '.awkd').strip()
if eventWise_path:
eventWise = Components.EventWise.from_file(eventWise_path)
jets = FormJets.get_jet_names(eventWise)
repeat = True
barrel_radius2 = np.max((eventWise.Track_OuterX**2 +
eventWise.Track_OuterY**2)))
barrel_radius = np.sqrt(barrel_radius2)
half_barrel_length = np.max(np.abs(eventWise.Track_OuterZ.flatten()))
while repeat:
eventWise.selected_event = int(input("Event number: "))
jet_name = InputTools.list_complete("Jet name (empty for none): ", jets).strip()
if not jet_name:
outer_pos, tower_pos, barrel_radius, halfbarrel_length\
= plot_tracks_towers(eventWise, barrel_radius=barrel_radius,
half_barrel_length=half_barrel_length)
else:
jets_here = getattr(eventWise, jet_name + "_Label")
print(f"Number of jets = {len(jets_here)}")
source_idx = eventWise.JetInputs_SourceIdx
n_source = len(source_idx)
jets_here = [source_idx[jet[jet < n_source]] for jet in jets_here]
all_jet_particles = set(np.concatenate(jets_here))
assert all_jet_particles == set(source_idx)
results = plot_tracks_towers(eventWise, particle_jets=jets_here,
barrel_radius=barrel_radius,
half_barrel_length=half_barrel_length)
outer_pos, tower_pos, barrel_radius, halfbarrel_length = results
plot_beamline(halfbarrel_length*3)
print(f"Barrel_radius = {barrel_radius}, half barrel length = {halfbarrel_length}")
mlab.show()
repeat = InputTools.yesNo_question("Repeat? ")
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
"Using the class method EventWise.from_file;")
print(" > from_disk = EventWise.from_file(new_path)")
from_disk = EventWise.from_file(new_path)
input("press enter to continue\n...")
print(
"There are a number of other useful methods avalible from the EventWise class. "
"fragment, split_unfinished, combine and recursive_combine all exist to "
"break up large file for processing on multiple cores. "
"append, remove, rename all exist to add and manipulate content. "
"append_hyperparameters allows adding content that is common to all events in the set. "
"match_indices is useful for doing more complex selections. "
"See the doc strings for a full explanation. ")
input("press enter to continue\n...")
if InputTools.yesNo_question(
"Do you want to try adding a root file for detector content? "):
print("~~ Example of reading a root file ~~")
print(
".root files produced by delphes can also be read into eventWise objects, "
"other root files could too in principle, "
"however, the code included in jetTools makes assumptions about the "
"structure of the root file that corrispond to a root file written by delphes. "
"This code also fixes various idiosyncrasies of delphes root files.")
root_path = "./mini_data/mini.root"
if os.path.exists(root_path):
print(f"Found a file to read in at {root_path}")
else:
message = "Please give the location of a root file (can tab complete): "
root_path = InputTools.get_file_name(message, "root").strip()
if not os.path.exists(root_path):
print("Not a vaild file path")
def main():
""" """
from jet_tools import Components, DrawBarrel
repeat = True
eventWise = Components.EventWise.from_file(
"megaIgnore/deltaRp4_akt.awkd")
jet_name = "HomeJet"
while repeat:
eventWise.selected_event = int(input("Event num: "))
outer_pos, tower_pos = DrawBarrel.plot_tracks_towers(eventWise)
tags_by_jet = from_hard_interaction(eventWise, jet_name)
tag_particle_idxs = []
tag_jet_idx = []
for j, tags in enumerate(tags_by_jet):
tag_particle_idxs += tags
tag_jet_idx += [j for _ in tags]
# start by putting the tag particles on the image
tag_distance = np.max(np.abs(tower_pos)) * 1.2
tag_colours = DrawBarrel.colour_set(len(tag_particle_idxs))
# set fat jets to share a colour
for i, tag in enumerate(tag_jet_idx):
locations = [
j for j, other_idx in enumerate(tag_jet_idx)
if other_idx == tag
]
for l in locations:
tag_colours[l] = tag_colours[i]
for colour, idx in zip(tag_colours, tag_particle_idxs):
pos = np.array([
eventWise.X[idx], eventWise.Y[idx], eventWise.Z[idx],
eventWise.Px[idx], eventWise.Py[idx], eventWise.Pz[idx]
])
# check the position is off the origin
if np.sum(np.abs(pos[:3])) <= 0.:
pos[:3] = pos[3:] #make the momentum the position
pos *= tag_distance / np.linalg.norm(pos)
DrawBarrel.add_single(pos,
colour,
name=f'tag({eventWise.MCPID[idx]})',
scale=300)
# highlight the towers and tracks assocated with each tag
for colour, jet_idx in zip(tag_colours, tag_jet_idx):
external_jetidx = getattr(eventWise,
jet_name + "_Child1")[jet_idx] < 0
input_jetidx = getattr(eventWise, jet_name +
"_Label")[jet_idx][external_jetidx]
particle_idx = eventWise.JetInputs_SourceIdx[input_jetidx]
tower_idx = eventWise.Particle_Tower[particle_idx]
tower_idx = tower_idx[tower_idx > 0]
track_idx = eventWise.Particle_Track[particle_idx]
track_idx = track_idx[track_idx > 0]
DrawBarrel.highlight_indices(tower_pos,
tower_idx,
colours=colour,
colourmap=False)
DrawBarrel.highlight_indices(outer_pos,
track_idx,
colours=colour,
colourmap=False)
repeat = InputTools.yesNo_question("Again? ")
def main():
"""Launch file, makes and saves a dot graph"""
repeat = True
eventWise_path = InputTools.get_file_name("Name the eventWise; ", '.awkd')
eventWise = Components.EventWise.from_file(eventWise_path)
while repeat:
from jet_tools import FormShower
eventWise.selected_event = int(input("Event number: "))
showers = FormShower.get_showers(eventWise)
jet_name = "HomeJet"
chosen_showers = []
for i, shower in enumerate(showers):
shower_roots = [shower.labels[i] for i in shower.root_local_idxs]
if 'b' not in shower_roots and 'bbar' not in shower_roots:
continue
chosen_showers.append(shower)
print(f"Shower roots {shower.root_idxs}")
max_children = max([len(d) for d in shower.children])
end_ids = shower.ends
print(
f"Drawing shower {i}, has {max_children} max children. Daughters to particles ratio = {max_children/len(shower.children)}"
)
# pick the jet with largest overlap
largest_overlap = 0
picked_jet = 0
for i in range(len(eventWise.HomeJet_Parent)):
is_external = getattr(eventWise, jet_name + "_Child1")[i] < 0
input_idx = getattr(eventWise,
jet_name + "_Label")[i][is_external]
jet_particles = eventWise.JetInputs_SourceIdx[input_idx]
matches_here = sum([p in end_ids for p in jet_particles])
if matches_here > largest_overlap:
largest_overlap = matches_here
picked_jet = i
print(
f"A jet contains {largest_overlap} out of {len(end_ids)} end products"
)
graph = DotGraph(shower=shower,
eventWise=eventWise,
jet_name=jet_name,
jet_num=picked_jet,
use_TracksTowers=True)
base_name = f"event{eventWise.selected_event}_plot"
dotName = base_name + str(i) + ".dot"
legendName = base_name + str(i) + "_ledg.dot"
with open(dotName, 'w') as dotFile:
dotFile.write(str(graph))
with open(legendName, 'w') as dotFile:
dotFile.write(graph.legend)
#amalgam_shower = chosen_showers[0]
#if len(chosen_showers)>1:
# for shower in chosen_showers[1:]:
# amalgam_shower.amalgamate(shower)
#print("Drawing the amalgam of all b showers")
#graph = DotGraph(shower=amalgam_shower, observables=obs)
#dotName = f"event{eventWise.selected_event}_mixing_plot.dot"
#legendName ="mixing_ledg.dot"
#with open(dotName, 'w') as dotFile:
# dotFile.write(str(graph))
#with open(legendName, 'w') as dotFile:
# dotFile.write(graph.legend)
repeat = InputTools.yesNo_question("Again? ")
elif header_line[0] == 'C':
i = 1
self.prepared_contents["Cross_section_pb"][event_index] = float(
header_line[i])
i += 1
self.prepared_contents["Cross_section_error_pb"][
event_index] = float(header_line[i])
def main():
""" """
input_file = InputTools.get_file_name("Input file? ", 'hepmc').strip()
new_name = InputTools.get_file_name(
"Output file? ", Components.EventWise.FILE_EXTENTION).strip()
dir_name, file_name = os.path.split(input_file)
new_dir_name, new_file_name = os.path.split(new_name)
print(f"Reading {file_name}")
eventWise = Hepmc(dir_name, file_name, 0, 10000)
eventWise.dir_name = new_dir_name
eventWise.file_name = new_file_name
print(f"Writing {eventWise.file_name}")
eventWise.write()
if __name__ == '__main__':
if InputTools.yesNo_question(
"Do you want to read that hepmc file and write it? "):
if InputTools.yesNo_question(
"Really? Did you close everything else? "):
main()
