Coverage for C: \ Users \ peaco \ OneDrive \ Documents \ GitHub \ mth5 \ mth5 \ processing \ spectre \ prewhitening.py: 61%

1""" 

2 This module has methods for pre-whitening time series to reduce spectral leakage before FFT. 

3""" 

4from typing import Literal 

5 

6import numpy as np 

7import xarray as xr 

8from loguru import logger 

9 

10 

11def apply_prewhitening( 

12 prewhitening_type: Literal["first difference", ""], 

13 run_xrds_input: xr.Dataset, 

14) -> xr.Dataset: 

15 """ 

16 Applies pre-whitening to time series to avoid spectral leakage when FFT is applied. 

17 

18 Parameters 

19 ---------- 

20 prewhitening_type: Literal["first difference", ] 

21 Placeholder to allow for multiple methods of pre-whitening. Currently only 

22 "first difference" is supported. 

23 run_xrds_input : xr.Dataset 

24 Time series to be pre-whitened (can be multivariate). 

25 

26 Returns 

27 ------- 

28 run_xrds : xr..Dataset 

29 pre-whitened time series 

30 

31 """ 

32 if not prewhitening_type: 

33 msg = "No prewhitening specified - skipping this step" 

34 logger.info(msg) 

35 return run_xrds_input 

36 

37 if prewhitening_type == "first difference": 

38 run_xrds = run_xrds_input.differentiate("time") 

39 else: 

40 msg = f"{prewhitening_type} pre-whitening not implemented" 

41 logger.exception(msg) 

42 raise NotImplementedError(msg) 

43 return run_xrds 

44 

45 

46def apply_recoloring( 

47 prewhitening_type: Literal["first difference",], 

48 stft_obj: xr.Dataset, 

49) -> xr.Dataset: 

50 """ 

51 Inverts the pre-whitening operation in frequency domain. Modifies the input xarray in-place. 

52 

53 Parameters 

54 ---------- 

55 prewhitening_type: Literal["first difference", ] 

56 Placeholder to allow for multiple methods of pre-whitening. Currently only 

57 "first difference" is supported. 

58 stft_obj : xarray.core.dataset.Dataset 

59 Time series of Fourier coefficients to be recoloured 

60 

61 

62 Returns 

63 ------- 

64 stft_obj : xarray.core.dataset.Dataset 

65 Recolored time series of Fourier coefficients. 

66 """ 

67 

68 # No recoloring needed if prewhitening not appiled, or recoloring set to False 

69 if not prewhitening_type: 

70 msg = "No prewhitening was defined -- skipping recoloring" 

71 logger.debug(msg) 

72 return stft_obj 

73 

74 if prewhitening_type == "first difference": 

75 freqs = stft_obj.frequency.data 

76 jw = 1.0j * 2 * np.pi * freqs 

77 stft_obj /= jw 

78 

79 # suppress nan and inf to mute later warnings 

80 if jw[0] == 0.0: 

81 cond = stft_obj.frequency != 0.0 

82 stft_obj = stft_obj.where(cond, complex(0.0)) 

83 # elif prewhitening_type == "ARMA": 

84 # from statsmodels.tsa.arima.model import ARIMA 

85 # AR = 3 # TODO: add this to processing config 

86 # MA = 4 # TODO: add this to processing config 

87 

88 else: 

89 msg = f"{prewhitening_type} recoloring not yet implemented" 

90 logger.error(msg) 

91 raise NotImplementedError(msg) 

92 

93 return stft_obj