Add constant and simple linear background subtraction

xsocs/gui/process/FitWidget.py

@@ -40,7 +40,7 @@ from ..widgets.Containers import GroupBox
 from ..widgets.RoiAxisWidget import RoiAxisWidget
 from ..widgets.Input import StyledLineEdit
 
-from ...process.fit.peak_fit import PeakFitter, FitTypes
+from ...process.fit.peak_fit import PeakFitter, FitTypes, BackgroundTypes
 
 
 _logger = logging.getLogger(__name__)
@@ -188,6 +188,13 @@ class FitWidget(Qt.QWidget):
         self.__fileEdit = StyledLineEdit(readOnly=True)
         layout.addRow('File:', self.__fileEdit)
 
+        self.__bgComboBox = Qt.QComboBox()
+        self.__bgComboBox.addItem('-', BackgroundTypes.NONE)
+        self.__bgComboBox.addItem('Constant', BackgroundTypes.CONSTANT)
+        self.__bgComboBox.addItem('Linear', BackgroundTypes.LINEAR)
+        self.__bgComboBox.setCurrentIndex(0)
+        layout.addRow('Background:', self.__bgComboBox)
+
         self.__fitTypeCb = Qt.QComboBox()
         self.__fitTypeCb.addItem('Gaussian', FitTypes.GAUSSIAN)
         self.__fitTypeCb.addItem('Centroid', FitTypes.CENTROID)
@@ -311,10 +318,14 @@ class FitWidget(Qt.QWidget):
         else:
             roiIndices = None
 
+        background = self.__bgComboBox.itemData(
+            self.__bgComboBox.currentIndex())
+
         self.__fitter = PeakFitter(self.__qspaceH5.filename,
                                    fit_type=fitType,
                                    roi_indices=roiIndices,
-                                   n_peaks=self.__nPeaks)
+                                   n_peaks=self.__nPeaks,
+                                   background=background)
         self.__statusLabel.setText('Running...')
 
         self.__progTimer = Qt.QTimer()
@@ -340,6 +351,7 @@ class FitWidget(Qt.QWidget):
     def __lock(self, lock):
         enable = not lock
         self.roiWidget().setEnabled(enable)
+        self.__bgComboBox.setEnabled(enable)
         self.__fitTypeCb.setEnabled(enable)
         self.__runButton.setEnabled(enable)
 

xsocs/process/fit/peak_fit.py

@@ -49,6 +49,11 @@ from .sharedresults import FitTypes
 disp_times = False
 
 
+class BackgroundTypes(object):
+    ALLOWED = range(3)
+    NONE, CONSTANT, LINEAR = ALLOWED
+
+
 class PeakFitter(Thread):
     """
     :param qspace_f: path to the HDF5 file containing the qspace cubes
@@ -64,6 +69,8 @@ class PeakFitter(Thread):
 
     :param Union[int,None] n_proc:
         Number of process to use. If None, the config value is used.
+
+    :param BackgroundTypes background: The background subtraction to perform
     """
 
     READY, RUNNING, DONE, ERROR, CANCELED = __STATUSES = range(5)
@@ -74,7 +81,8 @@ class PeakFitter(Thread):
                  n_peaks=1,
                  indices=None,
                  n_proc=None,
-                 roi_indices=None):
+                 roi_indices=None,
+                 background=None):
         super(PeakFitter, self).__init__()
 
         self.__results = None
@@ -88,6 +96,7 @@ class PeakFitter(Thread):
 
         self.__qspace_f = qspace_f
         self.__fit_type = fit_type
+        self.__background = background
         self.__n_peaks = n_peaks
 
         if n_proc:
@@ -105,7 +114,11 @@ class PeakFitter(Thread):
 
         if fit_type not in FitTypes.ALLOWED:
             self.__set_status(self.ERROR)
-            raise ValueError('Unknown fit type : {0}'.format(fit_type))
+            raise ValueError('Unknown fit type: {0}'.format(fit_type))
+
+        if background not in BackgroundTypes.ALLOWED:
+            self.__set_status(self.ERROR)
+            raise ValueError('Unknown background type: {}'.format(background))
 
         try:
             with QSpaceH5.QSpaceH5(qspace_f) as qspace_h5:
@@ -245,7 +258,8 @@ class PeakFitter(Thread):
                                  (n_indices, 9),
                                  idx_queue,
                                  qspace_f,
-                                 read_lock))
+                                 read_lock,
+                                 self.__background))
 
         if disp_times:
             class myTimes(object):
@@ -330,14 +344,16 @@ def _init_thread(shared_res_,
                  result_shape_,
                  idx_queue_,
                  qspace_f_,
-                 read_lock_):
+                 read_lock_,
+                 background_):
     global shared_res, \
         shared_progress, \
         fit_class, \
         result_shape, \
         idx_queue, \
         qspace_f, \
-        read_lock
+        read_lock, \
+        fit_background
 
     shared_res = shared_res_
     shared_progress = shared_prog_
@@ -346,6 +362,7 @@ def _init_thread(shared_res_,
     idx_queue = idx_queue_
     qspace_f = qspace_f_
     read_lock = read_lock_
+    fit_background = background_
 
 
 def _fit_process(th_idx, roiIndices=None):
@@ -430,6 +447,21 @@ def _fit_process(th_idx, roiIndices=None):
             y_sum = cube_sum_z.sum(axis=0)
             x_sum = cube_sum_z.sum(axis=1)
 
+            # Background subtraction
+            if fit_background == BackgroundTypes.CONSTANT:
+                # Shift data so that smallest value is 0
+                for array in (z_sum, y_sum, x_sum):
+                    array -= np.nanmin(array)
+
+            elif fit_background == BackgroundTypes.LINEAR:
+                # Simple linear background
+                for array in (z_sum, y_sum, x_sum):
+                    array -= np.linspace(
+                        array[0], array[-1], num=len(array), endpoint=True)
+
+            elif fit_background != BackgroundTypes.NONE:
+                raise RuntimeError("Unsupported background subtraction")
+
             fitter.fit(i_fit, i_cube, x_sum, y_sum, z_sum)
 
             t_fit += time.time() - t0
@@ -445,7 +477,3 @@ def _fit_process(th_idx, roiIndices=None):
     if disp_times:
         print('Thread {0} done ({1}).'.format(th_idx, times))
     return times
-
-
-if __name__ == '__main__':
-    pass