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
« 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."""
3import unittest
5import numpy as np
8class TestSpectra(unittest.TestCase):
9 """Ensure basic frequency transforms are working."""
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))
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]
23 def test_frequency_transform(self):
24 """Ensure basic frequency transforms are working."""
25 from ..spectra import frequency_transform
27 tol = 1e-3 # Relatively generous tol due to edge effects
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)
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)
43 def test_freq_from_phase(self):
44 """Ensure we get correct instantaneous frequency from phase."""
45 from ..spectra import freq_from_phase
47 tst = freq_from_phase(np.linspace(0, 2 * np.pi, 48), 47)
48 assert(np.allclose(tst, 1))
50 tst = freq_from_phase(np.linspace(0, 2 * np.pi * .5, 48), 47)
51 assert(np.allclose(tst, .5))
53 tst = freq_from_phase(np.linspace(0, 2 * np.pi * 2, 48), 47)
54 assert(np.allclose(tst, 2))
56 def test_phase_from_freq(self):
57 """Ensure we get correct instantaneous phase rom frequency."""
58 from ..spectra import phase_from_freq
60 tol = 1e-6
62 phs = phase_from_freq(np.ones((100,)), sample_rate=100)
63 assert(phs.max() - np.pi < tol)
65 def test_hilberthunang_1d(self):
66 """Ensure 1D Hilber-Huang Spectrum is working."""
67 from ..spectra import hilberthuang
69 IF = np.linspace(5, 15, 11)[:, None]
70 IA = np.ones_like(IF)
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]))
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]))
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]))
89 def test_hilberthuang(self):
90 """Ensure 2D Hilber-Huang Spectrum is working."""
91 from ..spectra import hilberthuang
93 IF = np.linspace(0, 12, 13)[:, None]
94 IA = np.ones_like(IF)
95 edges = np.linspace(0, 13, 3)
97 F, hht = hilberthuang(IF, IA, edges, sum_time=False)
99 # Check total amplitude is equal in HHT and IA
100 assert(hht.sum() == IA.sum())
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.])))
108class TestHistograms(unittest.TestCase):
109 """Ensure histogram binning is working."""
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
115 data = np.linspace(0, 1, 16)
116 edges, bins = define_hist_bins_from_data(data, tol=0)
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])))
121 def test_hist_bins(self):
122 """Check we get the right bins from user definition."""
123 from ..spectra import define_hist_bins
125 edges, bins = define_hist_bins(0, 1, 5)
127 edges = np.round(edges, 6) # Sometimes returns float errors 0.30000000000000004
128 bins = np.round(bins, 6) # Sometimes returns float errors 0.30000000000000004
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])))
134class TestHolospectrum(unittest.TestCase):
135 """Ensure Holospectrum is working."""
137 def test_holo(self):
138 """Ensure Holospectrum is working."""
139 from ..spectra import define_hist_bins, holospectrum
141 f_edges1, f_bins1 = define_hist_bins(0, 10, 5)
142 f_edges2, f_bins2 = define_hist_bins(0, 1, 5)
144 if1 = np.array([2, 6])[:, None]
145 if2 = np.array([.2, .3])[:, None, None]
146 ia2 = np.array([1, 2])[:, None, None]
148 fcarrier, fam, holo = holospectrum(if1, if2, ia2, f_edges1, f_edges2, sum_time=False)
150 assert(np.all(holo.shape == (5, 5, 2)))