Alignment: added 1d peak regions extraction along a given axis, from 2d images

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

nabu/preproc/alignment.py

@@ -42,7 +42,7 @@ class AlignmentBase(object):
             Estimated (vertical, horizontal) function max, according to the
             coordinates in fy and fx.
         """
-        if len(f_vals.shape) > 2:
+        if not (len(f_vals.shape) == 2):
             raise ValueError(
                 "The fitted values should form a 2-dimensional array. Array of shape: [%s] was given."
                 % (" ".join(("%d" % s for s in f_vals.shape)))
@@ -199,6 +199,56 @@ class AlignmentBase(object):
 
         return (f_vals, fv, fh)
 
+    @staticmethod
+    def extract_peak_regions_1d(cc, axis=-1, peak_radius=1, cc_coords=None):
+        """
+        Extracts a region around the maximum value.
+
+        Parameters
+        ----------
+        cc: numpy.ndarray
+            Correlation image.
+        axis: int, optional
+            Find the max values along the specified direction. The default is -1.
+        peak_radius: int, optional
+            The l_inf radius of the area to extract around the peak. The default is 1.
+        cc_coords: numpy.ndarray, optional
+            The coordinates of `cc` along the selected axis. The default is None.
+
+        Returns
+        -------
+        f_vals: numpy.ndarray
+            The extracted function values.
+        fc_ax: numpy.ndarray
+            The coordinates of the extracted values, along the selected axis.
+        """
+        img_shape = np.array(cc.shape)
+        if not (len(img_shape) == 2):
+            raise ValueError(
+                "The input image should be a 2-dimensional array. Array of shape: [%s] was given."
+                % (" ".join(("%d" % s for s in cc.shape)))
+            )
+        other_axis = (axis + 1) % 2
+        # get pixel having the maximum value of the correlation array
+        pix_max = np.argmax(cc, axis=axis)
+
+        # select a n neighborhood for the many 1D sub-pixel fittings (with wrapping)
+        p_ax_range = np.arange(- peak_radius, + peak_radius + 1)
+        p_ax = (pix_max[None, :] + p_ax_range[:, None]) % img_shape[axis]
+
+        p_ln = np.tile(np.arange(0, img_shape[axis])[None, :], [2 * peak_radius + 1, 1])
+
+        # extract the pixel coordinates along the axis
+        fc_ax = None if cc_coords is None else cc_coords[p_ax.flatten()].reshape(p_ax.shape)
+
+        # extract the correlation values
+        if other_axis == 0:
+            f_vals = cc[p_ln, p_ax]
+        else:
+            f_vals = cc[p_ax, p_ln]
+
+        return (f_vals, fc_ax)
+
     @staticmethod
     def _determine_roi(img_shape, roi_yxhw, do_truncate_horz_pow2):
         if roi_yxhw is None:

nabu/preproc/tests/test_alignment.py

@@ -100,6 +100,20 @@ class TestAlignmentBase(object):
         )
         assert "positions are outide the input margins" in str(ex.value), message
 
+    def test_extract_peak_regions_1d(self):
+        img = np.random.randint(0, 10, size=(8, 8))
+
+        peaks_pos = np.argmax(img, axis=-1)
+        peaks_val = np.max(img, axis=-1)
+
+        cc_coords = np.arange(0, 8)
+
+        (found_peaks_val, found_peaks_pos) = AlignmentBase.extract_peak_regions_1d(img, axis=-1, cc_coords=cc_coords)
+        message = "The found peak positions do not correspond to the expected peak positions:\n  Expected: %s\n  Found: %s" % (
+            peaks_pos, found_peaks_pos[1, :]
+        )
+        assert np.all(peaks_val == found_peaks_val[1, :]), message
+
 
 @pytest.mark.usefixtures("bootstrap_cor")
 class TestCor(object):
@@ -204,7 +218,6 @@ class TestDetectorTranslation(object):
     def test_alignxc(self):
         T_calc = DetectorTranslationAlongBeam()
 
-        found_shifts_list = []
         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" % (