Skip to content

Commit

Permalink
Compare with fixed now returns objects. Improved seeding in fit funct…
Browse files Browse the repository at this point in the history 
…ion.
  • Loading branch information
Cristovao Fernandes Vilela committed Jan 31, 2022
1 parent c59ffbe commit 26d6a0e
Show file tree
Hide file tree
Showing 2 changed files with 111 additions and 39 deletions.
74 changes: 52 additions & 22 deletions compare_with_fixed.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,13 +7,17 @@

import torch

import pickle

particle_name = ["gamma", "e", "mu"]
pmts_to_study_xy = [[22, 75], [22, 90], [22, 110]]
pmts_to_study_xy = [[25, 75], [25, 90], [25, 110]]
pmt_names = ["Center", "Edge", "Outside"]

n_bins_q = 100
n_bins_t = 100

mm = 0.1/2.54

def compare_with_fixed(args) :
# Get and initialize model
print("Loading model: "+args.model)
Expand All @@ -39,6 +43,12 @@ def compare_with_fixed(args) :
# Initialize model
network = model_module.model(**model_args_dict)

fname_suffix = "charge"
if network.use_time :
fname_suffix += "time"
if network.use_corr :
fname_suffix += "_corr"

# Load weigths
network.load_state_dict(torch.load(args.model_weights_path, map_location=network.device))

Expand Down Expand Up @@ -88,46 +98,55 @@ def compare_with_fixed(args) :
hit_prob = (1./(1+np.exp(prediction[0][0]))).reshape((51, 150))

# Hit probability
plt.figure()
plt.figure(figsize = (6.4, 6.4))
plt.figure(figsize = (180*mm, 180*mm))
plt.subplot(3, 1, 1)
plt.imshow(event_p, origin = "lower", vmin = 0, vmax = 1)
plt.title("Fixed events hit probability")
plt.title("Hit probability simulation")
plt.colorbar()
plt.subplot(3, 1, 2)
plt.imshow(hit_prob, origin = "lower", vmin = 0, vmax = 1)
plt.title("Hit probability prediction")
plt.title("Hit probability neural network output")
plt.colorbar()
plt.subplot(3, 1, 3)
plt.imshow(event_p - hit_prob, origin = "lower", vmin = -0.5, vmax = 0.5)
plt.title("Hit probability prediction - fixed events hit probability")
plt.title("Difference between simulation and neural network output")
plt.colorbar()
plt.tight_layout()
plt.savefig("compare_fixed_event_{0}_hit_prob_NGAUS_{1}.png".format(particle_name[i_flavour], network.N_GAUS))
plt.savefig("compare_fixed_event_{0}_hit_prob_{1}_NGAUS_{2}.png".format(particle_name[i_flavour], fname_suffix, network.N_GAUS))
plt.savefig("compare_fixed_event_{0}_hit_prob_{1}_NGAUS_{2}.pdf".format(particle_name[i_flavour], fname_suffix, network.N_GAUS))
with open("compare_fixed_event_{0}_hit_prob_{1}_NGAUS_{2}.pickle".format(particle_name[i_flavour], fname_suffix, network.N_GAUS), "wb") as f :
pickle.dump({"event_p" : event_p, "hit_prob" : hit_prob}, f)

# PMTs
if not network.use_time :
plt.figure(figsize = (6.4, 6.4))
plt.figure(figsize = (180*mm, 180*mm))
for i_pmt in range(len(pmts_to_study)) :

this_pmt_prediction = torch.tensor(prediction.reshape((-1, prediction.shape[1], 51,150))[:, :, pmts_to_study_xy[i_pmt][0], pmts_to_study_xy[i_pmt][1]]).unsqueeze(dim=2)

plt.subplot(len(pmts_to_study), 1, i_pmt+1)
contents, bins, _ = plt.hist(pmts_to_study[i_pmt][:,0], bins = n_bins_q, density = True)
contents, bins, _ = plt.hist(pmts_to_study[i_pmt][:,0], bins = n_bins_q, density = True, label = "Simulation")

xx = np.linspace(min(bins), max(bins), 500)
pdf_values = []
for x in xx :
loss = network.multiGausLoss(this_pmt_prediction, torch.tensor([x/network.charge_scale]).unsqueeze(dim=0))
pdf_values.append(loss['qt_loss'])
pdf_values = np.array(pdf_values)

plt.plot(xx, np.exp(-pdf_values)/network.charge_scale)
pdf_values = np.exp(-pdf_values)/network.charge_scale

plt.plot(xx, pdf_values, label = "Neural network output: {0} Gaussians".format(network.N_GAUS))
plt.title(pmt_names[i_pmt])
plt.xlabel("Charge [p.e.]")
if i_pmt == 0 :
plt.legend()
with open("compare_fixed_event_{0}_pdf_{1}_NGAUS_{2}_ipmt_{3}.pickle".format(particle_name[i_flavour], fname_suffix, network.N_GAUS, i_pmt), "wb") as f :
pickle.dump({"contents" : contents, "bins" : bins, "xx" : xx, "pdf_values" : pdf_values}, f)

plt.tight_layout()
plt.savefig("compare_fixed_event_{0}_charge_NGAUS_{1}.png".format(particle_name[i_flavour], network.N_GAUS))
plt.savefig("compare_fixed_event_{0}_pdf_{1}_NGAUS_{2}.png".format(particle_name[i_flavour], fname_suffix, network.N_GAUS))
plt.savefig("compare_fixed_event_{0}_pdf_{1}_NGAUS_{2}.pdf".format(particle_name[i_flavour], fname_suffix, network.N_GAUS))


else :
for i_pmt in range(len(pmts_to_study)) :
Expand Down Expand Up @@ -164,7 +183,7 @@ def compare_with_fixed(args) :
pdf_values = np.array(pdf_values).reshape(qq.shape)
pdf_values = np.exp(-pdf_values)/(network.charge_scale*network.time_scale)

fig = plt.figure()
fig = plt.figure(figsize = (1.5*85*mm, 1.5*90*mm))

left, width = 0.1, 0.65
bottom, height = 0.1, 0.65
Expand All @@ -178,24 +197,35 @@ def compare_with_fixed(args) :
ax_histx = fig.add_axes(rect_histx, sharex=ax)
ax_histy = fig.add_axes(rect_histy, sharey=ax)

ax.hist2d(pmts_to_study[i_pmt][:,1], pmts_to_study[i_pmt][:,0], bins = (n_bins_t, n_bins_q), range = (t_plot_limits, q_plot_limits), cmap = "Blues")
contents_2d, xedges, yedges, _ = ax.hist2d(pmts_to_study[i_pmt][:,1], pmts_to_study[i_pmt][:,0], bins = (n_bins_t, n_bins_q), range = (t_plot_limits, q_plot_limits), cmap = "Blues")
ax.contour(tt, qq, pdf_values, cmap = "Oranges", levels = 4)

ax_histx.hist(pmts_to_study[i_pmt][:,1], bins = n_bins_t, density = True, range = t_plot_limits)
ax_histx.plot(t_1D, (pdf_values*q_plot_range/n_bins_q).sum(axis = 1))
contents_t, bins_t, _ = ax_histx.hist(pmts_to_study[i_pmt][:,1], bins = n_bins_t, density = True, range = t_plot_limits, label = "Simulation")
pdf_values_t = (pdf_values*q_plot_range/n_bins_q).sum(axis = 1)
ax_histx.plot(t_1D, pdf_values_t, label = "Neural network output: {0} Gaussians".format(network.N_GAUS))

ax_histy.hist(pmts_to_study[i_pmt][:,0], bins = n_bins_q, density = True, orientation = 'horizontal', range = q_plot_limits)
ax_histy.plot((pdf_values*t_plot_range/n_bins_t).sum(axis = 0), q_1D) #, orientation = 'horizontal')
contents_q, bins_q, _ = ax_histy.hist(pmts_to_study[i_pmt][:,0], bins = n_bins_q, density = True, orientation = 'horizontal', range = q_plot_limits)
pdf_values_q = (pdf_values*t_plot_range/n_bins_t).sum(axis = 0)
ax_histy.plot(pdf_values_q, q_1D) #, orientation = 'horizontal')

ax.set_title(pmt_names[i_pmt])
ax_histx.set_title(pmt_names[i_pmt])
ax_histx.legend()
ax.set_xlabel("Time [ns]")
ax.set_ylabel("Charge [p.e.]")

ax_histx.tick_params(axis="x", labelbottom=False)
ax_histy.tick_params(axis="y", labelleft=False)

# plt.tight_layout()
plt.savefig("compare_fixed_event_{0}_charge_NGAUS_{1}_{2}.png".format(particle_name[i_flavour], network.N_GAUS, i_pmt))

plt.savefig("compare_fixed_event_{0}_pdf_{1}_NGAUS_{2}_i_pmt_{3}.png".format(particle_name[i_flavour], fname_suffix, network.N_GAUS, i_pmt))
plt.savefig("compare_fixed_event_{0}_pdf_{1}_NGAUS_{2}_i_pmt_{3}.pdf".format(particle_name[i_flavour], fname_suffix, network.N_GAUS, i_pmt))
with open("compare_fixed_event_{0}_pdf_{1}_NGAUS_{2}_ipmt_{3}.pickle".format(particle_name[i_flavour], fname_suffix, network.N_GAUS, i_pmt), "wb") as f :
pickle.dump({"contents_t" : contents_t, "bins_t" : bins_t,
"t_1D" : t_1D, "pdf_values_t" : pdf_values_t,
"contents_q" : contents_q, "bins_q" : bins_q,
"q_1D" : q_1D, "pdf_values_q" : pdf_values_q,
"contents_2d" : contents_2d, "xedges" : xedges, "yedges": yedges,
"tt" : tt, "qq" : qq, "pdf_values" : pdf_values}, f)



if __name__ == "__main__" :
Expand Down
76 changes: 59 additions & 17 deletions fit_event.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,30 +30,72 @@ def function(x, model, PID) :

def pre_fit(model, event) :
# For now this is just a placeholder. Using truth information, pick a random point, reasonably close to the truth.
if model.use_time :
seed = np.concatenate([model.data[0,3:6], model.data[0,6:9]*model.data[0,9], [0]])
else :
seed = np.concatenate([model.data[0,3:6], model.data[0,6:9]*model.data[0,9]])

print("TRUTH")
print(seed)

# Randomize truth for seed. To mimick a pre-fit algorithm.
# For position use 5/sqrt(3) m sigma on each coordinate
# For direction / energy: mutiply direction cosines by total energy, and then smear each component with 50%/sqrt(3).
# For time use sigma = 15 ns

sigma_space = 500. # cm
sigma_time = 15. # ns
sigma_E = 0.5 # Fractional
sigma_angle = 20 # degrees

# Randomize truth for seed. To mimick a pre-fit algorithm.
# Energy:
ran_E = model.data[0,9] * np.random.normal(loc = 1., scale = sigma_E*model.data[0,9])

# Random theta
ran_theta = np.abs(np.random.normal(0., sigma_angle*np.pi/180))
ran_phi = np.random.uniform(0, 2*np.pi)

original_dir = model.data[0,6:9]
# Just to make sure
original_dir = original_dir/np.linalg.norm(original_dir)

unit_vects = [np.array([1, 0, 0]),
np.array([0, 1, 0]),
np.array([0, 0, 1])]

# Find which coordinate is most orthogonal to original direction
i_coord = np.argmin(np.dot(original_dir, unit_vects))

# Get one orthogonal vector:
v_T_1 = np.cross(original_dir, unit_vects[i_coord])
# And get another orthogonal vector to complete the basis:
v_T_2 = np.cross(original_dir, v_T_1)

# Normalize basis
v_T_1 = v_T_1/np.linalg.norm(v_T_1)
v_T_2 = v_T_2/np.linalg.norm(v_T_2)

# Build rotation matrix
R = np.transpose(np.array([v_T_1,
v_T_2,
original_dir]))

# Random direction in new basis
ran_vec = np.array([np.sin(ran_theta)*np.cos(ran_phi),
np.sin(ran_theta)*np.sin(ran_phi),
np.cos(ran_theta)])

# Random direction in detector coordinates:
ran_dir = np.matmul(R, ran_vec)

# Now randomize position:
ran_pos = np.random.normal(loc = model.data[0,3:6], scale = [sigma_space/model.xy_scale/3**0.5, sigma_space/model.xy_scale/3**0.5, sigma_space/model.z_scale/3**0.5])

# And randomize time
if model.use_time :
ran_t = np.random.normal(loc = 0., scale = sigma_time)

print(ran_dir)
print(ran_dir*ran_E)

# Seed:
if model.use_time :
seed = np.concatenate([np.random.normal(loc = seed[0:3], scale = [sigma_space/model.xy_scale/3**0.5, sigma_space/model.xy_scale/3**0.5, sigma_space/model.z_scale/3**0.5]),
np.random.normal(loc = seed[3:6], scale = np.abs(seed[3:6]*sigma_E/3**0.5)),
np.expand_dims(np.random.normal(loc = 0, scale = sigma_time/model.time_scale), axis = 0)])
seed = np.concatenate([ran_pos, ran_dir*ran_E, [ran_t]])
else :
seed = np.concatenate([np.random.normal(loc = seed[0:3], scale = [sigma_space/model.xy_scale/3**0.5, sigma_space/model.xy_scale/3**0.5, sigma_space/model.z_scale/3**0.5]),
np.random.normal(loc = seed[3:6], scale = np.abs(seed[3:6]*sigma_E/3**0.5))])
seed = np.concatenate([ran_pos, ran_dir*ran_E])

print("TRUTH")
print(np.concatenate([model.data[0,3:6], model.data[0,6:9]*model.data[0,9]]))

print("SEED")
print(seed)

Expand Down

0 comments on commit 26d6a0e

Please sign in to comment.