From 5d18b8636c089b98edd7b083354615938df2da80 Mon Sep 17 00:00:00 2001
From: Jono Targett <jono.targett@gmail.com>
Date: Sat, 25 May 2024 16:14:15 +0930
Subject: [PATCH] Updated selcal-detect.py to have a nice pretty graph

---
 scripts/selcal-detect.py | 54 +++++++++++++++++-----------------------
 1 file changed, 23 insertions(+), 31 deletions(-)

diff --git a/scripts/selcal-detect.py b/scripts/selcal-detect.py
index d12b29b..8041239 100755
--- a/scripts/selcal-detect.py
+++ b/scripts/selcal-detect.py
@@ -4,12 +4,9 @@ import csv
 import sys
 import numpy as np
 from scipy import signal
-import matplotlib.pyplot as plt
 from scipy.io import wavfile
 from scipy.signal import butter, lfilter
-from math import log10
-import pyqtgraph as pg
-from PyQt5 import QtWidgets
+
 
 tones = {}
 with open('tones.csv', newline='') as csvfile:
@@ -55,9 +52,9 @@ def decimate_from_sample_rate(sample_rate):
         raise ValueError("Sample rate not supported")
 
 
-# analyze wav file
 if __name__ == '__main__':
-    FLT_LEN = 2000  # Samples
+    # TODO JMT: What is this?
+    FLT_LEN = 2000
 
     file_name = sys.argv[1]
     sample_rate, data = wavfile.read(file_name)
@@ -76,41 +73,36 @@ if __name__ == '__main__':
     pure_signals = {tone:note(freq, FLT_LEN, rate=sample_rate) for tone,freq in tones.items()}
     correlations = {tone:np.abs(signal.correlate(data, pure, mode='same')) for tone,pure in pure_signals.items()}
 
+    N = FLT_LEN # Rolling average length
+    cumsum_convolution = np.ones(N)/N
+    massaged = {tone:np.convolve(correlation, cumsum_convolution, mode='valid') for tone,correlation in correlations.items()}
+
+    # Only import if we're actually plotting, these imports are pretty heavy.
+    import pyqtgraph as pg
+    from PyQt5 import QtWidgets
 
     app = QtWidgets.QApplication([])
-    win = pg.GraphicsLayoutWidget(show=True, title="SELCAL Tone Correlation")
+    layout = pg.GraphicsLayoutWidget(show=True, title="SELCAL Tone Correlation")
+    layout.setGeometry(0, 0, 1600, 960)
 
-    plot = win.addPlot(title=file_name)
-    plot.addLegend()
+    plot = layout.addPlot(title=file_name)
+    legend_view = layout.addViewBox()
+
+    legend = pg.LegendItem(offset=(0, 0))
+    legend.setParentItem(legend_view)
 
     color_map = pg.colormap.get('CET-C6s')
     colors = color_map.getLookupTable(nPts=len(tones))
 
-    for (tone, correlation), color in zip(correlations.items(), colors):
-        plot.plot(correlation, pen=pg.mkPen(color=color), fillLevel=0, name=tone)
+    for (tone, correlation), color in zip(massaged.items(), colors):
+        line = plot.plot(correlation, pen=pg.mkPen(color=color), fillLevel=0.1, name=tone)
+        legend.addItem(line, tone)
 
     plot.setLabel('left', 'Signal Correlation')
     plot.setLabel('bottom', 'Time (samples)')
-
     plot.showGrid(x=True, y=True)
 
-    app.exec_()
+    legend_view.setFixedWidth(80)
+    layout.ci.layout.setColumnFixedWidth(1, 80)
 
-    # matplotlib is too slow
-    '''
-    fig, ax = plt.subplots()
-
-    # Step 4: Plot the data
-    for tone,correlation in correlations.items():
-        ax.plot(correlation, label=tone)
-
-    # Step 5: Customize the plot
-    ax.set_title('Multiple Lines Sharing the Same X Data')
-    ax.set_xlabel('Time (samples)')
-    ax.set_ylabel('Signal Correlation')
-    ax.legend()  # Add a legend
-    ax.grid(True)  # Add a grid
-
-    # Step 6: Display the plot
-    plt.show()
-    '''
\ No newline at end of file
+    app.exec_()
\ No newline at end of file