diff --git a/src/dataprovider/pumpcalibprovider.py b/src/dataprovider/pumpcalibprovider.py
index 07d7d17ecc99a53ba8a8500393a5e4840e31989f..65bcf5c9b6f8ab7ba5047af596f87a30ec71d8a7 100644
--- a/src/dataprovider/pumpcalibprovider.py
+++ b/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')
diff --git a/src/uim/pumpcalibuim.py b/src/uim/pumpcalibuim.py
index 77c2df04f049920c04e49f619643408af119318a..46038ab148fa0581194212fb20112af8ebae13e5 100644
--- a/src/uim/pumpcalibuim.py
+++ b/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)
diff --git a/src/utils.py b/src/utils.py
index f57f3dbb2aee1105f29a746a63eaad1da409ac79..d7218e70507bb595816260fa9080b2b29c473425 100644
--- a/src/utils.py
+++ b/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: