import cmath import ctypes import math import os import unittest from protocan.spectrum import Filter, NativeSpectrum, Window @unittest.skipUnless(os.environ.get("SETPROTOCOL_LIBRARY"), "Set SETPROTOCOL_LIBRARY to the built C library") class SpectrumTests(unittest.TestCase): @classmethod def setUpClass(cls): cls.core = NativeSpectrum(ctypes.CDLL(os.environ["SETPROTOCOL_LIBRARY"])) def sample(self, n=4096, fs=1024, frequencies=(64,), amplitude=1): times = [i / fs for i in range(n)] return times, [amplitude * sum(math.sin(2 * math.pi * f * t) for f in frequencies) for t in times] def test_fft_matches_independent_direct_dft(self): times, _ = self.sample(32) values = [math.sin(i * 1.37) + 0.1 * i for i in range(32)] result = self.core.analyze(times, values, window=Window.RECT, remove_mean=False) for k, amplitude in enumerate(result.amplitudes): direct = abs(sum(v * cmath.exp(-2j * math.pi * k * i / 32) for i, v in enumerate(values))) / 32 if k not in (0, 16): direct *= 2 self.assertAlmostEqual(direct, amplitude, places=11) def test_all_windows_preserve_bin_centered_peak_amplitude(self): times, values = self.sample(amplitude=3.25) for window in Window: with self.subTest(window=window): result = self.core.analyze(times, values, window=window) peak = max(range(len(result.amplitudes)), key=result.amplitudes.__getitem__) self.assertEqual(64, result.frequencies[peak]) self.assertAlmostEqual(3.25, result.amplitudes[peak], places=9) detected = self.core.dominant_peak(result) self.assertAlmostEqual(64, detected.frequency_hz, places=9) self.assertAlmostEqual(3.25, detected.amplitude, places=9) def test_dc_and_nyquist_are_not_doubled(self): times, _ = self.sample() result = self.core.analyze(times, [2.5] * len(times), remove_mean=False, window=Window.RECT) self.assertAlmostEqual(2.5, result.amplitudes[0], places=10) result = self.core.analyze(times, [3 * (-1) ** i for i in range(len(times))], window=Window.RECT) self.assertAlmostEqual(3, result.amplitudes[-1], places=10) result = self.core.analyze(times, [2.5] * len(times)) self.assertLess(max(result.amplitudes), 1e-12) def test_windows_suppress_far_leakage_and_flattop_recovers_off_bin_amplitude(self): times, values = self.sample(frequencies=(64.13,)) rect = self.core.analyze(times, values, window=Window.RECT) hann = self.core.analyze(times, values, window=Window.HANN) flat = self.core.analyze(times, values, window=Window.FLATTOP) self.assertLess(hann.amplitudes[400], rect.amplitudes[400] / 100) self.assertAlmostEqual(1, max(flat.amplitudes), delta=0.002) def test_filters_attenuate_expected_bands(self): times, values = self.sample(frequencies=(16, 64, 256)) low = self.core.analyze(times, values, filter=Filter.LOW_PASS, high_hz=64) high = self.core.analyze(times, values, filter=Filter.HIGH_PASS, low_hz=64) band = self.core.analyze(times, values, filter=Filter.BAND_PASS, low_hz=32, high_hz=128) notch = self.core.analyze(times, values, filter=Filter.NOTCH, low_hz=64) at = lambda result, hz: result.amplitudes[int(hz / (result.sample_rate / result.size))] self.assertGreater(at(low, 16), 0.99) self.assertLess(at(low, 256), 0.05) self.assertAlmostEqual(1 / math.sqrt(2), at(low, 64), delta=0.001) self.assertLess(at(high, 16), 0.07) self.assertGreater(at(high, 256), 0.99) self.assertGreater(at(band, 64), 0.93) self.assertLess(at(band, 16), 0.25) self.assertLess(at(band, 256), 0.2) self.assertLess(at(notch, 64), 0.02) self.assertGreater(at(notch, 16), 0.99) def test_timestamp_rate_not_requested_rate_and_jitter_interpolation(self): n, fs = 1024, 200 times = [(i + (0.05 if i % 2 else 0)) / fs for i in range(n)] values = [2 * math.sin(2 * math.pi * (16 * fs / n) * t) for t in times] result = self.core.analyze(times, values) self.assertAlmostEqual((n - 1) / (times[-1] - times[0]), result.sample_rate) self.assertGreater(result.jitter, 0.04) self.assertAlmostEqual(2, result.amplitudes[16], delta=0.005) def test_rejects_gaps_duplicates_bad_values_and_cutoffs(self): times, values = self.sample(128) for broken in ([0.0] * 128, times[:64] + [t + 1 for t in times[64:]], list(reversed(times))): with self.assertRaisesRegex(ValueError, "timestamps"): self.core.analyze(broken, values) for value in (float("nan"), float("inf")): with self.assertRaises(ValueError): self.core.analyze(times, values[:-1] + [value]) for cutoff in (0, -1, 512, 1000, float("nan")): with self.assertRaisesRegex(ValueError, "Filter frequencies"): self.core.analyze(times, values, filter=Filter.LOW_PASS, high_hz=cutoff) with self.assertRaises(ValueError): self.core.analyze(times, values, filter=Filter.BAND_PASS, low_hz=100, high_hz=50) def test_size_limits_tail_selection_and_inputs_unchanged(self): times, values = self.sample(1000) original = values[:] result = self.core.analyze(times, values) self.assertEqual(512, result.size) self.assertEqual(original, values) self.assertEqual(256, self.core.analyze(times, values, max_size=256).size) for n in (0, 1, 15): with self.assertRaisesRegex(ValueError, "16 samples"): self.core.analyze(times[:n], values[:n]) for size in (0, 15, 1000, 32768): with self.assertRaises(ValueError): self.core.analyze(times, values, max_size=size) with self.assertRaises(ValueError): self.core.analyze(times[:-1], values) times, values = self.sample(20000) self.assertEqual(16384, self.core.analyze(times, values, max_size=16384).size) if __name__ == "__main__": unittest.main()