use tree colors in fit result for sub-funcs (#253)

closes #201

Co-authored-by: Dominik Demuth <dominik.demuth@physik.tu-darmstadt.de>
Reviewed-on: #253

src/gui_qt/fit/result.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from math import isnan
 
-from pyqtgraph import mkBrush, mkPen
+from pyqtgraph import mkBrush, mkPen, mkColor
 from numpy import abs as np_abs, isfinite as np_isfinite
 
 from nmreval.utils.text import convert
@@ -18,7 +18,7 @@ class QFitResult(QtWidgets.QDialog, Ui_Dialog):
     closed = QtCore.pyqtSignal(dict, list, str, bool, bool, list)
     redoFit = QtCore.pyqtSignal(dict)
 
-    def __init__(self, results: list, management, parent=None):
+    def __init__(self, results: list, sub_colors: dict, management, parent=None):
         super().__init__(parent=parent)
         self.setupUi(self)
 
@@ -37,6 +37,7 @@ class QFitResult(QtWidgets.QDialog, Ui_Dialog):
         self._results = {}
         self.graph_opts = {}
         self.last_idx = None
+        self.func_colors = sub_colors
 
         self.fit_plot = self.graphicsView.addPlot(row=1, col=0, title='Fit')
         self.resid_plot = self.graphicsView.addPlot(row=0, col=0, title='Residual')
@@ -48,21 +49,29 @@ class QFitResult(QtWidgets.QDialog, Ui_Dialog):
         self.graphicsView.ci.layout.setRowStretchFactor(0, 1)
         self.graphicsView.ci.layout.setRowStretchFactor(1, 2)
 
-        self.resid_graph = PlotItem(x=[], y=[],
+        self.resid_graph = PlotItem(
+            x=[], y=[],
             symbol='o', symbolPen=None, symbolBrush=mkBrush(color=(31, 119, 180)),
-                                    pen=None)
-        self.resid_graph_imag = PlotItem(x=[], y=[],
+            pen=None
+        )
+        self.resid_graph_imag = PlotItem(
+            x=[], y=[],
             symbol='s', symbolPen=None, symbolBrush=mkBrush(color=(255, 127, 14)),
-                                         pen=None)
+            pen=None
+        )
         self.resid_plot.addItem(self.resid_graph)
         self.resid_plot.addItem(self.resid_graph_imag)
 
-        self.data_graph = PlotItem(x=[], y=[],
+        self.data_graph = PlotItem(
+            x=[], y=[],
             symbol='o', symbolPen=None, symbolBrush=mkBrush(color=(31, 119, 180)),
-                                   pen=None)
-        self.data_graph_imag = PlotItem(x=[], y=[],
+            pen=None
+        )
+        self.data_graph_imag = PlotItem(
+            x=[], y=[],
             symbol='s', symbolPen=None, symbolBrush=mkBrush(color=(255, 127, 14)),
-                                        pen=None)
+            pen=None
+        )
         self.fit_plot.addItem(self.data_graph)
         self.fit_plot.addItem(self.data_graph_imag)
 
@@ -85,13 +94,14 @@ class QFitResult(QtWidgets.QDialog, Ui_Dialog):
 
         self.set_results(results)
 
-    def __call__(self, results: list):
+    def __call__(self, results: list, sub_colors: dict):
         self._previous_fits = {}
         self.sets_comboBox.blockSignals(True)
         self.sets_comboBox.clear()
         self.sets_comboBox.blockSignals(False)
         self._results = {}
         self._opts = {}
+        self.func_colors = sub_colors
 
         self.set_results(results)
 
@@ -193,10 +203,11 @@ class QFitResult(QtWidgets.QDialog, Ui_Dialog):
             self.fit_graph.setData(x=res.x, y=res.y.real)
             self.fit_graph_imag.setData(x=res.x, y=res.y.imag)
 
-            for i, f in enumerate(sub_funcs):
-                item = PlotItem(x=f.x, y=f.y.real, pen=mkPen({'color': i, 'style': 2}))
+            for f, c in zip(sub_funcs, self.func_colors[idx]):
+                col = mkColor(*[c_i*255 for c_i in c])
+                item = PlotItem(x=f.x, y=f.y.real, pen=mkPen({'color': col, 'style': 2}))
                 self.fit_plot.addItem(item)
-                item = PlotItem(x=f.x, y=f.y.imag, pen=mkPen({'color': i, 'style': 2}))
+                item = PlotItem(x=f.x, y=f.y.imag, pen=mkPen({'color': col, 'style': 2}))
                 self.fit_plot.addItem(item)
 
         else:
@@ -205,8 +216,8 @@ class QFitResult(QtWidgets.QDialog, Ui_Dialog):
             self.fit_graph.setData(x=res.x, y=res.y)
             self.fit_graph_imag.setData(x=[], y=[])
 
-            for i, f in enumerate(sub_funcs):
-                item = PlotItem(x=f.x, y=f.y, pen=mkPen({'color': i, 'style': 2}))
+            for f, c in zip(sub_funcs, self.func_colors[idx]):
+                item = PlotItem(x=f.x, y=f.y, pen=mkPen({'color': mkColor(*[c_i*255 for c_i in c]), 'style': 2}))
                 self.fit_plot.addItem(item)
 
         self._plot_residuals(idx)

src/gui_qt/main/mainwindow.py

@@ -952,19 +952,19 @@ class NMRMainWindow(QtWidgets.QMainWindow, Ui_BaseWindow):
             for item in self.fit_dialog.preview_lines:
                 g.add_external(item)
 
-    @QtCore.pyqtSlot(list)
-    def show_fit_results(self, results: list):
+    @QtCore.pyqtSlot(list, dict)
+    def show_fit_results(self, results: list, sub_colors: dict[str, tuple[float, float, float]]):
         self.fit_dialog.fit_button.setEnabled(True)
         self.fit_timer.stop()
         self.status.setText('')
         if results:
             if self.fitresult_dialog is None:
-                self.fitresult_dialog = QFitResult(results, self.management, parent=self)
+                self.fitresult_dialog = QFitResult(results, sub_colors, self.management, parent=self)
                 self.fitresult_dialog.add_graphs(self.management.graphs.list())
                 self.fitresult_dialog.closed.connect(self.accepts_fit)
                 self.fitresult_dialog.redoFit.connect(self.management.redo_fits)
             else:
-                self.fitresult_dialog(results)
+                self.fitresult_dialog(results, sub_colors)
                 self.fitresult_dialog.add_graphs(self.management.graphs.list())
             self.fitresult_dialog.show()
 

src/gui_qt/main/management.py

@@ -85,7 +85,7 @@ class UpperManagement(QtCore.QObject):
     newData = QtCore.pyqtSignal([list, str], [list, str, bool])
     deleteData = QtCore.pyqtSignal(list)
     dataChanged = QtCore.pyqtSignal(str)
-    fitFinished = QtCore.pyqtSignal(list)
+    fitFinished = QtCore.pyqtSignal(list, dict)
     stopFit = QtCore.pyqtSignal()
     properties_collected = QtCore.pyqtSignal(dict)
     unset_state = QtCore.pyqtSignal(list)
@@ -586,13 +586,23 @@ class UpperManagement(QtCore.QObject):
     def end_fit(self, result: list, success: bool):
         if success:
             logger.info('Successful fit')
-            self.fitFinished.emit(result)
+
+            sub_colors = {}
+            for k, v in self.__fit_options[0].items():
+                sub_colors.update({set_id: v['color'] for set_id in v['data_parameter']})
+
+            self.fitFinished.emit(result, sub_colors)
+
         else:
             e = result[0]
             logger.exception(e, exc_info=True)
-            QtWidgets.QMessageBox.warning(QtWidgets.QWidget(), 'Fit failed',
-                                          f'Fit kaput with exception: \n\n{e!r}')
-            self.fitFinished.emit([])
+            QtWidgets.QMessageBox.warning(
+                QtWidgets.QWidget(),
+                'Fit failed',
+                f'Fit kaput with exception: \n\n{e!r}'
+            )
+            self.fitFinished.emit([], {})
+
         self._fit_active = False
 
     @QtCore.pyqtSlot(dict)