Coverage for C: \ Users \ peaco \ OneDrive \ Documents \ GitHub \ mt_metadata \ mt_metadata \ processing \ window.py: 100%

66 statements  

« prev     ^ index     » next       coverage.py v7.13.1, created at 2026-01-10 00:11 -0800

1""" 

2Updated 2025-01-02: kkappler, adding methods to generate taper values. In future this class 

3 can replace ApodizationWindow in aurora. 

4""" 

5 

6# ===================================================== 

7# Imports 

8# ===================================================== 

9from typing import Annotated 

10 

11import numpy as np 

12import pandas as pd 

13import scipy.signal as ssig 

14from pydantic import AliasChoices, computed_field, Field, field_validator, PrivateAttr 

15 

16from mt_metadata.base import MetadataBase 

17from mt_metadata.common.enumerations import StrEnumerationBase 

18from mt_metadata.common.mttime import MTime 

19 

20 

21# ===================================================== 

22class TypeEnum(StrEnumerationBase): 

23 boxcar = "boxcar" 

24 triang = "triang" 

25 blackman = "blackman" 

26 hamming = "hamming" 

27 hann = "hann" 

28 bartlett = "bartlett" 

29 flattop = "flattop" 

30 parzen = "parzen" 

31 bohman = "bohman" 

32 blackmanharris = "blackmanharris" 

33 nuttall = "nuttall" 

34 barthann = "barthann" 

35 kaiser = "kaiser" 

36 gaussian = "gaussian" 

37 general_gaussian = "general_gaussian" 

38 slepian = "slepian" 

39 chebwin = "chebwin" 

40 dpss = "dpss" 

41 

42 

43class ClockZeroTypeEnum(StrEnumerationBase): 

44 user_specified = "user specified" 

45 data_start = "data start" 

46 ignore = "ignore" 

47 

48 

49class Window(MetadataBase): 

50 _taper: np.ndarray | None = PrivateAttr(None) 

51 num_samples: Annotated[ 

52 int, 

53 Field( 

54 default=256, 

55 description="Number of samples in a single window", 

56 alias=None, 

57 json_schema_extra={ 

58 "units": "samples", 

59 "required": True, 

60 "examples": ["256"], 

61 }, 

62 ), 

63 ] 

64 

65 overlap: Annotated[ 

66 int, 

67 Field( 

68 default=32, 

69 description="Number of samples overlapped by adjacent windows", 

70 alias=None, 

71 json_schema_extra={ 

72 "units": "samples", 

73 "required": True, 

74 "examples": ["32"], 

75 }, 

76 ), 

77 ] 

78 

79 type: Annotated[ 

80 TypeEnum, 

81 Field( 

82 default=TypeEnum.boxcar, 

83 description="name of the window type", 

84 alias=None, 

85 json_schema_extra={ 

86 "units": None, 

87 "required": True, 

88 "examples": ["hamming"], 

89 }, 

90 ), 

91 ] 

92 

93 clock_zero_type: Annotated[ 

94 ClockZeroTypeEnum, 

95 Field( 

96 default=ClockZeroTypeEnum.ignore, 

97 description="how the clock-zero is specified", 

98 alias=None, 

99 json_schema_extra={ 

100 "units": None, 

101 "required": True, 

102 "examples": ["user specified"], 

103 }, 

104 ), 

105 ] 

106 

107 clock_zero: Annotated[ 

108 MTime | str | float | int | np.datetime64 | pd.Timestamp | None, 

109 Field( 

110 default_factory=lambda: MTime(time_stamp=None), 

111 description="Start date and time of the first data window", 

112 alias=None, 

113 json_schema_extra={ 

114 "units": None, 

115 "required": False, 

116 "examples": ["2020-02-01T09:23:45.453670+00:00"], 

117 }, 

118 ), 

119 ] 

120 

121 normalized: Annotated[ 

122 bool, 

123 Field( 

124 default=True, 

125 description="True if the window shall be normalized so the sum of the coefficients is 1", 

126 validation_alias=AliasChoices("normalised", "normalized"), 

127 json_schema_extra={ 

128 "units": None, 

129 "required": True, 

130 "examples": [False], 

131 }, 

132 ), 

133 ] 

134 

135 additional_args: Annotated[ 

136 dict, 

137 Field( 

138 default_factory=dict, 

139 description="Additional arguments for the window function", 

140 json_schema_extra={ 

141 "units": None, 

142 "required": False, 

143 "examples": [{"param": "value"}], 

144 }, 

145 ), 

146 ] 

147 

148 @field_validator("clock_zero", mode="before") 

149 @classmethod 

150 def validate_clock_zero( 

151 cls, field_value: MTime | float | int | np.datetime64 | pd.Timestamp | str 

152 ): 

153 return MTime(time_stamp=field_value) 

154 

155 @computed_field 

156 @property 

157 def num_samples_advance(self) -> int: 

158 return self.num_samples - self.overlap 

159 

160 def fft_harmonics(self, sample_rate: float) -> np.ndarray: 

161 """ 

162 Returns the frequencies for an fft.. 

163 :param sample_rate: 

164 :return: 

165 """ 

166 return get_fft_harmonics( 

167 samples_per_window=self.num_samples, sample_rate=sample_rate 

168 ) 

169 

170 def taper(self) -> np.ndarray: 

171 """ 

172 Get's the window coeffcients. via wrapper call to scipy.signal 

173 

174 Note: see scipy.signal.get_window for a description of what is expected in args[1:]. http://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.get_window.html 

175 

176 Returns 

177 ------- 

178 

179 """ 

180 if self._taper is None: 

181 # Repackaging the args so that scipy.signal.get_window() accepts all cases 

182 window_args = [v for k, v in self.additional_args.items()] 

183 window_args.insert(0, self.type) 

184 window_args = tuple(window_args) 

185 

186 taper = ssig.get_window(window_args, self.num_samples) 

187 

188 if self.normalized: 

189 taper /= np.sum(taper) 

190 

191 self._taper = taper 

192 

193 return self._taper 

194 

195 

196def get_fft_harmonics(samples_per_window: int, sample_rate: float) -> np.ndarray: 

197 """ 

198 Works for odd and even number of points. 

199 

200 Development notes: 

201 - Could be modified with arguments to support one_sided, two_sided, ignore_dc 

202 ignore_nyquist, and etc. Consider taking FrequencyBands as an argument. 

203 - This function was in decimation_level, but there were circular import issues. 

204 The function needs only a window length and sample rate, so putting it here for now. 

205 - TODO: switch to using np.fft.rfftfreq 

206 

207 Parameters 

208 ---------- 

209 samples_per_window: int 

210 Number of samples in a window that will be Fourier transformed. 

211 sample_rate: float 

212 Inverse of time step between samples; Samples per second in Hz. 

213 

214 Returns 

215 ------- 

216 harmonic_frequencies: numpy array 

217 The frequencies that the fft will be computed. 

218 These are one-sided (positive frequencies only) 

219 Does _not_ return Nyquist 

220 Does return DC component 

221 """ 

222 delta_t = 1.0 / sample_rate 

223 harmonic_frequencies = np.fft.fftfreq(samples_per_window, d=delta_t) 

224 n_fft_harmonics = int(samples_per_window / 2) # no bin at Nyquist, 

225 harmonic_frequencies = harmonic_frequencies[0:n_fft_harmonics] 

226 return harmonic_frequencies