from itertools import cycle import pyqtgraph as pg from numpy import nanmax, nanmin, inf, argsort, where try: # numpy > 1.19 renamed some integration functions from scipy.integrate import cumulative_trapezoid except ImportError: from scipy.integrate import cumtrapz as cumulative_trapezoid from ..Qt import QtWidgets, QtCore, QtGui from .._py.integral_widget import Ui_Form class IntegralWidget(QtWidgets.QWidget, Ui_Form): colors = cycle(['red', 'green', 'blue', 'cyan', 'magenta', 'darkRed', 'darkGreen', 'darkBlue', 'darkCyan', 'darkMagenta']) requestData = QtCore.pyqtSignal(str) item_deleted = QtCore.pyqtSignal(pg.GraphicsObject) newData = QtCore.pyqtSignal(str, list) def __init__(self, parent=None): super().__init__(parent=parent) self.setupUi(self) self.connected_figure = '' self.graph_shown = None self.shown_set = None self._data = None self.ranges = [] self.lines = [] self.max_area = 0 self.max_y = inf self.min_y = -inf def __call__(self, graph_name, items): self.label_2.setText(f'Connected to {graph_name}') self.clear() self.set_combobox.blockSignals(True) self.set_combobox.clear() for sid, name in items: self.set_combobox.addItem(name, userData=sid) self.set_combobox.blockSignals(False) self.set_combobox.setCurrentIndex(0) self.set_combobox.currentIndexChanged.emit(0) return self def keyPressEvent(self, e): if e.key() == QtCore.Qt.Key_Delete: self.remove_integral() else: super().keyPressEvent(e) @QtCore.pyqtSlot(int, name='on_set_combobox_currentIndexChanged') def change_set(self, idx: int): key = self.set_combobox.itemData(idx) self.requestData.emit(key) def set_data(self, ptr): self._data = ptr self.max_y = nanmax(self._data.y.real) self.min_y = nanmin(self._data.y.real) for idx, rnge in enumerate(self.ranges): self._update_values(idx, rnge) def add(self, pos): x = pos[0] self.ranges.append((x, x * 1.1)) c = next(IntegralWidget.colors) qc = QtGui.QColor(c) qc.setAlpha(40) region = pg.LinearRegionItem(values=[x, x*1.1], brush=QtGui.QBrush(qc), pen=pg.mkPen(QtGui.QColor(c))) integral_plot = pg.PlotDataItem(x=[], y=[]) region.sigRegionChanged.connect(self._update_integral) self.lines.append((region, integral_plot)) self.areas.append(0) self._make_entry(c) return region, integral_plot def _make_entry(self, c): item = QtWidgets.QTreeWidgetItem() item.setText(0, f'Integral {len(self.ranges)}') item.setForeground(0, QtGui.QBrush(QtGui.QColor(c))) pts_i = self.ranges[-1] item_list = [] for text, val in [('Start', pts_i[0]), ('Stop', pts_i[1]), ('Areas', 0), ('Ratio', 1.)]: child = QtWidgets.QTreeWidgetItem() child.setFlags(QtCore.Qt.NoItemFlags) child.setText(0, f'{text}: {val:.5g}') child.setForeground(0, QtGui.QBrush(QtGui.QColor('black'))) item_list.append(child) item.addChildren(item_list) self.treeWidget.addTopLevelItem(item) self.treeWidget.expandToDepth(1) self._update_values(len(self.ranges) - 1, pts_i) def _update_integral(self): idx = None sender = self.sender() for i, (reg, _) in enumerate(self.lines): if sender == reg: idx = i break if idx is None: return self._update_values(idx, sender.getRegion()) def _update_values(self, idx, new_range): self.ranges[idx] = new_range area = self.calc_integral(idx, *new_range) item = self.treeWidget.topLevelItem(idx) item.child(0).setText(0, f'Start: {new_range[0]:.5g}') item.child(1).setText(0, f'Stop: {new_range[1]:.5g}') if area is not None: self.areas[idx] = area item.child(2).setText(0, f'Area: {area:.5g}') if self.max_area > 0: self._set_ratios(idx, self.max_area) curr_max = max(self.areas) if curr_max != self.max_area: if curr_max > 0: root = self.treeWidget.invisibleRootItem() for i in range(root.childCount()): self._set_ratios(i, curr_max) self.max_area = curr_max def _set_ratios(self, idx, max_value): item = self.treeWidget.invisibleRootItem().child(idx) area_i = self.areas[idx] item.child(3).setText(0, f'Ratio: {area_i / max_value:.3g}') integral_line = self.lines[idx][1] x_i, y_i = integral_line.getData() scale = (self.max_y - self.min_y) / y_i[-1] * (area_i / max_value) integral_line.setData(x=x_i, y=y_i * scale) def calc_integral(self, idx, x_min, x_max): int_range = where((self._data.x >= x_min) & (self._data.x <= x_max))[0] if len(int_range) > 1: x_int = self._data.x[int_range] y_int = self._data.y[int_range].real order = argsort(x_int) integral = cumulative_trapezoid(y=y_int[order], x=x_int[order], initial=0) scale = (self.max_y-self.min_y) / integral[-1] self.lines[idx][1].setData(x=x_int[order], y=integral*scale + self.min_y) return integral[-1] else: self.lines[idx][1].setData(x=[], y=[]) return None def remove_integral(self): root = self.treeWidget.invisibleRootItem() for item in self.treeWidget.selectedItems(): idx = root.indexOfChild(item) self.ranges.pop(idx) self.item_deleted.emit(self.lines[idx][0]) self.item_deleted.emit(self.lines[idx][1]) self.lines.pop(idx) self.areas.pop(idx) self.treeWidget.takeTopLevelItem(idx) @QtCore.pyqtSlot(name='on_pushButton_clicked') def convert_to_datasets(self): set_id = self.set_combobox.currentData() values = [] for i in range(len(self.ranges)): x_i, y_i = self.lines[i][1].getData() start_i, stop_i = self.ranges[i] area_i = self.areas[i] values.append((x_i, y_i, start_i, stop_i, area_i)) self.newData.emit(set_id, values) def clear(self): self.connected_figure = '' self.graph_shown = None self.shown_set = None self._data = None self.ranges = [] self.lines = [] self.max_area = 0 self.max_y = inf self.min_y = -inf