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

1""" 

2This module has methods for applying the short-time-Fourier-transform. 

3 

4 

5""" 

6 

7from typing import Union 

8 

9import numpy as np 

10import scipy.signal as ssig 

11import xarray as xr 

12from mt_metadata.processing.aurora.decimation_level import ( 

13 DecimationLevel as AuroraDecimationLevel, 

14) 

15from mt_metadata.processing.fourier_coefficients import Decimation as FCDecimation 

16 

17from mth5.timeseries.spectre.spectrogram import Spectrogram 

18 

19from .prewhitening import apply_prewhitening, apply_recoloring 

20 

21 

22def run_ts_to_stft_scipy( 

23 decimation_obj: Union[AuroraDecimationLevel, FCDecimation], 

24 run_xrds_orig: xr.Dataset, 

25) -> Spectrogram: 

26 """ 

27 Converts a runts object into a time series of Fourier coefficients. 

28 This method uses scipy.signal.spectrogram. 

29 

30 TODO: consider making this a method of RunTS; runts.to_spectrogram(decimation_obj) 

31 

32 Parameters 

33 ---------- 

34 decimation_obj : Union[AuroraDecimationLevel, FCDecimation] 

35 Information about how the decimation level is to be processed 

36 run_xrds_orig : : xarray.core.dataset.Dataset 

37 Time series to be processed 

38 

39 Returns 

40 ------- 

41 stft_obj : xarray.core.dataset.Dataset 

42 Time series of Fourier coefficients 

43 """ 

44 run_xrds = apply_prewhitening(decimation_obj.stft.prewhitening_type, run_xrds_orig) 

45 

46 stft_obj = xr.Dataset() 

47 for channel_id in run_xrds.data_vars: 

48 ff, tt, specgm = ssig.spectrogram( 

49 run_xrds[channel_id].data, 

50 fs=decimation_obj.decimation.sample_rate, 

51 window=decimation_obj.stft.window.taper(), 

52 nperseg=decimation_obj.stft.window.num_samples, 

53 noverlap=decimation_obj.stft.window.overlap, 

54 detrend="linear", 

55 scaling="density", 

56 mode="complex", 

57 ) 

58 

59 # drop Nyquist 

60 ff = ff[:-1] 

61 specgm = specgm[:-1, :] 

62 specgm *= np.sqrt( 

63 2 

64 ) # compensate energy for keeping only positive harmonics (keep PSDs accurate) 

65 

66 # make time_axis 

67 tt = tt - tt[0] 

68 tt *= decimation_obj.decimation.sample_rate 

69 time_axis = run_xrds.time.data[tt.astype(int)] 

70 

71 xrd = xr.DataArray( 

72 specgm.T, 

73 dims=["time", "frequency"], 

74 coords={"frequency": ff, "time": time_axis}, 

75 ) 

76 stft_obj.update({channel_id: xrd}) 

77 

78 if decimation_obj.stft.recoloring: 

79 stft_obj = apply_recoloring(decimation_obj.stft.prewhitening_type, stft_obj) 

80 

81 return Spectrogram(dataset=stft_obj)