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._getIdealTimings()
        self._init_ui()
        self.settings_button.clicked.connect(show_settings_func)
        

    def update_data(self, 
                    system_config : dict,
                    operator_config: dict,
                    opt: OptAlgorithm,
                    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.theor_mode = mode
        self.scaler = int(system_config['UML_time_scaler'])
        self.WeldTime = operator_config['time_wielding'] #[sec]
        self._getIdealTimings()
        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 self.theor_mode:
            timings = np.arange(20000)/10
        else:
            timings = self._plotRealData()
        self._plotIdealData(timings)

    def _getIdealTimings(self):
        data = self.opt.Ts
        self.idealTime = [data['tclose'], data['tgrow'], self.opt.getMarkOpen(), data["tmovement"]]

    def _form_idealdatGraph(self, times):

        if self.theor_mode:
            closure_start = 0 
            compression_start = self.idealTime[0]*self.scaler 
            welding_start = sum(self.idealTime[:2])*self.scaler 
            opening_start = (sum(self.idealTime[:2])+ self.WeldTime)*self.scaler 
            opening_end = (sum(self.idealTime[:3])+ self.WeldTime)*self.scaler
        else:
            closure_start = self.timings_dict["closure"][0]*self.scaler 
            compression_start = self.timings_dict["compression"][0]*self.scaler 
            welding_start = self.timings_dict["welding"][0]*self.scaler 
            opening_start = self.timings_dict["opening"][0]*self.scaler 
            opening_end = self.timings_dict["opening"][1]*self.scaler 

        self.idealPhase0 = closure_start + (self.idealTime[0])*self.scaler                                                          #Подъезд
        self.idealPhase1 = compression_start + (self.idealTime[1])*self.scaler                                       #Сжатие
        self.idealPhase2 = welding_start + (self.WeldTime)*self.scaler                       #Сварка
        self.idealPhase3 = opening_start + (self.idealTime[2])*self.scaler   #Разъезд
        self.idealPhase4 = opening_end + (self.idealTime[3])*self.scaler                                    #Последнее смыкание  
        
        self.x_splitter = np.array([[closure_start, self.idealPhase0],
             [compression_start, self.idealPhase1],
             [welding_start, self.idealPhase2],
             [opening_start, self.idealPhase3],
             [opening_end, self.idealPhase4]])
        data = []
        for time in times:
            if time >= closure_start and time <= self.idealPhase0:
                x_fe, x_me, v_fe, v_me, f = self.opt.calcPhaseClose((time-closure_start)/self.scaler)
            elif time >= compression_start and time <= self.idealPhase1:
                x_fe, x_me, v_fe, v_me, f = self.opt.calcPhaseGrow((time-compression_start)/self.scaler)
            elif time >= welding_start and time <= self.idealPhase2:
                x_fe, x_me, v_fe, v_me, f = data[-1]
            elif time >= opening_start and time <= self.idealPhase3:
                x_fe, x_me, v_fe, v_me, f = self.opt.calcPhaseOpen((time-opening_start)/self.scaler)
            elif time >= opening_end and time <= self.idealPhase4:
                x_fe, x_me, v_fe, v_me, f = self.opt.calcPhaseMovement((time-opening_end)/self.scaler)
            else:
                if data:
                    x_fe, x_me, v_fe, v_me, f = data[-1]
                else: 
                    x_fe, x_me, v_fe, v_me, f = 0,0,0,0,0
            data.append([x_fe, x_me, v_fe, v_me, f])
        data = np.array(data).T
        return data
        

    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, times):
        data = self._form_idealdatGraph(times)

        x_fe = pg.PlotDataItem(times, data[0]*1000, pen=pg.mkPen(color=qts.colors[8], width=2), name='x_fe', autoDownsample=True, downsample=True)
        x_me = pg.PlotDataItem(times, data[1]*1000, pen=pg.mkPen(color=qts.colors[9], width=2), name='x_me', autoDownsample=True, downsample=True)
        v_fe = pg.PlotDataItem(times, data[2]*1000, pen=pg.mkPen(color=qts.colors[8], width=2), name='v_fe', autoDownsample=True, downsample=True)
        v_me = pg.PlotDataItem(times, data[3]*1000, pen=pg.mkPen(color=qts.colors[9], width=2), name='v_me', autoDownsample=True, downsample=True)
        f = pg.PlotDataItem(times, 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(times, data)

    def _addBackgroundSplitter(self):
        alpha = self.alpha
        y01 = np.array([10000, 10000])
        y0_1 = np.array([-10000, -10000])
        for i, _ in enumerate(self.x_splitter):
            a01 = pg.PlotDataItem(_, y01, pen=pg.mkPen(color=qts.colors[8], width=2), name=' ')
            a0_1 = pg.PlotDataItem(_, y0_1, pen=pg.mkPen(color=qts.colors[8], width=2), name=' ')
            bg1 = pg.FillBetweenItem(a01, a0_1, qts.RGBA[i]+(alpha,))
            bg2 = pg.FillBetweenItem(a01, a0_1, qts.RGBA[i]+(alpha,))
            bg3 = pg.FillBetweenItem(a01, a0_1, qts.RGBA[i]+(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 _addEquidistances(self, times, data):
        a1, b1, c1 = self._calculate_equidistant('fill', max(data[4]), 2.5, self.x_splitter[1], 3)
        self.p11.addItem(a1)
        self.p11.addItem(b1)
        self.p11.addItem(c1)

        a1, b1, c1 = self._calculate_equidistant(times, data[4], 0.75, self.x_splitter[2], 3)
        self.p11.addItem(a1)
        self.p11.addItem(b1)
        self.p11.addItem(c1)

        


    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, splitter, color):
        if str(x) == 'fill':
            x = [splitter[0]+0.001, splitter[0]+0.002, splitter[1]-0.002, splitter[1]-0.001]
            y = [y, y, y, y]
        
        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):
            if splitter[0]<=  x[i] and x[i] <= splitter[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)