def get_series(eventWise, fixed_params, vary_param, tag_index, name_start=""): all_cols, var_cols, score_cols, table = CompareClusters.tabulate_scores( eventWise) name_mask = [ TypeTools.restring(name).startswith(name_start) for name in table[:, all_cols.index("jet_name")] ] table = table[name_mask] mask = CompareClusters.filter_matching(all_cols, table, approx=fixed_params) table = table[mask] vary_values = table[:, all_cols.index(vary_param)] if not len(set(vary_values)) == len(vary_values): import jet_tools jet_tools.st() values_dict = {v: i for i, v in enumerate(vary_values)} sorted_values = CompareClusters.generic_sort(vary_values) order = [values_dict[v] for v in sorted_values] jet_names = table[order, all_cols.index("jet_name")] jet_indices = get_jet_index(eventWise, jet_names, tag_index) labels = [] for value in sorted_values: if isinstance(value, float): labels.append(f"{vary_param}={value:.1}") else: labels.append(f"{vary_param}={value}") return jet_names, jet_indices, labels
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 test_get_particle_jet_labels(): # will need JetInputs_Px jetName_Label jet_name = "GoodDog" params = {} # an empty event will have no labels params["JetInputs_Px"] = [[]] params[jet_name + "_Label"] = [[]] expected = [[]] # an event with one object and one jet params["JetInputs_Px"] += [[0.]] params[jet_name + "_Label"] += [[[0]]] expected += [[0]] # an event with two jets params["JetInputs_Px"] += [[0., 0.5, 0.1]] params[jet_name + "_Label"] += [[[2], [0, 1]]] expected += [[1, 1, 0]] params = {key: ak.from_iter(value) for key, value in params.items()} with TempTestDir("tst") as dir_name: # this will raise a value error if given an empty eventWise file_name = "test.parquet" ew = Components.EventWise(os.path.join(dir_name, file_name)) ew.append(**params) labels = CompareClusters.get_particle_jet_labels(ew, jet_name) assert len(labels) == len(expected) for l, e in zip(labels, expected): tst.assert_allclose(l, e)
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 test_add_bg_mass(): # will need TagIndex JetInputs_SourceIdx, DetectableTag_Leaves, # Energy Px Py Pz params = {} # an empty event should have 0 mass params["TagIndex"] = [ak.from_iter([3, 4])] params["JetInputs_SourceIdx"] = [ak.from_iter([])] params["DetectableTag_Leaves"] = [ak.from_iter([])] params["Energy"] = [ak.from_iter([])] params["Px"] = [ak.from_iter([])] params["Py"] = [ak.from_iter([])] params["Pz"] = [ak.from_iter([])] # an event with no jet_inputs should also have no mass params["TagIndex"] += [ak.from_iter([1, 2, 3, 4])] params["JetInputs_SourceIdx"] += [ak.from_iter([])] params["DetectableTag_Leaves"] += [ak.from_iter([[5, 6], [7, 8]])] params["Energy"] += [ak.from_iter([0, 0, 3, 0, 0, 0, 1, 10, 9])] params["Px"] += [ak.from_iter([0, 0, 1, 0, 0, 0, 1, 0, -3])] params["Py"] += [ak.from_iter([0, 0, 1, 1, 0, 0, -1, 0, 3])] params["Pz"] += [ak.from_iter([0, 0, 0, 0, 0, 0, 1, 1, 3])] # a standard event can have mass params["TagIndex"] += [ak.from_iter([1, 2, 3, 4])] params["JetInputs_SourceIdx"] += [ak.from_iter([2, 6, 7])] params["DetectableTag_Leaves"] += [ak.from_iter([[5, 6], [7, 8]])] params["Energy"] += [ak.from_iter([0, 0, 3, 0, 0, 0, 1, 10, 9])] params["Px"] += [ak.from_iter([0, 0, 1, 0, 0, 0, 1, 0, -3])] params["Py"] += [ak.from_iter([0, 0, 1, 1, 0, 0, -1, 0, 3])] params["Pz"] += [ak.from_iter([0, 0, 0, 0, 0, 0, 1, 1, 3])] expected = [0, 0, np.sqrt(9 - 2)] with TempTestDir("tst") as dir_name: # this will raise a value error if given an empty eventWise file_name = "test.parquet" ew = Components.EventWise(os.path.join(dir_name, file_name)) ew.append(**params) # settign a selected index should not screw things up ew.selected_event = 1 CompareClusters.add_bg_mass(ew) ew.selected_event = None tst.assert_allclose(ew.DetectableBG_Mass, expected)
def test_tabulate_scores(): params1 = {} params2 = {} params1["DogJet_QualityWidth"] = 4 params1["DogJet_DeltaR"] = .5 params2["CatJet_QualityWidth"] = np.inf with TempTestDir("tst") as dir_name: # this will raise a value error if given an empty eventWise file_name1 = "test1.parquet" file_name2 = "test2.parquet" ew1 = Components.EventWise(os.path.join(dir_name, file_name1)) ew1.append_hyperparameters(**params1) ew1.append(DogJet_Label=ak.from_iter([[]])) path1 = os.path.join(dir_name, file_name1) ew2 = Components.EventWise(os.path.join(dir_name, file_name2)) ew2.append_hyperparameters(**params2) ew2.append(CatJet_Label=ak.from_iter([[]])) path2 = os.path.join(dir_name, file_name2) all_cols, variable_cols, score_cols, table = CompareClusters.tabulate_scores( [path1, path2]) assert len(all_cols) == len(table[0]) assert len(all_cols) == len(variable_cols) + len(score_cols) + 3 assert len(table) == 2 dog_row = next( i for i, name in enumerate(table[:, all_cols.index("jet_name")]) if name == "DogJet") cat_row = next( i for i, name in enumerate(table[:, all_cols.index("jet_name")]) if name == "CatJet") use_cols = [ all_cols.index("QualityWidth"), all_cols.index("DeltaR"), all_cols.index("AveBGMassRatio") ] assert np.all( TypeTools.soft_equality(table[dog_row, use_cols], [4, .5, "Undefined"])) assert np.all( TypeTools.soft_equality(table[cat_row, use_cols], [np.inf, "Undefined", "Undefined"]))
def test_get_best(): with TempTestDir("tst") as dir_name: # this will raise a value error if given an empty eventWise file_name = "test.parquet" ew = Components.EventWise(os.path.join(dir_name, file_name)) with pytest.raises(ValueError): CompareClusters.get_best(ew, "Duck") # it should also raise a value error if told to get the best of a jet not present params = dict(DogJet1_AveSignalMassRatio=0.2, DogJet1_AveBGMassRatio=0.5, PigJet_AveSignalMassRatio=np.nan, PigJet_AveBGMassRatio=1., DogJet2_AveSignalMassRatio=0.5, DogJet2_AveBGMassRatio=0.1) ew.append_hyperparameters(**params) with pytest.raises(ValueError): CompareClusters.get_best(ew, "Duck") # if asked for something wiht only nan scores it shoudl return one of them found = CompareClusters.get_best(ew, "Pig") assert found == "PigJet" # if asked for something with multiple vaild scores, chose the one with the # best ratio found = CompareClusters.get_best(ew, "Dog") assert found == "DogJet2"
def test_make_scale(): # this has two modes, ordinal and float # either way it should return somehting sensible # given an empty list it should not choke empty = np.array([]) positions, scale_positions, scale_labels = CompareClusters.make_scale( empty) assert len(positions) == 0 assert len(scale_positions) == 0 assert len(scale_labels) == 0 # a list with just one number should return it one = np.array([1]) positions, scale_positions, scale_labels = CompareClusters.make_scale(one) assert len(positions) == 1 tst.assert_allclose(positions, one) assert len(scale_positions) == 1 tst.assert_allclose(scale_positions, one) assert len(scale_labels) == 1 assert scale_labels[0] == '1' # ditto for one string one = np.array(['frog']) positions, scale_positions, scale_labels = CompareClusters.make_scale(one) assert len(positions) == 1 assert len(scale_positions) == 1 assert len(scale_labels) == 1 assert scale_labels[0] == one[0] # one nan or inf should return that, but it may also add 0 as a scale pos one = np.array([np.nan]) positions, scale_positions, scale_labels = CompareClusters.make_scale(one) assert len(positions) == 1 scale_loc = np.argmin(np.abs(scale_positions - positions[0])) assert 'nan' in scale_labels[scale_loc] or 'None' in scale_labels[scale_loc] one = np.array([-np.inf]) positions, scale_positions, scale_labels = CompareClusters.make_scale(one) assert len(positions) == 1 scale_loc = np.argmin(np.abs(scale_positions - positions[0])) assert 'inf' in scale_labels[scale_loc] and '-' in scale_labels[scale_loc] one = np.array([np.inf]) positions, scale_positions, scale_labels = CompareClusters.make_scale(one) assert len(positions) == 1 scale_loc = np.argmin(np.abs(scale_positions - positions[0])) assert 'inf' in scale_labels[scale_loc] and '-' not in scale_labels[ scale_loc] # now try some more complex arrangements inputs = [[4, 3.2], [4, 3.2, np.nan], [4, 3.2, 0., np.nan], [-4, 3.2, np.nan], [4, 3.2, np.inf], [4, 3.2, -np.inf], [-4, 3.2, 0, np.nan, np.inf], [None, 'dog', 'cat'], [ np.inf, 1, 3, 5, 3.4, 10, np.nan, 2.2, 7, 3, 5, 9, 2, 2, -np.inf ]] closest_contains = [['4', '3'], ['4', '3', 'nan'], ['4', '3', '0', 'nan'], ['-4', '3', 'nan'], ['4', '3', 'inf'], ['4', '3', '-'], ['-4', '3', '0', 'nan', 'inf'], ['none', 'dog', 'cat'], [ 'inf', '1', '3', '5', '3', '10', 'nan', '2', '7', '3', '5', '9', '2', '2', '-' ]] for inp, close in zip(inputs, closest_contains): # make the test scale inp = np.array(inp) positions, scale_positions, scale_labels = CompareClusters.make_scale( inp) # check that for each of the values in the test scale for i, pos in enumerate(positions): s_position = np.argmin(np.abs(scale_positions - pos)) # the closest label contains the expected string assert close[i] in scale_labels[s_position].lower( ), f"inputs={inp}, expected={close[i]}, found={scale_labels[s_position]}" # finally check the behavior with tuples inp = [None, ('knn', 5), ('knn', 4), ('distance', 0), ('distance', 2)] # make the test scale inp = ak.from_iter(inp) positions, scale_positions, scale_labels = CompareClusters.make_scale(inp) assert positions[1] > positions[2] assert positions[4] > positions[3] label0 = scale_labels[np.argmin(np.abs(scale_positions - positions[0]))].lower() assert 'none' in label0 label1 = scale_labels[np.argmin(np.abs(scale_positions - positions[1]))].lower() assert 'knn' in label1 and '5' in label1 label2 = scale_labels[np.argmin(np.abs(scale_positions - positions[2]))].lower() assert 'knn' in label2 and '4' in label2 label3 = scale_labels[np.argmin(np.abs(scale_positions - positions[3]))].lower() assert 'distance' in label3 and '0' in label3 label4 = scale_labels[np.argmin(np.abs(scale_positions - positions[4]))].lower() assert 'distance' in label4 and '2' in label4
def test_append_scores(): # set up as for filter jets params = {} # event 0 params['Jet_Label'] = [ak.from_iter([])] params['Jet_Parent'] = [ak.from_iter([])] params['Jet_Child1'] = [ak.from_iter([])] params['Jet_PT'] = [ak.from_iter([])] params['Energy'] = [ak.from_iter([])] params['Px'] = [ak.from_iter([])] params['Py'] = [ak.from_iter([])] params['Pz'] = [ak.from_iter([])] params['Children'] = [ak.from_iter([])] params['Parents'] = [ak.from_iter([])] params['MCPID'] = [ak.from_iter([])] params['JetInputs_SourceIdx'] = [ak.from_iter([])] params['TagIndex'] = [ak.from_iter([])] # event 1 params['Jet_Label'] = [ak.from_iter([[0], [1, 2, 3, 4, 5]])] params['Jet_Parent'] += [ak.from_iter([[-1], [-1, 1, 1, 2, 2]])] params['Jet_Child1'] += [ak.from_iter([[-1], [1, 2, -1, -1, -1]])] params['Jet_PT'] += [ak.from_iter([[ 50., ], [0.2, 0.1, 0., 0., .1]])] params['JetInputs_SourceIdx'] += [ak.from_iter(np.arange(6))] params['Energy'] += [ ak.from_iter([30., 10., 20., 70., 20., 10., 45., 56., 40., 25.]) ] params['Px'] += [ak.from_iter([3., 0., 2., 1., 2., -1., 0., 3., -1., 0.])] params['Py'] += [ak.from_iter([3., 0., 2., 2., 2., 1., -1., -3., 0., -1.])] params['Pz'] += [ak.from_iter([3., 0., 2., 0., 2., 2., -5., -2., 1., 0.])] params['Children'] += [ ak.from_iter([[], [3], [], [5], [], [], [2, 7, 8, 9], [], [], [], []]) ] params['Parents'] += [ ak.from_iter([[], [], [6], [1], [], [3], [], [6], [6], [6], []]) ] params['MCPID'] += [ak.from_iter([4, -5, 5, 3, 2, 1, -5, -1, 7, 11, 12])] params['TagIndex'] += [ak.from_iter([1, 6])] # event 2 params['Jet_Label'] = [ak.from_iter([[0], [1, 2, 3]])] params['Jet_Parent'] += [ak.from_iter([[-1], [-1, 1, 1]])] params['Jet_Child1'] += [ak.from_iter([[-1], [1, -1, -1]])] params['Jet_PT'] += [ak.from_iter([[ 50., ], [21., 0.1, 0.]])] params['JetInputs_SourceIdx'] += [ak.from_iter(np.arange(6))] params['Energy'] += [ ak.from_iter([30., 10., 20., 70., 20., 10., 45., 56., 40., 25.]) ] params['Px'] += [ak.from_iter([3., 0., 2., 1., 2., -1., 0., 3., -1., 0.])] params['Py'] += [ak.from_iter([3., 0., 2., 2., 2., 1., -1., -3., 0., -1.])] params['Pz'] += [ak.from_iter([3., 0., 2., 0., 2., 2., -5., -2., 1., 0.])] params['Children'] += [ ak.from_iter([[], [3], [], [5], [], [], [2, 7, 8, 9], [], [], [], []]) ] params['Parents'] += [ ak.from_iter([[], [], [6], [1], [], [3], [], [6], [6], [6], []]) ] params['MCPID'] += [ak.from_iter([4, -5, 5, 3, 2, 1, -5, -1, 7, 11, 12])] params['TagIndex'] += [ak.from_iter([1, 6])] params = {key: ak.from_iter(val) for key, val in params.items()} with TempTestDir("tst") as dir_name: # this will raise a value error if given an empty eventWise file_name = "test.parquet" ew = Components.EventWise(os.path.join(dir_name, file_name)) ew.append(**params) # mock JetQuality.quality_width_fracton and get_detectable_comparisons with unittest.mock.patch('jet_tools.JetQuality.quality_width_fracton', new=fake_quality_width_fraction): with unittest.mock.patch( 'jet_tools.CompareClusters.get_detectable_comparisons', new=fake_detectable_comparisons): with unittest.mock.patch( 'jet_tools.TrueTag.add_detectable_fourvector', new=fake_empty): import jet_tools CompareClusters.append_scores(ew) tst.assert_allclose(ew.Jet_Bork, [np.inf, np.inf, np.inf]) tst.assert_allclose(ew.Jet_Quack[0], [0.5, np.nan]) tst.assert_allclose(ew.Jet_Quack[1], [0.5, -np.inf]) tst.assert_allclose(ew.Jet_Quack[2], []) assert np.isnan(ew.Jet_AveBork) tst.assert_allclose(ew.Jet_AveQuack, 0.5) assert np.isnan(ew.Jet_QualityWidth) tst.assert_allclose(ew.Jet_QualityFraction, 3 / Constants.dijet_mass)
def test_get_detectable_comparisons(): # will need TagIndex JetInputs_SourceIdx, DetectableTag_Leaves, DetectableTag_Roots # DetectableTag_Energy, DetectableTag_Px, DetectableTag_Py, DetectableTag_Pz # Energy Px Py Pz # Jet_Parent, Jet_TagMass, Jet_Label # Jet_Energy, Jet_Px, Jet_Py, Jet_Pz # Jet_Tags # will create Jet_DistanceRapidity Jet_DistancePT Jet_DistancePhi # Jet_SignalMassRatio Jet_BGMassRatio Jet_PercentFound # an empty event should be nan params = {} params["TagIndex"] = [ak.from_iter([])] params["JetInputs_SourceIdx"] = [ak.from_iter([])] params["DetectableTag_Roots"] = [ak.from_iter([])] params["DetectableTag_Leaves"] = [ak.from_iter([])] params["DetectableTag_Energy"] = [ak.from_iter([])] params["DetectableTag_Px"] = [ak.from_iter([])] params["DetectableTag_Py"] = [ak.from_iter([])] params["DetectableTag_Pz"] = [ak.from_iter([])] params["Energy"] = [ak.from_iter([])] params["Px"] = [ak.from_iter([])] params["Py"] = [ak.from_iter([])] params["Pz"] = [ak.from_iter([])] params["Jet_Energy"] = [ak.from_iter([])] params["Jet_Px"] = [ak.from_iter([])] params["Jet_Py"] = [ak.from_iter([])] params["Jet_Pz"] = [ak.from_iter([])] params["Jet_Parent"] = [ak.from_iter([])] params["Jet_TagMass"] = [ak.from_iter([])] params["Jet_Label"] = [ak.from_iter([])] params["Jet_Tags"] = [ak.from_iter([])] expected_distances = [[[], [], []]] expected_sig = [[]] expected_bg = [[]] expected_percent = [np.nan] # an event with tags but no jets params["TagIndex"] += [ak.from_iter([3, 4])] params["JetInputs_SourceIdx"] += [ak.from_iter([])] params["DetectableTag_Roots"] += [ak.from_iter([[3], [4]])] params["DetectableTag_Leaves"] += [ak.from_iter([[], []])] params["DetectableTag_Energy"] += [ak.from_iter([0., 0.])] params["DetectableTag_Px"] += [ak.from_iter([0., 0.])] params["DetectableTag_Py"] += [ak.from_iter([0., 0.])] params["DetectableTag_Pz"] += [ak.from_iter([0., 0.])] params["Energy"] += [ak.from_iter([])] params["Px"] += [ak.from_iter([])] params["Py"] += [ak.from_iter([])] params["Pz"] += [ak.from_iter([])] params["Jet_Energy"] += [ak.from_iter([])] params["Jet_Px"] += [ak.from_iter([])] params["Jet_Py"] += [ak.from_iter([])] params["Jet_Pz"] += [ak.from_iter([])] params["Jet_Parent"] += [ak.from_iter([])] params["Jet_TagMass"] += [ak.from_iter([])] params["Jet_Label"] += [ak.from_iter([])] params["Jet_Tags"] += [ak.from_iter([])] expected_distances += [[[np.nan, np.nan], [np.nan, np.nan], [np.nan, np.nan]]] expected_sig += [[np.nan, np.nan]] expected_bg += [[np.nan, np.nan]] expected_percent += [0.] # the only thing that changes in the next 3 events is the tags repeat = 4 params["TagIndex"] += [ak.from_iter([3, 4])] * repeat params["JetInputs_SourceIdx"] += [ak.from_iter([1, 2, 4, 5, 6, 8]) ] * repeat params["DetectableTag_Roots"] += [ak.from_iter([[3], [4]])] * repeat params["DetectableTag_Leaves"] += [ak.from_iter([[2, 4], [5]])] * repeat params["DetectableTag_Energy"] += [ak.from_iter([5., 10.])] * repeat params["DetectableTag_Px"] += [ak.from_iter([0., 0.])] * repeat params["DetectableTag_Py"] += [ak.from_iter([1., -1.])] * repeat params["DetectableTag_Pz"] += [ak.from_iter([3., 4.])] * repeat # Label 0 1 2 3 4 5 6 7 8 params["Energy"] += [ak.from_iter([3., 4., 4., 2., 1., 10., 6., 9., 6.]) ] * repeat params["Px"] += [ak.from_iter([0., 1., -1., 2., 1., 0., -1., 0., 5.]) ] * repeat params["Py"] += [ak.from_iter([0., 0., 3., 0., -2., -1., 0., 0., 0.]) ] * repeat params["Pz"] += [ak.from_iter([1., 3., 2., -1., 1., 4., 1., 0., 0.]) ] * repeat params["Jet_Energy"] = [ ak.from_iter([[4., 1., 5.], [10.], [12., 6., 6.]]) ] * repeat params["Jet_Px"] = [ak.from_iter([[-1., 1., 0.], [0.], [4., -1., 5.]]) ] * repeat params["Jet_Py"] = [ak.from_iter([[3., -2., 1.], [-1.], [0., 0., 0.]]) ] * repeat params["Jet_Pz"] = [ak.from_iter([[2., 1., 3.], [4.], [1., 1., 0.]]) ] * repeat params["Jet_Parent"] += [ak.from_iter([[10, 10, -1], [-1], [-1, 40, 40]]) ] * repeat params["Jet_TagMass"] += [ak.from_iter([[1., 0.], [0., 1.]])] * repeat params["Jet_Label"] += [ak.from_iter([[1, 2, 10], [3], [40, 4, 5]]) ] * repeat # a perfect event should return a perfect score params["Jet_Tags"] += [ak.from_iter([[3], [4], []])] expected_distances += [[[0, 0], [0, 0], [0, 0]]] expected_sig += [[1, 1]] expected_bg += [[0., 0.]] expected_percent += [1.] # a perfectly imperfect event should return the worst score params["Jet_Tags"] += [ak.from_iter([[], [], [3]])] # detectable tag is at [5, 0, 1, 3], matched jet is at [12, 4, 0, 1] (tag_phi, jet_phi), (tag_pt, jet_pt) = Components.pxpy_to_phipt(np.array([0., 4.]), np.array([1., 0.])) tag_rapidity, jet_rapidity = Components.ptpze_to_rapidity( np.array([tag_pt, jet_pt]), np.array([3., 1.]), np.array([5., 12.])) phi_distance = Components.angular_distance(tag_phi, jet_phi) expected_distances += [[[phi_distance, np.nan], [abs(tag_pt - jet_pt), np.nan], [abs(jet_rapidity - tag_rapidity), np.nan]]] expected_sig += [[0, 0]] # the total bg mass is made from indices 1, 6 and 8 total_bg = np.sqrt(16**2 - 25 - 16) # the bg in tagged jets is divided by this mass_jet3 = np.sqrt(12**2 - 17) expected_bg += [[mass_jet3 / total_bg, 0.]] expected_percent += [0.5] # a mixed event returns a score that can be calculated params["Jet_Tags"] += [ak.from_iter([[], [4], [3]])] expected_distances += [[[phi_distance, 0], [abs(tag_pt - jet_pt), 0], [abs(jet_rapidity - tag_rapidity), 0]]] expected_sig += [[0, 1]] expected_bg += [[mass_jet3 / total_bg, 0.]] expected_percent += [1.] # in an event where a jet has tags from two diferent shower clusters it should # be assocated witht he one where it has greates inheritance params["Jet_Tags"] += [ak.from_iter([[], [4, 3], []])] # the other one will simply be unfound expected_distances += [[[np.nan, 0], [np.nan, 0], [np.nan, 0]]] expected_sig += [[0, 1]] expected_bg += [[0., 0.]] expected_percent += [.5] params = {key: ak.from_iter(val) for key, val in params.items()} with TempTestDir("tst") as dir_name: # this will raise a value error if given an empty eventWise file_name = "test.parquet" ew = Components.EventWise(os.path.join(dir_name, file_name)) ew.append(**params) jet_idxs = ak.from_iter([[], [], [0, 1, 2], [0, 1, 2], [0, 1, 2], [0, 1, 2]]) CompareClusters.get_detectable_comparisons(ew, "Jet", jet_idxs, True) for event_n in range(len(expected_percent)): ew.selected_event = event_n err_msg = f"Event {event_n}; kinematics differ from expected" found = np.array([ ew.Jet_DistancePhi.tolist(), ew.Jet_DistancePT.tolist(), ew.Jet_DistanceRapidity.tolist() ]) #print() #print((found, expected_distances[event_n])) tst.assert_allclose(found, expected_distances[event_n], err_msg=err_msg) err_msg = f"Event {event_n}; signal ratios differs from expected" #print([ew.Jet_SignalMassRatio, expected_sig[event_n]]) tst.assert_allclose(ew.Jet_SignalMassRatio, expected_sig[event_n], err_msg=err_msg) err_msg = f"Event {event_n}; background ratios differs from expected" #print((ew.Jet_BGMassRatio, expected_bg[event_n])) tst.assert_allclose(ew.Jet_BGMassRatio, expected_bg[event_n], err_msg=err_msg) err_msg = f"Event {event_n}; percentage found differs from expected" #print([ew.Jet_PercentFound, expected_percent[event_n]]) tst.assert_allclose(ew.Jet_PercentFound, expected_percent[event_n], err_msg=err_msg)
def plot_all(eventWise, prefix, jet_names=None): if isinstance(eventWise, str): eventWise = Components.EventWise.from_file(eventWise) if jet_names is None: jet_names = FormJets.get_jet_names(eventWise) criteria_names = [ "$\\Delta R$ offset", "BG free $\\Delta R$ offset", "$p_T$ offset", "BG free $p_T$ offset", "width correlation", "BG free width correlation", "1/average multiplicity", "quality width", "quality fraction" ] values = np.empty((len(jet_names), len(criteria_names)), dtype=float) for i, jet_name in enumerate(jet_names): print(f"{i/len(jet_names):.1%} {jet_name}") 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) dr_off, sig_dr_off = plot_rapidity_phi_offset(jet_name, eventWise, rapidity_in, phi_in, tag_rapidity_in, tag_phi_in) values[i, criteria_names.index("$\\Delta R$ offset")] = dr_off values[i, criteria_names.index("BG free $\\Delta R$ offset")] = sig_dr_off plt.savefig(prefix + "_" + jet_name + "_deltaRoffset.png") plt.close() pt_off, sig_pt_off = plot_PT(jet_name, eventWise, pt_in, tag_pt_in) values[i, criteria_names.index("$p_T$ offset")] = pt_off values[i, criteria_names.index("BG free $p_T$ offset")] = sig_pt_off plt.savefig(prefix + "_" + jet_name + "_PToffset.png") plt.close() width, sig_width = plot_Width(jet_name, eventWise) values[i, criteria_names.index("width correlation")] = width values[i, criteria_names.index("BG free width correlation")] = sig_width plt.savefig(prefix + "_" + jet_name + "_Width.png") plt.close() mul = plot_Multiplicity(jet_name, eventWise) values[i, criteria_names.index("1/average multiplicity")] = mul plt.savefig(prefix + "_" + jet_name + "_Multiplicity.png") plt.close() plot_MassPeaks(jet_name, eventWise) plt.savefig(prefix + "_" + jet_name + "_MassPeaks.png") plt.close() # then jet quality measures values[i, criteria_names.index("quality width")] = getattr( eventWise, jet_name + "_QualityWidth") values[i, criteria_names.index("quality fraction")] = getattr( eventWise, jet_name + "_QualityFraction") # now plot and save the summary name plot_Summary(jet_names, criteria_names, values) plt.savefig(prefix + "_Summary.png") #st() plot_Summary(jet_names, criteria_names, values) plt.show() input()