Coverage for dynasor / core / time_averager.py: 100%

1import numpy as np 

2from numpy.typing import NDArray 

3 

4 

5class TimeAverager: 

6 """Naive special purpose averager class used in dynasor to collect and time-average arrays 

7 obtained from sliding time-window averaging. 

8 

9 It assists with keeping track of how many data samples have been added to each slot. 

10 

11 It will time-average arrays of shape ``(Nq, time_window)`` where ``Ǹq`` is the 

12 number of q-points and ``time_window`` is the size of the time window. 

13 

14 Parameters 

15 ---------- 

16 time_window 

17 Size of the time window in which the time-average happens. 

18 array_length 

19 Length of the array to be averaged for each time-lag, i.e., number of q-points. 

20 """ 

21 

22 def __init__(self, time_window: int, array_length: int): 

23 assert time_window >= 1 

24 self._time_window = time_window 

25 self._array_length = array_length 

26 

27 self._counts = np.zeros(time_window, dtype=int) 

28 self._arrays = [np.zeros(array_length) for _ in range(time_window)] 

29 

30 def add_sample(self, time_lag: int, sample: np.ndarray): 

31 assert len(sample) == self._array_length 

32 self._counts[time_lag] += 1 

33 self._arrays[time_lag] += sample 

34 

35 def get_average_at_timelag(self, time_lag: int): 

36 if self._counts[time_lag] == 0: 

37 array = np.full((self._array_length, ), np.nan) 

38 return array 

39 return self._arrays[time_lag] / self._counts[time_lag] 

40 

41 def get_average_all(self) -> NDArray[float]: 

42 """ 

43 Returns an averaged array of shape ``(array_length, time_window)``. 

44 """ 

45 return np.array([self.get_average_at_timelag(t) for t in range(self._time_window)]).T