| Geant4 Cross Reference |
1 from dataclasses import dataclass
2 from enum import Enum
3
4 import numpy as np
5
6 from core.constants import N_CELLS_Z, N_CELLS_R, SIZE_Z, SIZE_R
7
8
9 @dataclass
10 class Observable:
11 """ An abstract class defining interface of all observables.
12
13 Do not use this class directly.
14
15 Attributes:
16 _input: A numpy array with shape = (NE, R, PHI, Z), where NE stays for number of events.
17 """
18 _input: np.ndarray
19
20
21 class ProfileType(Enum):
22 """ Enum class of various profile types.
23
24 """
25 LONGITUDINAL = 0
26 LATERAL = 1
27
28
29 @dataclass
30 class Profile(Observable):
31 """ An abstract class describing behaviour of LongitudinalProfile and LateralProfile.
32
33 Do not use this class directly. Use LongitudinalProfile or LateralProfile instead.
34
35 """
36
37 def calc_profile(self) -> np.ndarray:
38 pass
39
40 def calc_first_moment(self) -> np.ndarray:
41 pass
42
43 def calc_second_moment(self) -> np.ndarray:
44 pass
45
46
47 @dataclass
48 class LongitudinalProfile(Profile):
49 """ A class defining observables related to LongitudinalProfile.
50
51 Attributes:
52 _energies_per_event: A numpy array with shape = (NE, Z) where NE stays for a number of events. An
53 element [i, j] is a sum of energies detected in all cells located in a jth layer for an ith event.
54 _total_energy_per_event: A numpy array with shape = (NE, ). An element [i] is a sum of energies detected in all
55 cells for an ith event.
56 _w: A numpy array = [0, 1, ..., Z - 1] which represents weights used in computation of first and second moment.
57
58 """
59
60 def __post_init__(self):
61 self._energies_per_event = np.sum(self._input, axis=(1, 2))
62 self._total_energy_per_event = np.sum(self._energies_per_event, axis=1)
63 self._w = np.arange(N_CELLS_Z)
64
65 def calc_profile(self) -> np.ndarray:
66 """ Calculates a longitudinal profile.
67
68 A longitudinal profile for a given layer l (l = 0, ..., Z - 1) is defined as:
69 sum_{i = 0}^{NE - 1} energy_per_event[i, l].
70
71 Returns:
72 A numpy array of longitudinal profiles for each layer with a shape = (Z, ).
73
74 """
75 return np.sum(self._energies_per_event, axis=0)
76
77 def calc_first_moment(self) -> np.ndarray:
78 """ Calculates a first moment of profile.
79
80 A first moment of a longitudinal profile for a given event e (e = 0, ..., NE - 1) is defined as:
81 FM[e] = alpha * (sum_{i = 0}^{Z - 1} energies_per_event[e, i] * w[i]) / total_energy_per_event[e], where
82 w = [0, 1, 2, ..., Z - 1],
83 alpha = SIZE_Z defined in core/constants.py.
84
85 Returns:
86 A numpy array of first moments of longitudinal profiles for each event with a shape = (NE, ).
87
88 """
89 return SIZE_Z * np.dot(self._energies_per_event, self._w) / self._total_energy_per_event
90
91 def calc_second_moment(self) -> np.ndarray:
92 """ Calculates a second moment of a longitudinal profile.
93
94 A second moment of a longitudinal profile for a given event e (e = 0, ..., NE - 1) is defined as:
95 SM[e] = (sum_{i = 0}^{Z - 1} (w[i] - alpha - FM[e])^2 * energies_per_event[e, i]) total_energy_per_event[e],
96 where
97 w = [0, 1, 2, ..., Z - 1],
98 alpha = SIZE_Z defined in ochre/constants.py
99
100 Returns:
101 A numpy array of second moments of longitudinal profiles for each event with a shape = (NE, ).
102 """
103 first_moment = self.calc_first_moment()
104 first_moment = np.expand_dims(first_moment, axis=1)
105 w = np.expand_dims(self._w, axis=0)
106 # w has now a shape = [1, Z] and first moment has a shape = [NE, 1]. There is a broadcasting in the line
107 # below how that one create an array with a shape = [NE, Z]
108 return np.sum(np.multiply(np.power(w * SIZE_Z - first_moment, 2), self._energies_per_event),
109 axis=1) / self._total_energy_per_event
110
111
112 @dataclass
113 class LateralProfile(Profile):
114 """ A class defining observables related to LateralProfile.
115
116 Attributes:
117 _energies_per_event: A numpy array with shape = (NE, R) where NE stays for a number of events. An
118 element [i, j] is a sum of energies detected in all cells located in a jth layer for an ith event.
119 _total_energy_per_event: A numpy array with shape = (NE, ). An element [i] is a sum of energies detected in all
120 cells for an ith event.
121 _w: A numpy array = [0, 1, ..., R - 1] which represents weights used in computation of first and second moment.
122
123 """
124
125 def __post_init__(self):
126 self._energies_per_event = np.sum(self._input, axis=(2, 3))
127 self._total_energy_per_event = np.sum(self._energies_per_event, axis=1)
128 self._w = np.arange(N_CELLS_R)
129
130 def calc_profile(self) -> np.ndarray:
131 """ Calculates a lateral profile.
132
133 A lateral profile for a given layer l (l = 0, ..., R - 1) is defined as:
134 sum_{i = 0}^{NE - 1} energy_per_event[i, l].
135
136 Returns:
137 A numpy array of longitudinal profiles for each layer with a shape = (R, ).
138
139 """
140 return np.sum(self._energies_per_event, axis=0)
141
142 def calc_first_moment(self) -> np.ndarray:
143 """ Calculates a first moment of profile.
144
145 A first moment of a lateral profile for a given event e (e = 0, ..., NE - 1) is defined as:
146 FM[e] = alpha * (sum_{i = 0}^{R - 1} energies_per_event[e, i] * w[i]) / total_energy_per_event[e], where
147 w = [0, 1, 2, ..., R - 1],
148 alpha = SIZE_R defined in core/constants.py.
149
150 Returns:
151 A numpy array of first moments of lateral profiles for each event with a shape = (NE, ).
152
153 """
154 return SIZE_R * np.dot(self._energies_per_event, self._w) / self._total_energy_per_event
155
156 def calc_second_moment(self) -> np.ndarray:
157 """ Calculates a second moment of a lateral profile.
158
159 A second moment of a lateral profile for a given event e (e = 0, ..., NE - 1) is defined as:
160 SM[e] = (sum_{i = 0}^{R - 1} (w[i] - alpha - FM[e])^2 * energies_per_event[e, i]) total_energy_per_event[e],
161 where
162 w = [0, 1, 2, ..., R - 1],
163 alpha = SIZE_R defined in ochre/constants.py
164
165 Returns:
166 A numpy array of second moments of lateral profiles for each event with a shape = (NE, ).
167 """
168 first_moment = self.calc_first_moment()
169 first_moment = np.expand_dims(first_moment, axis=1)
170 w = np.expand_dims(self._w, axis=0)
171 # w has now a shape = [1, R] and first moment has a shape = [NE, 1]. There is a broadcasting in the line
172 # below how that one create an array with a shape = [NE, R]
173 return np.sum(np.multiply(np.power(w * SIZE_R - first_moment, 2), self._energies_per_event),
174 axis=1) / self._total_energy_per_event
175
176
177 @dataclass
178 class Energy(Observable):
179 """ A class defining observables total energy per event and cell energy.
180
181 """
182
183 def calc_total_energy(self):
184 """ Calculates total energy detected in an event.
185
186 Total energy for a given event e (e = 0, ..., NE - 1) is defined as a sum of energies detected in all cells
187 for this event.
188
189 Returns:
190 A numpy array of total energy values with shape = (NE, ).
191 """
192 return np.sum(self._input, axis=(1, 2, 3))
193
194 def calc_cell_energy(self):
195 """ Calculates cell energy.
196
197 Cell energy for a given event (e = 0, ..., NE - 1) is defined by an array with shape (R * PHI * Z) storing
198 values of energy in particular cells.
199
200 Returns:
201 A numpy array of cell energy values with shape = (NE * R * PHI * Z, ).
202
203 """
204 return np.copy(self._input).reshape(-1)
205
206 def calc_energy_per_layer(self):
207 """ Calculates total energy detected in a particular layer.
208
209 Energy per layer for a given event (e = 0, ..., NE - 1) is defined by an array with shape (Z, ) storing
210 values of total energy detected in a particular layer
211
212 Returns:
213 A numpy array of cell energy values with shape = (NE, Z).
214
215 """
216 return np.sum(self._input, axis=(1, 2))