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_get_dir_name():
with TempTestDir("tst") as dir_name:
dir1 = "meow"
dir2 = "woof"
file1 = "maple.syrup"
new_file1 = os.path.join(dir_name, file1)
open(new_file1, 'w').close() # create the file
new_dir1 = os.path.join(dir_name, dir1)
os.mkdir(new_dir1)
new_dir2 = os.path.join(dir_name, dir2)
os.mkdir(new_dir2)
# chcek get dir name allows arbitary input
arbitary = "jflsafhdkas ;lasjdf"
with replace_stdin(io.StringIO(arbitary)):
found = InputTools.get_dir_name("Msg ")
assert found.strip() == arbitary
# tab completing a dir should get the dir back
complete = os.path.join(dir_name, 'm') + '\t'
with replace_stdin(io.StringIO(complete)):
with unittest.mock.patch('jet_tools.InputTools.tab_complete',
new=fake_tab_complete):
found = InputTools.get_dir_name("Msg ")
new_dir1 += '/'
assert found.strip(
) == new_dir1, f"Found {found} from {complete}, expected {new_dir1}"
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 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 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()
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 test_get_literal():
inputs = ['1', '1.', 'True', '[1,2, 3]', '(3,4)']
expecteds = [1, 1., True, [1, 2, 3], (3, 4)]
for inp, exp in zip(inputs, expecteds):
with replace_stdin(io.StringIO(inp)):
found = InputTools.get_literal("Msg ")
assert type(found) == type(exp)
tst.assert_allclose(found, exp)
inputs = ['1', '1.', 'True', '[1,2, 3]', '(3,4)']
expecteds = [1., 1, 1, (1, 2, 3), [3, 4]]
converters = [float, int, int, tuple, list]
for inp, exp, conv in zip(inputs, expecteds, converters):
with replace_stdin(io.StringIO(inp)):
found = InputTools.get_literal("Msg ", conv)
assert type(found) == type(exp)
tst.assert_allclose(found, exp)
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 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 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_select_value():
default = 4.
inputs = ['1', '1.', '-2', ' ']
expecteds = [1., 1., -2., default]
for inp, exp in zip(inputs, expecteds):
with replace_stdin(io.StringIO(inp)):
found = InputTools.select_value("dog ", default)
assert type(found) == type(exp)
assert found == exp
inputs = ['1', '1.', ' ']
expecteds = [1., 1, default]
converters = [float, int, int]
for inp, exp, conv in zip(inputs, expecteds, converters):
with replace_stdin(io.StringIO(inp)):
found = InputTools.select_value("dog ", default, conv)
assert type(found) == type(exp)
assert found == exp
def test_select_values():
default = [4., 12.]
inputs = ['1 4', '1., 10', '-2', ' ']
expecteds = [[1., 4.], [1., 10.], [-2.], default]
for inp, exp in zip(inputs, expecteds):
with replace_stdin(io.StringIO(inp)):
found = InputTools.select_values("friends", ["dog ", "frog"],
default)
assert isinstance(found[0], type(exp[0]))
tst.assert_allclose(found, exp)
inputs = ['1 4', '1.3, 10.2', '-2']
expecteds = [[1., 4.], [1, 10], [-2]]
converters = [float, int, int]
for inp, exp, conv in zip(inputs, expecteds, converters):
with replace_stdin(io.StringIO(inp)):
found = InputTools.select_values("animals", ["dog ", "frog"],
default, conv)
tst.assert_allclose(found, exp)
def test_get_time():
hours = 1
mins = 1
seconds = 1
expected = hours * 60**2 + mins * 60 + seconds
to_input = f'{hours}\n{mins}\n{seconds}'
with replace_stdin(io.StringIO(to_input)):
found = InputTools.get_time("Msg ")
assert found == expected
def check_problem():
""" """
eventWise = get_data_file()
BATCH_LENGTH = InputTools.get_literal("How long should the batch be (-1 for all events)? ", int)
if BATCH_LENGTH == -1:
BATCH_LENGTH = np.inf
eventWise = define_inputs(eventWise)
params = {'DeltaR': 0.4, 'ExponentMultiplier': 0.1, 'NumEigenvectors': 2, 'Laplacien': 'symmetric', "AffinityType": 'linear', "WithLaplacienScaling": False, "CutoffDistance": 5.2, "Invarient": 'normed'}
jet_name = "ProblemJet"
params["jet_name"] = jet_name
FormJets.cluster_multiapply(eventWise, FormJets.Spectral, params, batch_length=BATCH_LENGTH)
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_data_file():
""" """
# to start with we know nothing
root_file = hepmc_file = eventWise_file = False
start_file = InputTools.get_file_name("Name the data file to start from; ").strip()
if start_file.endswith('.root'):
root_file = start_file
hepmc_file = InputTools.get_file_name("Please give the corripsonding .hepmc file; ", '.hepmc')
elif start_file.endswith('.hepmc'):
hepmc_file = start_file
root_file = InputTools.get_file_name("If desired enter a .root file (blank to skip); ", '.root').strip()
if hepmc_file and root_file:
eventWise = JoinHepMCRoot.marry(hepmc_file, root_file)
elif hepmc_file:
dir_name, file_name = os.path.split(hepmc_file)
eventWise = ReadHepmc.Hepmc(dir_name, file_name)
elif start_file.endswith('.awkd'):
eventWise = Components.EventWise.from_file(start_file)
else:
raise ValueError(f"'{start_file}' not of recognised type,"
" please give a '.hepmc' (and optionally a '.root') or an '.awkd' file")
return eventWise
def test_get_file_name():
with TempTestDir("tst") as dir_name:
file1 = "meow.cat"
file2 = "maple.syrup"
new_name1 = os.path.join(dir_name, file1)
open(new_name1, 'w').close() # create the file
new_name2 = os.path.join(dir_name, file2)
open(new_name2, 'w').close() # create the file
# chcek get file name allows arbitary input
arbitary = "jflsafhdkas ;lasjdf"
with replace_stdin(io.StringIO(arbitary)):
found = InputTools.get_file_name("Msg ")
assert found.strip() == arbitary
with replace_stdin(io.StringIO(arbitary)):
found = InputTools.get_file_name("Msg ", "cat")
assert found.strip() == arbitary
# tab completing a file that ends in cat should get the first name
complete = os.path.join(dir_name, 'm') + '\t'
with replace_stdin(io.StringIO(complete)):
with unittest.mock.patch('jet_tools.InputTools.tab_complete',
new=fake_tab_complete):
found = InputTools.get_file_name("Msg ", "cat")
assert found.strip(
) == new_name1, f"Found {found} from {complete}, expected {new_name1}"
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 check_problem():
""" """
eventWise = get_data_file()
BATCH_LENGTH = InputTools.get_literal("How long should the batch be (-1 for all events)? ", int)
if BATCH_LENGTH == -1:
BATCH_LENGTH = np.inf
eventWise = define_inputs(eventWise)
params = {'DeltaR': 0.4, 'ExponentMultiplier': 0.1, 'NumEigenvectors': 2, 'Laplacien': 'symmetric', "AffinityType": 'linear', "WithLaplacienScaling": False, "CutoffDistance": 5.2, "Invarient": 'normed'}
jet_name = "ProblemJet"
params["jet_name"] = jet_name
FormJets.cluster_multiapply(eventWise, FormJets.Spectral, params, batch_length=BATCH_LENGTH)
if __name__ == '__main__':
eventWise = get_data_file()
BATCH_LENGTH = InputTools.get_literal("How long should the batch be (-1 for all events)? ", int)
if BATCH_LENGTH == -1:
BATCH_LENGTH = np.inf
eventWise = define_inputs(eventWise)
jet_name = get_existing_clusters(eventWise)
if not jet_name:
jet_name = make_new_cluster(eventWise)
# now do tagging
pretag_jet_pt_cut = get_make_tags(eventWise, jet_name)
# now calculate shapes and mass peaks
img_base = InputTools.get_file_name("Give a name to save the plots (empty for no save); ")
plot_results(eventWise, jet_name, pretag_jet_pt_cut, img_base)
"""Plot what changes the number of dimensions in the embedding space"""
from jet_tools import FormJets, InputTools, Components
import matplotlib
from matplotlib import pyplot as plt
from ipdb import set_trace as st
import numpy as np
import warnings
warnings.filterwarnings('ignore')
eventWise_path = InputTools.get_file_name(
"Where is the eventwise of collection fo eventWise? ", '.awkd').strip()
c_class = FormJets.SpectralFull
spectral_jet_params = dict(
ExpofPTMultiplier=0,
ExpofPTPosition='input',
ExpofPTFormat='Luclus',
NumEigenvectors=np.inf,
StoppingCondition='meandistance',
EigNormFactor=1.4,
EigenvalueLimit=0.4,
Laplacien='symmetric',
DeltaR=np.inf, # don't stop!
AffinityType='exponent',
AffinityExp=2.,
Sigma=.15,
CombineSize='sum',
EigDistance='abscos',
CutoffKNN=None,
CutoffDistance=None,
PhyDistance='angular')
max_particles = 140
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.)
eventWise.selected_event = None
if idealised_jet_clusters is None:
idealised_jet_clusters = eventWise.Solution
if jet_name is None:
jet_name = "SolJet"
# updated_dict will be replaced in the first batch
updated_dict = None
n_events = len(idealised_jet_clusters)
jet_list = []
for event_n, jets_gids in enumerate(idealised_jet_clusters):
if event_n % 100 == 0:
print(f"{event_n/n_events:.1%}", end='\r', flush=True)
eventWise.selected_event = event_n
jet = FormJets.Partitional(eventWise)
source_idx = eventWise.JetInputs_SourceIdx.tolist()
for gids in jets_gids:
if len(gids):
labels = [source_idx.index(i) for i in gids]
jet.create_jet(labels)
jet_list.append(jet)
updated_dict = FormJets.create_jet_contents(jet_list)
eventWise.append(**updated_dict)
if __name__ == '__main__':
path = InputTools.get_file_name("Name eventWise file: ", '.awkd').strip()
if path:
eventWise = Components.EventWise.from_file(path)
append_solution_and_weights(eventWise)
solution_to_jet(eventWise)
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):
if len(ew.Event_n) <= segment_length:
return [ew.path_name]
fragments = ew.fragment("Event_n", fragment_length=segment_length)
return fragments
if __name__ == '__main__':
from jet_tools import InputTools
import sys
args = sys.argv
if len(args) > 1:
if len(args) < 3 or len(args) > 6:
print("Format of arguments is;")
print("<input_data> <data_type> <use_pileup> (<pileup_file>)")
sys.exit()
use_pileup = InputTools.answer_to_bool(args[3])
input_data = args[1]
if input_data.endswith(Components.EventWise.FILE_EXTENTION):
print(f"Preparing files from {input_data}")
input_data = prepare_data(input_data, remove_excess=use_pileup)
else:
print(f"Using iles tarting with {input_data}")
input_data = find_varients(input_data)
data_type = args[2]
if len(args) == 5:
pileup_data = args[4]
else:
pileup_data = default_pileup
print(f"Running on {len(input_data)} files, type {data_type}," +
(" with {pileup_data}" if use_pileup else "without pileup"))
serial_run(input_data, data_type, use_pileup, pileup_data)
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? ")
try: # if the user presses ctrl-c we want to go back to the script
set_trace()
# now have a play with the objects created :)
# press c-enter when you're done
except:
pass
print("...")
print("~~ Example of reading a hepmc file ~~")
hepmc_path = "./mini_data/mini.hepmc"
if os.path.exists(hepmc_path):
print(f"Found a file to read in at {hepmc_path}")
else:
message = "Please give the location of a hepmc file (can tab complete): "
hepmc_path = InputTools.get_file_name(message, "hepmc").strip()
if not os.path.exists(hepmc_path):
print("Not a vaild file path")
exit
print("This hepmc file will be read into an awkward array")
input("press enter to continue\n...")
print(" > from jet_tools import ReadHepmc")
from jet_tools import ReadHepmc
print(" > required_size = 10")
required_size = 10
print(
" > eventWise = ReadHepmc.Hepmc(*os.path.split(hepmc_path), start=0, stop=required_size)"
)
eventWise = ReadHepmc.Hepmc(hepmc_path, start=0, stop=required_size)
print(
"If the chosen stop point is beyond the end of the hepmc (or equal to np.inf) "
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? ")