from dataclasses import dataclass from typing import Tuple import numpy as np from matplotlib import pyplot as plt from scipy.optimize import curve_fit from core.constants import N_CELLS_Z, N_CELLS_R, VALID_DIR, SIZE_Z, SIZE_R, HISTOGRAM_TYPE, FULL_SIM_HISTOGRAM_COLOR, \ ML_SIM_HISTOGRAM_COLOR, FULL_SIM_GAUSSIAN_COLOR, ML_SIM_GAUSSIAN_COLOR from utils.observables import LongitudinalProfile, ProfileType, Profile, Energy plt.rcParams.update({"font.size": 22}) @dataclass class Plotter: """ An abstract class defining interface of all plotters. Do not use this class directly. Use ProfilePlotter or EnergyPlotter instead. Attributes: _particle_energy: An integer which is energy of the primary particle in GeV units. _particle_angle: An integer which is an angle of the primary particle in degrees. _geometry: A string which is a name of the calorimeter geometry (e.g. SiW, SciPb). """ _particle_energy: int _particle_angle: int _geometry: str def plot_and_save(self): pass def _gaussian(x: np.ndarray, a: float, mu: float, sigma: float) -> np.ndarray: """ Computes a value of a Gaussian. Args: x: An argument of a function. a: A scaling parameter. mu: A mean. sigma: A variance. Returns: A value of a function for given arguments. """ return a * np.exp(-((x - mu)**2 / (2 * sigma**2))) def _best_fit(data: np.ndarray, bins: np.ndarray, hist: bool = False) -> Tuple[np.ndarray, np.ndarray]: """ Finds estimated shape of a Gaussian using Use non-linear least squares. Args: data: A numpy array with values of observables from multiple events. bins: A numpy array specifying histogram bins. hist: If histogram is calculated. Then data is the frequencies. Returns: A tuple of two lists. Xs and Ys of predicted curve. """ # Calculate histogram. if not hist: hist, _ = np.histogram(data, bins) else: hist = data # Choose only those bins which are nonzero. Nonzero() return a tuple of arrays. In this case it has a length = 1, # hence we are interested in its first element. indices = hist.nonzero()[0] # Based on previously chosen nonzero bin, calculate position of xs and ys_bar (true values) which will be used in # fitting procedure. Len(bins) == len(hist + 1), so we choose middles of bins as xs. bins_middles = (bins[:-1] + bins[1:]) / 2 xs = bins_middles[indices] ys_bar = hist[indices] # Set initial parameters for curve fitter. a0 = np.max(ys_bar) mu0 = np.mean(xs) sigma0 = np.var(xs) # Fit a Gaussian to the prepared data. (a, mu, sigma), _ = curve_fit(f=_gaussian, xdata=xs, ydata=ys_bar, p0=[a0, mu0, sigma0], method="trf", maxfev=1000) # Calculate values of an approximation in given points and return values. ys = _gaussian(xs, a, mu, sigma) return xs, ys @dataclass class ProfilePlotter(Plotter): """ Plotter responsible for preparing plots of profiles and their first and second moments. Attributes: _full_simulation: A numpy array representing a profile of data generated by Geant4. _ml_simulation: A numpy array representing a profile of data generated by ML model. _plot_gaussian: A boolean. Decides whether first and second moment should be plotted as a histogram or a fitted gaussian. _profile_type: An enum. A profile can be either lateral or longitudinal. """ _full_simulation: Profile _ml_simulation: Profile _plot_gaussian: bool = False def __post_init__(self): # Check if profiles are either both longitudinal or lateral. full_simulation_type = type(self._full_simulation) ml_generation_type = type(self._ml_simulation) assert full_simulation_type == ml_generation_type, "Both profiles within a ProfilePlotter must be the same " \ "type." # Set an attribute with profile type. if full_simulation_type == LongitudinalProfile: self._profile_type = ProfileType.LONGITUDINAL else: self._profile_type = ProfileType.LATERAL def _plot_and_save_customizable_histogram( self, full_simulation: np.ndarray, ml_simulation: np.ndarray, bins: np.ndarray, xlabel: str, observable_name: str, plot_profile: bool = False, y_log_scale: bool = False) -> None: """ Prepares and saves a histogram for a given pair of observables. Args: full_simulation: A numpy array of observables coming from full simulation. ml_simulation: A numpy array of observables coming from ML simulation. bins: A numpy array specifying histogram bins. xlabel: A string. Name of x-axis on the plot. observable_name: A string. Name of plotted observable. plot_profile: A boolean. If set to True, full_simulation and ml_simulation are histogram weights while x is defined by the number of layers. This means that in order to plot histogram (and gaussian), one first need to create a data repeating each layer or R index appropriate number of times. Should be set to True only while plotting profiles not first or second moments. y_log_scale: A boolean. Used log scale on y-axis is set to True. Returns: None. """ fig, axes = plt.subplots(2, 1, figsize=(15, 10), clear=True, sharex="all") # Plot histograms. if plot_profile: # We already have the bins (layers) and freqencies (energies), # therefore directly plotting a step plot + lines instead of a hist plot. axes[0].step(bins[:-1], full_simulation, label="FullSim", color=FULL_SIM_HISTOGRAM_COLOR) axes[0].step(bins[:-1], ml_simulation, label="MLSim", color=ML_SIM_HISTOGRAM_COLOR) axes[0].vlines(x=bins[0], ymin=0, ymax=full_simulation[0], color=FULL_SIM_HISTOGRAM_COLOR) axes[0].vlines(x=bins[-2], ymin=0, ymax=full_simulation[-1], color=FULL_SIM_HISTOGRAM_COLOR) axes[0].vlines(x=bins[0], ymin=0, ymax=ml_simulation[0], color=ML_SIM_HISTOGRAM_COLOR) axes[0].vlines(x=bins[-2], ymin=0, ymax=ml_simulation[-1], color=ML_SIM_HISTOGRAM_COLOR) axes[0].set_ylim(0, None) # For using it later for the ratios. energy_full_sim, energy_ml_sim = full_simulation, ml_simulation else: energy_full_sim, _, _ = axes[0].hist( x=full_simulation, bins=bins, label="FullSim", histtype=HISTOGRAM_TYPE, color=FULL_SIM_HISTOGRAM_COLOR) energy_ml_sim, _, _ = axes[0].hist(x=ml_simulation, bins=bins, label="MLSim", histtype=HISTOGRAM_TYPE, color=ML_SIM_HISTOGRAM_COLOR) # Plot Gaussians if needed. if self._plot_gaussian: if plot_profile: (xs_full_sim, ys_full_sim) = _best_fit(full_simulation, bins, hist=True) (xs_ml_sim, ys_ml_sim) = _best_fit(ml_simulation, bins, hist=True) else: (xs_full_sim, ys_full_sim) = _best_fit(full_simulation, bins) (xs_ml_sim, ys_ml_sim) = _best_fit(ml_simulation, bins) axes[0].plot(xs_full_sim, ys_full_sim, color=FULL_SIM_GAUSSIAN_COLOR, label="FullSim") axes[0].plot(xs_ml_sim, ys_ml_sim, color=ML_SIM_GAUSSIAN_COLOR, label="MLSim") if y_log_scale: axes[0].set_yscale("log") axes[0].legend(loc="best") axes[0].set_xlabel(xlabel) axes[0].set_ylabel("Energy [Mev]") axes[0].set_title( f" $e^-$, {self._particle_energy} [GeV], {self._particle_angle}$^{{\circ}}$, {self._geometry}" ) # Calculate ratios. ratio = np.divide(energy_ml_sim, energy_full_sim, out=np.ones_like(energy_ml_sim), where=(energy_full_sim != 0)) # Since len(bins) == 1 + data, we calculate middles of bins as xs. bins_middles = (bins[:-1] + bins[1:]) / 2 axes[1].plot(bins_middles, ratio, "-o") axes[1].set_xlabel(xlabel) axes[1].set_ylabel("MLSim/FullSim") axes[1].axhline(y=1, color="black") plt.savefig( f"{VALID_DIR}/{observable_name}_Geo_{self._geometry}_E_{self._particle_energy}_" + f"Angle_{self._particle_angle}.png") plt.clf() def _plot_profile(self) -> None: """ Plots profile of an observable. Returns: None. """ full_simulation_profile = self._full_simulation.calc_profile() ml_simulation_profile = self._ml_simulation.calc_profile() if self._profile_type == ProfileType.LONGITUDINAL: # matplotlib will include the right-limit for the last bar, # hence extending by 1. bins = np.linspace(0, N_CELLS_Z, N_CELLS_Z + 1) observable_name = "LongProf" xlabel = "Layer index" else: bins = np.linspace(0, N_CELLS_R, N_CELLS_R + 1) observable_name = "LatProf" xlabel = "R index" self._plot_and_save_customizable_histogram(full_simulation_profile, ml_simulation_profile, bins, xlabel, observable_name, plot_profile=True) def _plot_first_moment(self) -> None: """ Plots and saves a first moment of an observable's profile. Returns: None. """ full_simulation_first_moment = self._full_simulation.calc_first_moment( ) ml_simulation_first_moment = self._ml_simulation.calc_first_moment() if self._profile_type == ProfileType.LONGITUDINAL: xlabel = "$<\lambda> [mm]$" observable_name = "LongFirstMoment" bins = np.linspace(0, 0.4 * N_CELLS_Z * SIZE_Z, 128) else: xlabel = "$ [mm]$" observable_name = "LatFirstMoment" bins = np.linspace(0, 0.75 * N_CELLS_R * SIZE_R, 128) self._plot_and_save_customizable_histogram( full_simulation_first_moment, ml_simulation_first_moment, bins, xlabel, observable_name) def _plot_second_moment(self) -> None: """ Plots and saves a second moment of an observable's profile. Returns: None. """ full_simulation_second_moment = self._full_simulation.calc_second_moment( ) ml_simulation_second_moment = self._ml_simulation.calc_second_moment() if self._profile_type == ProfileType.LONGITUDINAL: xlabel = "$<\lambda^{2}> [mm^{2}]$" observable_name = "LongSecondMoment" bins = np.linspace(0, pow(N_CELLS_Z * SIZE_Z, 2) / 35., 128) else: xlabel = "$ [mm^{2}]$" observable_name = "LatSecondMoment" bins = np.linspace(0, pow(N_CELLS_R * SIZE_R, 2) / 8., 128) self._plot_and_save_customizable_histogram( full_simulation_second_moment, ml_simulation_second_moment, bins, xlabel, observable_name) def plot_and_save(self) -> None: """ Main plotting function. Calls private methods and prints the information about progress. Returns: None. """ if self._profile_type == ProfileType.LONGITUDINAL: profile_type_name = "longitudinal" else: profile_type_name = "lateral" print(f"Plotting the {profile_type_name} profile...") self._plot_profile() print(f"Plotting the first moment of {profile_type_name} profile...") self._plot_first_moment() print(f"Plotting the second moment of {profile_type_name} profile...") self._plot_second_moment() @dataclass class EnergyPlotter(Plotter): """ Plotter responsible for preparing plots of profiles and their first and second moments. Attributes: _full_simulation: A numpy array representing a profile of data generated by Geant4. _ml_simulation: A numpy array representing a profile of data generated by ML model. """ _full_simulation: Energy _ml_simulation: Energy def _plot_total_energy(self, y_log_scale=True) -> None: """ Plots and saves a histogram with total energy detected in an event. Args: y_log_scale: A boolean. Used log scale on y-axis is set to True. Returns: None. """ full_simulation_total_energy = self._full_simulation.calc_total_energy( ) ml_simulation_total_energy = self._ml_simulation.calc_total_energy() plt.figure(figsize=(12, 8)) bins = np.linspace( np.min(full_simulation_total_energy) - np.min(full_simulation_total_energy) * 0.05, np.max(full_simulation_total_energy) + np.max(full_simulation_total_energy) * 0.05, 50) plt.hist(x=full_simulation_total_energy, histtype=HISTOGRAM_TYPE, label="FullSim", bins=bins, color=FULL_SIM_HISTOGRAM_COLOR) plt.hist(x=ml_simulation_total_energy, histtype=HISTOGRAM_TYPE, label="MLSim", bins=bins, color=ML_SIM_HISTOGRAM_COLOR) plt.legend(loc="upper left") if y_log_scale: plt.yscale("log") plt.xlabel("Energy [MeV]") plt.ylabel("# events") plt.title( f" $e^-$, {self._particle_energy} [GeV], {self._particle_angle}$^{{\circ}}$, {self._geometry} " ) plt.savefig( f"{VALID_DIR}/E_tot_Geo_{self._geometry}_E_{self._particle_energy}_Angle_{self._particle_angle}.png" ) plt.clf() def _plot_cell_energy(self) -> None: """ Plots and saves a histogram with number of detector's cells across whole calorimeter with particular energy detected. Returns: None. """ full_simulation_cell_energy = self._full_simulation.calc_cell_energy() ml_simulation_cell_energy = self._ml_simulation.calc_cell_energy() log_full_simulation_cell_energy = np.log10( full_simulation_cell_energy, out=np.zeros_like(full_simulation_cell_energy), where=(full_simulation_cell_energy != 0)) log_ml_simulation_cell_energy = np.log10( ml_simulation_cell_energy, out=np.zeros_like(ml_simulation_cell_energy), where=(ml_simulation_cell_energy != 0)) plt.figure(figsize=(12, 8)) bins = np.linspace(-4, 1, 1000) plt.hist(x=log_full_simulation_cell_energy, bins=bins, histtype=HISTOGRAM_TYPE, label="FullSim", color=FULL_SIM_HISTOGRAM_COLOR) plt.hist(x=log_ml_simulation_cell_energy, bins=bins, histtype=HISTOGRAM_TYPE, label="MLSim", color=ML_SIM_HISTOGRAM_COLOR) plt.xlabel("log10(E/MeV)") plt.ylim(bottom=1) plt.yscale("log") plt.ylim(bottom=1) plt.ylabel("# entries") plt.title( f" $e^-$, {self._particle_energy} [GeV], {self._particle_angle}$^{{\circ}}$, {self._geometry} " ) plt.grid(True) plt.legend(loc="upper left") plt.savefig( f"{VALID_DIR}/E_cell_Geo_{self._geometry}_E_{self._particle_energy}_Angle_{self._particle_angle}.png" ) plt.clf() def _plot_energy_per_layer(self): """ Plots and saves N_CELLS_Z histograms with total energy detected in particular layers. Returns: None. """ full_simulation_energy_per_layer = self._full_simulation.calc_energy_per_layer( ) ml_simulation_energy_per_layer = self._ml_simulation.calc_energy_per_layer( ) number_of_plots_in_row = 9 number_of_plots_in_column = 5 bins = np.linspace(np.min(full_simulation_energy_per_layer - 10), np.max(full_simulation_energy_per_layer + 10), 25) fig, ax = plt.subplots(number_of_plots_in_column, number_of_plots_in_row, figsize=(20, 15), sharex="all", sharey="all", constrained_layout=True) for layer_nb in range(N_CELLS_Z): i = layer_nb // number_of_plots_in_row j = layer_nb % number_of_plots_in_row ax[i][j].hist(full_simulation_energy_per_layer[:, layer_nb], histtype=HISTOGRAM_TYPE, label="FullSim", bins=bins, color=FULL_SIM_HISTOGRAM_COLOR) ax[i][j].hist(ml_simulation_energy_per_layer[:, layer_nb], histtype=HISTOGRAM_TYPE, label="MLSim", bins=bins, color=ML_SIM_HISTOGRAM_COLOR) ax[i][j].set_title(f"Layer {layer_nb}", fontsize=13) ax[i][j].set_yscale("log") ax[i][j].tick_params(axis='both', which='major', labelsize=10) fig.supxlabel("Energy [MeV]", fontsize=14) fig.supylabel("# entries", fontsize=14) fig.suptitle( f" $e^-$, {self._particle_energy} [GeV], {self._particle_angle}$^{{\circ}}$, {self._geometry} " ) # Take legend from one plot and make it a global legend. handles, labels = ax[0][0].get_legend_handles_labels() fig.legend(handles, labels, bbox_to_anchor=(1.15, 0.5)) plt.savefig( f"{VALID_DIR}/E_layer_Geo_{self._geometry}_E_{self._particle_energy}_Angle_{self._particle_angle}.png", bbox_inches="tight") plt.clf() def plot_and_save(self): """ Main plotting function. Calls private methods and prints the information about progress. Returns: None. """ print("Plotting total energy...") self._plot_total_energy() print("Plotting cell energy...") self._plot_cell_energy() print("Plotting energy per layer...") self._plot_energy_per_layer()