nmreval/src/gui_qt/data/integral_widget.py

from itertools import cycle

import pyqtgraph as pg
from numpy import nanmax, nanmin, inf, argsort, where

from nmreval.lib.colors import Tab10
from ..lib.pg_objects import PlotItem, RegionItem

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(Tab10)

    requestData = QtCore.pyqtSignal(str)
    item_deleted = QtCore.pyqtSignal(pg.GraphicsObject)

    def __init__(self, parent=None):
        super().__init__(parent=parent)
        self.setupUi(self)

        self.connected_figure = ''
        self.management = None

        self.graph_shown = None
        self.shown_set = None

        self.ranges = []
        self.lines = []
        self.areas = []

        self._x = None
        self._y = None
        self.max_area = 0
        self.max_y = inf
        self.min_y = -inf

        self.treeWidget.itemChanged.connect(self._update_by_tree)

    def __call__(self, graph_name, items):
        self.label_2.setText(f'Connected to {graph_name}\nChanging tab will remove all integration limits.')

        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.change_set(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, _: int):
        # key = self.set_combobox.itemData(idx)
        key = self.set_combobox.currentData()

        self._data = self.management[key]
        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.hex())
        qc.setAlpha(40)
        pen = pg.mkPen({f'color': c.rgb()})

        region = RegionItem(values=[x, x*1.1], mode='mid', brush=QtGui.QBrush(qc), pen=pen)
        integral_plot = PlotItem(x=[], y=[], pen=pen)

        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.hex())))

        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()
            if text.startswith('S'):
                child.setFlags(child.flags() | QtCore.Qt.ItemIsEditable)
            else:
                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_by_tree(self, item: QtWidgets.QTreeWidgetItem) -> None:
        parent_item = item.parent()
        idx = self.treeWidget.invisibleRootItem().indexOfChild(parent_item)

        is_left_border = parent_item.indexOfChild(item) == 0
        current_region = self.lines[idx][0]
        current_limits = current_region.getRegion()

        new_value = item.text(0)
        try:
            new_value = float(new_value)
            if is_left_border:
                current_region.setRegion((new_value, current_limits[1]))
            else:
                current_region.setRegion((current_limits[0], new_value))
        except ValueError:
            self._update_values(idx, current_limits)

    def _update_integral(self):
        idx = None
        reg = 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, reg.getRegion())

    def _update_values(self, idx, new_range):
        self.ranges[idx] = new_range

        area = self.make_integral(idx, *new_range)

        self.treeWidget.blockSignals(True)

        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:
            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

        self.treeWidget.blockSignals(False)

    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 make_integral(self, idx, x_min, x_max):
        if self._data is None:
            self.change_set(0)

        integral = self._data.data.integrate(limits=(x_min, x_max), asarray=True)

        if integral.size != 0:
            area = integral[-1, 1]
            scale = (self.max_y-self.min_y) / area
            self.lines[idx][1].setData(x=integral[:, 0], y=integral[:, 1]*scale + self.min_y)

            self.areas[idx] = area

            return area

        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.item_deleted.emit(self.lines[idx][0])
            self.item_deleted.emit(self.lines[idx][1])

            self.ranges.pop(idx)
            self.lines.pop(idx)
            self.areas.pop(idx)

            self.treeWidget.takeTopLevelItem(idx)

    @QtCore.pyqtSlot(name='on_pushButton_clicked')
    def convert_to_datasets(self):
        set_values = []
        areas = []

        new_sets = True

        for r in self.ranges:
            area_i = []
            for idx in range(self.set_combobox.count()):
                set_id = self.set_combobox.itemData(idx)
                d = self.management[set_id]
                if new_sets:
                    set_values.append(d.value)
                integration = d.data.integrate(limits=r, asarray=True)
                area_i.append(integration[-1, 1])

            areas.append(area_i)
            new_sets = False

        self.management.integral_datasets(self.ranges, set_values, areas)

    def clear(self):
        self.connected_figure = ''
        self.graph_shown = None
        self.shown_set = None

        self._data = None
        self.areas = []
        self.ranges = []
        self.lines = []

        self.max_area = 0
        self.max_y = inf
        self.min_y = -inf

        self.treeWidget.clear()