Skip to content

Commit 85ca77b

Browse files
committed
whenxuan: create the unittest for stnbmd
1 parent b894707 commit 85ca77b

2 files changed

Lines changed: 249 additions & 0 deletions

File tree

‎pysdkit/__init__.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,9 @@
109109
# Adaptive Variational Nonlinear Chirp Mode Decomposition
110110
from ._vncmd import AVNCMD
111111

112+
# Short-Time Narrow-Banded Mode Decomposition
113+
from ._vncmd import STNBMD, stnbmd
114+
112115
# Adaptive Local Iterative Filtering (time iterative decomposition)
113116
from ._tid import ALIF
114117

@@ -208,6 +211,7 @@ def print_functions():
208211
Variational Nonlinear Chirp Mode Decomposition | VNCMD
209212
Iterative Nonlinear Chirp Mode Decomposition | INCMD
210213
Adaptive Variational Nonlinear Chirp Mode Dec. | AVNCMD
214+
Short-Time Narrow-Banded Mode Decomposition | STNBMD
211215
Adaptive Local Iterative Filtering | ALIF
212216
Variational TF Mode Tracking Decomposition | VTFMTD
213217
Adaptive Polymorphic Mode Decomposition | APMD
@@ -276,6 +280,8 @@ def print_functions():
276280
"VNCMD",
277281
"INCMD",
278282
"AVNCMD",
283+
"STNBMD",
284+
"stnbmd",
279285
"ALIF",
280286
"VTFMTD",
281287
"vtfmtd",

‎pysdkit/tests/test_stnbmd.py‎

Lines changed: 243 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,243 @@
1+
# -*- coding: utf-8 -*-
2+
"""
3+
Automated tests for Short-Time Narrow-Banded Mode Decomposition (STNBMD).
4+
"""
5+
6+
from __future__ import annotations
7+
8+
import unittest
9+
10+
import numpy as np
11+
12+
from pysdkit import STNBMD, stnbmd
13+
from pysdkit._vncmd.stnbmd import (
14+
ps90,
15+
ps90f,
16+
analytic_signal,
17+
first_difference_matrix,
18+
second_difference_matrix,
19+
build_smoothing_filters,
20+
schedule_index,
21+
fft_two_to_one,
22+
instantaneous_frequency,
23+
make_order_tracking_demo,
24+
constant_frequency_init,
25+
stnbm_decomp_ig,
26+
)
27+
28+
29+
class HilbertHelperTest(unittest.TestCase):
30+
def test_ps90f_shape_even_odd(self) -> None:
31+
x = np.random.randn(64, 2)
32+
y = ps90f(x)
33+
self.assertEqual(y.shape, x.shape)
34+
y2 = ps90f(x[:, 0])
35+
self.assertEqual(y2.shape, (64,))
36+
37+
x_odd = np.random.randn(63)
38+
y_odd = ps90f(x_odd)
39+
self.assertEqual(y_odd.shape, (63,))
40+
41+
def test_ps90f_custom_n(self) -> None:
42+
x = np.random.randn(50)
43+
y = ps90f(x, n=64)
44+
self.assertEqual(y.shape, (50,))
45+
46+
def test_ps90_sine_to_neg_cosine(self) -> None:
47+
n = 512
48+
t = np.arange(n) / float(n)
49+
x = np.sin(2 * np.pi * 8 * t)
50+
y = ps90(x)
51+
c = np.cos(2 * np.pi * 8 * t)
52+
corr = np.corrcoef(y, -c)[0, 1]
53+
self.assertGreater(corr, 0.99)
54+
55+
def test_analytic_signal(self) -> None:
56+
x = np.random.randn(100)
57+
z = analytic_signal(x)
58+
self.assertTrue(np.iscomplexobj(z))
59+
np.testing.assert_allclose(np.real(z), x)
60+
61+
62+
class OperatorHelperTest(unittest.TestCase):
63+
def test_first_difference_matrix(self) -> None:
64+
nt, fs = 10, 100.0
65+
d1 = first_difference_matrix(nt, fs)
66+
self.assertEqual(d1.shape, (nt - 1, nt))
67+
x = np.arange(nt, dtype=float)
68+
dx = d1 @ x
69+
np.testing.assert_allclose(dx, fs * np.diff(x))
70+
71+
def test_second_difference_matrix(self) -> None:
72+
nt, fs = 12, 50.0
73+
d2 = second_difference_matrix(nt, fs)
74+
self.assertEqual(d2.shape, (nt - 2, nt))
75+
x = np.arange(nt, dtype=float) ** 2
76+
# forward second difference: x[i] - 2 x[i+1] + x[i+2]
77+
expected = (fs**2) * (x[:-2] - 2 * x[1:-1] + x[2:])
78+
np.testing.assert_allclose(d2 @ x, expected)
79+
80+
def test_build_smoothing_filters(self) -> None:
81+
f1, f2 = build_smoothing_filters(32, 100.0, [0.1, 0.01], [1.0, 0.1])
82+
self.assertEqual(len(f1), 2)
83+
self.assertEqual(f1[0].shape, (32, 32))
84+
self.assertEqual(f2[1].shape, (32, 32))
85+
# F should be symmetric SPD-ish (positive eigenvalues)
86+
eig = np.linalg.eigvalsh(f1[0])
87+
self.assertTrue(np.all(eig > 0))
88+
89+
def test_build_smoothing_filters_mismatch(self) -> None:
90+
with self.assertRaises(ValueError):
91+
build_smoothing_filters(16, 10.0, [0.1], [0.1, 0.01])
92+
93+
def test_schedule_index(self) -> None:
94+
abitr = np.array([20.0, 50.0, 200.0])
95+
self.assertEqual(schedule_index(1, abitr), 0)
96+
self.assertEqual(schedule_index(20, abitr), 0)
97+
self.assertEqual(schedule_index(21, abitr), 1)
98+
self.assertEqual(schedule_index(50, abitr), 1)
99+
self.assertEqual(schedule_index(51, abitr), 2)
100+
self.assertEqual(schedule_index(500, abitr), 2)
101+
102+
103+
class SpectrumHelperTest(unittest.TestCase):
104+
def test_fft_two_to_one_even(self) -> None:
105+
nt = 128
106+
fs = 100.0
107+
x = np.sin(2 * np.pi * 5 * np.arange(nt) / fs)
108+
xf2 = np.fft.fft(x)
109+
xf1, f1 = fft_two_to_one(xf2, fs, nt)
110+
self.assertEqual(xf1.size, f1.size)
111+
self.assertEqual(f1[0], 0.0)
112+
peak = f1[np.argmax(np.abs(xf1))]
113+
self.assertAlmostEqual(peak, 5.0, delta=fs / nt)
114+
115+
def test_fft_two_to_one_odd_and_matrix(self) -> None:
116+
x = np.random.randn(65, 2)
117+
xf2 = np.fft.fft(x, axis=0)
118+
xf1, f1 = fft_two_to_one(xf2, 50.0, 65)
119+
self.assertEqual(xf1.shape[1], 2)
120+
self.assertEqual(xf1.shape[0], f1.size)
121+
122+
def test_instantaneous_frequency(self) -> None:
123+
fs = 100.0
124+
t = np.arange(200) / fs
125+
phz = 2 * np.pi * 7 * t
126+
ifrq = instantaneous_frequency(phz, fs)
127+
np.testing.assert_allclose(ifrq, 7.0, atol=1e-10)
128+
129+
phz2 = np.column_stack([phz, 2 * phz])
130+
ifrq2 = instantaneous_frequency(phz2, fs)
131+
self.assertEqual(ifrq2.shape, (199, 2))
132+
133+
134+
class DemoFactoryTest(unittest.TestCase):
135+
def test_make_order_tracking_demo(self) -> None:
136+
demo = make_order_tracking_demo(fs=100.0, nt=500)
137+
self.assertEqual(demo["signal"].shape, (500,))
138+
self.assertEqual(demo["modes"].shape, (500, 3))
139+
self.assertEqual(demo["true_if"].shape, (500, 3))
140+
141+
def test_constant_frequency_init(self) -> None:
142+
ampg, phzg = constant_frequency_init(100, 50.0, [2.0, 4.0])
143+
self.assertEqual(ampg.shape, (100, 2))
144+
self.assertEqual(phzg.shape, (100, 2))
145+
np.testing.assert_allclose(np.abs(ampg), 1.0)
146+
147+
148+
class STNBMDAlgoTest(unittest.TestCase):
149+
def test_str_and_call(self) -> None:
150+
model = STNBMD(fs=100.0, abitr=[5, 10, 15], tol=1e-4)
151+
self.assertIn("STNBMD", str(model))
152+
demo = make_order_tracking_demo(nt=300)
153+
modes = model(demo["signal"], frequencies=[2.0, 4.0, 6.0])
154+
self.assertEqual(modes.shape[0], 3)
155+
self.assertEqual(modes.shape[1], 300)
156+
157+
def test_functional_matches_class(self) -> None:
158+
demo = make_order_tracking_demo(nt=400)
159+
params = dict(
160+
fs=float(demo["fs"]),
161+
frequencies=[2.0, 4.0, 6.0],
162+
alpha=[0.1, 0.01, 0.01],
163+
beta=[1.0, 0.1, 0.001],
164+
abitr=[10, 20, 40],
165+
tol=1e-5,
166+
)
167+
xnb1, err1, amp1, phz1 = stnbmd(demo["signal"], **params)
168+
model = STNBMD(
169+
fs=params["fs"],
170+
alpha=params["alpha"],
171+
beta=params["beta"],
172+
abitr=params["abitr"],
173+
tol=params["tol"],
174+
)
175+
modes, xnb2, err2, amp2, phz2 = model.fit_transform(
176+
demo["signal"], frequencies=params["frequencies"], return_all=True
177+
)
178+
np.testing.assert_allclose(xnb1, xnb2)
179+
np.testing.assert_allclose(err1, err2)
180+
np.testing.assert_allclose(modes, np.real(xnb2).T)
181+
182+
def test_stnbm_decomp_ig_validation(self) -> None:
183+
x = np.random.randn(50)
184+
ampg = np.ones((50, 2), dtype=complex)
185+
phzg = np.zeros((50, 2))
186+
with self.assertRaises(ValueError):
187+
stnbm_decomp_ig(x, 100.0, ampg[:40], phzg, [0.1], [1.0], [10], 1e-6)
188+
with self.assertRaises(ValueError):
189+
stnbm_decomp_ig(x, 100.0, ampg, phzg, [0.1], [1.0, 0.1], [10], 1e-6)
190+
191+
def test_single_mode_tone(self) -> None:
192+
fs = 200.0
193+
t = np.arange(400) / fs
194+
x = np.sin(2 * np.pi * 5 * t)
195+
model = STNBMD(
196+
fs=fs,
197+
alpha=[0.1, 0.01],
198+
beta=[1.0, 0.01],
199+
abitr=[20, 60],
200+
tol=1e-6,
201+
)
202+
modes, xnb, err, amp, phz = model.fit_transform(
203+
x, frequencies=[5.0], return_all=True
204+
)
205+
self.assertEqual(modes.shape, (1, 400))
206+
# reconstruction should be close
207+
recon = np.real(xnb[:, 0])
208+
self.assertTrue(np.allclose(recon, x, atol=0.15))
209+
ifrq = model.instantaneous_frequency_hz()
210+
self.assertLess(np.abs(np.mean(ifrq) - 5.0), 0.25)
211+
212+
def test_order_tracking_demo_recon_and_if(self) -> None:
213+
demo = make_order_tracking_demo()
214+
model = STNBMD(
215+
fs=float(demo["fs"]),
216+
alpha=[1e-1, 1e-2, 1e-2],
217+
beta=[1.0, 1e-1, 1e-3],
218+
abitr=[20, 50, 120],
219+
tol=1e-6,
220+
)
221+
modes = model.fit_transform(demo["signal"], frequencies=[2.0, 4.0, 6.0])
222+
recon = np.sum(modes, axis=0)
223+
rel = np.linalg.norm(recon - demo["signal"]) / np.linalg.norm(demo["signal"])
224+
self.assertLess(rel, 0.05)
225+
226+
ifrq = model.instantaneous_frequency_hz()
227+
# mean IF should be near mean true upsweep harmonics
228+
true_mean = demo["true_if"].mean(axis=0)
229+
est_mean = ifrq.mean(axis=0)
230+
# allow mode order match by sorting
231+
np.testing.assert_allclose(np.sort(est_mean), np.sort(true_mean), rtol=0.15)
232+
233+
def test_if_before_fit_raises(self) -> None:
234+
with self.assertRaises(ValueError):
235+
STNBMD().instantaneous_frequency_hz()
236+
237+
def test_stnbmd_requires_init(self) -> None:
238+
with self.assertRaises(ValueError):
239+
stnbmd(np.random.randn(100), fs=50.0)
240+
241+
242+
if __name__ == "__main__":
243+
unittest.main()

0 commit comments

Comments
 (0)