214 lines
6.7 KiB
Python
214 lines
6.7 KiB
Python
|
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 <b>{graph_name}</b>')
|
||
|
|
||
|
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
|
||
|
|