restructured out the global search nesting

nabu/preproc/alignment.py

@@ -494,57 +494,6 @@ class AlignmentBase(object):
 
 
 class CenterOfRotation(AlignmentBase):
-    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.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 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 find_shift(
         self,
         img_1: np.ndarray,
@@ -635,15 +584,93 @@ class CenterOfRotation(AlignmentBase):
 
         >>> cor_position = CoR_calc.find_shift(radio1, radio2, median_filt_shape=(3, 3))
         """
+        if not global_search:
+            return 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:
-            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 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 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)