Alignment: detector translation: simplified function

Signed-off-by: Nicola VIGANÒ <nicola.vigano@esrf.fr>

nabu/preproc/alignment.py

@@ -481,27 +481,6 @@ class DetectorTranslationAlongBeam(AlignmentBase):
                 % (" ".join(("%d" % x for x in shape_stack)), " ".join(("%d" % x for x in shape_pos)),)
             )
 
-    @staticmethod
-    def _check_equispaced(img_pos, supposed_equispaced):
-        img_pos_incrs = np.diff(img_pos)
-        detected_equispaced = np.all(np.isclose(img_pos_incrs[0], img_pos_incrs[1:]))
-
-        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."
-            )
-        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."
-            )
-        return detected_equispaced
-
     def find_shift(
         self,
         img_stack: np.ndarray,
@@ -514,6 +493,7 @@ class DetectorTranslationAlongBeam(AlignmentBase):
         low_pass=None,
         equispaced_increments=True,
         return_shifts=False,
+        use_adjacent_imgs=False,
     ):
         """Find the deviation of the translation axis of the area detector
         along the beam propagation direction.
@@ -551,19 +531,12 @@ class DetectorTranslationAlongBeam(AlignmentBase):
             Low-pass filter properties, as described in `nabu.misc.fourier_filters`.
         high_pass: float or sequence of two floats
             High-pass filter properties, as described in `nabu.misc.fourier_filters`.
-        equispaced_increments: boolean, optional
-            Tells whether the position increments are equispaced or not. If
-            equispaced increments are used, we have to compute the correlation
-            images with respect to the first image, otherwise we can do it
-            against adjacent images.
-            The advantage of doing it between adjacent images is that we do not
-            build up large displacements in the correlation.
-            However, if this is done for equispaced images, the linear fitting
-            becomes unstable.
         found_shifts_list: list, optional
             if a list is given in input, it will be populated
             with the found shifts. For each given image of the stack a tuple of shifts will be
             appended to the list. This slot is intended for introspection.
+        use_adjacent_imgs: boolean, optional
+            Compute correlation between adjacent images. It is better when dealing with large shifts.
 
         Returns
         -------
@@ -576,7 +549,6 @@ class DetectorTranslationAlongBeam(AlignmentBase):
         TODO: Add examples here!
         """
         self._check_img_sizes(img_stack, img_pos)
-        equispaced_increments = self._check_equispaced(img_pos, equispaced_increments)
 
         if peak_fit_radius < 1:
             logging.getLogger(__name__).warning(
@@ -596,7 +568,7 @@ class DetectorTranslationAlongBeam(AlignmentBase):
         # do correlations
         ccs = [
             self._compute_correlation_fft(
-                img_stack[0 if equispaced_increments else ii - 1, ...],
+                img_stack[ii - 1 if use_adjacent_imgs else 0, ...],
                 img_stack[ii, ...],
                 padding_mode,
                 high_pass=high_pass,
@@ -610,10 +582,9 @@ class DetectorTranslationAlongBeam(AlignmentBase):
             (f_vals, fv, fh) = self.extract_peak_region_2d(cc, peak_radius=peak_fit_radius, cc_vs=cc_vs, cc_hs=cc_hs)
             shifts_vh[ii, :] = self.refine_max_position_2d(f_vals, fv, fh)
 
-        if equispaced_increments:
-            img_pos_increments = img_pos[1:] - img_pos[0]
-        else:
-            img_pos_increments = np.diff(img_pos)
+        if use_adjacent_imgs:
+            shifts_vh = np.cumsum(shifts_vh, axis=0)
+        img_pos_increments = img_pos[1:] - img_pos[0]
 
         # Polynomial.fit is supposed to be more numerically stable than polyfit
         # (according to numpy)
@@ -621,9 +592,6 @@ class DetectorTranslationAlongBeam(AlignmentBase):
         coeffs_h = Polynomial.fit(img_pos_increments, shifts_vh[:, 1], deg=1).convert().coef
 
         if return_shifts:
-            r_shifts = []
-            for vh in shifts_vh:
-                r_shifts.append(vh)
-            return coeffs_v[1], coeffs_h[1], r_shifts
+            return coeffs_v[1], coeffs_h[1], shifts_vh
         else:
             return coeffs_v[1], coeffs_h[1]