snr gets multi-tone

except intermodulation is a mess

barker/snr.py

@@ -18,6 +18,8 @@ from matplotlib import pyplot as plt
 LOW_CUT = 100  # ignore anything below 100Hz
 HIGH_CUT = 4500
 
+log = getLogger("snr")
+
 
 def rms_partition(
     seq, partition_key=lambda x: False, ignore=lambda x: False
@@ -41,7 +43,11 @@ def rms_partition(
         else:
             sum1 += v ** 2
             cnt1 += 1
-    return (np.sqrt(sum1 / cnt1), np.sqrt(sum2 / cnt2))
+
+    def rm(sum_, cnt):
+        return np.sqrt(sum_ / cnt)
+
+    return (rm(sum1, cnt1), rm(sum2, cnt2))
 
 
 class ToneAnalyzer:
@@ -74,17 +80,20 @@ class ToneAnalyzer:
         # x = np.arange(HIGH_CUT)
         # plt.plot(x, r)
         # plt.show()
-        return r
-
-    def analyze(self, freq2=None) -> float:
-        seq = self.get_wave()
-        r = self.get_fft(seq)
-
 
         def to_int(val):
             return int(np.absolute(val))
 
         fft = [to_int(val) for val in r]
+        return fft
+
+    def is_signal(self, f):
+        return abs(f - self.freq) < 5
+
+    def analyze(self, freq2=None) -> float:
+        seq = self.get_wave()
+        fft = self.get_fft(seq)
+
         if freq2 is not None:
             freq_val = fft[self.freq]
             second_best_val = fft[freq2]
@@ -94,7 +103,7 @@ class ToneAnalyzer:
 
         noise, signal = rms_partition(
             seq,
-            partition_key=lambda f: abs(f - self.freq) < 5,
+            partition_key=self.is_signal,
             ignore=lambda f: f < LOW_CUT,
         )
         self.log.debug("noise=%.2f signal=%.2f", noise, signal)
@@ -104,10 +113,18 @@ class ToneAnalyzer:
         # self.log.debug("RMS %.2f", rms)
 
 
+class MultiToneAnalyzer(ToneAnalyzer):
+    def is_signal(self, f):
+        for freq_i in self.freq:
+            if abs(f - freq_i) < 30:
+                return True
+        return False
+
+
 def get_parser():
     p = argparse.ArgumentParser()
     p.add_argument("fname")
-    p.add_argument("--freq", type=int, default=440)
+    p.add_argument("--freq", type=int, default=[440], nargs="*")
     p.add_argument("--freq2", type=int)
     p.add_argument("--offset", type=int, default=0, help="In milliseconds")
     p.add_argument("--duration", type=int, default=100, help="In milliseconds")
@@ -130,7 +147,7 @@ def main():
         raise ValueError("audio not mono maybe?")
         data = data[:, 0]
 
-    analyzer = ToneAnalyzer(
+    analyzer = MultiToneAnalyzer(
         args.freq,
         args.offset,
         args.duration,