Spaces:
Sleeping
Sleeping
File size: 4,071 Bytes
e75a247 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 |
from tqdm import tqdm
import numpy as np
from src.dataset.dataset import EventDataset, get_batch_bounds
def compute_f1_score(dataset, dataset_cap=-1):
R = 0.8
counts = {
"n_matched_dark_quarks": 0,
"n_jets": 0,
"n_dark_quarks": 0
} # Array of [n_relevant_retrieved, all_retrieved, all_relevant], or in our language, [n_matched_dark_quarks, n_jets, n_dark_quarks]
n = 0
for x in tqdm(range(len(dataset))):
data = dataset[x]
jets_object = data.model_jets
n += 1
if dataset_cap != -1 and n >= dataset_cap:
break
jets = [jets_object.eta, jets_object.phi]
dq = [data.matrix_element_gen_particles.eta, data.matrix_element_gen_particles.phi]
# calculate deltaR between each jet and each quark
distance_matrix = np.zeros((len(jets_object), len(data.matrix_element_gen_particles)))
for i in range(len(jets_object)):
for j in range(len(data.matrix_element_gen_particles)):
deta = jets[0][i] - dq[0][j]
dphi = jets[1][i] - dq[1][j]
distance_matrix[i, j] = np.sqrt(deta**2 + dphi**2)
# row-wise argmin
distance_matrix = distance_matrix.T
#min_distance = np.min(distance_matrix, axis=1)
n_jets = len(jets_object)
counts["n_jets"] += n_jets
counts["n_dark_quarks"] += len(data.matrix_element_gen_particles)
if len(jets_object):
quark_to_jet = np.min(distance_matrix, axis=1)
quark_to_jet[quark_to_jet > R] = -1
counts["n_matched_dark_quarks"] += np.sum(quark_to_jet != -1)
if counts["n_jets"] != 0:
precision = counts["n_matched_dark_quarks"] / counts["n_jets"]
else:
precision = 0
if counts["n_dark_quarks"] != 0:
recall = counts["n_matched_dark_quarks"] / counts["n_dark_quarks"]
else:
recall = 0
if precision == 0 or recall == 0:
return 0
f1_score = 2 * precision * recall / (precision + recall)
return f1_score
def compute_f1_score_from_result(result, dataset):
R = 0.8
counts = {
"n_matched_dark_quarks": 0,
"n_jets": 0,
"n_dark_quarks": 0
} # Array of [n_relevant_retrieved, all_retrieved, all_relevant], or in our language, [n_matched_dark_quarks, n_jets, n_dark_quarks]
result["event_idx_bounds"] = get_batch_bounds(result["event_idx"])
l = result["event_idx"].max().int().item()
for x in tqdm(range(len(dataset))):
if x >= l:
break
data = dataset[x]
jets_object = EventDataset.get_model_jets_static(x, data.pfcands, result, result["model_cluster"])
if jets_object is None:
continue
jets = [jets_object.eta, jets_object.phi]
dq = [data.matrix_element_gen_particles.eta, data.matrix_element_gen_particles.phi]
# calculate deltaR between each jet and each quark
distance_matrix = np.zeros((len(jets_object), len(data.matrix_element_gen_particles)))
for i in range(len(jets_object)):
for j in range(len(data.matrix_element_gen_particles)):
deta = jets[0][i] - dq[0][j]
dphi = jets[1][i] - dq[1][j]
distance_matrix[i, j] = np.sqrt(deta**2 + dphi**2)
# row-wise argmin
distance_matrix = distance_matrix.T
#min_distance = np.min(distance_matrix, axis=1)
n_jets = len(jets_object)
counts["n_jets"] += n_jets
counts["n_dark_quarks"] += len(data.matrix_element_gen_particles)
if len(jets_object):
quark_to_jet = np.min(distance_matrix, axis=1)
quark_to_jet[quark_to_jet > R] = -1
counts["n_matched_dark_quarks"] += np.sum(quark_to_jet != -1)
if counts["n_jets"] == 0 or counts["n_dark_quarks"] == 0:
return 0
precision = counts["n_matched_dark_quarks"] / counts["n_jets"]
recall = counts["n_matched_dark_quarks"] / counts["n_dark_quarks"]
f1_score = 2 * precision * recall / (precision + recall)
return f1_score
|