Spaces:
Sleeping
Sleeping
File size: 4,884 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 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 |
import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib
matplotlib.rc("font", size=35)
import matplotlib.pyplot as plt
import torch
from torch_scatter import scatter_sum, scatter_mean
def calculate_event_energy_resolution(df, pandora=False, full_vector=False):
true_e = torch.Tensor(df.true_showers_E.values)
mask_nan_true = np.isnan(df.true_showers_E.values)
true_e[mask_nan_true] = 0
batch_idx = df.number_batch
if pandora:
pred_E = df.pandora_calibrated_pfo.values
nan_mask = np.isnan(df.pandora_calibrated_pfo.values)
pred_E[nan_mask] = 0
pred_e1 = torch.tensor(pred_E).unsqueeze(1).repeat(1, 3)
pred_vect = torch.tensor(np.array(df.pandora_calibrated_pos.values.tolist()))
pred_vect[nan_mask] = 0
true_vect = torch.tensor(np.array(df.true_pos.values.tolist()))
true_vect[mask_nan_true] = 0
else:
pred_E = df.calibrated_E.values
nan_mask = np.isnan(df.calibrated_E.values)
# print(np.sum(nan_mask))
pred_E[nan_mask] = 0
pred_e1 = torch.tensor(pred_E).unsqueeze(1).repeat(1, 3)
pred_vect = torch.tensor(np.array(df.pred_pos_matched.values.tolist()))
pred_vect[nan_mask] = 0
true_vect = torch.tensor(np.array(df.true_pos.values.tolist()))
true_vect[mask_nan_true] = 0
batch_idx = torch.tensor(batch_idx.values).long()
pred_E = torch.tensor(pred_E)
true_jet_vect = scatter_sum(true_vect, batch_idx, dim=0)
pred_jet_vect = scatter_sum(pred_vect, batch_idx, dim=0)
true_E_jet = scatter_sum(torch.tensor(true_e), batch_idx)
pred_E_jet = scatter_sum(torch.tensor(pred_E), batch_idx)
true_jet_p = torch.norm(
true_jet_vect, dim=1
) # This is actually momentum resolution
pred_jet_p = torch.norm(pred_jet_vect, dim=1)
mass_true = torch.sqrt(torch.abs(true_E_jet**2 - true_jet_p**2))
mass_pred = torch.sqrt(torch.abs(pred_E_jet**2 - pred_jet_p**2))
mass_over_true = mass_pred / mass_true
return mass_over_true
def get_response_for_event_energy(matched_pandora, matched_):
mass_over_true_pandora = calculate_event_energy_resolution(
matched_pandora, True, True
)
decay_type = get_decay_type(matched_pandora)
mass_over_true_model = calculate_event_energy_resolution(matched_, False, True)
dic = {}
dic["mass_over_true_model"] = mass_over_true_model
dic["mass_over_true_pandora"] = mass_over_true_pandora
dic["decay_type"] = decay_type
return dic
def get_decay_type(sd_hgb1):
batch_number = sd_hgb1.number_batch.values
decay_type_list = []
for batch_id in range(0, int(np.max(batch_number)) + 1):
decay_type = determine_decay_type(sd_hgb1, batch_id)
decay_type_list.append(decay_type)
return torch.cat(decay_type_list)
def determine_decay_type(sd_hgb1, i):
pid_values = np.abs(sd_hgb1[sd_hgb1.number_batch == i].pid.values)
if len(pid_values) == 2:
decay_type = 0
charged = np.prod(pid_values == [211.0, 211])
elif len(pid_values) == 4 and np.count_nonzero(pid_values == 22.0) == 4:
decay_type = 1
neutral = np.prod(pid_values == [22.0, 22.0, 22.0, 22.0])
else:
decay_type = 2
return torch.Tensor([decay_type])
def plot_mass_resolution(event_res_dic, PATH_store):
mask_decay_charged = event_res_dic["decay_type"] == 0
fig, ax = plt.subplots()
ax.set_xlabel("M_pred/M_true")
ax.hist(
event_res_dic["mass_over_true_model"][mask_decay_charged],
bins=100,
histtype="step",
label="ML",
color="red",
density=True,
)
ax.hist(
event_res_dic["mass_over_true_pandora"][mask_decay_charged],
bins=100,
histtype="step",
label="Pandora",
color="blue",
density=True,
)
ax.grid()
ax.legend()
ax.set_xlim([0, 10])
fig.tight_layout()
fig.savefig(PATH_store + "mass_resolution_charged.pdf", bbox_inches="tight")
mask_decay_neutral = event_res_dic["decay_type"] == 1
fig, ax = plt.subplots()
ax.set_xlabel("M_pred/M_true")
ax.hist(
event_res_dic["mass_over_true_model"][mask_decay_neutral],
bins=100,
histtype="step",
label="ML",
color="red",
density=True,
)
ax.hist(
event_res_dic["mass_over_true_pandora"][mask_decay_neutral],
bins=100,
histtype="step",
label="Pandora",
color="blue",
density=True,
)
ax.grid()
ax.legend()
ax.set_xlim([0, 10])
fig.tight_layout()
fig.savefig(PATH_store + "mass_resolution_neutral.pdf", bbox_inches="tight")
def mass_Ks(matched_pandora, matched_, PATH_store):
event_res_dic = get_response_for_event_energy(matched_pandora, matched_)
plot_mass_resolution(event_res_dic, PATH_store)
|