|
| 1 | +from PyQt5.QtWidgets import QDialog, QVBoxLayout, QHBoxLayout, QPushButton, QLabel, QFormLayout, QRadioButton, QMessageBox |
| 2 | + |
| 3 | +import matplotlib.pyplot as plt |
| 4 | +from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas |
| 5 | +from matplotlib.widgets import SpanSelector |
| 6 | + |
| 7 | +import numpy as np |
| 8 | + |
| 9 | +from data_extractor import getInputOutputData |
| 10 | + |
| 11 | +class DataSelectionWindow(QDialog): |
| 12 | + def __init__(self, filename): |
| 13 | + QDialog.__init__(self) |
| 14 | + |
| 15 | + self.file_name = filename |
| 16 | + |
| 17 | + self.figure = plt.figure(1) |
| 18 | + self.canvas = FigureCanvas(self.figure) |
| 19 | + |
| 20 | + layout_v = QVBoxLayout() |
| 21 | + |
| 22 | + layout_v.addWidget(self.canvas) |
| 23 | + |
| 24 | + xyz_group = QHBoxLayout() |
| 25 | + r_x = QRadioButton("x") |
| 26 | + r_x.setChecked(True) |
| 27 | + r_y = QRadioButton("y") |
| 28 | + r_z = QRadioButton("z") |
| 29 | + xyz_group.addWidget(QLabel("Axis")) |
| 30 | + xyz_group.addWidget(r_x) |
| 31 | + xyz_group.addWidget(r_y) |
| 32 | + xyz_group.addWidget(r_z) |
| 33 | + r_x.clicked.connect(self.loadXData) |
| 34 | + r_y.clicked.connect(self.loadYData) |
| 35 | + r_z.clicked.connect(self.loadZData) |
| 36 | + |
| 37 | + layout_v.addLayout(xyz_group) |
| 38 | + |
| 39 | + btn_ok = QPushButton("Load selection") |
| 40 | + btn_ok.clicked.connect(self.loadLog) |
| 41 | + layout_v.addWidget(btn_ok) |
| 42 | + |
| 43 | + self.setLayout(layout_v) |
| 44 | + |
| 45 | + self.refreshInputOutputData() |
| 46 | + |
| 47 | + def loadLog(self): |
| 48 | + if self.t_stop > self.t_start: |
| 49 | + (self.t, self.u, self.y) = getInputOutputData(self.file_name, self.axis, self.t_start, self.t_stop) |
| 50 | + self.accept() |
| 51 | + else: |
| 52 | + self.printRangeError() |
| 53 | + |
| 54 | + def printRangeError(self): |
| 55 | + msg = QMessageBox() |
| 56 | + msg.setIcon(QMessageBox.Critical) |
| 57 | + msg.setWindowTitle("Error") |
| 58 | + msg.setText("Range is invalid") |
| 59 | + msg.exec_() |
| 60 | + |
| 61 | + def loadXData(self): |
| 62 | + if self.file_name: |
| 63 | + self.refreshInputOutputData(0) |
| 64 | + |
| 65 | + def loadYData(self): |
| 66 | + if self.file_name: |
| 67 | + self.refreshInputOutputData(1) |
| 68 | + |
| 69 | + def loadZData(self): |
| 70 | + if self.file_name: |
| 71 | + self.refreshInputOutputData(2) |
| 72 | + |
| 73 | + def refreshInputOutputData(self, axis=0): |
| 74 | + if self.file_name: |
| 75 | + self.axis = axis |
| 76 | + (self.t, self.u, self.y) = getInputOutputData(self.file_name, axis) |
| 77 | + self.plotInputOutput(redraw=True) |
| 78 | + |
| 79 | + def plotInputOutput(self, redraw=False): |
| 80 | + self.figure.clear() |
| 81 | + self.ax = self.figure.add_subplot(1,1,1) |
| 82 | + self.ax.plot(self.t, self.u, self.t, self.y) |
| 83 | + self.ax.set_title("Click and drag to select data range") |
| 84 | + self.ax.set_xlabel("Time (s)") |
| 85 | + self.ax.set_ylabel("Amplitude") |
| 86 | + self.ax.legend(["Input", "Output"]) |
| 87 | + |
| 88 | + self.span = SpanSelector(self.ax, self.onselect, 'horizontal', useblit=False, |
| 89 | + props=dict(alpha=0.2, facecolor='green'), interactive=True) |
| 90 | + |
| 91 | + self.t_start = self.t[0] |
| 92 | + self.t_stop = self.t[-1] |
| 93 | + |
| 94 | + self.canvas.mpl_connect('scroll_event', self.zoom_fun) |
| 95 | + |
| 96 | + self.canvas.draw() |
| 97 | + |
| 98 | + def onselect(self, xmin, xmax): |
| 99 | + indmin, indmax = np.searchsorted(self.t, (xmin, xmax)) |
| 100 | + indmax = min(len(self.t) - 1, indmax) |
| 101 | + indmin = min(indmin, indmax) |
| 102 | + |
| 103 | + self.t_start = self.t[indmin] |
| 104 | + self.t_stop = self.t[indmax] |
| 105 | + self.ax.set_xlim(self.t_start - 1.0, self.t_stop + 1.0) |
| 106 | + self.canvas.draw() |
| 107 | + |
| 108 | + def zoom_fun(self, event): |
| 109 | + base_scale = 1.1 |
| 110 | + # get the current x and y limits |
| 111 | + cur_xlim = self.ax.get_xlim() |
| 112 | + cur_xrange = cur_xlim[1] - cur_xlim[0] |
| 113 | + xdata = event.xdata # get event x location |
| 114 | + if xdata is None or xdata < cur_xlim[0] or xdata > cur_xlim[1]: |
| 115 | + return |
| 116 | + |
| 117 | + if event.button == 'up': |
| 118 | + # deal with zoom in |
| 119 | + scale_factor = 1/base_scale |
| 120 | + elif event.button == 'down': |
| 121 | + # deal with zoom out |
| 122 | + scale_factor = base_scale |
| 123 | + else: |
| 124 | + # deal with something that should never happen |
| 125 | + scale_factor = 1 |
| 126 | + # set new limits |
| 127 | + new_x_min = xdata - (xdata - cur_xlim[0])*scale_factor |
| 128 | + new_x_max = xdata + (xdata - new_x_min) / (xdata - cur_xlim[0]) * (cur_xlim[1] - xdata) |
| 129 | + |
| 130 | + new_x_min = max(new_x_min, self.t[0] - 1.0) |
| 131 | + new_x_max = min(new_x_max, self.t[-1] + 1.0) |
| 132 | + self.ax.set_xlim([new_x_min, new_x_max]) |
| 133 | + self.canvas.draw() |
0 commit comments