Coverage for emd/tests/test_spectra.py: 100%

86 statements  

« prev     ^ index     » next       coverage.py v7.6.10, created at 2025-01-09 10:07 +0000

1"""Tests for instantaneous frequency and power spectra in emd.spectra.""" 

2 

3import unittest 

4 

5import numpy as np 

6 

7 

8class TestSpectra(unittest.TestCase): 

9 """Ensure basic frequency transforms are working.""" 

10 

11 def setUp(self): 

12 """Set up data for testing.""" 

13 # Create core signal 

14 seconds = 10 

15 self.sample_rate = 2000 

16 self.f1 = 5 

17 self.f2 = 18 

18 time_vect = np.linspace(0, seconds, int(seconds * self.sample_rate)) 

19 

20 self.x1 = np.cos(2 * np.pi * self.f1 * time_vect)[:, None] 

21 self.x2 = 2 * np.cos(2 * np.pi * self.f2 * time_vect)[:, None] 

22 

23 def test_frequency_transform(self): 

24 """Ensure basic frequency transforms are working.""" 

25 from ..spectra import frequency_transform 

26 

27 tol = 1e-3 # Relatively generous tol due to edge effects 

28 

29 # Check first signal 

30 IP, IF, IA = frequency_transform(self.x1, self.sample_rate, 'hilbert') 

31 assert(IP.max() - (2 * np.pi) < tol) 

32 assert(IP.min() < tol) 

33 assert(IA.mean() - 1 < tol) 

34 assert(IF.mean() - self.f1 < tol) 

35 

36 # Check second signal 

37 IP, IF, IA = frequency_transform(self.x2, self.sample_rate, 'hilbert') 

38 assert(IP.max() - (2 * np.pi) < tol) 

39 assert(IP.min() < tol) 

40 assert(IA.mean() - 2 < tol) 

41 assert(IF.mean() - self.f2 < tol) 

42 

43 def test_freq_from_phase(self): 

44 """Ensure we get correct instantaneous frequency from phase.""" 

45 from ..spectra import freq_from_phase 

46 

47 tst = freq_from_phase(np.linspace(0, 2 * np.pi, 48), 47) 

48 assert(np.allclose(tst, 1)) 

49 

50 tst = freq_from_phase(np.linspace(0, 2 * np.pi * .5, 48), 47) 

51 assert(np.allclose(tst, .5)) 

52 

53 tst = freq_from_phase(np.linspace(0, 2 * np.pi * 2, 48), 47) 

54 assert(np.allclose(tst, 2)) 

55 

56 def test_phase_from_freq(self): 

57 """Ensure we get correct instantaneous phase rom frequency.""" 

58 from ..spectra import phase_from_freq 

59 

60 tol = 1e-6 

61 

62 phs = phase_from_freq(np.ones((100,)), sample_rate=100) 

63 assert(phs.max() - np.pi < tol) 

64 

65 def test_hilberthunang_1d(self): 

66 """Ensure 1D Hilber-Huang Spectrum is working.""" 

67 from ..spectra import hilberthuang 

68 

69 IF = np.linspace(5, 15, 11)[:, None] 

70 IA = np.ones_like(IF) 

71 

72 # We should 2 bins with 5 frequencies in each bin, top IF value is dropped. 

73 edges = np.linspace(5, 15, 3) 

74 f, spec = hilberthuang(IF, IA, edges, mode='amplitude', sum_imfs=False) 

75 assert(np.all(spec[:, 0] == [5, 5])) 

76 

77 # We should 4 bins with 2 or 3 frequencies in each bin. 

78 edges = np.linspace(5, 15, 5) 

79 f, spec = hilberthuang(IF, IA, edges, mode='amplitude', sum_imfs=False) 

80 assert(np.all(spec[:, 0] == [3, 2, 3, 2])) 

81 

82 IA = IA * 2 

83 # We should 2 bins with 5frequencies in each bin, energy should be 20 

84 # per bin (5*(2**2)) 

85 edges = np.linspace(5, 15, 3) 

86 f, spec = hilberthuang(IF, IA, edges, mode='power', sum_imfs=False) 

87 assert(np.all(spec[:, 0] == [20, 20])) 

88 

89 def test_hilberthuang(self): 

90 """Ensure 2D Hilber-Huang Spectrum is working.""" 

91 from ..spectra import hilberthuang 

92 

93 IF = np.linspace(0, 12, 13)[:, None] 

94 IA = np.ones_like(IF) 

95 edges = np.linspace(0, 13, 3) 

96 

97 F, hht = hilberthuang(IF, IA, edges, sum_time=False) 

98 

99 # Check total amplitude is equal in HHT and IA 

100 assert(hht.sum() == IA.sum()) 

101 

102 assert(np.all(hht[0, :7] == np.array([1., 1., 1., 1., 1., 1., 1.]))) 

103 assert(np.all(hht[1, :7] == np.array([0., 0., 0., 0., 0., 0., 0.]))) 

104 assert(np.all(hht[1, 7:] == np.array([1., 1., 1., 1., 1., 1.]))) 

105 assert(np.all(hht[0, 7:] == np.array([0., 0., 0., 0., 0., 0.]))) 

106 

107 

108class TestHistograms(unittest.TestCase): 

109 """Ensure histogram binning is working.""" 

110 

111 def test_hist_bins_from_data(self): 

112 """Check we get the right bins from data.""" 

113 from ..spectra import define_hist_bins_from_data 

114 

115 data = np.linspace(0, 1, 16) 

116 edges, bins = define_hist_bins_from_data(data, tol=0) 

117 

118 assert(np.all(edges == np.array([0., .25, .5, .75, 1.]))) 

119 assert(np.all(bins == np.array([0.125, 0.375, 0.625, 0.875]))) 

120 

121 def test_hist_bins(self): 

122 """Check we get the right bins from user definition.""" 

123 from ..spectra import define_hist_bins 

124 

125 edges, bins = define_hist_bins(0, 1, 5) 

126 

127 edges = np.round(edges, 6) # Sometimes returns float errors 0.30000000000000004 

128 bins = np.round(bins, 6) # Sometimes returns float errors 0.30000000000000004 

129 

130 assert(np.all(edges == np.array([0., 0.2, 0.4, 0.6, 0.8, 1.]))) 

131 assert(np.all(bins == np.array([0.1, 0.3, 0.5, 0.7, 0.9]))) 

132 

133 

134class TestHolospectrum(unittest.TestCase): 

135 """Ensure Holospectrum is working.""" 

136 

137 def test_holo(self): 

138 """Ensure Holospectrum is working.""" 

139 from ..spectra import define_hist_bins, holospectrum 

140 

141 f_edges1, f_bins1 = define_hist_bins(0, 10, 5) 

142 f_edges2, f_bins2 = define_hist_bins(0, 1, 5) 

143 

144 if1 = np.array([2, 6])[:, None] 

145 if2 = np.array([.2, .3])[:, None, None] 

146 ia2 = np.array([1, 2])[:, None, None] 

147 

148 fcarrier, fam, holo = holospectrum(if1, if2, ia2, f_edges1, f_edges2, sum_time=False) 

149 

150 assert(np.all(holo.shape == (5, 5, 2)))