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 basic_run():
#data_dir = "/home/henry/Programs/pythia8304/reshowered_restricted/"
data_dir = "/home/henry/Programs/pythia8304/reshowered_restricted_smalldr/"
suffix = ".lhe"
data_types = [
name.replace(suffix, "") for name in os.listdir(data_dir)
if name.endswith(suffix)
]
print(data_types)
data_type = InputTools.list_complete("Which data type? ",
data_types).strip()
if "hz_to_bbee" in data_type:
#events = [6589, 675, 5028]
events = [4657, 9785, 4081]
elif "gz_to_bbee" in data_type:
#events = [5334, 4210, 1394]
events = [40036, 2422, 35096]
else:
events = [0, 1, 2, 3]
fig, ax_arr = plt.subplots(2, 2, figsize=(10, 10))
ax_list = ax_arr.flatten()
calculate_plot(data_dir, data_type, events, ax_list)
matplotlib.rc('text', usetex=False)
fig.suptitle(f"{data_type} reshowered a thousand times")
fig.set_tight_layout(True)
def plot_one(eventWise, jet_name=None):
if isinstance(eventWise, str):
eventWise = Components.EventWise.from_file(eventWise)
if jet_name is None:
jet_names = FormJets.get_jet_names(eventWise)
jet_name = InputTools.list_complete("Chose a jet; ", jet_names).strip()
jet_idxs = FormJets.filter_jets(eventWise, jet_name)
[
tag_mass2_in, all_mass2_in, bg_mass2_in, rapidity_in, phi_in, pt_in,
mask, tag_rapidity_in, tag_phi_in, tag_pt_in, percent_found,
seperate_jets
] = CompareClusters.per_event_detectables(eventWise, jet_name, jet_idxs)
plot_rapidity_phi_offset(jet_name, eventWise, rapidity_in, phi_in,
tag_rapidity_in, tag_phi_in)
plt.show()
input()
plot_PT(jet_name, eventWise, pt_in, tag_pt_in)
plt.show()
input()
plot_Width(jet_name, eventWise)
plt.show()
input()
plot_Multiplicity(jet_name, eventWise)
plt.show()
input()
plot_MassPeaks(jet_name, eventWise)
plt.show()
input()
def make_new_cluster(eventWise):
"""
Parameters
----------
eventWise :
Returns
-------
"""
cluster_name, cluster_function, chosen_parameters = pick_class_params()
# now we have parameters, apply them
jet_name = InputTools.list_complete("Name this cluster (empty for autoname); ", [''])
if jet_name == '':
found = [name.split('_', 1)[0] for name in eventWise.hyperparameters
if name.startswith(cluster_name)]
i = 0
jet_name = cluster_name + "Jet" + str(i)
while jet_name in found:
i += 1
jet_name = cluster_name + "Jet" + str(i)
print(f"Naming this {jet_name}")
FormJets.cluster_multiapply(eventWise, cluster_function, chosen_parameters, batch_length=BATCH_LENGTH, jet_name=jet_name)
return jet_name
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 test_list_complete():
select_from = ['aaa', 'abb', 'bbb']
# chcek get list complete allows arbitary input
arbitary = "jflsafhdkas ;lasjdf"
with replace_stdin(io.StringIO(arbitary)):
found = InputTools.list_complete("Msg ", select_from)
assert found.strip() == arbitary
# tab completing
complete = 'b\t'
with replace_stdin(io.StringIO(complete)):
with unittest.mock.patch('jet_tools.InputTools.tab_complete',
new=fake_tab_complete):
found = InputTools.list_complete("Msg ", select_from)
expected = select_from[-1]
assert found.strip(
) == expected, f"Found {found} from {complete}, expected {expected}"
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_PreSelections():
# manually clear the last sugestions
InputTools.last_selections = InputTools.PreSelections()
messages1 = ["Blarfle", "good dog", "bork bork", "clip clop"]
constant_lenghts = [-1, 5, 9, -1]
functions = [
InputTools.get_file_name,
lambda msg, clen: InputTools.list_complete(msg, [], clen),
lambda msg, clen: InputTools.select_values(msg, ['bog', 'fog'], [2, 3],
consistant_length=clen),
lambda msg, clen: InputTools.select_value(
msg, 1, consistant_length=clen)
]
inputs = ["fog.meow", "boop", '3, 4', '5']
expected = ["fog.meow", "boop", [3., 4.], 5.]
for i, func in enumerate(functions):
with replace_stdin(io.StringIO(inputs[i])):
func(messages1[i], constant_lenghts[i])
# this should have filled the last_selections
for i, message in enumerate(messages1):
# there may be stuff printed besides the message
assert message in InputTools.last_selections.questions[i]
assert constant_lenghts[
i] == InputTools.last_selections.consistant_length[i]
try:
assert expected[i] == InputTools.last_selections.answers[i]
except ValueError:
tst.assert_allclose(expected[i],
InputTools.last_selections.answers[i])
# write and read
with TempTestDir("tst") as dir_name:
file_name = os.path.join(dir_name, "last_sel.dat")
InputTools.last_selections.write(file_name)
InputTools.pre_selections = InputTools.PreSelections(file_name)
messages2 = ["Blarfle", "good cat", "bork bork!"]
for i, msg in enumerate(messages2):
found = functions[i](msg, constant_lenghts[i])
try:
assert expected[i] == found
except ValueError:
tst.assert_allclose(expected[i], found)
def get_existing_clusters(eventWise):
"""
Parameters
----------
eventWise :
Returns
-------
"""
clusters = {name.split('_', 1)[0] for name in eventWise.columns
if name.endswith('Parent')}
if not clusters:
return False
choice = InputTools.list_complete("Do you want to use an existing cluster (blank for make new)? ", list(clusters)).strip()
if choice == '':
return False
return choice
def plot_phys_event(eventWise, event_num, metric_names=None, jet_names=None):
"""
Parameters
----------
eventWise :
event_num :
*jet_names :
Returns
-------
"""
if jet_names is None:
jet_names = [
name.split('_')[0] for name in eventWise.columns
if name.endswith("_PhysDistance")
]
jet_names = jet_names[::2]
if metric_names is None:
metric_names = []
name = True
while name:
name = InputTools.list_complete("Chose a metric (empty to stop); ",
metrics.keys())
name = name.strip()
metric_names.append(name)
del metric_names[-1]
num_jets = len(jet_names)
# get global data
eventWise.selected_event = event_num
phis = eventWise.JetInputs_Phi
y_lims = np.min(phis), np.max(phis)
rapidities = eventWise.JetInputs_Rapidity
x_lims = np.min(rapidities), np.max(rapidities)
same_mask = np.array(eventWise.JetInputs_PairLabels.tolist())
cross_mask = np.array(eventWise.JetInputs_PairCrossings.tolist())
colours = np.zeros((len(same_mask), len(same_mask[0]), 4), dtype=float)
colours += 0.3
colours[same_mask] = [0.1, 1., 0.1, 0.]
colours[cross_mask] = [1., 0.1, 0., 0.]
colours[:, :, -1] = same_mask.astype(float) * 0.5 + cross_mask.astype(
float) * 0.5 + 0.2
# make a grid of axis for each jet name and each metric
num_metrics = len(metric_names)
fig, ax_arr = plt.subplots(num_jets, num_metrics, sharex=True, sharey=True)
ax_arr = ax_arr.reshape((num_jets, num_metrics))
# now the other axis should contain the plots
metric_order = list(metrics.keys())
for jet_n, jet_name in enumerate(jet_names):
distances = getattr(eventWise, jet_name + "_PhysDistance")
# normalise the distances
distances = distances / np.nanmean(distances.tolist(), axis=(1, 2))
ratios = getattr(eventWise, jet_name + "_DifferencePhysDistance")
for metric_n, metric in enumerate(metric_names):
metric_pos = metric_order.index(metric)
ax = ax_arr[jet_n, metric_n]
for i1, (p1, r1) in enumerate(zip(phis, rapidities)):
for i2, (p2, r2) in enumerate(zip(phis, rapidities)):
width = distances[metric_pos, i1, i2]
line = matplotlib.lines.Line2D([r1, r2], [p1, p2],
c=colours[i1, i1],
lw=width)
ax.add_line(line)
if jet_n == num_jets - 1:
ax.set_xlabel(metric)
if metric_n == 0:
ax.set_ylabel(jet_name)
ax.set_xlim(*x_lims)
ax.set_ylim(*y_lims)
fig.set_size_inches(num_metrics * 3.5, num_jets * 1.8)
#fig.tight_layout()
fig.subplots_adjust(hspace=0.0, wspace=0., right=1., top=1.)
import numpy as np
from matplotlib import pyplot as plt
from jet_tools import CompareDatasets, InputTools
irc_information = awkward1.load("../megaIgnore/IRC_shapes.awkd")
jet_names = irc_information["jet_names"]
jet_classes = [names[0].split("Jet", 1)[0] for names in jet_names]
orders = irc_information["orders"]
kinematic_names = irc_information["kinematic_names"]
shape_names = irc_information["shape_names"]
# chose the variable
chosen_var = InputTools.list_complete("Which variable? ",
kinematic_names + shape_names)
chosen_var = chosen_var.strip()
is_kinematic = chosen_var in kinematic_names
js_scores = []
if is_kinematic:
data = irc_information["kinematics"]
var_index = kinematic_names.index(chosen_var)
subsampled_js_scores = []
else:
data = irc_information["shapes"]
var_index = shape_names.index(chosen_var)
# run a loop computing the scores
for c, class_name in enumerate(jet_classes):
def show_2d(eventWise, jet_name_base, show_params=None, take_best_slice=True, ax_arr=None):
if isinstance(eventWise, str):
eventWise = Components.EventWise.from_file("../megaIgnore/iridis_ScanAffinity_SpectralFull.awkd")
jet_names = [name for name in FormJets.get_jet_names(eventWise)
if name.startswith(jet_name_base)]
score_sg = np.array([getattr(eventWise, jet_name + "_SeperateAveDistanceSignal") for jet_name in jet_names])
score_bg = np.array([getattr(eventWise, jet_name + "_SeperateAveDistanceBG") for jet_name in jet_names])
possible_parameters = {getattr(FormJets, cluster_class).default_params.keys()
for cluster_class in FormJets.cluster_classes}
varying_names = []
varying_params = []
mappings = []
fixed_params = {}
for name in possible_parameters:
values = [getattr(eventWise, f"{jet_name}_{name}")
for jet_name in jet_names]
set_values = set(values)
if len(set_values) == 1:
fixed_params[name] = values[0]
else:
varying_names.append(name)
varying_params.append(values)
mappings.append({val: i for i, val in
enumerate(sorted(set_values))})
varying_params = np.array(varying_params)
num_configurations = varying_params.shape[1]
num_params = varying_params.shape[0]
image_dimensions = tuple(len(m) for m in mappings)
image_sg = np.empty(image_dimensions)
image_bg = np.empty(image_dimensions)
for i in range(num_configurations):
coordinates = tuple(mapping[values[i]] for mapping, values in
zip(mappings, varying_params))
image_bg[coordinates] = score_bg[i]
image_sg[coordinates] = score_sg[i]
combined = np.sqrt(image_bg**2 + image_sg**2)
# now turn this 2d
if show_params is None:
if len(varying_params) == 2:
show_params = varying_params
else:
first = InputTools.list_complete("Chose first parameter; ", varying_params)
second = InputTools.list_complete("Chose second parameter; ", varying_params)
show_params = [first.strip(), second.strip()]
show_params = [varying_names.index(s) for s in show_params]
assert len(show_params) == 2
if len(show_params) == len(varying_params):
pass
elif take_best_slice:
indices = [slice(None) for _ in varying_params]
for i in range(num_params):
if i not in show_params:
other_axis = tuple(j for j in range(num_params) if i!=j)
best = np.argmin(np.sum(combined, axis=other_axis))
indices[i] = best
image_sg = image_sg[indices]
image_bg = image_bg[indices]
combined = combined[indices]
else:
other_axis = tuple(i for i in range(num_params)
if i not in show_params)
image_sg = np.mean(image_sg, axis=other_axis)
image_bg = np.mean(image_bg, axis=other_axis)
combined = np.mean(combined, axis=other_axis)
if ax_arr is None:
fig, ax_arr = plt.subplots(1, 3)
ax1, ax2, ax3 = ax_arr
ax1.imshow(image_sg[indices].T, origin='lower')
ax3.imshow(image_bg[indices].T, origin='lower')
img = ax2.imshow(combined[indices].T, origin='lower')
plt.colorbar(img, ax=ax2)
x_mapping = mappings[show_params[0]]
y_mapping = mappings[show_params[1]]
for ax in ax_arr:
ax.set_xticks(sorted(x_mapping.values()))
ax.set_xticklabels(sorted(x_mapping.keys()))
ax.set_yticks(sorted(y_mapping.values()))
ax.set_yticklabels(sorted(y_mapping.keys()))
ax.set_xlabel(varying_params[show_params[0]])
ax.set_ylabel(varying_params[show_params[1]])
ax1.set_title("Signal lost")
ax3.set_title("Background contamination")
ax2.set_title("Euclidien combination")