nmreval/src/gui_qt/nmr/t1_from_tau.py

from __future__ import annotations

from nmreval.nmr.coupling import *
from nmreval.distributions import ColeCole, ColeDavidson, HavriliakNegami, KWW, LogGaussian
from nmreval.utils import pi
from nmreval.utils.text import convert

from ..Qt import QtGui, QtCore, QtWidgets
from ..lib.forms import SelectionWidget, Widget
from .._py.t1_calc_dialog import Ui_Dialog


class QRelaxCalc(QtWidgets.QDialog, Ui_Dialog):
    newData = QtCore.pyqtSignal(dict)

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

        self.graphs = {}

        self.specdens = [ColeCole, ColeDavidson, HavriliakNegami, KWW]
        self.coupling = [Quadrupolar, HomoDipolar, Czjzek]
        self.tau_parameter = []

        for line_edit in [self.ea_lineEdit, self.tau0_lineEdit, self.start_lineEdit, self.stop_lineEdit,
                          self.tau0_vft_lineEdit, self.b_vft_lineEdit, self.t0_vft_lineEdit]:
            line_edit.setValidator(QtGui.QDoubleValidator())

        for s in self.specdens:
            self.specdens_combobox.addItem(s.name)

        for c in self.coupling:
            self.coupling_combobox.addItem(c.name)

        self.specdens_combobox.currentIndexChanged.connect(self.update_specdens_model)
        self.coupling_combobox.currentIndexChanged.connect(self.update_coupling_model)
        self.buttonGroup.buttonClicked.connect(self.change_axis)

        self.update_specdens_model(0)
        self.update_coupling_model(0)
        self.temp_combo_change(0)
        self.tau_combo_changed(0)
        self.change_axis(self.radioButton)

    def set_graphs(self, graphs: dict):
        self.graph_combobox.clear()
        self.tau_graph_combobox.clear()

        self.tau_graph_combobox.blockSignals(True)
        for key, (name, _) in graphs.items():
            self.graph_combobox.addItem(name, userData=key)
            self.tau_graph_combobox.addItem(name, userData=key)
        self.tau_graph_combobox.blockSignals(False)

        self.graphs = graphs

        self.update_specdens_model(self.specdens_combobox.currentIndex())
        self.update_coupling_model(self.coupling_combobox.currentIndex())
        self.tau_graph_changed(self.tau_graph_combobox.currentIndex())

    def update_graphs(self, graphs: list[tuple[str, str]]):
        current_id = self.graph_combobox.currentData()
        self.graph_combobox.clear()
        for (gid, name) in graphs:
            self.graph_combobox.addItem(name, userData=gid)

        self.graph_combobox.setCurrentIndex(self.graph_combobox.findData(current_id, QtCore.Qt.UserRole))

    @QtCore.pyqtSlot(int, name='on_graph_checkbox_stateChanged')
    def changed_state(self, checked):
        self.graph_combobox.setEnabled(checked != QtCore.Qt.Checked)

    @QtCore.pyqtSlot(int, name='on_tau_graph_combobox_currentIndexChanged')
    def tau_graph_changed(self, idx: int):
        key = self.tau_graph_combobox.itemData(idx)
        self.tau_set_combobox.clear()
        if self.graphs:
            for (set_key, set_name) in self.graphs[key][1]:
                self.tau_set_combobox.addItem(set_name, userData=set_key)

    @QtCore.pyqtSlot(QtWidgets.QAbstractButton)
    def change_axis(self, bttn: QtWidgets.QRadioButton):
        self.temp_widget.setVisible(bttn in [self.radioButton_3, self.radioButton_4])
        if bttn == self.radioButton_2:
            self.label_7.setText('\u03c4 / s')
        else:
            self.label_7.setText('\u03c9 / Hz')

    @QtCore.pyqtSlot(int, name='on_temp_combobox_currentIndexChanged')
    def temp_combo_change(self, idx: int):
        self.arr_widget.setVisible(idx == 0)
        self.vft_widget.setVisible(idx == 1)

    @QtCore.pyqtSlot(int, name='on_x_input_combobox_currentIndexChanged')
    def tau_combo_changed(self, idx: int):
        self.range_widget.setVisible(idx == 0)
        self.data_widget.setVisible(idx == 1)

    @QtCore.pyqtSlot(int)
    def update_coupling_model(self, idx: int):
        while self.verticalLayout_4.count():
            item = self.verticalLayout_4.takeAt(0)
            try:
                item.widget().deleteLater()
            except AttributeError:
                pass

        if self.coupling[idx].parameter is not None:
            for p in self.coupling[idx].parameter:
                _temp = Widget(convert(p), tree=self.graphs, collapsing=True, parent=self)
                self.verticalLayout_4.addWidget(_temp)

        if self.coupling[idx].choice is not None:
            for c in self.coupling[idx].choice:
                _temp = SelectionWidget(*c, parent=self)
                self.verticalLayout_4.addWidget(_temp)

    @QtCore.pyqtSlot(int)
    def update_specdens_model(self, idx):
        while self.verticalLayout_3.count():
            item = self.verticalLayout_3.takeAt(0)
            try:
                item.widget().deleteLater()
            except AttributeError:
                pass

        if self.specdens[idx].parameter is not None:
            for param in self.specdens[idx].parameter:
                _temp = Widget(convert(param), tree=self.graphs, collapsing=True, parent=self)
                self.verticalLayout_3.addWidget(_temp)

    def get_taus(self, dic: dict):
        dic['tau_type'] = {0: 'raw', 1: 'mean', 2: 'peak', 3: 'logmean'}[self.xtype_combobox.currentIndex()]
        dic['axis1'] = {self.radioButton: 'tau', self.radioButton_2: 'omega',
                        self.radioButton_3: 't', self.radioButton_4: 'invt1000'}[self.buttonGroup.checkedButton()]

        dic['val2'] = float(self.second_x_lineEdit.text())
        if dic['axis1'] != 'omega':
            dic['val2'] *= 2*pi

        idx = self.x_input_combobox.currentIndex()
        if idx == 0:
            dic['pts'] = (float(self.start_lineEdit.text()), float(self.stop_lineEdit.text()),
                          self.spinBox.value(), self.checkBox.isChecked())
            if self.buttonGroup.checkedButton() in [self.radioButton_3, self.radioButton_4]:
                if self.temp_combobox.currentIndex() == 0:
                    dic['t_param'] = (float(self.tau0_lineEdit.text()), float(self.ea_lineEdit.text()))
                else:
                    dic['t_param'] = (float(self.tau0_vft_lineEdit.text()), float(self.b_vft_lineEdit.text()),
                                      float(self.t0_vft_lineEdit.text()))
        else:
            dic['pts'] = (self.tau_set_combobox.currentData(), self.y_radioButton.isChecked())

    def get_coupling(self, dic):
        dic['coup'] = self.coupling[self.coupling_combobox.currentIndex()]

        parameter = []
        kwargs = {}
        for i in range(self.verticalLayout_4.count()):
            p = self.verticalLayout_4.itemAt(i).widget()
            if isinstance(p, Widget):
                parameter.append(p.value)
            elif isinstance(p, SelectionWidget):
                kwargs.update(p.value)
            else:
                raise TypeError('WTF: Unknown widget', p)

        dic['cp_param'] = (parameter, kwargs)

    def get_specdens(self, dic):
        dic['spec_dens'] = self.specdens[self.specdens_combobox.currentIndex()]

        parameter = []
        kwargs = {}
        for i in range(self.verticalLayout_3.count()):
            p = self.verticalLayout_3.itemAt(i).widget()
            if isinstance(p, Widget):
                parameter.append(p.value)
            elif isinstance(p, SelectionWidget):
                kwargs[p.argname] = p.value
            else:
                raise TypeError('WTF?: Unknown widget', p)

        dic['sd_param'] = (parameter, kwargs)

    def calc_relaxation(self):
        opts = {}

        self.get_taus(opts)
        self.get_coupling(opts)
        self.get_specdens(opts)

        opts['out'] = {0: 't1', 1: 't2'}[self.relax_combox.currentIndex()]
        opts['graph'] = '' if self.graph_checkbox.isChecked() else self.graph_combobox.currentData()

        self.newData.emit(opts)

    def accept(self):
        self.calc_relaxation()
        super().accept()