from __future__ import annotations from functools import reduce from itertools import count, cycle from operator import add from string import ascii_letters import numpy as np from pyqtgraph import mkPen from nmreval.fit._meta import MultiModel, ModelFactory from nmreval.fit.result import FitResult from .fit_forms import FitTableWidget from .fit_parameter import QFitParameterWidget from ..lib.pg_objects import PlotItem from ..Qt import QtGui, QtCore, QtWidgets from .._py.fitdialog import Ui_FitDialog class QFitDialog(QtWidgets.QWidget, Ui_FitDialog): func_cnt = count() model_cnt = cycle(ascii_letters) preview_num = 201 preview_emit = QtCore.pyqtSignal(dict, int, bool) fitStartSig = QtCore.pyqtSignal(dict, list, dict) abortFit = QtCore.pyqtSignal() def __init__(self, mgmt=None, parent=None): super().__init__(parent=parent) self.setupUi(self) self.parameters = {} self.preview_lines = [] self._current_function = None self.param_widgets = {} self._management = mgmt self._current_model = next(QFitDialog.model_cnt) self.show_combobox.setItemData(0, self._current_model, QtCore.Qt.UserRole) self.default_combobox.setItemData(0, self._current_model, QtCore.Qt.UserRole) self.data_table = FitTableWidget(self.data_widget) self.data_widget.addWidget(self.data_table) self.data_widget.setText('Data') self.models = {} self._func_list = {} self._complex = {} self.connected_figure = '' self.model_frame.hide() self.preview_button.hide() self.abort_button.clicked.connect(lambda: self.abortFit.emit()) self.functionwidget.newFunction.connect(self.add_function) self.functionwidget.showFunction.connect(self.show_function_parameter) self.functionwidget.itemRemoved.connect(self.remove_function) self.read_and_load_functions = self.functionwidget.read_and_load_functions @QtCore.pyqtSlot(int, int) def add_function(self, function_idx: int, function_id: int): self.show_function_parameter(function_id, function_idx) self.newmodel_button.setEnabled(True) @QtCore.pyqtSlot(int) def remove_function(self, idx: int): """ Remove function and children from tree and dictionary """ w = self.param_widgets[idx] self.stackedWidget.removeWidget(w) w.deleteLater() del self.param_widgets[idx] if len(self.functionwidget) == 0: # empty model self.newmodel_button.setEnabled(False) self.deletemodel_button.setEnabled(False) self._current_function = None else: f_tree = self.functionwidget.functree func_idx = f_tree.currentItem().data(0, f_tree.counterRole) self._current_function = self.functionwidget.functions[func_idx] @QtCore.pyqtSlot(int) def show_function_parameter(self, function_id: int, function_idx: int = None): """ Display parameter associated with selected function. """ if function_id in self.param_widgets: dialog = self.param_widgets[function_id] else: # create new widget for function if function_idx is not None: function = self.functionwidget.functions[function_idx] else: raise ValueError('No function index given') if function is None: return dialog = QFitParameterWidget(self.stackedWidget) data_names = self.data_table.data_list(include_name=True) dialog.set_function(function, function_idx) dialog.load(data_names) dialog.value_requested.connect(self.look_value) self.stackedWidget.addWidget(dialog) self.param_widgets[function_id] = dialog self.stackedWidget.setCurrentWidget(dialog) # collect parameter names etc. to allow linkage self._func_list[self._current_model] = self.functionwidget.get_parameter_list() dialog.set_links(self._func_list) # show same tab (general parameter/Data parameter) tab_idx = 0 if self._current_function is not None: tab_idx = self.param_widgets[self._current_function].tabWidget.currentIndex() dialog.tabWidget.setCurrentIndex(tab_idx) self._current_function = function_id def look_value(self, idx: int): func_widget = self.param_widgets[self._current_function] set_ids = [func_widget.comboBox.itemData(i) for i in range(func_widget.comboBox.count())] for s in set_ids: func_widget.data_values[s][idx] = self._management[s].value func_widget.change_data(func_widget.comboBox.currentIndex()) def get_functions(self): """ update functions, parameters""" self.models[self._current_model] = self.functionwidget.get_functions() self._complex[self._current_model] = self.functionwidget.get_complex_state() self._func_list[self._current_model] = self.functionwidget.get_parameter_list() def load(self, ids: list[str]): """ Add name and id of dataset to list. """ self.data_table.load(ids) if self.models: for m in self.models.keys(): self.data_table.add_model(m) else: self.data_table.add_model(self._current_model) for dialog in self.param_widgets.values(): dialog.load(ids) @QtCore.pyqtSlot(name='on_newmodel_button_clicked') def make_new_model(self): """ Save model with all its functions in dictionary and adjust gui. """ self.deletemodel_button.setEnabled(True) self.model_frame.show() idx = next(QFitDialog.model_cnt) self.data_table.add_model(idx) self.default_combobox.addItem('Model '+idx, userData=idx) self.show_combobox.addItem('Model '+idx, userData=idx) self.show_combobox.setItemData(self.show_combobox.count()-1, idx, QtCore.Qt.UserRole) self.show_combobox.setCurrentIndex(self.show_combobox.count()-1) self._current_model = idx self.stackedWidget.setCurrentIndex(0) @QtCore.pyqtSlot(int, name='on_show_combobox_currentIndexChanged') def change_model(self, idx: int): """ Save old model and display new model. """ self.get_functions() self.functionwidget.clear() self._current_model = self.show_combobox.itemData(idx, QtCore.Qt.UserRole) if self._current_model in self.models and len(self.models[self._current_model]): for el in self.models[self._current_model]: self.functionwidget.add_function(**el) self.functionwidget.set_complex_state(self._complex[self._current_model]) else: self.stackedWidget.setCurrentIndex(0) @QtCore.pyqtSlot(name='on_deletemodel_button_clicked') def remove_model(self): model_id = self._current_model self.show_combobox.removeItem(self.show_combobox.findData(model_id)) self.default_combobox.removeItem(self.default_combobox.findData(model_id)) for m in self.models[model_id]: func_id = m['cnt'] self.stackedWidget.removeWidget(self.param_widgets[func_id]) self.param_widgets.pop(func_id) self._complex.pop(model_id) self._func_list.pop(model_id) self.models.pop(model_id) self.data_table.remove_model(model_id) if len(self.models) == 1: self.model_frame.hide() def _prepare(self, model: list, function_use: list = None, parameter: dict = None, add_idx: bool = False, cnt: int = 0) -> tuple[dict, int]: if parameter is None: parameter = {'parameter': {}, 'lb': (), 'ub': (), 'var': [], 'glob': {'idx': [], 'p': [], 'var': [], 'lb': [], 'ub': []}, 'links': [], 'color': []} for i, f in enumerate(model): if not f['active']: continue try: p, lb, ub, var, glob, links = self.param_widgets[f['cnt']].get_parameter(function_use) except ValueError as e: _ = QtWidgets.QMessageBox().warning(self, 'Invalid value', str(e), QtWidgets.QMessageBox.Ok) return {}, -1 p_len = len(parameter['lb']) parameter['lb'] += lb parameter['ub'] += ub parameter['var'] += var parameter['links'] += links parameter['color'] += [f['color']] cnt = f['cnt'] for p_k, v_k in p.items(): if add_idx: kw_k = {f'{k}_{cnt}': v for k, v in v_k[1].items()} else: kw_k = v_k[1] if p_k in parameter['parameter']: params, kw = parameter['parameter'][p_k] params += v_k[0] kw.update(kw_k) else: parameter['parameter'][p_k] = (v_k[0], kw_k) for g_k, g_v in glob.items(): if g_k != 'idx': parameter['glob'][g_k] += g_v else: parameter['glob']['idx'] += [idx_i + p_len for idx_i in g_v] if add_idx: cnt += 1 if f['children']: # recurse for children child_parameter, cnt = self._prepare(f['children'], parameter=parameter, add_idx=add_idx, cnt=cnt) return parameter, cnt @QtCore.pyqtSlot(name='on_fit_button_clicked') def start_fit(self): self.get_functions() data = self.data_table.collect_data(default=self.default_combobox.currentData()) func_dict = {} for k, mod in self.models.items(): func, order, param_len = ModelFactory.create_from_list(mod) if func is None: continue if k in data: parameter, _ = self._prepare(mod, function_use=data[k], add_idx=isinstance(func, MultiModel)) if parameter is None: return parameter['func'] = func parameter['order'] = order parameter['len'] = param_len parameter['complex'] = self._complex[k] if self._complex[k] is not None: for p_k, p_v in parameter['parameter'].items(): p_v[1].update({'complex_mode': self._complex[k]}) parameter['parameter'][p_k] = p_v[0], p_v[1] func_dict[k] = parameter replaceable = [] for k, v in func_dict.items(): for i, link_i in enumerate(v['links']): if link_i is None: continue rep_model, rep_func, rep_pos = link_i try: f = func_dict[rep_model] except KeyError: QtWidgets.QMessageBox().warning(self, 'Invalid value', 'Parameter cannot be linked: Model is unused', QtWidgets.QMessageBox.Ok) return try: f_idx = f['order'].index(rep_func) except ValueError: QtWidgets.QMessageBox().warning(self, 'Invalid value', 'Parameter cannot be linked: ' 'Function is probably not checked or deleted', QtWidgets.QMessageBox.Ok) return repl_idx = sum(f['len'][:f_idx])+rep_pos if repl_idx not in f['glob']['idx']: _ = QtWidgets.QMessageBox().warning(self, 'Invalid value', 'Parameter cannot be linked: ' 'Destination is not a global parameter.', QtWidgets.QMessageBox.Ok) return replaceable.append((k, i, rep_model, repl_idx)) replace_value = None for p_k in f['parameter'].values(): replace_value = p_k[0][repl_idx] break if replace_value is not None: for p_k in v['parameter'].values(): p_k[0][i] = replace_value weight = ['None', 'y', 'y2', 'Deltay'][self.weight_combobox.currentIndex()] fit_args = {'we': weight} if func_dict: self.fitStartSig.emit(func_dict, replaceable, fit_args) return func_dict @QtCore.pyqtSlot(int, name='on_preview_checkbox_stateChanged') def show_preview(self, state: int): if state: self.preview_button.show() self.preview_checkbox.setText('') self._prepare_preview() else: self.preview_emit.emit({}, -1, False) self.preview_lines = [] self.preview_button.hide() self.preview_checkbox.setText('Preview') @QtCore.pyqtSlot(name='on_preview_button_clicked') def _prepare_preview(self): self.get_functions() default_model = self.default_combobox.currentData() data = self.data_table.collect_data(default=default_model) func_dict = {} for k, mod in self.models.items(): func, order, param_len = ModelFactory.create_from_list(mod) multiple_funcs = isinstance(func, MultiModel) if k in data: parameter, _ = self._prepare(mod, function_use=data[k], add_idx=multiple_funcs) parameter['func'] = func parameter['order'] = order parameter['len'] = param_len func_dict[k] = parameter for v in func_dict.values(): for i, link_i in enumerate(v['links']): if link_i is None: continue rep_model, rep_func, rep_pos = link_i f = func_dict[rep_model] f_idx = f['order'].index(rep_func) repl_idx = sum(f['len'][:f_idx]) + rep_pos replace_value = None for p_k in f['parameter'].values(): replace_value = p_k[0][repl_idx] break if replace_value is not None: for p_k in v['parameter'].values(): p_k[0][i] = replace_value self.preview_emit.emit(func_dict, QFitDialog.preview_num, True) def make_previews(self, x, models_parameters: dict): self.preview_lines = [] for k, model in models_parameters.items(): f = model['func'] is_complex = self._complex[k] parameters = model['parameter'] color = model['color'] seen_parameter = [] for p, kwargs in parameters.values(): if (p, kwargs) in seen_parameter: # plot only previews with different parameter continue seen_parameter.append((p, kwargs)) if is_complex is not None: y = f.func(x, *p, complex_mode=is_complex, **kwargs) if np.iscomplexobj(y): self.preview_lines.append(PlotItem(x=x, y=y.real, pen=mkPen(width=3))) self.preview_lines.append(PlotItem(x=x, y=y.imag, pen=mkPen(width=3))) else: self.preview_lines.append(PlotItem(x=x, y=y, pen=mkPen(width=3))) else: y = f.func(x, *p, **kwargs) self.preview_lines.append(PlotItem(x=x, y=y, pen=mkPen(width=3))) if isinstance(f, MultiModel): sub_kwargs = kwargs.copy() if is_complex is not None: sub_kwargs.update({'complex_mode': is_complex}) for i, s in enumerate(f.subs(x, *p, **sub_kwargs)): pen_i = mkPen(QtGui.QColor.fromRgbF(*color[i])) if np.iscomplexobj(s): self.preview_lines.append(PlotItem(x=x, y=s.real, pen=pen_i)) self.preview_lines.append(PlotItem(x=x, y=s.imag, pen=pen_i)) else: self.preview_lines.append(PlotItem(x=x, y=s, pen=pen_i)) return self.preview_lines def set_parameter(self, parameter: dict[str, FitResult]): # which data uses which model data = self.data_table.collect_data(default=self.default_combobox.currentData()) glob_fit_parameter = [] for fitted_model, fitted_data in data.items(): for fit_id, fit_curve in parameter.items(): if fit_id in fitted_data: fit_parameter = list(fit_curve.parameter.values()) glob_fit_parameter.append(fit_parameter) self.set_parameter_iter(fit_id, [p.value for p in fit_parameter], self.models[fitted_model]) mean_parameter = [reduce(add, p, 0)/len(p) for p in zip(*glob_fit_parameter)] self.set_parameter_iter(None, mean_parameter, self.models[fitted_model]) def set_parameter_iter(self, fit_id: str | None, param: list[float], functions: list, cnt: int = 0): for model_p in functions: if model_p['active']: cnt += self.param_widgets[model_p['cnt']].set_parameter(fit_id, param[cnt:]) if model_p['children']: cnt += self.set_parameter_iter(fit_id, param, model_p['children'], cnt=cnt) return cnt def closeEvent(self, evt: QtGui.QCloseEvent): self.preview_emit.emit({}, -1, False) self.preview_lines = [] super().closeEvent(evt)