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

1import numpy as np 

2 

3 

4class TimeAverager: 

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

6 obtained from sliding time-window averaging. 

7 

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

9 

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. 

12 

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 """ 

20 

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 

25 

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

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

28 

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 

33 

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] 

39 

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