remove debugging print

charon_vna/gui.py

@@ -1,13 +1,14 @@
 # %% imports
 import sys
 from pathlib import Path
-from typing import Callable, List, Tuple
+from typing import Callable, List, Literal, Tuple
 
 import matplotlib as mpl
 import numpy as np
 import xarray as xr
 from matplotlib import pyplot as plt
 from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg
+from matplotlib.lines import Line2D
 from matplotlib.ticker import EngFormatter
 from numpy import typing as npt
 from PySide6.QtGui import QAction, QKeySequence
@@ -34,20 +35,23 @@ DEFAULT_CONFIG = dict(
 
 
 class PlotWidget(QWidget):
-    enabled_ports: List[Tuple[int | str]]
+    traces: List[Tuple[int | str]]
+    lines: List[Line2D]
 
-    def __init__(self):
+    def __init__(self, type_: str = "logmag"):
         super().__init__()
 
-        self.enabled_ports = [(1, 1)]
+        self.traces = [(1, 1)]
 
         layout = QVBoxLayout()
         self.setLayout(layout)
 
         self.fig = plt.Figure(figsize=(5, 4), dpi=100, tight_layout=True)
         self.ax = self.fig.add_subplot(111)
-        self.setup_logmag()
-        # self.setup_smith()
+        self.set_plot_type(type_)
+        self.lines = [
+            self.ax.plot([np.nan], [np.nan], label="$S_{" + str(m) + str(n) + "}$")[0] for m, n in self.traces
+        ]
         self.ax.legend(loc="upper right")
 
         canvas = FigureCanvasQTAgg(self.fig)
@@ -57,6 +61,37 @@ class PlotWidget(QWidget):
         # toolbar.addAction("blah")
         # self.addToolBar(toolbar)
 
+    def set_plot_type(
+        self,
+        type_: Literal["logmag", "phase", "vswr", "smith"],
+        sweep_type: Literal["frequency", "time"] = "frequency",
+    ) -> None:
+        if sweep_type != "frequency":
+            raise NotImplementedError("Only frequency sweeps are currently supported")
+
+        if type_ == "logmag":
+            self.setup_logmag()
+        elif type_ == "phase":
+            self.setup_phase()
+        elif type_ == "vswr":
+            self.setup_vswr()
+        elif type_ == "smith":
+            self.setup_smith()
+        else:
+            raise ValueError(f"Unknown plot type: {type_}")
+
+        self._plot_type = type_
+
+    def update_plot(self, data: xr.DataArray):
+        if self._plot_type == "logmag":
+            self.update_logmag(data)
+        elif self._plot_type == "phase":
+            self.update_phase(data)
+        elif self._plot_type == "vswr":
+            self.update_vswr(data)
+        elif self._plot_type == "smith":
+            self.update_smith(data)
+
     def setup_rect(self) -> None:
         self.ax.grid(True)
         self.ax.xaxis.set_major_formatter(EngFormatter())
@@ -64,11 +99,11 @@ class PlotWidget(QWidget):
 
     def update_rect(self, data: xr.DataArray, func: Callable[[npt.ArrayLike], npt.ArrayLike]) -> None:
         self.ax.set_xlim(data["frequency"].min().data, data["frequency"].max().data)
-        # remove old lines
-        for line in self.ax.lines:
-            line.remove()
-        for m, n in self.enabled_ports:
-            self.ax.plot(data["frequency"], func(data.sel(m=m, n=n)))
+        for ii, (m, n) in enumerate(self.traces):
+            self.lines[ii].set_xdata(data["frequency"])
+            self.lines[ii].set_ydata(func(data.sel(m=m, n=n)))
+
+        self.fig.canvas.draw()
 
     def setup_logmag(self, ylim: List[float] = [-30, 30]) -> None:
         self.setup_rect()
@@ -103,12 +138,12 @@ class PlotWidget(QWidget):
         rf_plt.smith(ax=self.ax, smithR=1, chart_type="z", draw_vswr=None)
 
     def update_smith(self, data: xr.DataArray) -> None:
-        # remove old lines
-        for line in self.ax.lines:
-            line.remove()
-        for m, n in self.enabled_ports:
+        for ii, (m, n) in enumerate(self.traces):
             sel = data.sel(m=m, n=n)
-            self.ax.plot(sel.real, sel.imag)
+            self.lines[ii].set_xdata(sel.real)
+            self.lines[ii].set_ydata(sel.imag)
+
+        self.fig.canvas.draw()
 
 
 # Subclass QMainWindow to customize your application's main window
@@ -116,6 +151,8 @@ class MainWindow(QMainWindow):
     config_path: Path | None
     # device: Charon
 
+    plots: List[PlotWidget]
+
     def __init__(self):
         super().__init__()
 
@@ -157,22 +194,31 @@ class MainWindow(QMainWindow):
         menu_calibration = QMenu("&Calibration")
         menubar.addMenu(menu_calibration)
 
+        menu_simulation = QMenu("Si&mulation")
+        menubar.addMenu(menu_simulation)
+        action_generate_data = QAction("&Generate data", self)
+        menu_file.addAction(action_generate_data)
+        action_generate_data.triggered.connect(self.generate_sim_data)
+        action_generate_data.setShortcut(QKeySequence("Ctrl+G"))
+
         # Content
         window_layout = QVBoxLayout()
 
         prog_sweep = QProgressBar()
         prog_sweep.setMinimum(0)
         prog_sweep.setMaximum(100)
-        prog_sweep.setTextVisible(False)
+        prog_sweep.setFormat("%v / %m")
+        # prog_sweep.setTextVisible(False)
         prog_sweep.setValue(50)
         window_layout.addWidget(prog_sweep)
 
         # window_widget.se
         plot_layout = QVBoxLayout()
         # TODO: handle plots properly
-        for i in range(2):
-            plot0 = PlotWidget()
-            plot_layout.addWidget(plot0)
+        self.plots = []
+        for type_ in ["logmag", "phase", "vswr", "smith"]:
+            self.plots.append(PlotWidget(type_=type_))
+            plot_layout.addWidget(self.plots[-1])
         plot_widget = QWidget()
         plot_widget.setLayout(plot_layout)
         window_layout.addWidget(plot_widget)
@@ -200,6 +246,18 @@ class MainWindow(QMainWindow):
         print("loading config")
         # TODO: load config
 
+    def generate_sim_data(self) -> None:
+        coords = {"frequency": np.linspace(1e9, 2e9, 101), "m": [1], "n": [1]}
+        shape = tuple(len(v) for v in coords.values())
+        data = xr.DataArray(
+            ((-1 + 2 * np.random.rand(*shape)) + 1j * (-1 + 2 * np.random.rand(*shape))) / np.sqrt(2),
+            dims=list(coords.keys()),
+            coords=coords,
+        )
+
+        for plot in self.plots:
+            plot.update_plot(data)
+
 
 def main() -> None:
     app = QApplication(sys.argv)