| Geant4 Cross Reference |
1 from dataclasses import dataclass
2 from typing import Tuple
3
4 import numpy as np
5 from matplotlib import pyplot as plt
6 from scipy.optimize import curve_fit
7
8 from core.constants import N_CELLS_Z, N_CELLS_R, VALID_DIR, SIZE_Z, SIZE_R, HISTOGRAM_TYPE, FULL_SIM_HISTOGRAM_COLOR, \
9 ML_SIM_HISTOGRAM_COLOR, FULL_SIM_GAUSSIAN_COLOR, ML_SIM_GAUSSIAN_COLOR
10 from utils.observables import LongitudinalProfile, ProfileType, Profile, Energy
11
12 plt.rcParams.update({"font.size": 22})
13
14
15 @dataclass
16 class Plotter:
17 """ An abstract class defining interface of all plotters.
18
19 Do not use this class directly. Use ProfilePlotter or EnergyPlotter instead.
20
21 Attributes:
22 _particle_energy: An integer which is energy of the primary particle in GeV units.
23 _particle_angle: An integer which is an angle of the primary particle in degrees.
24 _geometry: A string which is a name of the calorimeter geometry (e.g. SiW, SciPb).
25
26 """
27 _particle_energy: int
28 _particle_angle: int
29 _geometry: str
30
31 def plot_and_save(self):
32 pass
33
34
35 def _gaussian(x: np.ndarray, a: float, mu: float, sigma: float) -> np.ndarray:
36 """ Computes a value of a Gaussian.
37
38 Args:
39 x: An argument of a function.
40 a: A scaling parameter.
41 mu: A mean.
42 sigma: A variance.
43
44 Returns:
45 A value of a function for given arguments.
46
47 """
48 return a * np.exp(-((x - mu)**2 / (2 * sigma**2)))
49
50
51 def _best_fit(data: np.ndarray,
52 bins: np.ndarray,
53 hist: bool = False) -> Tuple[np.ndarray, np.ndarray]:
54 """ Finds estimated shape of a Gaussian using Use non-linear least squares.
55
56 Args:
57 data: A numpy array with values of observables from multiple events.
58 bins: A numpy array specifying histogram bins.
59 hist: If histogram is calculated. Then data is the frequencies.
60
61 Returns:
62 A tuple of two lists. Xs and Ys of predicted curve.
63
64 """
65 # Calculate histogram.
66 if not hist:
67 hist, _ = np.histogram(data, bins)
68 else:
69 hist = data
70
71 # Choose only those bins which are nonzero. Nonzero() return a tuple of arrays. In this case it has a length = 1,
72 # hence we are interested in its first element.
73 indices = hist.nonzero()[0]
74
75 # Based on previously chosen nonzero bin, calculate position of xs and ys_bar (true values) which will be used in
76 # fitting procedure. Len(bins) == len(hist + 1), so we choose middles of bins as xs.
77 bins_middles = (bins[:-1] + bins[1:]) / 2
78 xs = bins_middles[indices]
79 ys_bar = hist[indices]
80
81 # Set initial parameters for curve fitter.
82 a0 = np.max(ys_bar)
83 mu0 = np.mean(xs)
84 sigma0 = np.var(xs)
85
86 # Fit a Gaussian to the prepared data.
87 (a, mu, sigma), _ = curve_fit(f=_gaussian,
88 xdata=xs,
89 ydata=ys_bar,
90 p0=[a0, mu0, sigma0],
91 method="trf",
92 maxfev=1000)
93
94 # Calculate values of an approximation in given points and return values.
95 ys = _gaussian(xs, a, mu, sigma)
96 return xs, ys
97
98
99 @dataclass
100 class ProfilePlotter(Plotter):
101 """ Plotter responsible for preparing plots of profiles and their first and second moments.
102
103 Attributes:
104 _full_simulation: A numpy array representing a profile of data generated by Geant4.
105 _ml_simulation: A numpy array representing a profile of data generated by ML model.
106 _plot_gaussian: A boolean. Decides whether first and second moment should be plotted as a histogram or
107 a fitted gaussian.
108 _profile_type: An enum. A profile can be either lateral or longitudinal.
109
110 """
111 _full_simulation: Profile
112 _ml_simulation: Profile
113 _plot_gaussian: bool = False
114
115 def __post_init__(self):
116 # Check if profiles are either both longitudinal or lateral.
117 full_simulation_type = type(self._full_simulation)
118 ml_generation_type = type(self._ml_simulation)
119 assert full_simulation_type == ml_generation_type, "Both profiles within a ProfilePlotter must be the same " \
120 "type."
121
122 # Set an attribute with profile type.
123 if full_simulation_type == LongitudinalProfile:
124 self._profile_type = ProfileType.LONGITUDINAL
125 else:
126 self._profile_type = ProfileType.LATERAL
127
128 def _plot_and_save_customizable_histogram(
129 self,
130 full_simulation: np.ndarray,
131 ml_simulation: np.ndarray,
132 bins: np.ndarray,
133 xlabel: str,
134 observable_name: str,
135 plot_profile: bool = False,
136 y_log_scale: bool = False) -> None:
137 """ Prepares and saves a histogram for a given pair of observables.
138
139 Args:
140 full_simulation: A numpy array of observables coming from full simulation.
141 ml_simulation: A numpy array of observables coming from ML simulation.
142 bins: A numpy array specifying histogram bins.
143 xlabel: A string. Name of x-axis on the plot.
144 observable_name: A string. Name of plotted observable.
145 plot_profile: A boolean. If set to True, full_simulation and ml_simulation are histogram weights while x is
146 defined by the number of layers. This means that in order to plot histogram (and gaussian), one first
147 need to create a data repeating each layer or R index appropriate number of times. Should be set to True
148 only while plotting profiles not first or second moments.
149 y_log_scale: A boolean. Used log scale on y-axis is set to True.
150
151 Returns:
152 None.
153
154 """
155 fig, axes = plt.subplots(2,
156 1,
157 figsize=(15, 10),
158 clear=True,
159 sharex="all")
160
161 # Plot histograms.
162 if plot_profile:
163 # We already have the bins (layers) and freqencies (energies),
164 # therefore directly plotting a step plot + lines instead of a hist plot.
165 axes[0].step(bins[:-1],
166 full_simulation,
167 label="FullSim",
168 color=FULL_SIM_HISTOGRAM_COLOR)
169 axes[0].step(bins[:-1],
170 ml_simulation,
171 label="MLSim",
172 color=ML_SIM_HISTOGRAM_COLOR)
173 axes[0].vlines(x=bins[0],
174 ymin=0,
175 ymax=full_simulation[0],
176 color=FULL_SIM_HISTOGRAM_COLOR)
177 axes[0].vlines(x=bins[-2],
178 ymin=0,
179 ymax=full_simulation[-1],
180 color=FULL_SIM_HISTOGRAM_COLOR)
181 axes[0].vlines(x=bins[0],
182 ymin=0,
183 ymax=ml_simulation[0],
184 color=ML_SIM_HISTOGRAM_COLOR)
185 axes[0].vlines(x=bins[-2],
186 ymin=0,
187 ymax=ml_simulation[-1],
188 color=ML_SIM_HISTOGRAM_COLOR)
189 axes[0].set_ylim(0, None)
190
191 # For using it later for the ratios.
192 energy_full_sim, energy_ml_sim = full_simulation, ml_simulation
193 else:
194 energy_full_sim, _, _ = axes[0].hist(
195 x=full_simulation,
196 bins=bins,
197 label="FullSim",
198 histtype=HISTOGRAM_TYPE,
199 color=FULL_SIM_HISTOGRAM_COLOR)
200 energy_ml_sim, _, _ = axes[0].hist(x=ml_simulation,
201 bins=bins,
202 label="MLSim",
203 histtype=HISTOGRAM_TYPE,
204 color=ML_SIM_HISTOGRAM_COLOR)
205
206 # Plot Gaussians if needed.
207 if self._plot_gaussian:
208 if plot_profile:
209 (xs_full_sim, ys_full_sim) = _best_fit(full_simulation,
210 bins,
211 hist=True)
212 (xs_ml_sim, ys_ml_sim) = _best_fit(ml_simulation,
213 bins,
214 hist=True)
215 else:
216 (xs_full_sim, ys_full_sim) = _best_fit(full_simulation, bins)
217 (xs_ml_sim, ys_ml_sim) = _best_fit(ml_simulation, bins)
218 axes[0].plot(xs_full_sim,
219 ys_full_sim,
220 color=FULL_SIM_GAUSSIAN_COLOR,
221 label="FullSim")
222 axes[0].plot(xs_ml_sim,
223 ys_ml_sim,
224 color=ML_SIM_GAUSSIAN_COLOR,
225 label="MLSim")
226
227 if y_log_scale:
228 axes[0].set_yscale("log")
229 axes[0].legend(loc="best")
230 axes[0].set_xlabel(xlabel)
231 axes[0].set_ylabel("Energy [Mev]")
232 axes[0].set_title(
233 f" $e^-$, {self._particle_energy} [GeV], {self._particle_angle}$^{{\circ}}$, {self._geometry}"
234 )
235
236 # Calculate ratios.
237 ratio = np.divide(energy_ml_sim,
238 energy_full_sim,
239 out=np.ones_like(energy_ml_sim),
240 where=(energy_full_sim != 0))
241 # Since len(bins) == 1 + data, we calculate middles of bins as xs.
242 bins_middles = (bins[:-1] + bins[1:]) / 2
243 axes[1].plot(bins_middles, ratio, "-o")
244 axes[1].set_xlabel(xlabel)
245 axes[1].set_ylabel("MLSim/FullSim")
246 axes[1].axhline(y=1, color="black")
247 plt.savefig(
248 f"{VALID_DIR}/{observable_name}_Geo_{self._geometry}_E_{self._particle_energy}_"
249 + f"Angle_{self._particle_angle}.png")
250 plt.clf()
251
252 def _plot_profile(self) -> None:
253 """ Plots profile of an observable.
254
255 Returns:
256 None.
257
258 """
259 full_simulation_profile = self._full_simulation.calc_profile()
260 ml_simulation_profile = self._ml_simulation.calc_profile()
261 if self._profile_type == ProfileType.LONGITUDINAL:
262 # matplotlib will include the right-limit for the last bar,
263 # hence extending by 1.
264 bins = np.linspace(0, N_CELLS_Z, N_CELLS_Z + 1)
265 observable_name = "LongProf"
266 xlabel = "Layer index"
267 else:
268 bins = np.linspace(0, N_CELLS_R, N_CELLS_R + 1)
269 observable_name = "LatProf"
270 xlabel = "R index"
271 self._plot_and_save_customizable_histogram(full_simulation_profile,
272 ml_simulation_profile,
273 bins,
274 xlabel,
275 observable_name,
276 plot_profile=True)
277
278 def _plot_first_moment(self) -> None:
279 """ Plots and saves a first moment of an observable's profile.
280
281 Returns:
282 None.
283
284 """
285 full_simulation_first_moment = self._full_simulation.calc_first_moment(
286 )
287 ml_simulation_first_moment = self._ml_simulation.calc_first_moment()
288 if self._profile_type == ProfileType.LONGITUDINAL:
289 xlabel = "$<\lambda> [mm]$"
290 observable_name = "LongFirstMoment"
291 bins = np.linspace(0, 0.4 * N_CELLS_Z * SIZE_Z, 128)
292 else:
293 xlabel = "$<r> [mm]$"
294 observable_name = "LatFirstMoment"
295 bins = np.linspace(0, 0.75 * N_CELLS_R * SIZE_R, 128)
296
297 self._plot_and_save_customizable_histogram(
298 full_simulation_first_moment, ml_simulation_first_moment, bins,
299 xlabel, observable_name)
300
301 def _plot_second_moment(self) -> None:
302 """ Plots and saves a second moment of an observable's profile.
303
304 Returns:
305 None.
306
307 """
308 full_simulation_second_moment = self._full_simulation.calc_second_moment(
309 )
310 ml_simulation_second_moment = self._ml_simulation.calc_second_moment()
311 if self._profile_type == ProfileType.LONGITUDINAL:
312 xlabel = "$<\lambda^{2}> [mm^{2}]$"
313 observable_name = "LongSecondMoment"
314 bins = np.linspace(0, pow(N_CELLS_Z * SIZE_Z, 2) / 35., 128)
315 else:
316 xlabel = "$<r^{2}> [mm^{2}]$"
317 observable_name = "LatSecondMoment"
318 bins = np.linspace(0, pow(N_CELLS_R * SIZE_R, 2) / 8., 128)
319
320 self._plot_and_save_customizable_histogram(
321 full_simulation_second_moment, ml_simulation_second_moment, bins,
322 xlabel, observable_name)
323
324 def plot_and_save(self) -> None:
325 """ Main plotting function.
326
327 Calls private methods and prints the information about progress.
328
329 Returns:
330 None.
331
332 """
333 if self._profile_type == ProfileType.LONGITUDINAL:
334 profile_type_name = "longitudinal"
335 else:
336 profile_type_name = "lateral"
337 print(f"Plotting the {profile_type_name} profile...")
338 self._plot_profile()
339 print(f"Plotting the first moment of {profile_type_name} profile...")
340 self._plot_first_moment()
341 print(f"Plotting the second moment of {profile_type_name} profile...")
342 self._plot_second_moment()
343
344
345 @dataclass
346 class EnergyPlotter(Plotter):
347 """ Plotter responsible for preparing plots of profiles and their first and second moments.
348
349 Attributes:
350 _full_simulation: A numpy array representing a profile of data generated by Geant4.
351 _ml_simulation: A numpy array representing a profile of data generated by ML model.
352
353 """
354 _full_simulation: Energy
355 _ml_simulation: Energy
356
357 def _plot_total_energy(self, y_log_scale=True) -> None:
358 """ Plots and saves a histogram with total energy detected in an event.
359
360 Args:
361 y_log_scale: A boolean. Used log scale on y-axis is set to True.
362
363 Returns:
364 None.
365
366 """
367 full_simulation_total_energy = self._full_simulation.calc_total_energy(
368 )
369 ml_simulation_total_energy = self._ml_simulation.calc_total_energy()
370
371 plt.figure(figsize=(12, 8))
372 bins = np.linspace(
373 np.min(full_simulation_total_energy) -
374 np.min(full_simulation_total_energy) * 0.05,
375 np.max(full_simulation_total_energy) +
376 np.max(full_simulation_total_energy) * 0.05, 50)
377 plt.hist(x=full_simulation_total_energy,
378 histtype=HISTOGRAM_TYPE,
379 label="FullSim",
380 bins=bins,
381 color=FULL_SIM_HISTOGRAM_COLOR)
382 plt.hist(x=ml_simulation_total_energy,
383 histtype=HISTOGRAM_TYPE,
384 label="MLSim",
385 bins=bins,
386 color=ML_SIM_HISTOGRAM_COLOR)
387 plt.legend(loc="upper left")
388 if y_log_scale:
389 plt.yscale("log")
390 plt.xlabel("Energy [MeV]")
391 plt.ylabel("# events")
392 plt.title(
393 f" $e^-$, {self._particle_energy} [GeV], {self._particle_angle}$^{{\circ}}$, {self._geometry} "
394 )
395 plt.savefig(
396 f"{VALID_DIR}/E_tot_Geo_{self._geometry}_E_{self._particle_energy}_Angle_{self._particle_angle}.png"
397 )
398 plt.clf()
399
400 def _plot_cell_energy(self) -> None:
401 """ Plots and saves a histogram with number of detector's cells across whole
402 calorimeter with particular energy detected.
403
404 Returns:
405 None.
406
407 """
408 full_simulation_cell_energy = self._full_simulation.calc_cell_energy()
409 ml_simulation_cell_energy = self._ml_simulation.calc_cell_energy()
410
411 log_full_simulation_cell_energy = np.log10(
412 full_simulation_cell_energy,
413 out=np.zeros_like(full_simulation_cell_energy),
414 where=(full_simulation_cell_energy != 0))
415 log_ml_simulation_cell_energy = np.log10(
416 ml_simulation_cell_energy,
417 out=np.zeros_like(ml_simulation_cell_energy),
418 where=(ml_simulation_cell_energy != 0))
419 plt.figure(figsize=(12, 8))
420 bins = np.linspace(-4, 1, 1000)
421 plt.hist(x=log_full_simulation_cell_energy,
422 bins=bins,
423 histtype=HISTOGRAM_TYPE,
424 label="FullSim",
425 color=FULL_SIM_HISTOGRAM_COLOR)
426 plt.hist(x=log_ml_simulation_cell_energy,
427 bins=bins,
428 histtype=HISTOGRAM_TYPE,
429 label="MLSim",
430 color=ML_SIM_HISTOGRAM_COLOR)
431 plt.xlabel("log10(E/MeV)")
432 plt.ylim(bottom=1)
433 plt.yscale("log")
434 plt.ylim(bottom=1)
435 plt.ylabel("# entries")
436 plt.title(
437 f" $e^-$, {self._particle_energy} [GeV], {self._particle_angle}$^{{\circ}}$, {self._geometry} "
438 )
439 plt.grid(True)
440 plt.legend(loc="upper left")
441 plt.savefig(
442 f"{VALID_DIR}/E_cell_Geo_{self._geometry}_E_{self._particle_energy}_Angle_{self._particle_angle}.png"
443 )
444 plt.clf()
445
446 def _plot_energy_per_layer(self):
447 """ Plots and saves N_CELLS_Z histograms with total energy detected in particular layers.
448
449 Returns:
450 None.
451
452 """
453 full_simulation_energy_per_layer = self._full_simulation.calc_energy_per_layer(
454 )
455 ml_simulation_energy_per_layer = self._ml_simulation.calc_energy_per_layer(
456 )
457
458 number_of_plots_in_row = 9
459 number_of_plots_in_column = 5
460
461 bins = np.linspace(np.min(full_simulation_energy_per_layer - 10),
462 np.max(full_simulation_energy_per_layer + 10), 25)
463
464 fig, ax = plt.subplots(number_of_plots_in_column,
465 number_of_plots_in_row,
466 figsize=(20, 15),
467 sharex="all",
468 sharey="all",
469 constrained_layout=True)
470
471 for layer_nb in range(N_CELLS_Z):
472 i = layer_nb // number_of_plots_in_row
473 j = layer_nb % number_of_plots_in_row
474
475 ax[i][j].hist(full_simulation_energy_per_layer[:, layer_nb],
476 histtype=HISTOGRAM_TYPE,
477 label="FullSim",
478 bins=bins,
479 color=FULL_SIM_HISTOGRAM_COLOR)
480 ax[i][j].hist(ml_simulation_energy_per_layer[:, layer_nb],
481 histtype=HISTOGRAM_TYPE,
482 label="MLSim",
483 bins=bins,
484 color=ML_SIM_HISTOGRAM_COLOR)
485 ax[i][j].set_title(f"Layer {layer_nb}", fontsize=13)
486 ax[i][j].set_yscale("log")
487 ax[i][j].tick_params(axis='both', which='major', labelsize=10)
488
489 fig.supxlabel("Energy [MeV]", fontsize=14)
490 fig.supylabel("# entries", fontsize=14)
491 fig.suptitle(
492 f" $e^-$, {self._particle_energy} [GeV], {self._particle_angle}$^{{\circ}}$, {self._geometry} "
493 )
494
495 # Take legend from one plot and make it a global legend.
496 handles, labels = ax[0][0].get_legend_handles_labels()
497 fig.legend(handles, labels, bbox_to_anchor=(1.15, 0.5))
498
499 plt.savefig(
500 f"{VALID_DIR}/E_layer_Geo_{self._geometry}_E_{self._particle_energy}_Angle_{self._particle_angle}.png",
501 bbox_inches="tight")
502 plt.clf()
503
504 def plot_and_save(self):
505 """ Main plotting function.
506
507 Calls private methods and prints the information about progress.
508
509 Returns:
510 None.
511
512 """
513 print("Plotting total energy...")
514 self._plot_total_energy()
515 print("Plotting cell energy...")
516 self._plot_cell_energy()
517 print("Plotting energy per layer...")
518 self._plot_energy_per_layer()