Расширить общие API графиков и GAS обмена

python/protocan/spectrum.py

@@ -37,6 +37,12 @@ class Spectrum:
         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
@@ -46,6 +52,11 @@ class NativeSpectrum:
             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:
@@ -69,3 +80,15 @@ class NativeSpectrum:
             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