forked from IPKM/nmreval
239 lines
8.2 KiB
Python
239 lines
8.2 KiB
Python
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 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()
|