optional return argument for shifts. Reference shifts given by bootstrap.

nabu/preproc/alignment.py

@@ -131,7 +131,7 @@ class AlignmentBase(object):
 
         # For a 1D parabola `f(x) = c + bx + ax^2`, the vertex position is:
         # x_v = -b / 2a.
-        vertex_x = - coeffs[1] / (2 * coeffs[2])
+        vertex_x = -coeffs[1] / (2 * coeffs[2])
 
         vertex_min_x = np.min(fx)
         vertex_max_x = np.max(fx)
@@ -202,9 +202,7 @@ class AlignmentBase(object):
         return roi_yxhw
 
     @staticmethod
-    def _prepare_image(
-        img, invalid_val=1e-5, roi_yxhw=None, median_filt_shape=None, high_pass=None, low_pass=None
-    ):
+    def _prepare_image(img, invalid_val=1e-5, roi_yxhw=None, median_filt_shape=None, high_pass=None, low_pass=None):
         """
         Prepare and returns  a cropped  and filtered image, or array of filtered images if the input is an  array of images.
 
@@ -241,7 +239,7 @@ class AlignmentBase(object):
 
         if roi_yxhw is not None:
             img = img[
-                ..., roi_yxhw[0]: roi_yxhw[0] + roi_yxhw[2], roi_yxhw[1]: roi_yxhw[1] + roi_yxhw[3],
+                ..., roi_yxhw[0] : roi_yxhw[0] + roi_yxhw[2], roi_yxhw[1] : roi_yxhw[1] + roi_yxhw[3],
             ]
 
         img = img.copy()
@@ -295,7 +293,7 @@ class AlignmentBase(object):
 
         if not do_circular_conv:
             cc = np.fft.fftshift(cc, axes=(-2, -1))
-            cc = cc[padding[0]: -padding[0], padding[1]: -padding[1]]
+            cc = cc[padding[0] : -padding[0], padding[1] : -padding[1]]
             cc = np.fft.ifftshift(cc, axes=(-2, -1))
 
         return cc
@@ -346,7 +344,7 @@ class CenterOfRotation(AlignmentBase):
         padding_mode=None,
         peak_fit_radius=1,
         high_pass=None,
-        low_pass=None
+        low_pass=None,
     ):
         """Find the Center of Rotation (CoR), given to images.
 
@@ -441,18 +439,10 @@ class CenterOfRotation(AlignmentBase):
         roi_yxhw = self._determine_roi(img_shape, roi_yxhw, self.truncate_horz_pow2)
 
         img_1 = self._prepare_image(
-            img_1,
-            roi_yxhw=roi_yxhw,
-            median_filt_shape=median_filt_shape,
-            high_pass=high_pass,
-            low_pass=low_pass
+            img_1, roi_yxhw=roi_yxhw, median_filt_shape=median_filt_shape, high_pass=high_pass, low_pass=low_pass
         )
         img_2 = self._prepare_image(
-            img_2,
-            roi_yxhw=roi_yxhw,
-            median_filt_shape=median_filt_shape,
-            high_pass=high_pass,
-            low_pass=low_pass
+            img_2, roi_yxhw=roi_yxhw, median_filt_shape=median_filt_shape, high_pass=high_pass, low_pass=low_pass
         )
 
         cc = self._compute_correlation_fft(img_1, img_2, padding_mode)
@@ -496,17 +486,17 @@ class DetectorTranslationAlongBeam(AlignmentBase):
 
         if detected_equispaced and not supposed_equispaced:
             logging.getLogger(__name__).warning(
-                "The image position increments were supposed to NOT be equispaced, " +
-                "but they seem to be equispaced: " +
-                " ".join(("%f" % x for x in img_pos)) +
-                ". Forcing behavior according to the detected condition."
+                "The image position increments were supposed to NOT be equispaced, "
+                + "but they seem to be equispaced: "
+                + " ".join(("%f" % x for x in img_pos))
+                + ". Forcing behavior according to the detected condition."
             )
         if not detected_equispaced and supposed_equispaced:
             logging.getLogger(__name__).warning(
-                "The image position increments were supposed to be equispaced, " +
-                "but they seem to NOT be equispaced: " +
-                " ".join(("%f" % x for x in img_pos)) +
-                ". Forcing behavior according to the detected condition."
+                "The image position increments were supposed to be equispaced, "
+                + "but they seem to NOT be equispaced: "
+                + " ".join(("%f" % x for x in img_pos))
+                + ". Forcing behavior according to the detected condition."
             )
         return detected_equispaced
 
@@ -519,7 +509,7 @@ class DetectorTranslationAlongBeam(AlignmentBase):
         padding_mode=None,
         peak_fit_radius=1,
         equispaced_increments=True,
-        found_shifts_list=[]
+        return_shifts=False,
     ):
         """Find the deviation of the translation axis of the area detector
         along the beam propagation direction.
@@ -618,11 +608,10 @@ class DetectorTranslationAlongBeam(AlignmentBase):
         coeffs_v = Polynomial.fit(img_pos_increments, shifts_vh[:, 0], deg=1).convert().coef
         coeffs_h = Polynomial.fit(img_pos_increments, shifts_vh[:, 1], deg=1).convert().coef
 
-        if found_shifts_list is not None:
-            if not isinstance(found_shifts_list, list):
-                raise ValueError("If optional parameter found_shifts_list is given it must be a list."
-                                 "You have given an object of type %s" % type(found_shifts_list))
+        if return_shifts:
+            r_shifts = []
             for vh in shifts_vh:
-                found_shifts_list.append(vh)
-
-        return coeffs_v[1], coeffs_h[1]
+                r_shifts.append(vh)
+            return coeffs_v[1], coeffs_h[1], r_shifts
+        else:
+            return coeffs_v[1], coeffs_h[1]

nabu/preproc/tests/test_alignment.py

@@ -37,6 +37,14 @@ def bootstrap_dtr(request):
 
     cls.expected_shifts_vh = np.array((2551.25, 308.9))
 
+    cls.reference_shifts_list = [
+        [-121.39034751, 11.29080154],
+        [-117.01864405, 3.81080018],
+        [-115.66542712, 9.73375546],
+        [-116.72357645, 16.71435394],
+        [6.024768, 20.28406386],
+    ]
+
 
 def get_data_h5(*dataset_path):
     """
@@ -197,15 +205,7 @@ class TestDetectorTranslation(object):
         T_calc = DetectorTranslationAlongBeam()
 
         found_shifts_list = []
-        shifts_v, shifts_h = T_calc.find_shift(self.align_images, self.img_pos, found_shifts_list=found_shifts_list,)
-
-        reference_shifts_list = [
-            [-121.39034751, 11.29080154],
-            [-117.01864405, 3.81080018],
-            [-115.66542712, 9.73375546],
-            [-116.72357645, 16.71435394],
-            [6.024768, 20.28406386],
-        ]
+        shifts_v, shifts_h, found_shifts_list = T_calc.find_shift(self.align_images, self.img_pos, return_shifts=True)
 
         message = "Computed shifts coefficients %s and found ones %s do not coincide" % (
             (shifts_v, shifts_h),
@@ -213,5 +213,5 @@ class TestDetectorTranslation(object):
         )
         assert np.all(np.isclose(self.expected_shifts_vh, [shifts_v, shifts_h], atol=self.abs_tol)), message
 
-        message = "Computed shifts %s and found ones %s do not coincide" % (found_shifts_list, reference_shifts_list)
-        assert np.all(np.isclose(found_shifts_list, reference_shifts_list, atol=self.abs_tol)), message
+        message = "Computed shifts %s and found ones %s do not coincide" % (found_shifts_list, self.reference_shifts_list)
+        assert np.all(np.isclose(found_shifts_list, self.reference_shifts_list, atol=self.abs_tol)), message