data_selection_window: cleanup plotCoherence

Author: mahima-yoga

autotune/data_selection_window.py

@@ -57,7 +57,9 @@ def __init__(self, filename):
 
         self.input_ref = None
         self.output_ref = None
-        self.figure = plt.figure(figsize=(8, 6))
+        self.coherence_ref = None 
+        self.coherence_info_text = None
+        self.figure = plt.figure(figsize=(8, 6), layout="constrained")
         self.canvas = FigureCanvas(self.figure)
         self.initPlots()
 
@@ -87,6 +89,10 @@ def __init__(self, filename):
         layout_v.addLayout(top_group)
         layout_v.addWidget(self.canvas)
 
+        self.label_warning = QLabel("")
+        self.label_warning.setStyleSheet("color: red; font-weight: bold;")
+        layout_v.addWidget(self.label_warning)
+
         btn_ok = QPushButton("Load selection")
         btn_ok.clicked.connect(self.loadSelection)
         layout_v.addWidget(btn_ok)
@@ -101,7 +107,7 @@ def __init__(self, filename):
             self.browseFiles()
 
     def loadSelection(self):
-        if (self.t_start is None and self.t_start is None) or (self.t_stop > self.t_start):
+        if (self.t_start is None and self.t_stop is None) or (self.t_stop > self.t_start):
             (self.t, self.u, self.y, self.v) = self.data_extractor.getInputOutputData(self.topics[self.index_u], self.topics[self.index_y], self.t_start, self.t_stop)
             self.accept()
         else:
@@ -213,12 +219,12 @@ def initPlots(self):
 
             # --- Coherence Plot (Bottom) --- 
             self.ax_coherence = self.figure.add_subplot(2, 1, 2)
+            color_coherence = 'tab:grey'
+            plot_refs = self.ax_coherence.plot([], [], color=color_coherence)
+            self.coherence_ref = plot_refs[0]
             self.ax_coherence.set_title("Coherence")
-            self.ax_coherence.set_xlabel("Frequency [Hz]")
+            self.ax_coherence.set_xlabel("Frequency (Hz)")
             self.ax_coherence.set_ylabel("Coherence")
-            self.coherence_ref, = self.ax_coherence.plot([], [])
-
-            self.figure.tight_layout()
 
             self.canvas.mpl_connect('scroll_event', self.zoom_fun)
             self.canvas.draw()
@@ -243,10 +249,22 @@ def plotY(self):
             self.ax_out.set_ylim([min_y, max_y])
         self.canvas.draw()
 
+    def onselect(self, xmin, xmax):
+        indmin, indmax = np.searchsorted(self.t, (xmin, xmax))
+        indmax = min(len(self.t) - 1, indmax)
+        indmin = min(indmin, indmax)
+
+        self.t_start = self.t[indmin]
+        self.t_stop = self.t[indmax]
+        self.ax.set_xlim(self.t_start - 1.0, self.t_stop + 1.0)
+        self.canvas.draw()
+
+        self.plotCoherence()
+
     def plotCoherence(self):
         if len(self.t) == 0 or len(self.u) == 0 or len(self.y) == 0:
             return
-
+    
         # Use getInputOutputData with selected range
         if self.t_start is not None and self.t_stop is not None and self.t_stop > self.t_start:
             t_sel, u_sel, y_sel, _ = self.data_extractor.getInputOutputData(
@@ -260,14 +278,16 @@ def plotCoherence(self):
             )   
 
         num_samples = len(t_sel)
-        if num_samples < 256: # I kind of made up this number -> maybe requires some research 
-            self.ax_coherence.clear()
-            self.ax_coherence.set_title("Coherence (Selection too short)")
-            self.ax_coherence.text(0.5, 0.5, f"Not enough data ({num_samples} samples).\nSelect a larger window.",
-                                ha='center', va='center', transform=self.ax_coherence.transAxes,
-                                fontsize=10, color='red')
-            self.canvas.draw()
-            return
+        
+        if num_samples < 256:
+            self.label_warning.setText(f"Not enough data ({num_samples} samples). Select a larger window.")
+            self.label_warning.show()
+            
+            self.coherence_ref.set_xdata([])
+            self.coherence_ref.set_ydata([])
+            return      
+        else:
+            self.label_warning.hide()
 
         # Estimate sampling frequency
         time_diffs = np.diff(t_sel)
@@ -277,45 +297,35 @@ def plotCoherence(self):
         fs = 1 / avg_time_diff
 
         # Choose segment size
-        nperseg = min(1024, max(256, num_samples // 8))
+        nperseg = min(1024, num_samples // 4)
 
         # Compute coherence
         freq, Cuy = signal.coherence(u_sel, y_sel, fs, nperseg=nperseg)
 
-        # Update coherence plot
-        self.ax_coherence.clear()
-        self.ax_coherence.plot(freq, Cuy, label='Coherence(u, y)')
-        self.ax_coherence.set_title("Coherence")
-        self.ax_coherence.set_xlabel("Frequency [Hz]")
-        self.ax_coherence.set_ylabel("Coherence")
+        # Update coherence plot 
+        self.coherence_ref.set_xdata(freq)
+        self.coherence_ref.set_ydata(Cuy)
         self.ax_coherence.set_xlim([min(freq), max(freq)])
         self.ax_coherence.set_ylim([0, 1])
-        self.ax_coherence.grid(True)
 
-        # Add informative annotation
+        # Remove previous annotation if it exists
+        if self.coherence_info_text is not None:
+            self.coherence_info_text.remove()
+            self.coherence_info_text = None
+
         duration = t_sel[-1] - t_sel[0]
         freq_res = fs / nperseg
         info_text = (f"Samples: {num_samples}, Duration: {duration:.2f}s, "
                     f"fs: {fs:.1f}Hz, nperseg: {nperseg}, Δf: {freq_res:.2f}Hz")
+        
+        self.coherence_info_text = self.ax_coherence.text(
+            0.98, 0.02, info_text,
+            ha='right', va='bottom',
+            transform=self.ax_coherence.transAxes,
+            fontsize=8, color='gray')
 
-        self.ax_coherence.text(0.98, 0.02, info_text,
-                            ha='right', va='bottom',
-                            transform=self.ax_coherence.transAxes,
-                            fontsize=8, color='gray')
-
-        self.canvas.draw()
-
-    def onselect(self, xmin, xmax):
-        indmin, indmax = np.searchsorted(self.t, (xmin, xmax))
-        indmax = min(len(self.t) - 1, indmax)
-        indmin = min(indmin, indmax)
-
-        self.t_start = self.t[indmin]
-        self.t_stop = self.t[indmax]
-        self.ax.set_xlim(self.t_start - 1.0, self.t_stop + 1.0)
         self.canvas.draw()
 
-        self.plotCoherence()
 
     def zoom_fun(self, event):
         base_scale = 1.1