WeldingSpotPerformance/src/gui/plot_window.py

import pyqtgraph as pg
from pyqtgraph.Qt import QtWidgets
import numpy as np

from src.gui import qt_settings as qts

from src.OptAlgorithm import OptAlgorithm


class PlotWindow:
    def __init__(self, opt: OptAlgorithm, show_settings_func):
        pg.setConfigOptions(antialias=True)
        self.alpha = 100 #[0-255 прозрачность фона]
        self.opt = opt
        self._init_ui()
        self.settings_button.clicked.connect(show_settings_func)
        

    def update_data(self, 
                    system_config : dict,
                    operator_config: dict,
                    opt: OptAlgorithm,
                    ideal_time: list[float],
                    bool_dict: dict,
                    float_dict: dict,
                    timings_dict: dict,
                    mode: bool):
        self.opt = opt
        self.bool_dict = bool_dict
        self.float_dict = float_dict
        self.timings_dict = timings_dict
        self.idealTime = ideal_time
        self.theor_mode = mode
        self.scaler = int(system_config['UML_time_scaler'])
        self.WeldTime = operator_config['time_wielding'] #[sec]
        self.WeldData = self.opt.calcPhaseGrow(self.idealTime[1])
        self._updatePlots()
    
    def _init_ui(self):
        self.widget = QtWidgets.QWidget()
        layout = QtWidgets.QVBoxLayout()
        self.widget.setLayout(layout)

        self.win = pg.GraphicsLayoutWidget(show=True, title="")
        self.win.resize(1000,600)
        self.win.setWindowTitle('')

        layout.addWidget(self.win)
        self.settings_button = QtWidgets.QPushButton("Show settings")
        self.settings_button.setFixedWidth(160)
        layout.addWidget(self.settings_button)

        self.p11, self.l11 = self._init_graph('Electrode force, closure', 'Force', 'N', 'Time', 'ms')
        #self.p21, _ = self._init_graph('Electrode force, compression', 'Force', 'N', 'Time', 'ms')
        #self.p31, _ = self._init_graph('Electrode force, compression', 'Force', 'N', 'Time', 'ms')
        self.win.nextRow()
        self.p12, self.l12 = self._init_graph('Rotor Position, closure', 'Posicion', 'mm', 'Time', 'ms')
        #self.p22, _ = self._init_graph('Rotor Position, compression', 'Posicion', 'mm', 'Time', 'ms')
        #self.p32, _ = self._init_graph('Rotor Position, compression', 'Posicion', 'mm', 'Time', 'ms')
        self.win.nextRow()
        self.p13, self.l13 = self._init_graph('Rotor Speed, closure', 'Speed', 'mm/s', 'Time', 'ms')
        #self.p23, _ = self._init_graph('Rotor Speed, compression', 'Speed', 'mm/s', 'Time', 'ms')
        #self.p33, _ = self._init_graph('Rotor Speed, compression', 'Speed', 'mm/s', 'Time', 'ms')
        self.win.nextRow()

        self.p12.setXLink(self.p11)
        self.p13.setXLink(self.p11)

        self.p11.setAutoVisible(x=False, y=True)
        self.p12.setAutoVisible(x=False, y=True)
        self.p13.setAutoVisible(x=False, y=True)
        self.widget.setStyleSheet(qts.dark_style)
        self.widget.show()

    def _init_graph(self, title, Yname, Yunits, Xname, Xunits):
        plot = self.win.addPlot(title = title)
        plot.showGrid(x=True, y=True)
        plot.setLabel('left', Yname, units=Yunits)
        plot.setLabel('bottom', Xname, units=Xunits)
        legend1 = pg.LegendItem((80,60), offset=(70,20))
        legend1.setParentItem(plot)
        return plot, legend1

    def _updatePlots(self):
        self.p11.clear()
        self.l11.clear()
        self.p12.clear()
        self.l12.clear()
        self.p13.clear()
        self.l13.clear()

        if not self.theor_mode:
            self._plotRealData()
        self._form_idealdatGraph()


    def _form_idealdatGraph(self):

        if self.theor_mode:
            self.timings_dict["closure"] = [[0, self.idealTime[0]]]
            self.timings_dict["compression"] = [[self.idealTime[0], sum(self.idealTime[:2])]]
            self.timings_dict["welding"] = [[sum(self.idealTime[:2]), sum(self.idealTime[:2])+self.WeldTime]]
            self.timings_dict["opening"] = [[sum(self.idealTime[:2])+self.WeldTime, sum(self.idealTime[:3])+self.WeldTime]]

        delta = 10 #points_per_ms
        for key, items in self.timings_dict.items():
            if key == 'closure': 
                ideal_time = self.idealTime[0]
                calc = self.opt.calcPhaseClose
                color = qts.RGBA[0]
            elif key == 'compression': 
                ideal_time = self.idealTime[1]
                calc = self.opt.calcPhaseGrow
                color = qts.RGBA[1]
            elif key == 'welding': 
                ideal_time = self.WeldTime
                calc = self._returnWeldData
                color = qts.RGBA[2]
            elif key == 'opening': 
                calc = self.opt.calcPhaseOpen
                ideal_time = self.idealTime[2]
                ideal_closure = self.idealTime[3]
                color = qts.RGBA[3]
                color_closure = qts.RGBA[4]

            for item in items:
                item_data = []
                time_data = []
                for i in range(0, int(ideal_time*self.scaler)*delta):
                    time = i/delta
                    item_data.append(calc(time/self.scaler))
                    time_data.append(time+item[0]*self.scaler)
                    #print (item_data[-1], time_data[-1])
                self._plotIdealData(np.array(time_data), np.array(item_data).T)
                self._addBackgroundSplitter([item[0]*self.scaler,item[0]*self.scaler + time], color)
                
                if key == 'opening':
                    item_data = []
                    time_data = []
                    for i in range(0, int(ideal_closure*self.scaler)*delta):
                        time = i/delta
                        item_data.append(self.opt.calcPhaseMovement(time/self.scaler))
                        time_data.append(time+item[1]*self.scaler)
                    self._plotIdealData(np.array(time_data), np.array(item_data).T)
                    self._addBackgroundSplitter([item[1]*self.scaler,item[1]*self.scaler + time], color_closure)
                
                elif key == 'welding':
                    x = [time_data[0], time_data[-1], time_data[-1]+0.0001]
                    y = [item_data[0][4], item_data[0][4], item_data[0][4]]
                    a1, b1, c1 = self._calculate_equidistant(x, y, 0.75, 3)
                    self.p11.addItem(a1)
                    self.p11.addItem(b1)
                    self.p11.addItem(c1)
                    
                elif key == 'compression':
                    temp = item_data[-1][4]
                    x = [time_data[0], time_data[-1], time_data[-1]+0.0001]
                    y = [temp, temp, temp]
                    a1, b1, c1 = self._calculate_equidistant(x, y, 2.5, 3)
                    self.p11.addItem(a1)
                    self.p11.addItem(b1)
                    self.p11.addItem(c1)     

    def _returnWeldData(self, _):
        return self.WeldData



    def _plotRealData(self):
        for i, (key, dat) in enumerate(self.float_dict.items()):
            dat = np.array(dat).T
            dat[0] = dat[0]*self.scaler
            curve = pg.PlotDataItem(dat[0], dat[1], pen=pg.mkPen(color=qts.colors[i], width=2), name=key, autoDownsample=True, downsample=True)
            if 'Electrode Force' in key:
                self.p11.addItem(curve)
                self.l11.addItem(curve, key)
            elif 'Rotor Position' in key:
                self.p12.addItem(curve)
                self.l12.addItem(curve, key)
            elif 'Rotor Speed' in key:
                self.p13.addItem(curve)
                self.l13.addItem(curve, key)
        return dat[0]

    def _plotIdealData(self, time, data):
        x_fe = pg.PlotDataItem(time, data[0]*1000, pen=pg.mkPen(color=qts.colors[8], width=2), name='x_fe', autoDownsample=True, downsample=True)
        x_me = pg.PlotDataItem(time, data[1]*1000, pen=pg.mkPen(color=qts.colors[9], width=2), name='x_me', autoDownsample=True, downsample=True)
        v_fe = pg.PlotDataItem(time, data[2]*1000, pen=pg.mkPen(color=qts.colors[8], width=2), name='v_fe', autoDownsample=True, downsample=True)
        v_me = pg.PlotDataItem(time, data[3]*1000, pen=pg.mkPen(color=qts.colors[9], width=2), name='v_me', autoDownsample=True, downsample=True)
        f = pg.PlotDataItem(time, data[4], pen=pg.mkPen(color=qts.colors[8], width=2), name='f', autoDownsample=True, downsample=True)

        self.p11.addItem(f)
        #self.l11.addItem(f, 'Ideal force')

        self.p12.addItem(x_fe)
        #self.l12.addItem(x_fe, 'FE POS')
        self.p12.addItem(x_me)
        #self.l12.addItem(x_me, 'ME POS')

        self.p13.addItem(v_fe)
        #self.l13.addItem(v_fe, 'FE VEL')
        self.p13.addItem(v_me)
        #self.l13.addItem(v_me, 'ME VEL')
        #self._addBackgroundSplitter()
        #self._addEquidistances(time, data)

    def _addBackgroundSplitter(self, x, color):
        alpha = self.alpha
        y01 = np.array([10000, 10000])
        y0_1 = np.array([-10000, -10000])
        a01 = pg.PlotDataItem(x, y01, pen=pg.mkPen(color=qts.colors[8], width=2), name=' ')
        a0_1 = pg.PlotDataItem(x, y0_1, pen=pg.mkPen(color=qts.colors[8], width=2), name=' ')
        bg1 = pg.FillBetweenItem(a01, a0_1, color+(alpha,))
        bg2 = pg.FillBetweenItem(a01, a0_1, color+(alpha,))
        bg3 = pg.FillBetweenItem(a01, a0_1, color+(alpha,))
        self.p11.addItem(bg1)
        self.p12.addItem(bg2)
        self.p13.addItem(bg3)

        self.p11.setYRange(-1000, 5000)
        self.p12.setYRange(-50, 250)
        self.p13.setYRange(-400, 400)

        
    def _makeFiller(self, x1, y1, x2, y2, color):
        alpha = self.alpha
        eq1 = pg.PlotDataItem(x1, y1, pen=pg.mkPen(color='#000000', width=1))
        eq2 = pg.PlotDataItem(x2, y2, pen=pg.mkPen(color='#000000', width=1))
        bg = pg.FillBetweenItem(eq1, eq2, qts.RGBA[color]+(alpha,))
        return eq1, eq2, bg


    def _calculate_equidistant(self, x, y, percent, color):
        if len(x) != len(y):
            raise ValueError("x и y должны быть одного размера")
        distance = max(y)/100*percent
        x_eq1 = []
        y_eq1 = []
        x_eq2 = []
        y_eq2 = []

        for i in range(0, len(x) - 1):
            dx = x[i + 1] - x[i]
            dy = y[i + 1] - y[i]
            length = np.sqrt(dx ** 2 + dy ** 2)
            sinA = dy/length
            sinB = dx/length
            
            nx = -sinA*distance
            ny = sinB*distance
            x_eq1.append(x[i] + nx)
            y_eq1.append(y[i] + ny)
            x_eq2.append(x[i] - nx)
            y_eq2.append(y[i] - ny)
        
        return self._makeFiller(np.array(x_eq1), np.array(y_eq1), np.array(x_eq2), np.array(y_eq2), color)