nmreval/src/gui_qt/data/plot_dialog.py

from random import randint

import numpy as np

from nmreval import models
from nmreval.lib.importer import find_models
from nmreval.lib.utils import valid_function

from ..Qt import QtGui, QtCore, QtWidgets
from .._py.setbyfunction_dialog import Ui_NewCurveDialog


class QPlotDialog(QtWidgets.QDialog, Ui_NewCurveDialog):
    line_created = QtCore.pyqtSignal(object)

    def __init__(self):
        super().__init__()
        self.setupUi(self)

        self._function = find_models(models)

        self.lineEdit_3.setValidator(QtGui.QDoubleValidator())
        self.lineEdit_4.setValidator(QtGui.QDoubleValidator())
        self.lineEdit_5.setValidator(QtGui.QIntValidator().setBottom(0))

        self.buttonBox.accepted.connect(self.make_line)

        for cb in [self.comboBox, self.comboBox_2, self.comboBox_4, self.comboBox_5]:
            cb.setCurrentIndex(randint(0, cb.count()))

        for cb in [self.comboBox_6, self.comboBox_7]:
            self.load_models(cb)

    def load_models(self, cb: QtWidgets.QComboBox):
        for f in self._function:
            cb.addItem(f'{f.name} ({f.type})', userData=f)

    @QtCore.pyqtSlot(name='on_pushButton_clicked')
    def check_input(self):
        err = []
        try:
            start = float(self.lineEdit_3.text())
        except ValueError:
            err.append(0)
            start = 1
        try:
            stop = float(self.lineEdit_4.text())
        except ValueError:
            err.append(1)
            stop = 10

        try:
            nums = int(self.lineEdit_5.text())
        except ValueError:
            err.append(2)
            nums = 10

        if self.checkBox.isChecked():
            if start <= 0 or stop <= 0:
                err.append(3)
                start, stop = abs(start)+1e-12, abs(stop)
            grid = np.geomspace(start, stop, num=nums)
        else:
            grid = np.linspace(start, stop, num=nums)

        x_func = self.lineEdit.text()
        x, isok = valid_function(x_func, extra_namespace={'i': grid})
        if not isok:
            err.append(4)
            x = grid

        y_func = self.lineEdit_2.text()
        y, isok = valid_function(y_func, extra_namespace={'i': grid, 'x':x})
        if not isok:
            err.append(5)

        msg_err = {0: 'Invalid value for grid start',
                   1: 'Invalid value for grid end',
                   2: 'Invalid number of grid steps',
                   3: 'Negative numbers in logarithmic grid',
                   4: 'Invalid expression for x',
                   5: 'Invalid expression for y'
                   }

        if err:
            m = '\n'.join([msg_err[e] for e in err])
            QtWidgets.QMessageBox().information(self, 'Error detected', m)
            return False

        return True

    def make_line(self):
        if not self.check_input():
            return

        start = float(self.lineEdit_3.text())
        stop = float(self.lineEdit_4.text())
        nums = int(self.lineEdit_5.text())
        if self.checkBox.isChecked():
            x = np.geomspace(start, stop, num=nums)
        else:
            x = np.linspace(start, stop, num=nums)
        x_func = self.lineEdit.text()
        y_func = self.lineEdit_2.text()

        sym = self.comboBox.currentText()
        lin = self.comboBox_2.currentText()

        name = self.lineEdit_6.text()
        if not name:
            name = 'self done'

        lw = self.doubleSpinBox.value()
        sw = self.spinBox.value()