got it to work as separate method

nabu/preproc/alignment.py

@@ -494,7 +494,7 @@ class AlignmentBase(object):
 
 
 class CenterOfRotation(AlignmentBase):
-    def find_shift(
+    def find_shift_original(
         self,
         img_1: np.ndarray,
         img_2: np.ndarray,
@@ -584,93 +584,15 @@ class CenterOfRotation(AlignmentBase):
 
         >>> cor_position = CoR_calc.find_shift(radio1, radio2, median_filt_shape=(3, 3))
         """
-        if not global_search:
-            return self.basic_find_shift_(
-                img_1,
-                img_2,
-                roi_yxhw,
-                median_filt_shape=median_filt_shape,
-                padding_mode=padding_mode,
-                peak_fit_radius=peak_fit_radius,
-                high_pass=high_pass,
-                low_pass=low_pass,
-                half_tomo_cor_guess=half_tomo_cor_guess,
-                half_tomo_return_stats=half_tomo_return_stats,
-            )
-
-        if global_search is True:
-            limits = None
-        else:
-            limits = global_search
-
-        return self.global_search_master_(
-            img_1, img_2, roi_yxhw, median_filt_shape, padding_mode, peak_fit_radius, high_pass, low_pass, limits=limits
-        )
-
-    def global_search_master_(
-        self,
-        img_1: np.ndarray,
-        img_2: np.ndarray,
-        roi_yxhw=None,
-        median_filt_shape=None,
-        padding_mode=None,
-        peak_fit_radius=1,
-        high_pass=None,
-        low_pass=None,
-        half_tomo_cor_guess=None,
-        half_tomo_return_stats=False,
-        limits=None,
-    ):
 
-        dim_radio = img_1.shape[1]
-
-        if limits is None:
-            lim1, lim2 = 10, dim_radio - 10
-        else:
-            lim1, lim2 = limits
-
-        if lim1 < 1:
-            lim1 = 1
-        if lim2 > dim_radio - 2:
-            lim2 = dim_radio - 2
-
-        assert isinstance(lim1, (int, float))
-        assert isinstance(lim2, (int, float))
-
-        found_centers = []
-        Xcor = lim1
-        while Xcor < lim2:
-            Xcor_rel = Xcor - (img_1.shape[1] // 2)
-            cor_position, merit, energy = self.basic_find_shift_(
-                img_1,
-                img_2,
-                low_pass=low_pass,
-                high_pass=high_pass,
-                half_tomo_cor_guess=Xcor_rel,
-                half_tomo_return_stats=True,
+        if global_search:
+            if global_search is True:
+                limits = None
+            else:
+                limits = global_search
+            return self.global_search_master_(
+                img_1, img_2, roi_yxhw, median_filt_shape, padding_mode, peak_fit_radius, high_pass, low_pass, limits=limits
             )
-            if not np.isnan(merit):
-                found_centers.append([merit, abs(Xcor_rel - cor_position), cor_position, energy])
-
-            Xcor = min(Xcor + Xcor / 6.0, Xcor + (dim_radio - Xcor) / 6.0)
-
-        found_centers.sort()
-        cor_position = found_centers[0][2]
-        return cor_position
-
-    def basic_find_shift_(
-        self,
-        img_1: np.ndarray,
-        img_2: np.ndarray,
-        roi_yxhw=None,
-        median_filt_shape=None,
-        padding_mode=None,
-        peak_fit_radius=1,
-        high_pass=None,
-        low_pass=None,
-        half_tomo_cor_guess=None,
-        half_tomo_return_stats=False,
-    ):
 
         self._check_img_pair_sizes(img_1, img_2)
 
@@ -742,6 +664,336 @@ class CenterOfRotation(AlignmentBase):
 
         return return_value
 
+    def find_shift(
+                self,
+                img_1: np.ndarray,
+                img_2: np.ndarray,
+                roi_yxhw=None,
+                median_filt_shape=None,
+                padding_mode=None,
+                peak_fit_radius=1,
+                high_pass=None,
+                low_pass=None,
+        ):
+        """Find the Center of Rotation (CoR), given two images.
+
+        This method finds the half-shift between two opposite images, by
+            means of correlation computed in Fourier space.
+
+        The output of this function, allows to compute motor movements for
+            aligning the sample rotation axis. Given the following values:
+
+        - L1: distance from source to motor
+            - L2: distance from source to detector
+            - ps: physical pixel size
+            - v: output of this function
+
+        displacement of motor = (L1 / L2 * ps) * v
+
+        Parameters
+            ----------
+            img_1: numpy.ndarray
+            First image
+            img_2: numpy.ndarray
+            Second image, it needs to have been flipped already (e.g. using numpy.fliplr).
+            roi_yxhw: (2, ) or (4, ) numpy.ndarray, tuple, or array, optional
+            4 elements vector containing: vertical and horizontal coordinates
+            of first pixel, plus height and width of the Region of Interest (RoI).
+            Or a 2 elements vector containing: plus height and width of the
+            centered Region of Interest (RoI).
+            Default is None -> deactivated.
+            median_filt_shape: (2, ) numpy.ndarray, tuple, or array, optional
+            Shape of the median filter window. Default is None -> deactivated.
+            padding_mode: str in numpy.pad's mode list, optional
+            Padding mode, which determines the type of convolution. If None or
+            'wrap' are passed, this resorts to the traditional circular convolution.
+            If 'edge' or 'constant' are passed, it results in a linear convolution.
+            Default is the circular convolution.
+            All options are:
+                None | 'constant' | 'edge' | 'linear_ramp' | 'maximum' | 'mean'
+                | 'median' | 'minimum' | 'reflect' | 'symmetric' |'wrap'
+            peak_fit_radius: int, optional
+            Radius size around the max correlation pixel, for sub-pixel fitting.
+            Minimum and default value is 1.
+            low_pass: float or sequence of two floats
+            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`
+
+        Raises
+            ------
+            ValueError
+            In case images are not 2-dimensional or have different sizes.
+
+        Returns
+            -------
+            float
+            Estimated center of rotation position from the center of the RoI in pixels.
+
+        Examples
+            --------
+            The following code computes the center of rotation position for two
+            given images in a tomography scan, where the second image is taken at
+            180 degrees from the first.
+
+        >>> radio1 = data[0, :, :]
+            ... radio2 = np.fliplr(data[1, :, :])
+            ... CoR_calc = CenterOfRotation()
+            ... cor_position = CoR_calc.find_shift(radio1, radio2)
+
+        Or for noisy images:
+
+        >>> cor_position = CoR_calc.find_shift(radio1, radio2, median_filt_shape=(3, 3))
+        """
+        self._check_img_pair_sizes(img_1, img_2)
+
+        if peak_fit_radius < 1:
+            logging.getLogger(__name__).warning(
+                "Parameter peak_fit_radius should be at least 1, given: %d instead." % peak_fit_radius
+            )
+            peak_fit_radius = 1
+
+        img_shape = img_2.shape
+        roi_yxhw = self._determine_roi(img_shape, roi_yxhw)
+
+        img_1 = self._prepare_image(img_1, roi_yxhw=roi_yxhw, median_filt_shape=median_filt_shape)
+        img_2 = self._prepare_image(img_2, roi_yxhw=roi_yxhw, median_filt_shape=median_filt_shape)
+
+        cc = self._compute_correlation_fft(img_1, img_2, padding_mode, high_pass=high_pass, low_pass=low_pass)
+        img_shape = img_2.shape
+        cc_vs = np.fft.fftfreq(img_shape[-2], 1 / img_shape[-2])
+        cc_hs = np.fft.fftfreq(img_shape[-1], 1 / img_shape[-1])
+
+        (f_vals, fv, fh) = self.extract_peak_region_2d(cc, peak_radius=peak_fit_radius, cc_vs=cc_vs, cc_hs=cc_hs)
+        fitted_shifts_vh = self.refine_max_position_2d(f_vals, fv, fh)
+
+        return fitted_shifts_vh[-1] / 2.0
+
+    def global_adaptive_find_shift(
+            self,
+            img_1: np.ndarray,
+            img_2: np.ndarray,
+            roi_yxhw=None,
+            median_filt_shape=None,
+            padding_mode=None,
+            peak_fit_radius=1,
+            high_pass=None,
+            low_pass=None,
+            margins=None,
+    ):
+        """Find the Center of Rotation (CoR), given two images.
+
+        This method finds the half-shift between two opposite images, by
+        means of correlation computed in Fourier space.
+        A global search is done on on the detector span (minus a margin) without assuming centered scan conditions.
+
+        The output of this function, allows to compute motor movements for
+        aligning the sample rotation axis. Given the following values:
+
+        - L1: distance from source to motor
+        - L2: distance from source to detector
+        - ps: physical pixel size
+        - v: output of this function
+
+        displacement of motor = (L1 / L2 * ps) * v
+
+        Parameters
+        ----------
+        img_1: numpy.ndarray
+            First image
+        img_2: numpy.ndarray
+            Second image, it needs to have been flipped already (e.g. using numpy.fliplr).
+        roi_yxhw: (2, ) or (4, ) numpy.ndarray, tuple, or array, optional
+            4 elements vector containing: vertical and horizontal coordinates
+            of first pixel, plus height and width of the Region of Interest (RoI).
+            Or a 2 elements vector containing: plus height and width of the
+            centered Region of Interest (RoI).
+            Default is None -> deactivated.
+        median_filt_shape: (2, ) numpy.ndarray, tuple, or array, optional
+            Shape of the median filter window. Default is None -> deactivated.
+        padding_mode: str in numpy.pad's mode list, optional
+            Padding mode, which determines the type of convolution. If None or
+            'wrap' are passed, this resorts to the traditional circular convolution.
+            If 'edge' or 'constant' are passed, it results in a linear convolution.
+            Default is the circular convolution.
+            All options are:
+                None | 'constant' | 'edge' | 'linear_ramp' | 'maximum' | 'mean'
+                | 'median' | 'minimum' | 'reflect' | 'symmetric' |'wrap'
+        peak_fit_radius: int, optional
+            Radius size around the max correlation pixel, for sub-pixel fitting.
+            Minimum and default value is 1.
+        low_pass: float or sequence of two floats
+            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`
+        half_tomo_cor_guess: float or None
+            The approximate position of the rotation axis from the image center. Optional.
+            When given a special algorithm is used which can work also in half-tomo conditions
+        margins:  None or a couple of floats or ints
+            margins is None or in the form of  (margin1,margin2) the search is done between margin1 and  dim_x-1-margin2.
+            If left to None then by default (margin1,margin2)  = ( 10, 10 )
+            
+
+        Raises
+        ------
+        ValueError
+            In case images are not 2-dimensional or have different sizes.
+
+        Returns
+        -------
+        float
+            Estimated center of rotation position from the center of the RoI in pixels.
+
+        Examples
+        --------
+        The following code computes the center of rotation position for two
+        given images in a tomography scan, where the second image is taken at
+        180 degrees from the first.
+
+        >>> radio1 = data[0, :, :]
+        ... radio2 = np.fliplr(data[1, :, :])
+        ... CoR_calc = CenterOfRotation()
+        ... cor_position = CoR_calc.find_shift(radio1, radio2)
+
+        Or for noisy images:
+
+        >>> cor_position = CoR_calc.find_shift(radio1, radio2, median_filt_shape=(3, 3))
+        """
+        
+        if peak_fit_radius < 1:
+            logging.getLogger(__name__).warning(
+                "Parameter peak_fit_radius should be at least 1, given: %d instead." % peak_fit_radius
+            )
+            peak_fit_radius = 1
+            
+        self._check_img_pair_sizes(img_1, img_2)
+
+        used_type = img_1.dtype
+
+
+        img_1_original = img_1
+        img_2_original = img_2
+        
+        if low_pass is not None or hig_hpass is not None:
+            pad_size = np.ceil(np.array(img_2.shape) / 2).astype(np.int)
+            pad_array = [(0,)] * len(img_2.shape)
+            for a in range(2):
+                pad_array[a] = (pad_size[a],)
+
+
+            print( " PADA ", pad_array)
+            print(" IMAS ", img_1.shape)
+            tmp_img_1 = np.pad(img_1, pad_array, mode="edge")
+            print(" tmp_img_1 S ", tmp_img_1.shape)
+
+            tmp_img_2 = np.pad(img_2, pad_array, mode="edge")
+            tmp_img_fft_1 = local_fftn(tmp_img_1)
+            tmp_img_fft_2 = local_fftn(tmp_img_2)
+
+            print(" SHAPES ",  tmp_img_1.shape ) 
+            print(" SHAPES ",  tmp_img_2.shape ) 
+
+            print(" LOCAL ", local_fftn)
+            print(" LOCAL use", __have_scipy__)
+            
+            filt = fourier_filters.get_bandpass_filter(
+                tmp_img_1.shape,
+                cutoff_lowpass=low_pass,
+                cutoff_highpass=high_pass,
+                use_rfft=__have_scipy__,
+                data_type=self.data_type,
+            )
+            print(" SHAPES F",  filt.shape ) 
+            print(" SHAPES F",  tmp_img_fft_1.shape  ) 
+            print(" SHAPES F",  tmp_img_1.shape  ) 
+
+            
+            tmp_img_1  = np.real(local_ifftn(   filt*tmp_img_fft_1 ))
+            tmp_img_2  = np.real(local_ifftn(   filt*tmp_img_fft_2 ))
+
+
+            
+            
+            img_1 = tmp_img_1[pad_array[0][0]:-pad_array[0][0]  , pad_array[1][0]:-pad_array[1][0]   ]
+            img_2 = tmp_img_2[pad_array[0][0]:-pad_array[0][0]  , pad_array[1][0]:-pad_array[1][0]   ]
+            
+
+        roi_yxhw = self._determine_roi(img_1.shape, roi_yxhw)
+
+        img_1 = self._prepare_image(img_1, roi_yxhw=roi_yxhw, median_filt_shape=median_filt_shape).astype(used_type)
+        img_2 = self._prepare_image(img_2, roi_yxhw=roi_yxhw, median_filt_shape=median_filt_shape).astype(used_type)
+
+        dim_radio = img_1.shape[1]
+
+        if margins is None:
+            lim1, lim2 = 10, dim_radio-1 - 10
+        else:
+            lim1, lim2 = margins
+            lim2 = dim_radio - 1 - lim2
+
+        if lim1 < 1:
+            lim1 = 1
+        if lim2 > dim_radio - 2:
+            lim2 = dim_radio - 2
+
+        if  lim2 <= lim1  :
+            message = (
+                "Image shape or cropped selection too small for global search. After removal of the margins the search limit collide. The cropped size is %d\n" % (dim_radio) 
+            )
+            raise ValueError(message)
+        
+        found_centers = []
+        Xcor = lim1
+        while Xcor < lim2:
+
+            tmpsigma = min(
+                (img_1.shape[1] - Xcor) / 4.0,
+                (Xcor) / 4.0,
+            )
+
+            tmpx = (np.arange(img_1.shape[1]) - Xcor) / tmpsigma
+            apodis = np.exp(-tmpx * tmpx / 2.0)
+            
+            Xcor_rel = Xcor - (img_1.shape[1] // 2)
+
+            img_1_apodised = img_1 * apodis
+            
+            cor_position  = self.find_shift_original(
+                img_1_original,
+                img_2_original,
+                low_pass=low_pass,
+                high_pass=high_pass,
+                half_tomo_cor_guess = Xcor_rel,
+                roi_yxhw=            roi_yxhw,
+            )
+
+            
+            cor_in_img = img_1.shape[1] // 2 + cor_position
+            tmpsigma = min(
+                (img_1.shape[1] - cor_in_img) / 4.0,
+                (cor_in_img) / 4.0,
+            )
+            
+            M1 = int(round(cor_position + img_1.shape[1] // 2)) - int(round(tmpsigma))
+            M2 = int(round(cor_position + img_1.shape[1] // 2)) + int(round(tmpsigma))
+
+            piece1 = img_1_original[:, M1:M2]
+            piece2 = img_2_original[:, img_1.shape[1] - M2 : img_1.shape[1] - M1]
+            energy = np.array(piece1 * piece1 + piece2 * piece2, "d").sum()
+            diff_energy = np.array((piece1 - piece2) * (piece1 - piece2), "d").sum()
+            cost = diff_energy / energy
+
+            print( cost, abs(Xcor_rel - cor_position), cor_position, energy  ) 
+            if not np.isnan(cost):
+                found_centers.append([cost, abs(Xcor_rel - cor_position), cor_position, energy])
+
+            Xcor = min(Xcor + Xcor / 6.0, Xcor + (dim_radio - Xcor) / 6.0)
+
+        found_centers.sort()
+        cor_position = found_centers[0][2]
+        return cor_position
+
+
     __call__ = find_shift
 
 

nabu/preproc/tests/test_alignment.py

@@ -351,7 +351,7 @@ class TestCor(object):
         cor_pos = self.ht_cor
 
         CoR_calc = alignment.CenterOfRotation()
-        cor_position = CoR_calc.find_shift(radio1, radio2, low_pass=1, high_pass=20, half_tomo_cor_guess=cor_pos - 10.0)
+        cor_position = CoR_calc.global_adaptive_find_shift(radio1, radio2, low_pass=1, high_pass=20)
 
         message = (
             "Computed CoR %f " % cor_position
@@ -360,7 +360,7 @@ class TestCor(object):
         assert np.isclose(cor_pos, cor_position, atol=self.abs_tol), message
 
     def test_half_tomo_cor_exp(self):
-        """test the hal_tomo algorithm on experimental data and global search"""
+        """test the hal_tomo algorithm on experimental data """
 
         radios = nabu_get_data("ha_autocor_radios.npz")
         radio1 = radios["radio1"]
@@ -372,7 +372,7 @@ class TestCor(object):
 
         CoR_calc = alignment.CenterOfRotation()
 
-        cor_position = CoR_calc.find_shift(radio1, radio2, low_pass=1, high_pass=20, global_search=True)
+        cor_position = CoR_calc.global_adaptive_find_shift(radio1, radio2, low_pass=1, high_pass=20)
 
         print("Found cor_position", cor_position)
 
@@ -388,10 +388,10 @@ class TestCor(object):
         radios = nabu_get_data("ha_autocor_radios.npz")
         radio1 = radios["radio1"]
         radio2 = radios["radio2"]
+        cor_pos =radios["cor_pos"]
 
         radio2 = np.fliplr(radio2)
 
-        cor_pos = 983.107
 
         CoR_calc = alignment.CenterOfRotation()