489 lines
18 KiB
Python
489 lines
18 KiB
Python
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 import Relations
|
|
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 = None
|
|
|
|
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]
|
|
|
|
self._current_function = None
|
|
if len(self.param_widgets) == 0:
|
|
# empty model
|
|
self.newmodel_button.setEnabled(False)
|
|
self.deletemodel_button.setEnabled(False)
|
|
|
|
@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)
|
|
|
|
# deselect all fit sets
|
|
for i in range(self.data_table.rowCount()):
|
|
data_id = self.data_table.item(i, 0).data(QtCore.Qt.UserRole+1)
|
|
if self._management[data_id].mode == 'fit' or self._management[data_id].has_relation(Relations.isFitPartOf):
|
|
self.data_table.item(i, 0).setCheckState(QtCore.Qt.Unchecked)
|
|
|
|
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)
|