training/utils/observables.py

0001 from dataclasses import dataclass
0002 from enum import Enum
0003
0004 import numpy as np
0005
0006 from core.constants import N_CELLS_Z, N_CELLS_R, SIZE_Z, SIZE_R
0007
0008
0009 @dataclass
0010 class Observable:
0011     """ An abstract class defining interface of all observables.
0012
0013     Do not use this class directly.
0014
0015     Attributes:
0016           _input: A numpy array with shape = (NE, R, PHI, Z), where NE stays for number of events.
0017     """
0018     _input: np.ndarray
0019
0020
0021 class ProfileType(Enum):
0022     """ Enum class of various profile types.
0023
0024     """
0025     LONGITUDINAL = 0
0026     LATERAL = 1
0027
0028
0029 @dataclass
0030 class Profile(Observable):
0031     """ An abstract class describing behaviour of LongitudinalProfile and LateralProfile.
0032
0033     Do not use this class directly. Use LongitudinalProfile or LateralProfile instead.
0034
0035     """
0036
0037     def calc_profile(self) -> np.ndarray:
0038         pass
0039
0040     def calc_first_moment(self) -> np.ndarray:
0041         pass
0042
0043     def calc_second_moment(self) -> np.ndarray:
0044         pass
0045
0046
0047 @dataclass
0048 class LongitudinalProfile(Profile):
0049     """ A class defining observables related to LongitudinalProfile.
0050
0051     Attributes:
0052         _energies_per_event: A numpy array with shape = (NE, Z) where NE stays for a number of events. An
0053             element [i, j] is a sum of energies detected in all cells located in a jth layer for an ith event.
0054         _total_energy_per_event: A numpy array with shape = (NE, ). An element [i] is a sum of energies detected in all
0055             cells for an ith event.
0056         _w: A numpy array = [0, 1, ..., Z - 1] which represents weights used in computation of first and second moment.
0057
0058     """
0059
0060     def __post_init__(self):
0061         self._energies_per_event = np.sum(self._input, axis=(1, 2))
0062         self._total_energy_per_event = np.sum(self._energies_per_event, axis=1)
0063         self._w = np.arange(N_CELLS_Z)
0064
0065     def calc_profile(self) -> np.ndarray:
0066         """ Calculates a longitudinal profile.
0067
0068         A longitudinal profile for a given layer l (l = 0, ..., Z - 1) is defined as:
0069         sum_{i = 0}^{NE - 1} energy_per_event[i, l].
0070
0071         Returns:
0072             A numpy array of longitudinal profiles for each layer with a shape = (Z, ).
0073
0074         """
0075         return np.sum(self._energies_per_event, axis=0)
0076
0077     def calc_first_moment(self) -> np.ndarray:
0078         """ Calculates a first moment of profile.
0079
0080         A first moment of a longitudinal profile for a given event e (e = 0, ..., NE - 1) is defined as:
0081         FM[e] = alpha * (sum_{i = 0}^{Z - 1} energies_per_event[e, i] * w[i]) / total_energy_per_event[e], where
0082         w = [0, 1, 2, ..., Z - 1],
0083         alpha = SIZE_Z defined in core/constants.py.
0084
0085         Returns:
0086             A numpy array of first moments of longitudinal profiles for each event with a shape = (NE, ).
0087
0088         """
0089         return SIZE_Z * np.dot(self._energies_per_event, self._w) / self._total_energy_per_event
0090
0091     def calc_second_moment(self) -> np.ndarray:
0092         """ Calculates a second moment of a longitudinal profile.
0093
0094         A second moment of a longitudinal profile for a given event e (e = 0, ..., NE - 1) is defined as:
0095         SM[e] = (sum_{i = 0}^{Z - 1} (w[i] - alpha - FM[e])^2 * energies_per_event[e, i]) total_energy_per_event[e],
0096         where
0097         w = [0, 1, 2, ..., Z - 1],
0098         alpha = SIZE_Z defined in ochre/constants.py
0099
0100         Returns:
0101             A numpy array of second moments of longitudinal profiles for each event with a shape = (NE, ).
0102         """
0103         first_moment = self.calc_first_moment()
0104         first_moment = np.expand_dims(first_moment, axis=1)
0105         w = np.expand_dims(self._w, axis=0)
0106         # w has now a shape = [1, Z] and first moment has a shape = [NE, 1]. There is a broadcasting in the line
0107         # below how that one create an array with a shape = [NE, Z]
0108         return np.sum(np.multiply(np.power(w * SIZE_Z - first_moment, 2), self._energies_per_event),
0109                       axis=1) / self._total_energy_per_event
0110
0111
0112 @dataclass
0113 class LateralProfile(Profile):
0114     """ A class defining observables related to LateralProfile.
0115
0116     Attributes:
0117         _energies_per_event: A numpy array with shape = (NE, R) where NE stays for a number of events. An
0118             element [i, j] is a sum of energies detected in all cells located in a jth layer for an ith event.
0119         _total_energy_per_event: A numpy array with shape = (NE, ). An element [i] is a sum of energies detected in all
0120             cells for an ith event.
0121         _w: A numpy array = [0, 1, ..., R - 1] which represents weights used in computation of first and second moment.
0122
0123     """
0124
0125     def __post_init__(self):
0126         self._energies_per_event = np.sum(self._input, axis=(2, 3))
0127         self._total_energy_per_event = np.sum(self._energies_per_event, axis=1)
0128         self._w = np.arange(N_CELLS_R)
0129
0130     def calc_profile(self) -> np.ndarray:
0131         """ Calculates a lateral profile.
0132
0133         A lateral profile for a given layer l (l = 0, ..., R - 1) is defined as:
0134         sum_{i = 0}^{NE - 1} energy_per_event[i, l].
0135
0136         Returns:
0137             A numpy array of longitudinal profiles for each layer with a shape = (R, ).
0138
0139         """
0140         return np.sum(self._energies_per_event, axis=0)
0141
0142     def calc_first_moment(self) -> np.ndarray:
0143         """ Calculates a first moment of profile.
0144
0145         A first moment of a lateral profile for a given event e (e = 0, ..., NE - 1) is defined as:
0146         FM[e] = alpha * (sum_{i = 0}^{R - 1} energies_per_event[e, i] * w[i]) / total_energy_per_event[e], where
0147         w = [0, 1, 2, ..., R - 1],
0148         alpha = SIZE_R defined in core/constants.py.
0149
0150         Returns:
0151             A numpy array of first moments of lateral profiles for each event with a shape = (NE, ).
0152
0153         """
0154         return SIZE_R * np.dot(self._energies_per_event, self._w) / self._total_energy_per_event
0155
0156     def calc_second_moment(self) -> np.ndarray:
0157         """ Calculates a second moment of a lateral profile.
0158
0159         A second moment of a lateral profile for a given event e (e = 0, ..., NE - 1) is defined as:
0160         SM[e] = (sum_{i = 0}^{R - 1} (w[i] - alpha - FM[e])^2 * energies_per_event[e, i]) total_energy_per_event[e],
0161         where
0162         w = [0, 1, 2, ..., R - 1],
0163         alpha = SIZE_R defined in ochre/constants.py
0164
0165         Returns:
0166             A numpy array of second moments of lateral profiles for each event with a shape = (NE, ).
0167         """
0168         first_moment = self.calc_first_moment()
0169         first_moment = np.expand_dims(first_moment, axis=1)
0170         w = np.expand_dims(self._w, axis=0)
0171         # w has now a shape = [1, R] and first moment has a shape = [NE, 1]. There is a broadcasting in the line
0172         # below how that one create an array with a shape = [NE, R]
0173         return np.sum(np.multiply(np.power(w * SIZE_R - first_moment, 2), self._energies_per_event),
0174                       axis=1) / self._total_energy_per_event
0175
0176
0177 @dataclass
0178 class Energy(Observable):
0179     """ A class defining observables total energy per event and cell energy.
0180
0181     """
0182
0183     def calc_total_energy(self):
0184         """ Calculates total energy detected in an event.
0185
0186         Total energy for a given event e (e = 0, ..., NE - 1) is defined as a sum of energies detected in all cells
0187         for this event.
0188
0189         Returns:
0190             A numpy array of total energy values with shape = (NE, ).
0191         """
0192         return np.sum(self._input, axis=(1, 2, 3))
0193
0194     def calc_cell_energy(self):
0195         """ Calculates cell energy.
0196
0197         Cell energy for a given event (e = 0, ..., NE - 1) is defined by an array with shape (R * PHI * Z) storing
0198         values of energy in particular cells.
0199
0200         Returns:
0201             A numpy array of cell energy values with shape = (NE * R * PHI * Z, ).
0202
0203         """
0204         return np.copy(self._input).reshape(-1)
0205
0206     def calc_energy_per_layer(self):
0207         """ Calculates total energy detected in a particular layer.
0208
0209         Energy per layer for a given event (e = 0, ..., NE - 1) is defined by an array with shape (Z, ) storing
0210         values of total energy detected in a particular layer
0211
0212         Returns:
0213             A numpy array of cell energy values with shape = (NE, Z).
0214
0215         """
0216         return np.sum(self._input, axis=(1, 2))