nmreval/src/gui_qt/fit/result.py

from math import isnan

from pyqtgraph import mkBrush

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, dict)
    redoFit = QtCore.pyqtSignal(dict)

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

        self._management = management

        self._prevs = {}
        self._models = {}

        for (res, parts) in results:
            idx = res.idx
            data_k = management.data[idx]

            if res.name not in self._models:
                self._models[res.name] = []

            self._models[res.name].append(idx)

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

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

        self.residplot = self.graphicsView.addPlot(row=0, col=0)
        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.residplot.addItem(self.resid_graph)
        self.residplot.addItem(self.resid_graph_imag)
        self.residplot.setLabel('left', 'Residual')

        self.fitplot = self.graphicsView.addPlot(row=1, col=0)
        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.fitplot.addItem(self.data_graph)
        self.fitplot.addItem(self.data_graph_imag)
        self.fitplot.setLabel('left', 'Function')

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

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

        self.sets_comboBox.blockSignals(True)
        for n in self._models.keys():
            self.sets_comboBox.addItem(n)
        self.sets_comboBox.blockSignals(False)

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

        self.param_tableWidget.itemClicked.connect(self.show_results)
        self.param_tableWidget.horizontalHeader().sectionClicked.connect(lambda i: self.show_results(None, idx=i))

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

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

    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):
        model_name = self.sets_comboBox.itemText(idx)
        sets = self._models[model_name]
        self.param_tableWidget.setColumnCount(len(sets))

        r = self._results[sets[0]]
        self.param_tableWidget.setRowCount(len(r.parameter))

        for i, pval in enumerate(r.parameter.values()):
            name = pval.full_name
            p_header = QtWidgets.QTableWidgetItem(convert(name, 'tex', 'html', brackets=False))
            self.param_tableWidget.setVerticalHeaderItem(i, p_header)

        for i, set_id in enumerate(sets):
            data_i = self._management[set_id]
            header_item = QtWidgets.QTableWidgetItem(data_i.name)
            header_item.setData(QtCore.Qt.UserRole, set_id)
            self.param_tableWidget.setHorizontalHeaderItem(i, header_item)

            res = self._results[set_id]
            for j, pvalue in enumerate(res.parameter.values()):
                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, i, QtWidgets.QTableWidgetItem('-'))
                else:
                    self.param_tableWidget.setItem(2*j+1, i, QtWidgets.QTableWidgetItem())
                item = QtWidgets.QTableWidgetItem(item_text)
                self.param_tableWidget.setItem(j, i, item)

        self.param_tableWidget.resizeColumnsToContents()
        self.param_tableWidget.selectColumn(0)
        self.show_results(None, idx=0)

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

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

    def show_results(self, item, idx=None):
        if item is not None:
            idx = self.param_tableWidget.indexFromItem(item).column()

        set_id = self.param_tableWidget.horizontalHeaderItem(idx).data(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):
        res = self._results[idx]
        iscomplex = res.iscomplex

        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)
        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=[])

        self.fitplot.setLogMode(x=res.islog)
        self.residplot.setLogMode(x=res.islog)

    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._prevs[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._prevs[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._prevs[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._prevs[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()

            if self.partial_checkBox.checkState() == QtCore.Qt.Checked:
                self.closed.emit(self._results, self._opts, graph, plot_fits, self._parts)
            else:
                self.closed.emit(self._results, self._opts, graph, plot_fits, {})

            self.accept()

        else:
            self.reject()