Draft: Added support for plotting result from many simulations. Solves #209

b_asic/gui_utils/plot_window.py

@@ -47,7 +47,7 @@ class PlotWindow(QWidget):
     def __init__(
         self,
         sim_result: Mapping[ResultKey, Sequence[Num]],
-        sfg_name: Optional[str] = None,
+        figure_name: Optional[str] = None,
     ):
         super().__init__()
         self.setWindowFlags(
@@ -57,8 +57,8 @@ class PlotWindow(QWidget):
             | Qt.WindowMaximizeButtonHint
         )
         title = (
-            f"Simulation results: {sfg_name}"
-            if sfg_name is not None
+            f"Simulation results: {figure_name}"
+            if figure_name is not None
             else "Simulation results"
         )
         self.setWindowTitle(title)
@@ -67,17 +67,17 @@ class PlotWindow(QWidget):
         ########### Flattening sim_result, if it is a list of results #######
         # take: sim_result (possibly on form ['name1', simres1, 'name2', simres2, ...]
         # generate: sim_result (dict)
-        if type(sim_result) == Simulation:
+        if isinstance(sim_result, Simulation):
             sim_result = sim_result._results
-            assert type(sim_result) == dict, TypeError(
+            assert isinstance(sim_result, dict), TypeError(
                 "Parsing sim_result as a Simulation, but the _result seems broken."
             )
-        elif type(sim_result) == list:
+        elif isinstance(sim_result, list):
             new_result = dict()
             nr = 0  # value number. Used for automatic names.
             name = None
             for element in sim_result:
-                if type(element) == str:
+                if isinstance(element, str):
                     assert not name, Exception(
                         "Parsing sim_result as a list. Did you provide two names after"
                         " each other?"
@@ -87,11 +87,11 @@ class PlotWindow(QWidget):
                     if not name:
                         nr = nr + 1
                         name = "(res" + str(nr) + ")"
-                    if type(element) == dict:
+                    if isinstance(element, dict):
                         res = element
-                    elif type(element) == Simulation:
+                    elif isinstance(element, Simulation):
                         res = element._results
-                        assert type(res) == dict, TypeError(
+                        assert isinstance(res, dict), TypeError(
                             f"Parsing sim_result as a list. Result '{name}' is a"
                             " Simulation, and its _result seems broken."
                         )
@@ -110,10 +110,10 @@ class PlotWindow(QWidget):
                 " in the list?"
             )
             sim_result = new_result
-        elif type(sim_result) != dict:
+        elif not isinstance(sim_result, dict):
             raise TypeError(
-                "sim_result must be a dict, Simulation or list. Found"
-                f" {type(sim_result)}"
+                "sim_result must be a dict, Simulation or list."
+                f" Found {type(sim_result)}"
             )
 
         ########### Categorizing/sorting/renaming sim_results: ##############
@@ -238,8 +238,8 @@ class PlotWindow(QWidget):
         # Add additional checkboxes
         self._legend = self._plot_axes.legend()
         self._legend_checkbox = QCheckBox("&Legend")
-        self._legend_checkbox.stateChanged.connect(self._legend_checkbox_change)
         self._legend_checkbox.setCheckState(Qt.CheckState.Checked)
+        self._legend_checkbox.stateChanged.connect(self._legend_checkbox_change)
         self._legend_checkbox.setIcon(get_icon('legend'))
         listlayout.addWidget(self._legend_checkbox)
         # self.ontop_checkbox = QCheckBox("&On top")
@@ -261,12 +261,63 @@ class PlotWindow(QWidget):
         # self.plotcanvas.draw()
         self._auto_redraw = True
 
+    def add_result(self, name, result) -> None:
+        """
+        Add another result to the plot.
+
+        The new signals are added in natural order without sorting them.
+        ---
+        name: str
+          The name of the result
+        result: dict
+          The result from one simulation.
+        """
+        markers = ".ov<^>s+*xd|_"
+        ix = 0
+        self._auto_redraw = False
+        for key, res in result.items():
+            key2 = name + "." + key
+            ischecked = Qt.CheckState.Unchecked
+            if re.fullmatch(r"[0-9]+", key):
+                key2 = name + '.out' + key
+                ischecked = Qt.CheckState.Checked
+
+            if len(res) <= 100:
+                fmt = markers[ix] + '-'
+                ix = (ix + 1) % len(markers)
+            else:
+                fmt = '-'
+
+            def addline(key, vector, checked):
+                line = self._plot_axes.plot(np.real(res), fmt, label=key)
+                self._lines[key] = line[0]
+                list_item = QListWidgetItem(key)
+                list_item.setCheckState(
+                    Qt.CheckState.Unchecked
+                )  # will add it if checked
+                self._checklist.addItem(list_item)
+                self._lines[
+                    key
+                ].remove()  # remove the line from plot. Keep it in _lines.
+                list_item.setCheckState(checked)  # will add it if checked
+
+            if all(np.imag(np.real_if_close(res)) == 0):
+                # real: add one line with corresponding checkbox
+                addline(key2, np.real(res), ischecked)
+            else:
+                # complex: add '_re', '_im', '_mag', '_ang'
+                addline(key2 + "_re", np.real(res), ischecked)
+                addline(key2 + "_im", np.imag(res), ischecked)
+                addline(key2 + "_mag", np.absolute(res), Qt.CheckState.Unchecked)
+                addline(key2 + "_ang", np.angle(res), Qt.CheckState.Unchecked)
+        self._auto_redraw = True
+        self._update_legend()
+        self._plot_canvas.draw()
+
     def _legend_checkbox_change(self, check_state):
         self._legend.set(visible=(check_state == Qt.CheckState.Checked))
-        if self._auto_redraw:
-            if check_state == Qt.CheckState.Checked:
-                self._legend = self._plot_axes.legend()
-            self._plot_canvas.draw()
+        self._update_legend()
+        self._plot_canvas.draw()
 
     # def _ontop_checkbox_change(self, checkState):
     # Bugg: It seems the window closes if you change the WindowStaysOnTopHint.
@@ -281,10 +332,12 @@ class PlotWindow(QWidget):
             self._checklist.item(x).setCheckState(Qt.CheckState.Checked)
         self._auto_redraw = True
         self._update_legend()
+        self._plot_canvas.draw()
 
     def _update_legend(self):
-        self._legend = self._plot_axes.legend()
-        self._plot_canvas.draw()
+        # if self._legend_checkbox.checkState == Qt.CheckState.Checked:
+        if self._legend_checkbox.isChecked():
+            self._legend = self._plot_axes.legend()
 
     def _button_none_click(self, event):
         self._auto_redraw = False
@@ -292,6 +345,7 @@ class PlotWindow(QWidget):
             self._checklist.item(x).setCheckState(Qt.CheckState.Unchecked)
         self._auto_redraw = True
         self._update_legend()
+        self._plot_canvas.draw()
 
     def _item_change(self, listitem):
         key = listitem.text()
@@ -301,6 +355,7 @@ class PlotWindow(QWidget):
             self._lines[key].remove()
         if self._auto_redraw:
             self._update_legend()
+            self._plot_canvas.draw()
 
     def _relim(self, event=None):
         self._plot_axes.relim(True)
@@ -312,6 +367,13 @@ class PlotWindow(QWidget):
 
 def start_simulation_dialog(
     sim_results: Dict[str, List[complex]], sfg_name: Optional[str] = None
+):
+    """Deprecated. Use `show_simulation_result` instead."""
+    show_simulation_result(sim_results, sfg_name)
+
+
+def show_simulation_result(
+    sim_results: Dict[str, List[complex]], figure_name: Optional[str] = None
 ):
     """
     Display the simulation results window.
@@ -328,14 +390,14 @@ def start_simulation_dialog(
         app = QApplication(sys.argv)
     else:
         app = QApplication.instance()
-    win = PlotWindow(sim_result=sim_results, sfg_name=sfg_name)
+    win = PlotWindow(sim_result=sim_results, figure_name=figure_name)
     win.show()
     app.exec_()
 
 
 # Simple test of the dialog
 if __name__ == "__main__":
-    sim_res1 = {
+    res1 = {
         '0': [1.5, 1.6, 1.5, 1],
         '1': [1.0, 2.0, 1.5, 1.1],
         'add1': [1.5, 1.5, 1, 1],
@@ -347,7 +409,7 @@ if __name__ == "__main__":
         't2': [1, 1, 2, 1],
         't3': [1, 1, 1, 2],
     }
-    sim_res2 = {
+    res2 = {
         '0': [0.5, 0.6, 0.5, 0],
         '1': [0.0, 1.0 + 0.3j, 0.5, 0.1j],
         'add1': [0.5, 0.5, 0, 0],
@@ -359,11 +421,19 @@ if __name__ == "__main__":
         't2': [0, 0, 1, 0],
         't3': [0, 0, 0, 1],
     }
-    sim_res3 = {
-        '0': np.random.rand(200).tolist(),
+    res3 = {
+        #'0': np.random.rand(50).tolist(),
+        '0': np.random.rand(50),
         '1': np.random.rand(200).tolist(),
     }
-    # start_simulation_dialog(sim_res3)
-    start_simulation_dialog(
-        ['simReal', sim_res1, 'simCpx', sim_res2, sim_res3], "Test data"
-    )
+    res4 = {
+        '0': np.random.rand(60).tolist(),
+        '1': np.random.rand(220).tolist(),
+        't4': np.random.rand(50).tolist(),
+    }
+    # start_simulation_dialog(res3)
+    app = QApplication(sys.argv)
+    win2 = PlotWindow(['Real', res1, 'Cpx', res2, res3], "Test data")
+    win2.add_result('res4', res4)
+    win2.show()
+    app.exec_()