nmreval/src/gui_qt/data/signaledit/baseline_dialog.py

import numpy as np
import pyqtgraph as pg

from scipy.interpolate import splrep, splev

from ...Qt import QtCore, QtWidgets
from ..._py.baseline_dialog import Ui_SignalEdit


class QBaselineDialog(QtWidgets.QDialog, Ui_SignalEdit):
    finished = QtCore.pyqtSignal(str, tuple)

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

        self.data = None

        self.graph = pg.PlotDataItem(x=[], y=[], pen=pg.mkPen({'color': 'b'}), name='Original')
        self.graph_corr = pg.PlotDataItem(x=[], y=[], pen=pg.mkPen({'color': 'r'}), name='Corrected')
        self.baseline = pg.PlotDataItem(x=[], y=[], name='Baseline')

        self.anchors = []
        self.anchor_lines = []
        self.spline = None

        self.legend = self.graphicsView.addLegend()

        self.graphicsView.scene().sigMouseClicked.connect(self.add_node)
        self.graphicsView.addItem(self.graph_corr)
        self.graphicsView.addItem(self.graph)
        self.graphicsView.addItem(self.baseline)

    def add_data(self, x, y):
        if self.data is not None:
            QtWidgets.QMessageBox().information(self, 'Invalid number of datasets',
                                                'Baseline correction is only working on one set at a time.')
            self.close()
        self.anchors.extend([np.min(x), np.max(x)])
        self.data = (x, y)
        self.graph.setData(x=x, y=y.real)
        self.graph_corr.setData(x=x, y=y.real)

    def accept(self):
        self.finished.emit('bls', (splev(self.data[0], self.spline),))
        self.close()

    def add_node(self, evt):
        vb = self.graphicsView.plotItem.vb

        if self.graphicsView.plotItem.sceneBoundingRect().contains(evt.scenePos()) and evt.button() == 1:
            pos = vb.mapSceneToView(evt.scenePos())
            x = pos.x()

            self.anchors.append(x)
            self.anchors.sort()
            row = self.anchors.index(x)
            self.listWidget.insertItem(row-1, QtWidgets.QListWidgetItem(str(x)))

            inf_line = pg.InfiniteLine(pos=x)
            self.anchor_lines.insert(row-1, inf_line)
            self.graphicsView.addItem(inf_line)

        self.change_baseline()

    def change_baseline(self):
        if self.data:
            x, y = self.data

            def mean(xx):
                return np.mean(y[max(0, np.argmin(abs(x-xx))-5):min(len(x), np.argmin(abs(x-xx))+6)].real)

            y_node = [mean(x_node) for x_node in self.anchors]
            try:
                self.spline = splrep(self.anchors, y_node, per=False)
            except TypeError:
                self.spline = splrep(self.anchors, y_node, per=False, k=1)

            bl = splev(x, self.spline)

            self.baseline.setData(x=x, y=bl)
            self.graph_corr.setData(x=x, y=y.real-bl)

    def keyPressEvent(self, evt):
        if self.listWidget.hasFocus() and evt.key() == QtCore.Qt.Key_Delete:
            r = self.listWidget.currentRow()
            self.anchors.pop(r+1)
            listitem = self.listWidget.takeItem(r)
            del listitem
            self.graphicsView.removeItem(self.anchor_lines.pop(r))
            self.change_baseline()

        else:
            super().keyPressEvent(evt)