Make a background_estimation method

xsocs/process/fit/peak_fit.py

@@ -62,6 +62,31 @@ class BackgroundTypes(object):
     ALLOWED = NONE, CONSTANT, LINEAR
 
 
+def background_estimation(mode, data):
+    """Estimates a background of mode kind for data
+
+    :param BackgroundTypes mode: The kind of background to compute
+    :param numpy.ndarray data: The data array to process
+    :return: The estimated background as an array with same shape as input
+    :rtype: numpy.ndarray
+    :raises ValueError: In case mode is not known
+    """
+    # Background subtraction
+    if mode == BackgroundTypes.CONSTANT:
+        # Shift data so that smallest value is 0
+        return np.ones_like(data) * np.nanmin(data)
+
+    elif mode == BackgroundTypes.LINEAR:
+        # Simple linear background
+        return np.linspace(data[0], data[-1], num=len(data), endpoint=True)
+
+    elif mode == BackgroundTypes.NONE:
+        return np.zeros_like(data)
+
+    else:
+        raise ValueError("Unsupported background mode")
+
+
 class PeakFitter(Thread):
     """
     :param str qspace_f: path to the HDF5 file containing the qspace cubes
@@ -456,19 +481,8 @@ def _fit_process(th_idx, roiIndices=None):
             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")
+            for array in (z_sum, y_sum, x_sum):
+                array -= background_estimation(fit_background, array)
 
             fitter.fit(i_fit, i_cube, x_sum, y_sum, z_sum)