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

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 psd_from_acf(acf, dt=1, angular=True, even=True): 

13 """Computes the power spectral density from auto-correlation function 

14 

15 Parameters 

16 ---------- 

17 acf 

18 The acf as an array 

19 dt 

20 The timestep between samples 

21 angular 

22 Whether to return normal or angular freqs 

23 even 

24 Whether to mirror the acf and force psd to be purely real 

25 """ 

26 assert acf.ndim == 1 

27 signal = np.array(acf) 

28 if even: 

29 signal = np.hstack((signal, signal[:0:-1])) 

30 fft = np.fft.fft(signal) 

31 if even: 

32 assert np.allclose(fft.imag, 0) 

33 fft = fft.real if even else fft 

34 

35 freqs = np.fft.fftfreq(len(signal), dt) 

36 freqs = 2*np.pi * freqs if angular else freqs 

37 

38 return freqs, fft 

39 

40 

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

42 r""" 

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

44 

45 .. math:: 

46 

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

48 

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

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

51 

52 Parameters 

53 ---------- 

54 Z 

55 complex time signal 

56 delta_t 

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

58 method 

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

60 """ 

61 

62 # keep only real part and normalize 

63 acf = _compute_correlation_function(Z, Z, method) 

64 acf = np.real(acf) 

65 acf /= acf[0] 

66 

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

68 return time_lags, acf 

69 

70 

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

72 N = len(Z1) 

73 assert len(Z1) == len(Z2) 

74 if method == 'scipy': 

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

76 elif method == 'numpy': 

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

78 else: 

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

80 return cf 

81 

82 

83# smoothing functions / FFT filters 

84# ------------------------------------- 

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

86 r""" 

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

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

89 

90 .. math:: 

91 

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

93 

94 Parameters 

95 ---------- 

96 t 

97 time array 

98 t_sigma 

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

100 """ 

101 

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

103 

104 

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

106 r""" 

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

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

109 the short-time correlations too much. 

110 

111 .. math:: 

112 

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

114 

115 Parameters 

116 ---------- 

117 t 

118 time array 

119 t_0 

120 starting time for decay 

121 t_width 

122 width of the decay 

123 

124 """ 

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

126 

127 

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

129 """ 

130 Smoothing function for 1D arrays. 

131 This functions employs pandas rolling window average 

132 

133 Parameters 

134 ---------- 

135 data 

136 1D data array 

137 window_size 

138 The size of smoothing/smearing window 

139 window_type 

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

141 (see pandas and scipy documentaiton for more details) 

142 

143 """ 

144 series = pd.Series(data) 

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

146 return np.array(new_data)