Coverage for src / tracekit / utils / windowing.py: 100%

1"""Window function support for spectral analysis. 

2 

3This module provides standard window functions for FFT and spectral 

4analysis, implementing the requirements for windowed spectral estimation. 

5 

6 

7Example: 

8 >>> from tracekit.utils.windowing import get_window, WINDOW_FUNCTIONS 

9 >>> window = get_window("hann", 1024) 

10 >>> print(f"Available windows: {list(WINDOW_FUNCTIONS.keys())}") 

11 

12References: 

13 Harris, F. J. (1978). "On the use of windows for harmonic analysis 

14 with the discrete Fourier transform." Proceedings of the IEEE, 66(1). 

15""" 

16 

17from __future__ import annotations 

18 

19from collections.abc import Callable 

20from typing import TYPE_CHECKING, Any, Literal 

21 

22import numpy as np 

23 

24if TYPE_CHECKING: 

25 from numpy.typing import NDArray 

26 

27# Type alias for window function (using string annotation for TYPE_CHECKING compatibility) 

28WindowFunction = Callable[[int], "NDArray[np.float64]"] 

29 

30 

31def rectangular(n: int) -> NDArray[np.float64]: 

32 """Rectangular (boxcar) window. 

33 

34 No tapering applied - all samples weighted equally. 

35 

36 Args: 

37 n: Window length in samples. 

38 

39 Returns: 

40 Window coefficients (all ones). 

41 

42 Example: 

43 >>> w = rectangular(64) 

44 >>> assert np.all(w == 1.0) 

45 """ 

46 return np.ones(n, dtype=np.float64) 

47 

48 

49def hann(n: int) -> NDArray[np.float64]: 

50 """Hann (raised cosine) window. 

51 

52 Also known as Hanning window. Provides good frequency resolution 

53 with moderate sidelobe suppression. 

54 

55 Args: 

56 n: Window length in samples. 

57 

58 Returns: 

59 Window coefficients. 

60 

61 Example: 

62 >>> w = hann(64) 

63 >>> assert w[0] == w[-1] # Symmetric 

64 

65 References: 

66 IEEE Std 1057-2017 Section 4.4.2 

67 """ 

68 return np.hanning(n).astype(np.float64) 

69 

70 

71def hamming(n: int) -> NDArray[np.float64]: 

72 """Hamming window. 

73 

74 Similar to Hann but with reduced first sidelobe at cost of 

75 slower rolloff. 

76 

77 Args: 

78 n: Window length in samples. 

79 

80 Returns: 

81 Window coefficients. 

82 

83 Example: 

84 >>> w = hamming(64) 

85 >>> assert w[32] > w[0] # Peak in center 

86 """ 

87 return np.hamming(n).astype(np.float64) 

88 

89 

90def blackman(n: int) -> NDArray[np.float64]: 

91 """Blackman window. 

92 

93 Three-term cosine window with excellent sidelobe suppression 

94 (-58 dB first sidelobe). 

95 

96 Args: 

97 n: Window length in samples. 

98 

99 Returns: 

100 Window coefficients. 

101 

102 Example: 

103 >>> w = blackman(64) 

104 >>> assert w[32] > w[0] 

105 """ 

106 return np.blackman(n).astype(np.float64) 

107 

108 

109def kaiser(n: int, beta: float = 8.6) -> NDArray[np.float64]: 

110 """Kaiser window with configurable shape parameter. 

111 

112 Provides adjustable tradeoff between main lobe width and 

113 sidelobe attenuation. 

114 

115 Args: 

116 n: Window length in samples. 

117 beta: Shape parameter (default 8.6 for ~60 dB sidelobe attenuation). 

118 - beta=0: Rectangular 

119 - beta=5: ~30 dB sidelobe attenuation 

120 - beta=8.6: ~60 dB sidelobe attenuation 

121 - beta=14: ~90 dB sidelobe attenuation 

122 

123 Returns: 

124 Window coefficients. 

125 

126 Example: 

127 >>> w = kaiser(64, beta=10) 

128 >>> assert 0 < w[0] < w[32] 

129 """ 

130 return np.kaiser(n, beta).astype(np.float64) 

131 

132 

133def flattop(n: int) -> NDArray[np.float64]: 

134 """Flat-top window for accurate amplitude measurements. 

135 

136 Provides minimal scalloping loss (<0.01 dB) at cost of 

137 wider main lobe. Best for amplitude accuracy when frequency 

138 resolution is not critical. 

139 

140 Args: 

141 n: Window length in samples. 

142 

143 Returns: 

144 Window coefficients. 

145 

146 Example: 

147 >>> w = flattop(64) 

148 >>> # Flat-top has characteristic near-zero values at edges 

149 

150 References: 

151 D'Antona, G. & Ferrero, A. (2006). "Digital Signal Processing 

152 for Measurement Systems." 

153 """ 

154 # Flat-top coefficients per HP/Agilent standard 

155 a0 = 0.21557895 

156 a1 = 0.41663158 

157 a2 = 0.277263158 

158 a3 = 0.083578947 

159 a4 = 0.006947368 

160 

161 k = np.arange(n, dtype=np.float64) 

162 w = ( 

163 a0 

164 - a1 * np.cos(2 * np.pi * k / (n - 1)) 

165 + a2 * np.cos(4 * np.pi * k / (n - 1)) 

166 - a3 * np.cos(6 * np.pi * k / (n - 1)) 

167 + a4 * np.cos(8 * np.pi * k / (n - 1)) 

168 ) 

169 return np.asarray(w, dtype=np.float64) 

170 

171 

172def bartlett(n: int) -> NDArray[np.float64]: 

173 """Bartlett (triangular) window. 

174 

175 Linear taper from zero at edges to maximum at center. 

176 

177 Args: 

178 n: Window length in samples. 

179 

180 Returns: 

181 Window coefficients. 

182 

183 Example: 

184 >>> w = bartlett(64) 

185 >>> assert w[32] == 1.0 # Maximum at center 

186 """ 

187 return np.bartlett(n).astype(np.float64) 

188 

189 

190def blackman_harris(n: int) -> NDArray[np.float64]: 

191 """Blackman-Harris window (4-term). 

192 

193 Four-term cosine window with excellent sidelobe suppression 

194 (-92 dB first sidelobe). 

195 

196 Args: 

197 n: Window length in samples. 

198 

199 Returns: 

200 Window coefficients. 

201 

202 Example: 

203 >>> w = blackman_harris(64) 

204 >>> assert w[32] > w[0] 

205 """ 

206 a0 = 0.35875 

207 a1 = 0.48829 

208 a2 = 0.14128 

209 a3 = 0.01168 

210 

211 k = np.arange(n, dtype=np.float64) 

212 w = ( 

213 a0 

214 - a1 * np.cos(2 * np.pi * k / (n - 1)) 

215 + a2 * np.cos(4 * np.pi * k / (n - 1)) 

216 - a3 * np.cos(6 * np.pi * k / (n - 1)) 

217 ) 

218 return np.asarray(w, dtype=np.float64) 

219 

220 

221# Window function registry 

222WINDOW_FUNCTIONS: dict[str, WindowFunction] = { 

223 "rectangular": rectangular, 

224 "boxcar": rectangular, 

225 "rect": rectangular, 

226 "hann": hann, 

227 "hanning": hann, 

228 "hamming": hamming, 

229 "blackman": blackman, 

230 "kaiser": lambda n: kaiser(n, beta=8.6), 

231 "flattop": flattop, 

232 "flat_top": flattop, 

233 "bartlett": bartlett, 

234 "triangular": bartlett, 

235 "blackman_harris": blackman_harris, 

236 "blackmanharris": blackman_harris, 

237} 

238 

239 

240# Type for window names 

241WindowName = Literal[ 

242 "rectangular", 

243 "boxcar", 

244 "rect", 

245 "hann", 

246 "hanning", 

247 "hamming", 

248 "blackman", 

249 "kaiser", 

250 "flattop", 

251 "flat_top", 

252 "bartlett", 

253 "triangular", 

254 "blackman_harris", 

255 "blackmanharris", 

256] 

257 

258 

259def get_window( 

260 window: str | WindowFunction | NDArray[np.floating[Any]], 

261 n: int, 

262 *, 

263 beta: float | None = None, 

264) -> NDArray[np.float64]: 

265 """Get window coefficients by name or callable. 

266 

267 Args: 

268 window: Window specification. Can be: 

269 - A string name from WINDOW_FUNCTIONS 

270 - A callable that takes length and returns coefficients 

271 - A pre-computed array of coefficients 

272 n: Window length in samples. 

273 beta: Optional beta parameter for Kaiser window. 

274 

275 Returns: 

276 Window coefficients array of length n. 

277 

278 Raises: 

279 ValueError: If window name is unknown. 

280 

281 Example: 

282 >>> w = get_window("hann", 1024) 

283 >>> w = get_window("kaiser", 1024, beta=10) 

284 >>> w = get_window(np.hamming, 1024) 

285 """ 

286 if isinstance(window, np.ndarray): 

287 if len(window) != n: 

288 raise ValueError(f"Window array length {len(window)} != requested {n}") 

289 return window.astype(np.float64) 

290 

291 if callable(window) and not isinstance(window, str): 

292 return np.asarray(window(n), dtype=np.float64) 

293 

294 window_name = window.lower() 

295 

296 if window_name == "kaiser" and beta is not None: 

297 return kaiser(n, beta) 

298 

299 if window_name not in WINDOW_FUNCTIONS: 

300 available = ", ".join(sorted(set(WINDOW_FUNCTIONS.keys()))) 

301 raise ValueError(f"Unknown window: {window}. Available: {available}") 

302 

303 return WINDOW_FUNCTIONS[window_name](n) 

304 

305 

306def window_properties(window: str | NDArray[np.floating[Any]], n: int = 1024) -> dict[str, Any]: 

307 """Compute window properties for analysis. 

308 

309 Args: 

310 window: Window name or coefficients. 

311 n: Window length for named windows. 

312 

313 Returns: 

314 Dictionary with window properties: 

315 - coherent_gain: Sum of window / length 

316 - noise_bandwidth: Normalized equivalent noise bandwidth 

317 - scalloping_loss: Peak amplitude error in dB 

318 

319 Example: 

320 >>> props = window_properties("hann") 

321 >>> print(f"ENBW: {props['noise_bandwidth']:.3f}") 

322 """ 

323 if isinstance(window, str): 

324 w = get_window(window, n) 

325 else: 

326 w = np.asarray(window, dtype=np.float64) 

327 n = len(w) 

328 

329 # Coherent gain (DC gain) 

330 coherent_gain = np.sum(w) / n 

331 

332 # Noise equivalent bandwidth 

333 # ENBW = N * sum(w^2) / (sum(w))^2 

334 noise_bandwidth = n * np.sum(w**2) / np.sum(w) ** 2 

335 

336 # Scalloping loss (worst-case amplitude error at bin edge) 

337 # Approximate by evaluating window at half-bin offset 

338 k = np.arange(n) 

339 w_shifted = w * np.exp(2j * np.pi * 0.5 * k / n) 

340 scalloping_loss = 20 * np.log10(np.abs(np.sum(w_shifted)) / np.abs(np.sum(w))) 

341 

342 return { 

343 "coherent_gain": float(coherent_gain), 

344 "noise_bandwidth": float(noise_bandwidth), 

345 "scalloping_loss": float(scalloping_loss), 

346 "length": n, 

347 } 

348 

349 

350__all__ = [ 

351 "WINDOW_FUNCTIONS", 

352 "bartlett", 

353 "blackman", 

354 "blackman_harris", 

355 "flattop", 

356 "get_window", 

357 "hamming", 

358 "hann", 

359 "kaiser", 

360 "rectangular", 

361 "window_properties", 

362]