Alignment: added plotting and moved properties in base class

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

nabu/preproc/alignment.py

 import numpy as np
 
 import logging
-from numpy.polynomial.polynomial import Polynomial
+from numpy.polynomial.polynomial import Polynomial, polyval
 
 from nabu.utils import previouspow2
 from nabu.misc import fourier_filters
@@ -15,8 +15,40 @@ except ImportError:
 
     __have_scipy__ = False
 
+try:
+    import matplotlib.pyplot as plt
+
+    __have_matplotlib__ = True
+except ImportError:
+    logging.getLogger(__name__).warning("Matplotlib not available. Plotting disabled")
+
+    __have_matplotlib__ = False
+
 
 class AlignmentBase(object):
+    def __init__(self, horz_fft_width=False, verbose=False):
+        """
+        Alignment basic functions.
+
+        Parameters
+        ----------
+        horz_fft_width: boolean, optional
+            If True, restrict the horizontal size to a power of 2:
+
+            >>> new_x_dim = 2 ** math.floor(math.log2(x_dim))
+        verbose: boolean, optional
+            When True it will produce verbose output, including plots.
+        """
+        self._init_parameters(horz_fft_width, verbose)
+
+    def _init_parameters(self, horz_fft_width, verbose):
+        self.truncate_horz_pow2 = horz_fft_width
+
+        if verbose and not __have_matplotlib__:
+            logging.getLogger(__name__).warning("Matplotlib not available. Plotting disabled, despite being activated by user")
+            verbose = False
+        self.verbose = verbose
+
     @staticmethod
     def refine_max_position_2d(f_vals: np.ndarray, fy=None, fx=None):
         """Computes the sub-pixel max position of the given function sampling.
@@ -83,7 +115,7 @@ class AlignmentBase(object):
         if np.any(vertex_yx < vertex_min_yx) or np.any(vertex_yx > vertex_max_yx):
             raise ValueError(
                 "Fitted (y: {}, x: {}) positions are outide the input margins y: [{}, {}], and x: [{}, {}]".format(
-                    vertex_yx[0], vertex_yx[1], vertex_min_yx[0], vertex_max_yx[0], vertex_min_yx[1], vertex_max_yx[1]
+                    vertex_yx[0], vertex_yx[1], vertex_min_yx[0], vertex_max_yx[0], vertex_min_yx[1], vertex_max_yx[1],
                 )
             )
         return vertex_yx
@@ -150,7 +182,7 @@ class AlignmentBase(object):
                 )
             else:
                 message = "Fitted positions outide the input margins [{}, {}]: %d below and %d above".format(
-                    vertex_min_x, vertex_max_x, np.sum(1 - lower_bound_ok), np.sum(1 - upper_bound_ok)
+                    vertex_min_x, vertex_max_x, np.sum(1 - lower_bound_ok), np.sum(1 - upper_bound_ok),
                 )
             raise ValueError(message)
         return vertex_x
@@ -233,7 +265,7 @@ class AlignmentBase(object):
         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_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])
@@ -249,13 +281,12 @@ class AlignmentBase(object):
 
         return (f_vals, fc_ax)
 
-    @staticmethod
-    def _determine_roi(img_shape, roi_yxhw, do_truncate_horz_pow2):
+    def _determine_roi(self, img_shape, roi_yxhw):
         if roi_yxhw is None:
             # vertical window size is reduced to a power of 2 to accelerate fft
             # same thing horizontal window - if requested. Default is not
             roi_yxhw = previouspow2(img_shape)
-            if not do_truncate_horz_pow2:
+            if not self.truncate_horz_pow2:
                 roi_yxhw[1] = img_shape[1]
 
         if len(roi_yxhw) == 2:  # Convert centered 2-element roi into 4-element
@@ -264,7 +295,9 @@ class AlignmentBase(object):
         return roi_yxhw
 
     @staticmethod
-    def _prepare_image(img, invalid_val=1e-5, roi_yxhw=None, median_filt_shape=None, low_pass=None, high_pass=None):
+    def _prepare_image(
+        img, invalid_val=1e-5, roi_yxhw=None, median_filt_shape=None, low_pass=None, high_pass=None,
+    ):
         """
         Prepare and returns  a cropped  and filtered image, or array of filtered images if the input is an  array of images.
 
@@ -367,23 +400,6 @@ class AlignmentBase(object):
 
 
 class CenterOfRotation(AlignmentBase):
-    def __init__(self, horz_fft_width=False):
-        """
-        Center of Rotation (CoR) computation object.
-        This class is used on radios.
-
-        Parameters
-        ----------
-        horz_fft_width: boolean, optional
-            If True, restrict the horizontal size to a power of 2:
-
-            >>> new_x_dim = 2 ** math.floor(math.log2(x_dim))
-        """
-        self._init_parameters(horz_fft_width)
-
-    def _init_parameters(self, horz_fft_width):
-        self.truncate_horz_pow2 = horz_fft_width
-
     @staticmethod
     def _check_img_sizes(img_1: np.ndarray, img_2: np.ndarray):
         shape_1 = np.squeeze(img_1).shape
@@ -492,7 +508,7 @@ class CenterOfRotation(AlignmentBase):
             peak_fit_radius = 1
 
         img_shape = img_2.shape
-        roi_yxhw = self._determine_roi(img_shape, roi_yxhw, self.truncate_horz_pow2)
+        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)
@@ -608,7 +624,7 @@ class DetectorTranslationAlongBeam(AlignmentBase):
 
         num_imgs = img_stack.shape[0]
         img_shape = img_stack.shape[-2:]
-        roi_yxhw = self._determine_roi(img_shape, roi_yxhw, False)
+        roi_yxhw = self._determine_roi(img_shape, roi_yxhw)
 
         img_stack = self._prepare_image(img_stack, roi_yxhw=roi_yxhw, median_filt_shape=median_filt_shape)
 
@@ -641,6 +657,16 @@ 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 self.verbose:
+            f, axs = plt.subplots(1, 2)
+            axs[0].scatter(img_pos_increments, shifts_vh[:, 0])
+            axs[0].plot(img_pos_increments, polyval(img_pos_increments, coeffs_v))
+            axs[0].set_title("Vertical shifts")
+            axs[1].scatter(img_pos_increments, shifts_vh[:, 1])
+            axs[1].plot(img_pos_increments, polyval(img_pos_increments, coeffs_h))
+            axs[1].set_title("Horizontal shifts")
+            plt.show(block=False)
+
         if return_shifts:
             return coeffs_v[1], coeffs_h[1], shifts_vh
         else: