Coverage for local_installation/dynasor/tools/acfs.py: 97%

1""" 

2A number of utility functions, for example for dealing with 

3autocorrelation functions, Fourier transforms, and smoothing. 

4""" 

5 

6import numpy as np 

7from scipy.signal import correlate 

8from numpy.typing import NDArray 

9import pandas as pd 

10 

11 

12def compute_acf(Z: NDArray[float], delta_t: float = 1.0, method='scipy'): 

13 r""" 

14 Computes the autocorrelation function (ACF) for a one-dimensional signal :math:`Z` in time as 

15 

16 .. math:: 

17 

18 ACF(\tau) = \frac{\left < Z(t) Z^*(t+\tau) \right >}{\left < Z(t) Z^*(t) \right >} 

19 

20 Here, only the real part of the ACF is returned since if :math:`Z` is complex 

21 the imaginary part should average out to zero for any stationary signal. 

22 

23 Parameters 

24 ---------- 

25 Z 

26 complex time signal 

27 delta_t 

28 spacing in time between two consecutive values in :math:`Z` 

29 method 

30 implementation to use; possible values: `numpy` and `scipy` (default and usually faster) 

31 """ 

32 

33 # keep only real part and normalize 

34 acf = _compute_correlation_function(Z, Z, method) 

35 acf = np.real(acf) 

36 acf /= acf[0] 

37 

38 time_lags = delta_t * np.arange(0, len(acf), 1) 

39 return time_lags, acf 

40 

41 

42def _compute_correlation_function(Z1, Z2, method='scipy'): 

43 N = len(Z1) 

44 assert len(Z1) == len(Z2) 

45 if method == 'scipy': 

46 cf = correlate(Z1, Z2, mode='full')[N - 1:] / np.arange(N, 0, -1) 

47 elif method == 'numpy': 

48 cf = np.correlate(Z1, Z2, mode='full')[N - 1:] / np.arange(N, 0, -1) 

49 else: 

50 raise ValueError('method must be either numpy or scipy') 

51 return cf 

52 

53 

54# smoothing functions / FFT filters 

55# ------------------------------------- 

56def gaussian_decay(t: NDArray[float], t_sigma: float): 

57 r""" 

58 Evaluates a gaussian distribution in time :math:`f(t)`, which can be applied to an ACF in time 

59 to artificially damp it, i.e., forcing it to go to zero for long times. 

60 

61 .. math:: 

62 

63 f(t) = \exp{\left [-\frac{1}{2} \left (\frac{t}{t_\mathrm{sigma}}\right )^2 \right ] } 

64 

65 Parameters 

66 ---------- 

67 t 

68 time array 

69 t_sigma 

70 width (standard deviation of the gaussian) of the decay 

71 """ 

72 

73 return np.exp(- 1 / 2 * (t / t_sigma) ** 2) 

74 

75 

76def fermi_dirac(t: NDArray[float], t_0: float, t_width: float): 

77 r""" 

78 Evaluates a Fermi-Dirac-like function in time :math:`f(t)`, which can be applied to an ACF in 

79 time to artificially damp it, i.e., forcing it to go to zero for long times without affecting 

80 the short-time correlations too much. 

81 

82 .. math:: 

83 

84 f(t) = \frac{1}{\exp{[(t-t_0)/t_\mathrm{width}}] + 1} 

85 

86 Parameters 

87 ---------- 

88 t 

89 time array 

90 t_0 

91 starting time for decay 

92 t_width 

93 width of the decay 

94 

95 """ 

96 return 1.0 / (np.exp((t - t_0) / t_width) + 1) 

97 

98 

99def smoothing_function(data: NDArray[float], window_size: int, window_type: str = 'hamming'): 

100 """ 

101 Smoothing function for 1D arrays. 

102 This functions employs pandas rolling window average 

103 

104 Parameters 

105 ---------- 

106 data 

107 1D data array 

108 window_size 

109 The size of smoothing/smearing window 

110 window_type 

111 What type of window-shape to use, e.g. ``'blackman'``, ``'hamming'``, ``'boxcar'`` 

112 (see pandas and scipy documentaiton for more details) 

113 

114 """ 

115 series = pd.Series(data) 

116 new_data = series.rolling(window_size, win_type=window_type, center=True, min_periods=1).mean() 

117 return np.array(new_data)