Merge pull request #267 from zhujun98/add_curve_fitting_in_histogram_window

Add curve fitting in histogram window

Author: zhujun98

extra_foam/algorithms/curve_fitting.py

@@ -83,11 +83,11 @@ def __init__(self):
             self.n = 4
 
         def __call__(self, x, a, b, c, d):
-            return a / ((x - c)**2 + b**2) + d
+            return a / ((x - b)**2 + c**2) + d
 
         @staticmethod
         def format(a, b, c, d):
-            return f"y = a / (4 * (x - c)^2 + b ^ 2) + d, \n" \
+            return f"y = a / (4 * (x - b)^2 + c ^ 2) + d, \n" \
                    f"a = {a:.4E}, b = {b:.4E}, c = {c:.4E}, d = {d:.4E}"
 
     class Erf(_BaseCurveFitting):

extra_foam/algorithms/tests/test_curve_fitting.py

@@ -72,10 +72,10 @@ def testLorentzian(self):
                       105.22065798, 78.23873672, 52.12398640, 40.89729233, 40.60966561,
                        35.44761773, 40.06147386, 39.26495213, 36.69235285, 32.29345426])
 
-        popt = algo.fit(x, y, p0=[100, 1, 100, 1])
+        popt = algo.fit(x, y, p0=[100, 100, 1, 1])
         assert abs(popt[0] - 100) < 10
-        assert abs(popt[1] - 1) < 0.2
-        assert abs(popt[2] - 100) < 1
+        assert abs(popt[1] - 100) < 1
+        assert abs(popt[2] - 1) < 0.2
         assert abs(popt[3] - 30) < 5
 
     def testERF(self):

extra_foam/gui/ctrl_widgets/correlation_ctrl_widget.py

@@ -17,7 +17,7 @@
     QLabel, QPushButton, QTableWidget
 )
 
-from .curve_fitting_ctrl_widget import _BaseFittingCtrlWidget, FittingType
+from .curve_fitting_ctrl_widget import _BaseFittingCtrlWidget
 from .base_ctrl_widgets import _AbstractCtrlWidget
 from .smart_widgets import SmartLineEdit
 from ..gui_helpers import invert_dict
@@ -29,7 +29,7 @@
 _DEFAULT_RESOLUTION = 0.0
 
 
-class FittingCtrlWidget(_BaseFittingCtrlWidget):
+class _FittingCtrlWidget(_BaseFittingCtrlWidget):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -121,7 +121,7 @@ def __init__(self, *args, **kwargs):
             }
         }
 
-        self._fitting = FittingCtrlWidget()
+        self._fitting = _FittingCtrlWidget()
 
         self.initParamTable()
         self.initUI()

extra_foam/gui/ctrl_widgets/curve_fitting_ctrl_widget.py

@@ -12,8 +12,7 @@
 import numpy as np
 
 from PyQt5.QtWidgets import (
-    QCheckBox, QComboBox, QGridLayout, QLabel, QPlainTextEdit, QPushButton,
-    QWidget
+    QCheckBox, QComboBox, QPlainTextEdit, QPushButton, QWidget
 )
 
 from .smart_widgets import SmartLineEdit

extra_foam/gui/ctrl_widgets/histogram_ctrl_widget.py

@@ -9,20 +9,56 @@
 """
 from collections import OrderedDict
 
-from PyQt5.QtCore import Qt
+from PyQt5.QtCore import pyqtSignal, Qt
 from PyQt5.QtGui import QIntValidator
 from PyQt5.QtWidgets import (
     QCheckBox, QComboBox, QFrame, QGridLayout, QHBoxLayout, QLabel,
-    QPushButton
+    QPushButton, QWidget
 )
 
 from .base_ctrl_widgets import _AbstractCtrlWidget
+from .curve_fitting_ctrl_widget import _BaseFittingCtrlWidget, FittingType
 from .smart_widgets import SmartBoundaryLineEdit, SmartLineEdit
 from ..gui_helpers import invert_dict
 from ...config import AnalysisType
 from ...database import Metadata as mt
 
 
+class _FittingCtrlWidget(_BaseFittingCtrlWidget):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.initUI()
+        self.initConnections()
+
+    def initUI(self):
+        AR = Qt.AlignRight
+
+        layout = QGridLayout()
+        layout.addWidget(QLabel("Param a0 = "), 1, 0, AR)
+        layout.addWidget(self._params[0], 1, 1)
+        layout.addWidget(QLabel("Param b0 = "), 1, 2, AR)
+        layout.addWidget(self._params[1], 1, 3)
+        layout.addWidget(QLabel("Param c0 = "), 1, 4, AR)
+        layout.addWidget(self._params[2], 1, 5)
+        layout.addWidget(QLabel("Param d0 = "), 2, 0, AR)
+        layout.addWidget(self._params[3], 2, 1)
+        layout.addWidget(QLabel("Param e0 = "), 2, 2, AR)
+        layout.addWidget(self._params[4], 2, 3)
+        layout.addWidget(QLabel("Param f0 = "), 2, 4, AR)
+        layout.addWidget(self._params[5], 2, 5)
+        layout.addWidget(self.fit_btn, 3, 0, 1, 2)
+        layout.addWidget(self.clear_btn, 3, 2, 1, 2)
+        layout.addWidget(QLabel("Fit type: "), 3, 4, AR)
+        layout.addWidget(self.fit_type_cb, 3, 5)
+        layout.addWidget(self._output, 4, 0, 1, 6)
+        layout.setContentsMargins(0, 0, 0, 0)
+        self.setLayout(layout)
+
+        self.setFixedWidth(self.minimumSizeHint().width())
+
+
 class HistogramCtrlWidget(_AbstractCtrlWidget):
     """Widget for setting up histogram analysis parameters."""
 
@@ -33,6 +69,9 @@ class HistogramCtrlWidget(_AbstractCtrlWidget):
     })
     _analysis_types_inv = invert_dict(_analysis_types)
 
+    fit_curve_sgn = pyqtSignal()
+    clear_fitting_sgn = pyqtSignal()
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
@@ -53,6 +92,8 @@ def __init__(self, *args, **kwargs):
 
         self._reset_btn = QPushButton("Reset")
 
+        self._fitting = _FittingCtrlWidget()
+
         self.initUI()
         self.initConnections()
 
@@ -61,27 +102,23 @@ def __init__(self, *args, **kwargs):
     def initUI(self):
         """Overload."""
         AR = Qt.AlignRight
-        AT = Qt.AlignTop
 
-        lwidget = QFrame()
+        ctrl_widget = QFrame()
         llayout = QGridLayout()
         llayout.addWidget(QLabel("Analysis type: "), 0, 0, AR)
         llayout.addWidget(self._analysis_type_cb, 0, 1)
-        llayout.addWidget(self._pulse_resolved_cb, 1, 1)
-        llayout.addWidget(self._reset_btn, 2, 0, 1, 2)
-        lwidget.setLayout(llayout)
-
-        rwidget = QFrame()
-        rlayout = QGridLayout()
-        rlayout.addWidget(QLabel("Bin range: "), 1, 0, AR)
-        rlayout.addWidget(self._bin_range_le, 1, 1)
-        rlayout.addWidget(QLabel("# of bins: "), 2, 0, AR)
-        rlayout.addWidget(self._n_bins_le, 2, 1)
-        rwidget.setLayout(rlayout)
+        llayout.addWidget(self._pulse_resolved_cb, 0, 3, AR)
+        llayout.addWidget(self._reset_btn, 0, 5)
+        llayout.addWidget(QLabel("Bin range: "), 1, 2, AR)
+        llayout.addWidget(self._bin_range_le, 1, 3)
+        llayout.addWidget(QLabel("# of bins: "), 1, 4, AR)
+        llayout.addWidget(self._n_bins_le, 1, 5)
+        llayout.setContentsMargins(0, 0, 0, 0)
+        ctrl_widget.setLayout(llayout)
 
         layout = QHBoxLayout()
-        layout.addWidget(lwidget, alignment=AT)
-        layout.addWidget(rwidget, alignment=AT)
+        layout.addWidget(ctrl_widget, alignment=Qt.AlignTop)
+        layout.addWidget(self._fitting)
         self.setLayout(layout)
 
     def initConnections(self):
@@ -100,6 +137,9 @@ def initConnections(self):
 
         self._reset_btn.clicked.connect(mediator.onHistReset)
 
+        self._fitting.fit_btn.clicked.connect(self.fit_curve_sgn)
+        self._fitting.clear_btn.clicked.connect(self.clear_fitting_sgn)
+
     def updateMetaData(self):
         """Overload."""
         self._analysis_type_cb.currentTextChanged.emit(
@@ -127,3 +167,6 @@ def loadMetaData(self):
     def resetAnalysisType(self):
         self._analysis_type_cb.setCurrentText(
             self._analysis_types_inv[AnalysisType.UNDEFINED])
+
+    def fit_curve(self, x, y):
+        return self._fitting.fit(x, y)

extra_foam/gui/windows/histogram_w.py

@@ -14,6 +14,7 @@
 
 from .base_window import _AbstractPlotWindow
 from ..ctrl_widgets import HistogramCtrlWidget
+from ..misc_widgets import FColor
 from ..plot_widgets import HistMixin, PlotWidgetF, TimedPlotWidgetF
 from ...config import config
 
@@ -51,6 +52,7 @@ def __init__(self, *, parent=None):
         super().__init__(parent=parent)
 
         self._plot = self.plotBar()
+        self._fitted = self.plotCurve(pen=FColor.mkPen('r'))
 
         self._title_template = Template(
             f"FOM Histogram (mean: $mean, median: $median, std: $std)")
@@ -68,6 +70,12 @@ def refresh(self):
             self._plot.setData(bin_centers, hist.hist)
             self.updateTitle(hist.mean, hist.median, hist.std)
 
+    def data(self):
+        return self._plot.data()
+
+    def setFitted(self, x, y):
+        self._fitted.setData(x, y)
+
 
 class HistogramWindow(_AbstractPlotWindow):
     """HistogramWindow class.
@@ -115,7 +123,17 @@ def initUI(self):
 
     def initConnections(self):
         """Override."""
-        pass
+        self._ctrl_widget.fit_curve_sgn.connect(self._onCurveFit)
+        self._ctrl_widget.clear_fitting_sgn.connect(self._onClearFitting)
+
+    def _onCurveFit(self):
+        data = self._fom_hist.data()
+
+        x, y = self._ctrl_widget.fit_curve(*data)
+        self._fom_hist.setFitted(x, y)
+
+    def _onClearFitting(self):
+        self._fom_hist.setFitted([], [])
 
     def closeEvent(self, QCloseEvent):
         self._ctrl_widget.resetAnalysisType()

extra_foam/gui/windows/tests/test_plot_widgets.py

@@ -1,13 +1,8 @@
 import unittest
 from unittest.mock import patch
 
-import numpy as np
-
 from extra_foam.logger import logger
 from extra_foam.gui import mkQApp
-from extra_foam.gui.ctrl_widgets.correlation_ctrl_widget import (
-    FittingType, FittingCtrlWidget
-)
 from extra_foam.gui.plot_widgets.plot_items import ScatterPlotItem
 from extra_foam.pipeline.data_model import ProcessedData
 from extra_foam.pipeline.tests import _TestDataMixin

extra_foam/gui/windows/tests/test_plot_windows.py

@@ -101,6 +101,7 @@ def testOpenCloseWindows(self):
         self.assertIsInstance(histogram_window, HistogramWindow)
         self._checkHistogramWindow(histogram_window)
         self._checkHistogramCtrlWidget(histogram_window)
+        self._checkHistogramCurveFitting(histogram_window)
 
         poi_window = self._check_open_window(self.poi_action)
         self.assertIsInstance(poi_window, PulseOfInterestWindow)
@@ -596,11 +597,11 @@ def _checkCorrelationCurveFitting(self, win):
                 QTest.mouseClick(fitting.fit_btn, Qt.LeftButton)
                 mocked_fit.assert_called_once_with(x1, y1, p0=[1.1, 2.2])
 
-                mocked_fit.side_effect=RuntimeError("runtime error")
+                mocked_fit.side_effect = RuntimeError("runtime error")
                 QTest.mouseClick(fitting.fit_btn, Qt.LeftButton)
                 self.assertEqual("RuntimeError('runtime error')", fitting._output.toPlainText())
 
-                mocked_fit.side_effect=ValueError("value error")
+                mocked_fit.side_effect = ValueError("value error")
                 QTest.mouseClick(fitting.fit_btn, Qt.LeftButton)
                 self.assertEqual("ValueError('value error')", fitting._output.toPlainText())
 
@@ -808,6 +809,46 @@ def _checkHistogramCtrlWidget(self, win):
         self.assertEqual("55", widget._n_bins_le.text())
         self.assertEqual(True, widget._pulse_resolved_cb.isChecked())
 
+    def _checkHistogramCurveFitting(self, win):
+        widget = win._ctrl_widget
+        fitting = widget._fitting
+
+        with patch.object(win._fom_hist, "setFitted") as mocked_set_fitted:
+            self.assertIsNone(fitting._algo)
+            QTest.mouseClick(fitting.fit_btn, Qt.LeftButton)
+            mocked_set_fitted.assert_called_with(None, None)
+            mocked_set_fitted.reset_mock()
+
+            fitting.fit_type_cb.setCurrentText("Gaussian")
+            with patch.object(fitting._algo, "fit") as mocked_fit:
+                x1, y1 = np.random.rand(1), np.random.rand(1)
+                win._fom_hist._plot.setData(x1, y1)
+                QTest.mouseClick(fitting.fit_btn, Qt.LeftButton)
+                mocked_fit.assert_not_called()
+                mocked_set_fitted.assert_called_with(None, None)
+                mocked_set_fitted.reset_mock()
+                self.assertEqual("Not enough data", fitting._output.toPlainText())
+
+                x1, y1 = np.random.rand(10), np.random.rand(10)
+                win._fom_hist._plot.setData(x1, y1)
+                QTest.mouseClick(fitting.fit_btn, Qt.LeftButton)
+                mocked_fit.assert_called_once_with(x1, y1, p0=[1.0, 1.0, 1.0, 1.0])
+                mocked_fit.reset_mock()
+                mocked_set_fitted.assert_called_once()
+                mocked_set_fitted.reset_mock()
+                fitting._params[2].setText("1.1")
+                fitting._params[3].setText("2.2")
+                QTest.mouseClick(fitting.fit_btn, Qt.LeftButton)
+                mocked_fit.assert_called_once_with(x1, y1, p0=[1.0, 1.0, 1.1, 2.2])
+
+                mocked_fit.side_effect = RuntimeError("runtime error")
+                QTest.mouseClick(fitting.fit_btn, Qt.LeftButton)
+                self.assertEqual("RuntimeError('runtime error')", fitting._output.toPlainText())
+
+                mocked_fit.side_effect = ValueError("value error")
+                QTest.mouseClick(fitting.fit_btn, Qt.LeftButton)
+                self.assertEqual("ValueError('value error')", fitting._output.toPlainText())
+
     def _checkHistogramCtrlWidgetTs(self, win):
         widget = win._ctrl_widget