95 lines
3.9 KiB
Python
95 lines
3.9 KiB
Python
"""Shared FFT adapter for desktop GUIs; no numpy/Qt dependency or duplicated DSP.
|
|
|
|
Pass NativeProtocol.lib (rebuilt with set_spectrum.c). Timestamp inputs are seconds.
|
|
See c/set-protocol/include/set_spectrum.h for filtering and normalization semantics.
|
|
"""
|
|
from __future__ import annotations
|
|
import ctypes
|
|
from dataclasses import dataclass
|
|
from enum import IntEnum
|
|
|
|
|
|
class Window(IntEnum):
|
|
RECT = 0
|
|
HANN = 1
|
|
HAMMING = 2
|
|
BLACKMAN = 3
|
|
FLATTOP = 4
|
|
|
|
|
|
class Filter(IntEnum):
|
|
NONE = 0
|
|
LOW_PASS = 1
|
|
HIGH_PASS = 2
|
|
BAND_PASS = 3
|
|
NOTCH = 4
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Spectrum:
|
|
size: int
|
|
sample_rate: float
|
|
jitter: float
|
|
amplitudes: tuple[float, ...]
|
|
|
|
@property
|
|
def frequencies(self) -> tuple[float, ...]:
|
|
return tuple(i * self.sample_rate / self.size for i in range(len(self.amplitudes)))
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class SpectrumPeak:
|
|
frequency_hz: float
|
|
amplitude: float
|
|
|
|
|
|
class NativeSpectrum:
|
|
def __init__(self, library: ctypes.CDLL):
|
|
self.lib = library
|
|
self._analyze = library.set_spectrum_analyze
|
|
pointer = ctypes.POINTER(ctypes.c_double)
|
|
self._analyze.argtypes = [pointer, pointer, ctypes.c_size_t, ctypes.c_size_t,
|
|
ctypes.c_int, ctypes.c_int, ctypes.c_double, ctypes.c_double, ctypes.c_int,
|
|
pointer, ctypes.c_size_t, pointer]
|
|
self._analyze.restype = ctypes.c_int
|
|
self._peak = library.set_spectrum_dominant_peak
|
|
self._peak.argtypes = [pointer, ctypes.c_size_t, ctypes.c_double,
|
|
ctypes.c_double, ctypes.c_double, pointer, ctypes.c_size_t,
|
|
pointer, ctypes.c_size_t]
|
|
self._peak.restype = ctypes.c_int
|
|
|
|
def analyze(self, times, values, *, max_size=4096, window=Window.HANN, filter=Filter.NONE,
|
|
low_hz=10.0, high_hz=100.0, remove_mean=True) -> Spectrum:
|
|
if type(max_size) is not int or not 16 <= max_size <= 16384 or max_size & (max_size - 1):
|
|
raise ValueError("FFT size must be a power of two in 16..16384")
|
|
if len(times) != len(values):
|
|
raise ValueError("Timestamp and value counts differ")
|
|
count = min(len(times), max_size)
|
|
origin = times[len(times) - count] if count else 0.0
|
|
t = (ctypes.c_double * count)(*(value - origin for value in times[-count:])) if count else (ctypes.c_double * 0)()
|
|
v = (ctypes.c_double * count)(*values[-count:]) if count else (ctypes.c_double * 0)()
|
|
capacity = max_size // 2 + 1
|
|
output, meta = (ctypes.c_double * capacity)(), (ctypes.c_double * 3)()
|
|
status = self._analyze(t, v, count, max_size, int(window), int(filter),
|
|
low_hz, high_hz, remove_mean, output, capacity, meta)
|
|
if status:
|
|
message = {1: "At least 16 samples are required", 2: "Invalid spectrum input",
|
|
3: "Gaps, duplicate or irregular timestamps (>50% interval deviation)",
|
|
4: f"Filter frequencies must be between 0 and Fs/2 ({meta[1] / 2:g} Hz)",
|
|
5: "Cannot allocate FFT workspace"}
|
|
raise ValueError(message.get(status, "FFT failed"))
|
|
n = int(meta[0])
|
|
return Spectrum(n, meta[1], meta[2], tuple(output[:n // 2 + 1]))
|
|
|
|
def dominant_peak(self, spectrum: Spectrum, *, relative_threshold: float = 3.0,
|
|
absolute_floor: float = 1e-6) -> SpectrumPeak | None:
|
|
amplitudes = (ctypes.c_double * len(spectrum.amplitudes))(*spectrum.amplitudes)
|
|
scratch = (ctypes.c_double * max(1, len(spectrum.amplitudes) - 1))()
|
|
output = (ctypes.c_double * 2)()
|
|
status = self._peak(amplitudes, len(spectrum.amplitudes),
|
|
spectrum.sample_rate / spectrum.size, relative_threshold, absolute_floor,
|
|
scratch, len(scratch), output, 2)
|
|
if status < 0:
|
|
raise ValueError("Invalid spectrum peak input")
|
|
return SpectrumPeak(output[0], output[1]) if status else None
|