diff --git a/src/base/base.py b/src/base/base.py index e2220a0..78a482e 100644 --- a/src/base/base.py +++ b/src/base/base.py @@ -101,6 +101,13 @@ class BasePlotWidget: "Welding" ] + self._stage_colors = { + "Closing": [208, 28, 31, 100], + "Squeeze": [45, 51, 89, 150], + "Welding": [247, 183, 24, 100], + "Relief": [0, 134, 88, 100] + } + self._plt_channels = { "Electrode Force, N & Welding Current, kA": { "Settings": { diff --git a/src/gui/widgets/plot.py b/src/gui/widgets/plot.py index 95941c8..8fad208 100644 --- a/src/gui/widgets/plot.py +++ b/src/gui/widgets/plot.py @@ -1,16 +1,27 @@ import pandas as pd from PyQt5.QtWidgets import QWidget, QVBoxLayout import pyqtgraph as pg +import numpy as np from base.base import BasePlotWidget class PlotWidget(BasePlotWidget): - def _draw_stages(self, dataframe: pd.DataFrame) -> None: - ... + def _create_stage_region(self, + stage: str, + times: pd.Series, + dataframe: pd.DataFrame) -> pg.LinearRegionItem: + stage_diff = np.diff(dataframe[stage]) + start_index = np.where(stage_diff == 1)[0] + finish_index = np.where(stage_diff == -1)[0] + start_timestamp = times[start_index[0]] + finish_timestamp = times[finish_index[0]] + region = pg.LinearRegionItem([start_timestamp, finish_timestamp], movable=False) + region.setBrush(pg.mkBrush(self._stage_colors[stage])) + return region - def _create_widget(self, dataframe: pd.DataFrame) -> QWidget: + def _build_widget(self, dataframe: pd.DataFrame) -> QWidget: widget = QWidget() layout = QVBoxLayout() @@ -24,6 +35,11 @@ class PlotWidget(BasePlotWidget): legend.setParentItem(plot_widget.graphicsItem()) settings = description["Settings"] + if settings["stages"] and all([stage in dataframe_headers for stage in self._stages]): + for stage in self._stages: + region = self._create_stage_region(stage, time_axis, dataframe) + plot_widget.addItem(region) + for signal in description["Signals"]: if signal["name"] in dataframe_headers: plot = plot_widget.plot(time_axis, dataframe[signal["name"]], pen=signal["pen"]) @@ -39,5 +55,5 @@ class PlotWidget(BasePlotWidget): return widget def build(self, data: list[pd.DataFrame]) -> None: - widgets = [self._create_widget(data_sample) for data_sample in data] + widgets = [self._build_widget(data_sample) for data_sample in data] self._mediator.notify(self, widgets)