nmreval/src/gui_qt/fit/result.py

from math import isnan

from pyqtgraph import mkBrush, mkPen

from nmreval.utils.text import convert

from ..lib.utils import RdBuCMap
from ..Qt import QtWidgets, QtGui, QtCore
from .._py.fitresult import Ui_Dialog
from ..lib.pg_objects import PlotItem


class QFitResult(QtWidgets.QDialog, Ui_Dialog):
    closed = QtCore.pyqtSignal(dict, list, str, bool, bool, list)
    redoFit = QtCore.pyqtSignal(dict)

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

        self._management = management

        self.maxx_line.setValidator(QtGui.QDoubleValidator())
        self.minx_line.setValidator(QtGui.QDoubleValidator())
        self.numx_line.setValidator(QtGui.QIntValidator())
        self.extrapolate_box.stateChanged.connect(lambda x: self.maxx_line.setEnabled(x))
        self.extrapolate_box.stateChanged.connect(lambda x: self.minx_line.setEnabled(x))
        self.extrapolate_box.stateChanged.connect(lambda x: self.numx_line.setEnabled(x))

        self._previous_fits = {}
        self._opts = []
        self._results = {}
        self.graph_opts = {}
        self.last_idx = None

        self.resid_plot = self.graphicsView.addPlot(row=0, col=0, title='Residual')
        self.fit_plot = self.graphicsView.addPlot(row=1, col=0, title='Fit')

        self.graphicsView.ci.layout.setRowStretchFactor(0, 1)
        self.graphicsView.ci.layout.setRowStretchFactor(1, 2)

        self.resid_graph = PlotItem(x=[], y=[],
                                    symbol='o', symbolPen=None, symbolBrush=mkBrush(color=(31, 119, 180)),
                                    pen=None)
        self.resid_graph_imag = PlotItem(x=[], y=[],
                                         symbol='s', symbolPen=None, symbolBrush=mkBrush(color=(255, 127, 14)),
                                         pen=None)
        self.resid_plot.addItem(self.resid_graph)
        self.resid_plot.addItem(self.resid_graph_imag)

        self.data_graph = PlotItem(x=[], y=[],
                                   symbol='o', symbolPen=None, symbolBrush=mkBrush(color=(31, 119, 180)),
                                   pen=None)
        self.data_graph_imag = PlotItem(x=[], y=[],
                                        symbol='s', symbolPen=None, symbolBrush=mkBrush(color=(255, 127, 14)),
                                        pen=None)
        self.fit_plot.addItem(self.data_graph)
        self.fit_plot.addItem(self.data_graph_imag)

        self.fit_graph = PlotItem(x=[], y=[])
        self.fit_graph_imag = PlotItem(x=[], y=[])
        self.fit_plot.addItem(self.fit_graph)
        self.fit_plot.addItem(self.fit_graph_imag)

        self.cmap = RdBuCMap(vmin=-1, vmax=1)

        self.buttonBox.accepted.connect(self.accept)

        self.graph_checkBox.stateChanged.connect(lambda x: self.graph_comboBox.setEnabled(x == QtCore.Qt.Unchecked))

        self.logy_box.stateChanged.connect(lambda x: self.fit_plot.setLogMode(y=bool(x)))
        self.logx_box.stateChanged.connect(lambda x: self.fit_plot.setLogMode(x=bool(x)))
        self.resid_plot.setXLink(self.fit_plot)

        self.set_results(results)

    def __call__(self, results: list):
        self._previous_fits = {}
        self.sets_comboBox.blockSignals(True)
        self.sets_comboBox.clear()
        self.sets_comboBox.blockSignals(False)
        self._results = {}
        self._opts = {}

        self.set_results(results)

    def set_results(self, results: list):
        self.sets_comboBox.blockSignals(True)
        for res in results:
            idx = res.idx
            data_k = self._management.data[idx]

            self._previous_fits[idx] = []
            for fit in data_k.get_fits():
                self._previous_fits[idx].append((fit.name, fit.statistics, fit.nobs - fit.nvar))

            self.sets_comboBox.addItem(data_k.name, userData=idx)
        self.sets_comboBox.blockSignals(False)

        self._results = {res.idx: res for res in results}
        self._opts = [(False, False) for _ in range(len(self._results))]

        self.set_parameter(0)

    def add_graphs(self, graphs: list):
        self.graph_comboBox.clear()
        for (graph_id, graph_name) in graphs:
            self.graph_comboBox.addItem(graph_name, userData=graph_id)

    @QtCore.pyqtSlot(int, name='on_sets_comboBox_currentIndexChanged')
    def set_parameter(self, idx: int):
        set_id = self.sets_comboBox.itemData(idx, QtCore.Qt.UserRole)

        res = self._results[set_id]
        self.param_tableWidget.setRowCount(len(res.parameter))
        for j, pvalue in enumerate(res.parameter.values()):
            name = pvalue.name
            p_header = QtWidgets.QTableWidgetItem(convert(name, 'tex', 'str', brackets=True))
            self.param_tableWidget.setVerticalHeaderItem(j, p_header)

            item_text = f'{pvalue.value:.4g}'
            if pvalue.error is not None:
                item_text += f' \u00b1 {pvalue.error:.4g}'
                self.param_tableWidget.setItem(2*j+1, 0, QtWidgets.QTableWidgetItem('-'))
            else:
                self.param_tableWidget.setItem(2*j+1, 0, QtWidgets.QTableWidgetItem())
            item = QtWidgets.QTableWidgetItem(item_text)
            self.param_tableWidget.setItem(j, 0, item)

        self.param_tableWidget.resizeColumnToContents(0)
        self.show_results(idx)

    @QtCore.pyqtSlot(int, name='on_reject_fit_checkBox_stateChanged')
    @QtCore.pyqtSlot(int, name='on_del_prev_checkBox_stateChanged')
    def change_opts(self, _):
        idx = self.sets_comboBox.currentIndex()

        self._opts[idx] = (self.reject_fit_checkBox.checkState() == QtCore.Qt.Checked,
                           self.del_prev_checkBox.checkState() == QtCore.Qt.Checked)

    def show_results(self, idx):
        set_id = self.sets_comboBox.itemData(idx, QtCore.Qt.UserRole)
        self.set_plot(set_id)
        self.set_correlation(set_id)
        self.set_statistics(set_id)
        self.reject_fit_checkBox.blockSignals(True)
        self.reject_fit_checkBox.setChecked(self._opts[idx][0])
        self.reject_fit_checkBox.blockSignals(False)
        self.del_prev_checkBox.blockSignals(True)
        self.del_prev_checkBox.setChecked(self._opts[idx][1])
        self.del_prev_checkBox.blockSignals(False)

    def set_plot(self, idx: str):
        if self.last_idx is not None:
            self.graph_opts[self.last_idx] = (
                self.fit_plot.viewRange(),
                self.logx_box.isChecked(),
                self.logy_box.isChecked(),
            )

        self.last_idx = idx
        res = self._results[idx]
        iscomplex = res.iscomplex

        sub_funcs = res.sub(res.x)
        for item in self.fit_plot.items[::-1]:
            if item not in [self.data_graph, self.data_graph_imag, self.fit_graph, self.fit_graph_imag]:
                self.fit_plot.removeItem(item)

        if iscomplex:
            self.data_graph.setData(x=res.x_data, y=res.y_data.real)
            self.data_graph_imag.setData(x=res.x_data, y=res.y_data.imag)
            self.fit_graph.setData(x=res.x, y=res.y.real)
            self.fit_graph_imag.setData(x=res.x, y=res.y.imag)
            self.resid_graph.setData(x=res.x_data, y=res.residual.real)
            self.resid_graph_imag.setData(x=res.x_data, y=res.residual.imag)

            for i, f in enumerate(sub_funcs):
                item = PlotItem(x=f.x, y=f.y.real, pen=mkPen({'color': i, 'style': 2}))
                self.fit_plot.addItem(item)
                item = PlotItem(x=f.x, y=f.y.imag, pen=mkPen({'color': i, 'style': 2}))
                self.fit_plot.addItem(item)

        else:
            self.resid_graph.setData(x=res.x_data, y=res.residual)
            self.resid_graph_imag.setData(x=[], y=[])
            self.data_graph.setData(x=res.x_data, y=res.y_data)
            self.data_graph_imag.setData(x=[], y=[])
            self.fit_graph.setData(x=res.x, y=res.y)
            self.fit_graph_imag.setData(x=[], y=[])

            for i, f in enumerate(sub_funcs):
                item = PlotItem(x=f.x, y=f.y, pen=mkPen({'color': i, 'style': 2}))
                self.fit_plot.addItem(item)

        self.logx_box.blockSignals(True)
        self.logx_box.setChecked(res.islog)
        self.logx_box.blockSignals(False)

        self.fit_plot.setLogMode(x=res.islog)
        self.resid_plot.setLogMode(x=res.islog)

        if idx in self.graph_opts:
            view_range, logx, logy = self.graph_opts[idx]
            self.fit_plot.setRange(xRange=view_range[0], yRange=view_range[1], padding=0)
            self.fit_plot.setLogMode(x=logx, y=logy)
            self.logx_box.blockSignals(True)
            self.logx_box.setChecked(logx)
            self.logx_box.blockSignals(False)
            self.logy_box.blockSignals(True)
            self.logy_box.setChecked(logy)
            self.logy_box.blockSignals(False)

    def set_correlation(self, idx: str):
        while self.corr_tableWidget.rowCount():
            self.corr_tableWidget.removeRow(0)

        res = self._results[idx]
        c = res.correlation_list()
        for pi, pj, corr, pcorr in c:
            cnt = self.corr_tableWidget.rowCount()
            self.corr_tableWidget.insertRow(cnt)
            self.corr_tableWidget.setItem(cnt, 0, QtWidgets.QTableWidgetItem(convert(pi, old='tex', new='str')))
            self.corr_tableWidget.setItem(cnt, 1, QtWidgets.QTableWidgetItem(convert(pj, old='tex', new='str')))

            for i, val in enumerate([corr, pcorr]):
                if isnan(val):
                    val = 1000.
                val_item = QtWidgets.QTableWidgetItem(f'{val:.4g}')
                val_item.setBackground(self.cmap.color(val))
                if abs(val) > 0.75:
                    val_item.setForeground(QtGui.QColor('white'))
                self.corr_tableWidget.setItem(cnt, i+2, val_item)

        self.corr_tableWidget.resizeColumnsToContents()

    def set_statistics(self, idx: str):
        while self.stats_tableWidget.rowCount():
            self.stats_tableWidget.removeRow(0)

        res = self._results[idx]

        self.stats_tableWidget.setColumnCount(1 + len(self._previous_fits[idx]))
        self.stats_tableWidget.setRowCount(len(res.statistics)+3)

        it = QtWidgets.QTableWidgetItem(f'{res.dof}')
        it.setFlags(it.flags() ^ QtCore.Qt.ItemIsEditable)
        self.stats_tableWidget.setVerticalHeaderItem(0, QtWidgets.QTableWidgetItem('DoF'))
        self.stats_tableWidget.setItem(0, 0, it)

        for col, (name, _, dof) in enumerate(self._previous_fits[idx], start=1):
            self.stats_tableWidget.setHorizontalHeaderItem(0, QtWidgets.QTableWidgetItem(name))
            it = QtWidgets.QTableWidgetItem(f'{dof}')
            it.setFlags(it.flags() ^ QtCore.Qt.ItemIsEditable)
            self.stats_tableWidget.setItem(0, col, it)

        for row, (k, v) in enumerate(res.statistics.items(), start=1):
            self.stats_tableWidget.setVerticalHeaderItem(row, QtWidgets.QTableWidgetItem(k))
            it = QtWidgets.QTableWidgetItem(f'{v:.4f}')
            it.setFlags(it.flags() ^ QtCore.Qt.ItemIsEditable)
            self.stats_tableWidget.setItem(row, 0, it)

            best_idx = -1
            best_val = v
            for col, (_, stats, _) in enumerate(self._previous_fits[idx], start=1):
                if k in ['adj. R^2', 'R^2']:
                    best_idx = col if best_val < stats[k] else max(0, best_idx)
                else:
                    best_idx = col if best_val > stats[k] else max(0, best_idx)
                it = QtWidgets.QTableWidgetItem(f'{stats[k]:.4f}')
                it.setFlags(it.flags() ^ QtCore.Qt.ItemIsEditable)
                self.stats_tableWidget.setItem(row, col, it)

            if best_idx > -1:
                self.stats_tableWidget.item(row, best_idx).setBackground(QtGui.QColor('green'))
                self.stats_tableWidget.item(row, best_idx).setForeground(QtGui.QColor('white'))

        row = self.stats_tableWidget.rowCount() - 2
        self.stats_tableWidget.setVerticalHeaderItem(row, QtWidgets.QTableWidgetItem('F'))
        self.stats_tableWidget.setItem(row, 0, QtWidgets.QTableWidgetItem('-'))

        self.stats_tableWidget.setVerticalHeaderItem(row+1, QtWidgets.QTableWidgetItem('Pr(>F)'))
        self.stats_tableWidget.setItem(row+1, 0, QtWidgets.QTableWidgetItem('-'))

        for col, (_, stats, dof) in enumerate(self._previous_fits[idx], start=1):
            f_value, prob_f = res.f_test(stats['chi^2'], dof)
            it = QtWidgets.QTableWidgetItem(f'{f_value:.4g}')
            it.setFlags(it.flags() ^ QtCore.Qt.ItemIsEditable)
            self.corr_tableWidget.setItem(row, col, it)

            it = QtWidgets.QTableWidgetItem(f'{prob_f:.4g}')
            it.setFlags(it.flags() ^ QtCore.Qt.ItemIsEditable)
            if prob_f < 0.05:
                it.setBackground(QtGui.QColor('green'))
                it.setForeground(QtGui.QColor('white'))
            self.stats_tableWidget.setItem(row+1, col, it)

    @QtCore.pyqtSlot(QtWidgets.QAbstractButton)
    def on_buttonBox_clicked(self, button: QtWidgets.QAbstractButton):
        button_type = self.buttonBox.standardButton(button)

        if button_type == self.buttonBox.Retry:
            self.redoFit.emit(self._results)

        elif button_type == self.buttonBox.Ok:
            graph = '-1'
            if self.parameter_checkbox.isChecked():
                if self.graph_checkBox.checkState() == QtCore.Qt.Checked:
                    graph = ''
                else:
                    graph = self.graph_comboBox.currentData()

            plot_fits = self.curve_checkbox.isChecked()

            parts = self.partial_checkBox.checkState() == QtCore.Qt.Checked

            extrapolate = [None, None, None]
            if self.extrapolate_box.isChecked():
                try:
                    extrapolate[0] = float(self.minx_line.text())
                except TypeError:
                    pass
                try:
                    extrapolate[1] = float(self.maxx_line.text())
                except TypeError:
                    pass
                try:
                    extrapolate[2] = int(self.numx_line.text())
                except TypeError:
                    pass

            self.closed.emit(self._results, self._opts, graph, plot_fits, parts, extrapolate)
            self.accept()

        else:
            self.reject()


class FitExtension(QtWidgets.QDialog):
    def __init__(self, parent=None):
        super().__init__(parent=parent)
        gridLayout = QtWidgets.QGridLayout(self)

        self.label = QtWidgets.QLabel('Minimum value')
        gridLayout.addWidget(self.label, 0, 0, 1, 1)

        self.min_line = QtWidgets.QLineEdit()
        self.min_line.setValidator(QtGui.QDoubleValidator())
        gridLayout.addWidget(self.min_line, 0, 1, 1, 1)

        self.label_2 = QtWidgets.QLabel('Maximum value')
        gridLayout.addWidget(self.label_2, 1, 0, 1, 1)

        self.max_line = QtWidgets.QLineEdit()
        self.max_line.setValidator(QtGui.QDoubleValidator())
        gridLayout.addWidget(self.max_line, 1, 1, 1, 1)

        self.label_3 = QtWidgets.QLabel('Number of pts.')
        gridLayout.addWidget(self.label_3, 2, 0, 1, 1)

        self.num_pts = QtWidgets.QLineEdit()
        self.num_pts.setValidator(QtGui.QIntValidator())
        gridLayout.addWidget(self.num_pts, 2, 1, 1, 1)

        self.buttonBox = QtWidgets.QDialogButtonBox()
        self.buttonBox.setOrientation(QtCore.Qt.Horizontal)
        self.buttonBox.setStandardButtons(QtWidgets.QDialogButtonBox.Cancel | QtWidgets.QDialogButtonBox.Ok)
        gridLayout.addWidget(self.buttonBox, 3, 0, 1, 2)

        self.setLayout(gridLayout)

        self.buttonBox.accepted.connect(self.accept)
        self.buttonBox.rejected.connect(self.reject)

    @property
    def values(self):
        try:
            xmin = float(self.min_line.text())
            xmax = float(self.max_line.text())
            nums = int(self.num_pts.text())
        except TypeError:
            return None

        return xmin, xmax, nums