from __future__ import annotations from itertools import cycle, count from nmreval.configs import config_paths from nmreval import models from nmreval.lib.importer import find_models from nmreval.lib.colors import BaseColor, Tab10 from nmreval.utils.text import convert from ..lib.iconloading import get_icon from .._py.fitfunctionwidget import Ui_Form from ..Qt import QtWidgets, QtCore, QtGui class QFunctionWidget(QtWidgets.QWidget, Ui_Form): func_cnt = count() func_colors = cycle(Tab10) op_names = ['plus_icon', 'mal_icon', 'minus_icon', 'geteilt_icon'] newFunction = QtCore.pyqtSignal(int, int) treeChanged = QtCore.pyqtSignal() itemRemoved = QtCore.pyqtSignal(int) showFunction = QtCore.pyqtSignal(int) def __init__(self, parent=None): super().__init__(parent=parent) self.setupUi(self) self._types = [] self.user_functions = {} self.functions = find_models(models) self.read_and_load_functions() self.functree.treeChanged.connect(lambda: self.treeChanged.emit()) self.functree.itemRemoved.connect(self.remove_function) self.iscomplex = False self.complex_widget.hide() for i, op_icon in enumerate(self.op_names): self.operator_combobox.setItemIcon(i, get_icon(op_icon)) def read_and_load_functions(self): """ Reads user-defined fit functions from usermodels.py and adds them to list of fit function. Previous class definitions are replaced. """ user_defined = [] try: user_defined = find_models(config_paths() / 'usermodels.py') except FileNotFoundError: pass for model in user_defined: name = model.__name__ if name not in self.user_functions: self.functions.append(model) self.user_functions[name] = model else: idx = self.functions.index(self.user_functions[name]) self.functions[idx] = model self.user_functions[name] = model self._types = [] for m in self.functions: try: m.type except AttributeError: m.type = 'Other' if m.type not in self._types: self._types.append(m.type) self.typecomboBox.blockSignals(True) self.typecomboBox.clear() self.typecomboBox.addItems(sorted(self._types)) self.typecomboBox.blockSignals(False) self.change_group(0) def __len__(self) -> int: num = 0 iterator = QtWidgets.QTreeWidgetItemIterator(self.functree) while iterator.value(): num += 1 iterator += 1 return num @QtCore.pyqtSlot(int, name='on_typecomboBox_currentIndexChanged') def change_group(self, idx: int): """ Change items in fitcombobox to new entries """ self.fitcomboBox.blockSignals(True) while self.fitcomboBox.count(): self.fitcomboBox.removeItem(0) selected_type = self.typecomboBox.itemText(idx) for m in self.functions: if m.type == selected_type: self.fitcomboBox.addItem(m.name, userData=self.functions.index(m)) self.fitcomboBox.blockSignals(False) self.on_fitcomboBox_currentIndexChanged(0) @QtCore.pyqtSlot(int, name='on_fitcomboBox_currentIndexChanged') def change_function(self, idx: int): """ Display new equation on changing function """ index = self.fitcomboBox.itemData(idx) if self.functions: fitfunc = self.functions[index] try: self.fitequation.setText(convert(fitfunc.equation)) except AttributeError: self.fitequation.setText('') @QtCore.pyqtSlot(name='on_use_function_button_clicked') def new_function(self): idx = self.fitcomboBox.itemData(self.fitcomboBox.currentIndex()) cnt = next(self.func_cnt) op = self.operator_combobox.currentIndex() name = self.functions[idx].name col = next(self.func_colors) self.newFunction.emit(idx, cnt) self.add_function(idx, cnt, op, name, col) def add_function(self, idx: int, cnt: int, op: int, name: str, color: str | tuple[float, float, float] | BaseColor, **kwargs): """ Add function to tree and dictionary of functions. """ if isinstance(color, BaseColor): qcolor = QtGui.QColor.fromRgbF(*color.rgb(normed=True)) elif isinstance(color, tuple): qcolor = QtGui.QColor.fromRgbF(*color) else: qcolor = QtGui.QColor(color) self.functree.add_function(idx, cnt, op, name, qcolor, **kwargs) f = self.functions[idx] if hasattr(f, 'iscomplex') and f.iscomplex: self.iscomplex = True self.complex_widget.show() @QtCore.pyqtSlot(int) def remove_function(self, idx: int): iterator = QtWidgets.QTreeWidgetItemIterator(self.functree) self.iscomplex = False while iterator.value(): item = iterator.value() f = self.functions[item.data(0, QtCore.Qt.UserRole)] if hasattr(f, 'iscomplex') and f.iscomplex: self.iscomplex = True break iterator += 1 self.complex_widget.setVisible(self.iscomplex) self.itemRemoved.emit(idx) @QtCore.pyqtSlot(QtWidgets.QTreeWidgetItem, int, name='on_functree_itemClicked') def show_parameter(self, item: QtWidgets.QTreeWidgetItem, _: int): self.showFunction.emit(item.data(0, self.functree.counterRole)) def get_selected(self): function_nr, idx = self.functree.get_selected() if function_nr is not None: return self.functions[function_nr], idx else: return None, None def get_functions(self, full: bool = True, clsname=False, include_all: bool = True): """ Create nested list of functions in tree. Parameters saved are idx (Index of function in list of all functions), cnt (counter of number to associate with functione values), ops (+, -, *, /), and maybe children. """ used_functions = self.functree.get_functions(full=full) self._prepare_function_for_model(used_functions, full=full, clsname=clsname, include_all=include_all) return used_functions def _prepare_function_for_model(self, func_list: list[dict], full: bool = True, clsname: bool = False, include_all: bool = True): for func_args in func_list: is_active = func_args.get('active') if (not is_active) and (not include_all): continue if not clsname: func_args['func'] = self.functions[func_args['idx']] else: func_args['func'] = self.functions[func_args['idx']].name if not full: func_args.pop('active') func_args.pop('idx') if func_args['children']: self._prepare_function_for_model(func_args['children'], full=full, clsname=clsname, include_all=include_all) def get_parameter_list(self): all_parameters = {} iterator = QtWidgets.QTreeWidgetItemIterator(self.functree) while iterator.value(): item = iterator.value() f = self.functions[item.data(0, QtCore.Qt.UserRole)] cnt = item.data(0, self.functree.counterRole) all_parameters[f'{f.name}_{cnt}'] = [(convert(pp, new='str'), (cnt, i)) for i, pp in enumerate(f.params)] iterator += 1 return all_parameters def get_complex_state(self): return self.complex_comboBox.currentIndex() if self.iscomplex else None def set_complex_state(self, state): if state is not None: self.complex_comboBox.setCurrentIndex(state) def clear(self): self.functree.blockSignals(True) self.functree.clear() self.functree.blockSignals(False) self.complex_comboBox.setCurrentIndex(0) self.complex_widget.hide()