forked from IPKM/nmreval
456 lines
17 KiB
Python
456 lines
17 KiB
Python
|
from itertools import count, cycle
|
||
|
|
from string import ascii_letters
|
||
|
|
|
||
|
|
from pyqtgraph import PlotDataItem, mkPen
|
||
|
|
|
||
|
|
from nmreval.gui_qt.lib.pg_objects import PlotItem
|
||
|
|
from ...fit._meta import MultiModel, ModelFactory
|
||
|
|
from ..Qt import QtGui, QtCore, QtWidgets
|
||
|
|
from .._py.fitdialog import Ui_FitDialog
|
||
|
|
from .fit_forms import FitTableWidget
|
||
|
|
from .fit_parameter import QFitParameterWidget
|
||
|
|
|
||
|
|
|
||
|
|
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.function_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)
|
||
|
|
|
||
|
|
@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.function_widgets[idx]
|
||
|
|
self.stackedWidget.removeWidget(w)
|
||
|
|
w.deleteLater()
|
||
|
|
del self.function_widgets[idx]
|
||
|
|
|
||
|
|
if len(self.functionwidget) == 0:
|
||
|
|
# empty model
|
||
|
|
self.newmodel_button.setEnabled(False)
|
||
|
|
self.deletemodel_button.setEnabled(False)
|
||
|
|
self._current_function = None
|
||
|
|
|
||
|
|
else:
|
||
|
|
self._current_function = self.functionwidget.use_combobox.currentData()
|
||
|
|
|
||
|
|
@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.function_widgets:
|
||
|
|
dialog = self.function_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()
|
||
|
|
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.function_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.function_widgets[self._current_function].tabWidget.currentIndex()
|
||
|
|
dialog.tabWidget.setCurrentIndex(tab_idx)
|
||
|
|
|
||
|
|
self._current_function = function_id
|
||
|
|
|
||
|
|
def look_value(self, idx):
|
||
|
|
func_widget = self.function_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):
|
||
|
|
"""
|
||
|
|
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.function_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.function_widgets[func_id])
|
||
|
|
|
||
|
|
self.function_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=None, parameter=None, add_idx=False, cnt=0):
|
||
|
|
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.function_widgets[f['cnt']].get_parameter(function_use)
|
||
|
|
except ValueError as e:
|
||
|
|
_ = QtWidgets.QMessageBox().warning(self, 'Invalid value', str(e),
|
||
|
|
QtWidgets.QMessageBox.Ok)
|
||
|
|
return None, -1
|
||
|
|
|
||
|
|
p_len = len(parameter['lb'])
|
||
|
|
|
||
|
|
parameter['lb'] += lb
|
||
|
|
parameter['ub'] += ub
|
||
|
|
parameter['var'] += var
|
||
|
|
parameter['links'] += links
|
||
|
|
parameter['color'] += [f['color']]
|
||
|
|
|
||
|
|
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
|
||
|
|
if self._complex[k] is None:
|
||
|
|
parameter['complex'] = self._complex[k]
|
||
|
|
else:
|
||
|
|
parameter['complex'] = ['complex', 'real', 'imag'][self._complex[k]]
|
||
|
|
|
||
|
|
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):
|
||
|
|
print('state', state)
|
||
|
|
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']
|
||
|
|
|
||
|
|
for p, kwargs in parameters.values():
|
||
|
|
y = f.func(x, *p, **kwargs)
|
||
|
|
if is_complex is None:
|
||
|
|
self.preview_lines.append(PlotItem(x=x, y=y, pen=mkPen(width=3)))
|
||
|
|
|
||
|
|
elif is_complex == 0:
|
||
|
|
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)))
|
||
|
|
elif is_complex == 1:
|
||
|
|
self.preview_lines.append(PlotItem(x=x, y=y.real, pen=mkPen(width=3)))
|
||
|
|
else:
|
||
|
|
self.preview_lines.append(PlotItem(x=x, y=y.imag, pen=mkPen(width=3)))
|
||
|
|
|
||
|
|
if isinstance(f, MultiModel):
|
||
|
|
for i, s in enumerate(f.subs(x, *p, **kwargs)):
|
||
|
|
pen_i = mkPen(QtGui.QColor.fromRgbF(*color[i]))
|
||
|
|
if is_complex is None:
|
||
|
|
self.preview_lines.append(PlotItem(x=x, y=s, pen=pen_i))
|
||
|
|
elif is_complex == 0:
|
||
|
|
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))
|
||
|
|
elif is_complex == 1:
|
||
|
|
self.preview_lines.append(PlotItem(x=x, y=s.real, pen=pen_i))
|
||
|
|
else:
|
||
|
|
self.preview_lines.append(PlotItem(x=x, y=s.imag, pen=pen_i))
|
||
|
|
|
||
|
|
return self.preview_lines
|
||
|
|
|
||
|
|
def closeEvent(self, evt: QtGui.QCloseEvent):
|
||
|
|
self.preview_emit.emit({}, -1, False)
|
||
|
|
self.preview_lines = []
|
||
|
|
|
||
|
|
super().closeEvent(evt)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
import sys
|
||
|
|
from numpy.random import choice
|
||
|
|
|
||
|
|
app = QtWidgets.QApplication(sys.argv)
|
||
|
|
# while qw == 'QtCurve':
|
||
|
|
qw = choice(QtWidgets.QStyleFactory.keys())
|
||
|
|
app.setStyle(QtWidgets.QStyleFactory.create(qw))
|
||
|
|
|
||
|
|
fd = QFitDialog()
|
||
|
|
|
||
|
|
fd.load([('fff', 'testtesttesttest'),
|
||
|
|
('ggg', 'testtesttesttesttest'),
|
||
|
|
('hhh', 'testtesttesttesttesttest')])
|
||
|
|
fd.show()
|
||
|
|
|
||
|
|
sys.exit(app.exec())
|