PumpCalib. Legend not fully working...

src/dataprovider/pumpcalibprovider.py

@@ -38,7 +38,7 @@ class PumpCalibProvider:
         return runs_df
 
     def get_coef_df(self) -> pd.DataFrame:
-        df = pd.read_csv(filepath_or_buffer=self.data_dir + "COLPUMP_calib_coef.csv",
+        df = pd.read_csv(filepath_or_buffer=self.data_dir + self.pump_id + "_calib_coef.csv",
                          sep=",",
                          parse_dates=["first_run_datetime", "apply_datetime"])
         df["first_run_datetime"] = df["first_run_datetime"].dt.tz_localize('UTC')

src/uim/pumpcalibuim.py

@@ -55,7 +55,6 @@ class PumpCalibUim:
         self.coef_df = self.pump_prvd.get_coef_df()
         self.__update_runs_plot__(self.runs_df, self.coef_df)
 
-
     ####################################################################
     # Plot
 
@@ -63,21 +62,31 @@ class PumpCalibUim:
 
         self.plot_item = pg.PlotItem()
         self.main_ui.pump_graphicsview_runs.setCentralItem(self.plot_item)
+
+        # Create legend
+        self.legend = pg.LegendItem(offset=(40, 10))
+        self.legend.setParentItem(self.plot_item.graphicsItem())
+
         self.runs_points_list = []
 
     def __update_runs_plot__(self, runs_df: pd.DataFrame, coef_df: pd.DataFrame) -> None:
+        self.plot_item.clear()
+        # for item in self.runs_points_list:
+        #     self.legend.removeItem(item)
+        self.legend = utils.clear_legend_items(self.legend)
+        # while self.legend.layout.count() > 0:
+        #     self.legend.layout.removeAt(0)
+        # grphic_item = self.plot_item.graphicsItem()
+        # grphic_item.removeItem(self.legend)
+        # self.legend = pg.LegendItem(offset=(40, 10))
+        # self.legend.setParentItem(self.plot_item)
+
+
         # Create the color scale
         color_map = pg.ColorMap(pos=[runs_df["calib_id_int"].min(),
                                      runs_df["calib_id_int"].max()],
                                 color=[[0, 0, 255], [0, 255, 0]])
 
-        # Add legend
-        # self.plot_item.addLegend((100,60), offset=(70,30))
-        # self.plot_item.plot([1, 3, 2, 4], pen='r', name='red plot')
-        # l = pg.LegendItem((100,60), offset=(70,30))  # args are (size, offset)
-        l = pg.LegendItem()  # args are (size, offset)
-        l.setParentItem(self.plot_item.graphicsItem())   # Note we do NOT call plt.addItem in this case
-
         # Add linear regression line
         for row_id, row in coef_df.iterrows():
             reg_line = pg.InfiniteLine(pos=QPointF(0, row["intercept"]),
@@ -89,26 +98,11 @@ class PumpCalibUim:
         # Add runs points
         for calib_id in runs_df["calib_id_int"].unique():
             run_df = runs_df[runs_df["calib_id_int"] == calib_id]
-            # brushes = [pg.mkBrush(color=color) for color in color_map.mapToQColor(runs_df["calib_id_int"])]
             run_points = pg.ScatterPlotItem(run_df["rpm"],
-                                             run_df["flow"],
-                                             size=6,
-                                             brush=color_map.mapToQColor(calib_id),
-                                             name=run_df["calib_id_str"].unique())
-            l.addItem(run_points,run_df["calib_id_str"].unique())
+                                            run_df["flow"],
+                                            size=6,
+                                            brush=color_map.mapToQColor(calib_id),
+                                            name=run_df["calib_id_str"].unique()[0])
+            self.legend.addItem(run_points, run_df["calib_id_str"].unique()[0])
             self.plot_item.addItem(run_points)
-
-
-
-
-        # Add runs text
-        # for row_id, row in runs_df.iterrows():
-        #     if row_id != 1:
-        #         continue
-        #     run_text = pg.TextItem(text=row["calib_id_str"],
-        #                            # color=color_map.mapToQColor(row["calib_id_int"]),
-        #                            anchor=QPointF(row["rpm"], row["flow"]))
-        #     # self.plot_item.getViewBox().disableAutoRange()
-        #     self.plot_item.addItem(run_text)
-        #     # self.plot_item.addItem(run_text, ignoreBounds=True)
-
+            self.runs_points_list.append(run_points)

src/utils.py

@@ -44,6 +44,24 @@ class TimeAxisItem(pg.AxisItem):
         return tick_strings
 
 
+def clear_legend_items(legend_item: pg.LegendItem) -> pg.LegendItem:
+    """Remove all the elements of the input `legend_item`.
+
+    Parameters
+    ----------
+    legend_item: pg.LegendItem
+        Legend item to be emptied.
+
+    Returns
+    -------
+    pg.LegendItem:
+        Same as input, but empty.
+    """
+    for sample, label in legend_item.items:
+        legend_item.removeItem(label.text)
+    return legend_item
+
+
 def highlight_row(table: QTableWidget, highlighted_row_id: int,
                   highlighted_row_color: QColor = QColor(173, 216, 230),
                   other_rows_color: QColor = QColor(255, 255, 255)) -> None: