Coverage for local_installation/dynasor/core/time_averager.py: 100%
19 statements
« prev ^ index » next coverage.py v7.3.2, created at 2024-08-05 09:53 +0000
« prev ^ index » next coverage.py v7.3.2, created at 2024-08-05 09:53 +0000
1import numpy as np
4class TimeAverager:
5 """Naive special purpose averager class used in dynasor to collect and time-average arrays
6 obtained from sliding time-window averaging.
8 It assists with keeping track of how many data samples have been added to each slot.
10 It will time-average arrays of shape ``(Nq, time_window)`` where ``Ǹq`` is the
11 number of q-points and ``time_window`` is the size of the time window.
13 Parameters
14 ----------
15 time_window
16 size of the time window in which the time-average happens
17 array_length
18 length of the array to be averaged for each time-lag, i.e. number of q-points
19 """
21 def __init__(self, time_window: int, array_length: int):
22 assert time_window >= 1
23 self._time_window = time_window
24 self._array_length = array_length
26 self._counts = np.zeros(time_window, dtype=int)
27 self._arrays = [np.zeros(array_length) for _ in range(time_window)]
29 def add_sample(self, time_lag: int, sample: np.ndarray):
30 assert len(sample) == self._array_length
31 self._counts[time_lag] += 1
32 self._arrays[time_lag] += sample
34 def get_average_at_timelag(self, time_lag: int):
35 if self._counts[time_lag] == 0:
36 array = np.full((self._array_length, ), np.nan)
37 return array
38 return self._arrays[time_lag] / self._counts[time_lag]
40 def get_average_all(self):
41 """
42 Returns an averaged array of shape ``(array_length, time_window)``
43 """
44 return np.array([self.get_average_at_timelag(t) for t in range(self._time_window)]).T