nabu 2024.1.1__py3-none-any.whl → 2024.1.2__py3-none-any.whl

build/lib/build/lib/nabu/estimation/alignment.py

@@ -1,578 +0,0 @@
-# import math
-import logging
-import numpy as np
-from numpy.polynomial.polynomial import Polynomial
-from silx.math.medianfilter import medfilt2d
-import scipy.fft  # pylint: disable=E0611
-from ..utils import previouspow2
-from ..misc import fourier_filters
-from ..resources.logger import LoggerOrPrint
-
-try:
-    import matplotlib.pyplot as plt
-
-    __have_matplotlib__ = True
-except ImportError:
-    logging.getLogger(__name__).warning("Matplotlib not available. Plotting disabled")
-    plt = None
-    __have_matplotlib__ = False
-
-try:
-    from tqdm import tqdm
-
-    def progress_bar(x, verbose=True):
-        if verbose:
-            return tqdm(x)
-        else:
-            return x
-
-except ImportError:
-
-    def progress_bar(x, verbose=True):
-        return x
-
-
-local_fftn = scipy.fft.rfftn
-local_ifftn = scipy.fft.irfftn
-
-
-class AlignmentBase:
-    default_extra_options = {"blocking_plots": False}
-    _default_cor_options = {}
-
-    def __init__(
-        self,
-        vert_fft_width=False,
-        horz_fft_width=False,
-        verbose=False,
-        logger=None,
-        data_type=np.float32,
-        extra_options=None,
-        cor_options=None,
-    ):
-        """
-        Alignment basic functions.
-
-        Parameters
-        ----------
-        vert_fft_width: boolean, optional
-            If True, restrict the vertical size to a power of 2:
-
-            >>> new_v_dim = 2 ** math.floor(math.log2(v_dim))
-
-        horz_fft_width: boolean, optional
-            If True, restrict the horizontal size to a power of 2:
-
-            >>> new_h_dim = 2 ** math.floor(math.log2(h_dim))
-
-        verbose: boolean, optional
-            When True it will produce verbose output, including plots.
-        data_type: `numpy.float32`
-            Computation data type.
-        """
-
-        self._init_parameters(vert_fft_width, horz_fft_width, verbose, logger, data_type, extra_options=extra_options)
-        self._plot_windows = {}
-        self._set_cor_options(cor_options)
-
-    def _set_cor_options(self, cor_options):
-        default_dict = self._default_cor_options.copy()
-        if cor_options is not None:
-            default_dict.update(cor_options)
-        self.cor_options = default_dict
-
-    def _init_parameters(self, vert_fft_width, horz_fft_width, verbose, logger, data_type, extra_options=None):
-        self.logger = LoggerOrPrint(logger)
-        self.truncate_vert_pow2 = vert_fft_width
-        self.truncate_horz_pow2 = horz_fft_width
-
-        if verbose and not __have_matplotlib__:
-            self.logger.warning("Matplotlib not available. Plotting disabled, despite being activated by user")
-            verbose = False
-        self.verbose = verbose
-        self.data_type = data_type
-        self.extra_options = self.default_extra_options.copy()
-        self.extra_options.update(extra_options or {})
-
-    @staticmethod
-    def _check_img_stack_size(img_stack: np.ndarray, img_pos: np.ndarray):
-        shape_stack = np.squeeze(img_stack).shape
-        shape_pos = np.squeeze(img_pos).shape
-        if not len(shape_stack) == 3:
-            raise ValueError(
-                "A stack of 2-dimensional images is required. Shape of stack: %s"
-                % (" ".join(("%d" % x for x in shape_stack)))
-            )
-        if not len(shape_pos) == 1:
-            raise ValueError(
-                "Positions need to be a 1-dimensional array. Shape of the positions variable: %s"
-                % (" ".join(("%d" % x for x in shape_pos)))
-            )
-        if not shape_stack[0] == shape_pos[0]:
-            raise ValueError(
-                "The same number of images and positions is required."
-                + " Shape of stack: %s, shape of positions variable: %s"
-                % (
-                    " ".join(("%d" % x for x in shape_stack)),
-                    " ".join(("%d" % x for x in shape_pos)),
-                )
-            )
-
-    @staticmethod
-    def _check_img_pair_sizes(img_1: np.ndarray, img_2: np.ndarray):
-        shape_1 = np.squeeze(img_1).shape
-        shape_2 = np.squeeze(img_2).shape
-        if not len(shape_1) == 2:
-            raise ValueError(
-                "Images need to be 2-dimensional. Shape of image #1: %s" % (" ".join(("%d" % x for x in shape_1)))
-            )
-        if not len(shape_2) == 2:
-            raise ValueError(
-                "Images need to be 2-dimensional. Shape of image #2: %s" % (" ".join(("%d" % x for x in shape_2)))
-            )
-        if not np.all(shape_1 == shape_2):
-            raise ValueError(
-                "Images need to be of the same shape. Shape of image #1: %s, image #2: %s"
-                % (
-                    " ".join(("%d" % x for x in shape_1)),
-                    " ".join(("%d" % x for x in shape_2)),
-                )
-            )
-
-    @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.
-
-        Parameters
-        ----------
-        f_vals: numpy.ndarray
-            Function values of the sampled points
-        fy: numpy.ndarray, optional
-            Vertical coordinates of the sampled points
-        fx: numpy.ndarray, optional
-            Horizontal coordinates of the sampled points
-
-        Raises
-        ------
-        ValueError
-            In case position and values do not have the same size, or in case
-            the fitted maximum is outside the fitting region.
-
-        Returns
-        -------
-        tuple(float, float)
-            Estimated (vertical, horizontal) function max, according to the
-            coordinates in fy and fx.
-        """
-        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)))
-            )
-        if fy is None:
-            fy_half_size = (f_vals.shape[0] - 1) / 2
-            fy = np.linspace(-fy_half_size, fy_half_size, f_vals.shape[0])
-        elif not (len(fy.shape) == 1 and np.all(fy.size == f_vals.shape[0])):
-            raise ValueError(
-                "Vertical coordinates should have the same length as values matrix. Sizes of fy: %d, f_vals: [%s]"
-                % (fy.size, " ".join(("%d" % s for s in f_vals.shape)))
-            )
-        if fx is None:
-            fx_half_size = (f_vals.shape[1] - 1) / 2
-            fx = np.linspace(-fx_half_size, fx_half_size, f_vals.shape[1])
-        elif not (len(fx.shape) == 1 and np.all(fx.size == f_vals.shape[1])):
-            raise ValueError(
-                "Horizontal coordinates should have the same length as values matrix. Sizes of fx: %d, f_vals: [%s]"
-                % (fx.size, " ".join(("%d" % s for s in f_vals.shape)))
-            )
-
-        fy, fx = np.meshgrid(fy, fx, indexing="ij")
-        fy = fy.flatten()
-        fx = fx.flatten()
-        coords = np.array([np.ones(f_vals.size), fy, fx, fy * fx, fy**2, fx**2])
-
-        coeffs = np.linalg.lstsq(coords.T, f_vals.flatten(), rcond=None)[0]
-
-        # For a 1D parabola `f(x) = ax^2 + bx + c`, the vertex position is:
-        # x_v = -b / 2a. For a 2D parabola, the vertex position is:
-        # (y, x)_v = - b / A, where:
-        A = [[2 * coeffs[4], coeffs[3]], [coeffs[3], 2 * coeffs[5]]]
-        b = coeffs[1:3]
-        vertex_yx = np.linalg.lstsq(A, -b, rcond=None)[0]
-
-        vertex_min_yx = [np.min(fy), np.min(fx)]
-        vertex_max_yx = [np.max(fy), np.max(fx)]
-        if np.any(vertex_yx < vertex_min_yx) or np.any(vertex_yx > vertex_max_yx):
-            raise ValueError(
-                "Fitted (y: {}, x: {}) positions are outside 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],
-                )
-            )
-        return vertex_yx
-
-    @staticmethod
-    def refine_max_position_1d(f_vals, fx=None, return_vertex_val=False, return_all_coeffs=False):
-        """Computes the sub-pixel max position of the given function sampling.
-
-        Parameters
-        ----------
-        f_vals: numpy.ndarray
-            Function values of the sampled points
-        fx: numpy.ndarray, optional
-            Coordinates of the sampled points
-        return_vertex_val: boolean, option
-            Enables returning the vertex values. Defaults to False.
-
-        Raises
-        ------
-        ValueError
-            In case position and values do not have the same size, or in case
-            the fitted maximum is outside the fitting region.
-
-        Returns
-        -------
-        float
-            Estimated function max, according to the coordinates in fx.
-        """
-        if not len(f_vals.shape) in (1, 2):
-            raise ValueError(
-                "The fitted values should be either one or a collection of 1-dimensional arrays. Array of shape: [%s] was given."
-                % (" ".join(("%d" % s for s in f_vals.shape)))
-            )
-        num_vals = f_vals.shape[0]
-
-        if fx is None:
-            fx_half_size = (num_vals - 1) / 2
-            fx = np.linspace(-fx_half_size, fx_half_size, num_vals)
-        else:
-            fx = np.squeeze(fx)
-            if not (len(fx.shape) == 1 and np.all(fx.size == num_vals)):
-                raise ValueError(
-                    "Base coordinates should have the same length as values array. Sizes of fx: %d, f_vals: %d"
-                    % (fx.size, num_vals)
-                )
-
-        if len(f_vals.shape) == 1:
-            # using Polynomial.fit, because supposed to be more numerically
-            # stable than previous solutions (according to numpy).
-            poly = Polynomial.fit(fx, f_vals, deg=2)
-            coeffs = poly.convert().coef
-        else:
-            coords = np.array([np.ones(num_vals), fx, fx**2])
-            coeffs = np.linalg.lstsq(coords.T, f_vals, rcond=None)[0]
-
-        # For a 1D parabola `f(x) = c + bx + ax^2`, the vertex position is:
-        # x_v = -b / 2a.
-        vertex_x = -coeffs[1, :] / (2 * coeffs[2, :])
-        if not return_all_coeffs:
-            vertex_x = vertex_x[0]
-
-        vertex_min_x = np.min(fx)
-        vertex_max_x = np.max(fx)
-        lower_bound_ok = vertex_min_x < vertex_x
-        upper_bound_ok = vertex_x < vertex_max_x
-        if not np.all(lower_bound_ok * upper_bound_ok):
-            if len(f_vals.shape) == 1:
-                message = "Fitted position {} is outide the input margins [{}, {}]".format(
-                    vertex_x, vertex_min_x, vertex_max_x
-                )
-            else:
-                message = "Fitted positions outside the input margins [{}, {}]: {} below and {} above".format(
-                    vertex_min_x,
-                    vertex_max_x,
-                    np.sum(1 - lower_bound_ok),
-                    np.sum(1 - upper_bound_ok),
-                )
-            raise ValueError(message)
-        if return_vertex_val:
-            vertex_val = coeffs[0, :] + vertex_x * coeffs[1, :] / 2
-            return vertex_x, vertex_val
-        else:
-            return vertex_x
-
-    @staticmethod
-    def extract_peak_region_2d(cc, peak_radius=1, cc_vs=None, cc_hs=None):
-        """
-        Extracts a region around the maximum value.
-
-        Parameters
-        ----------
-        cc: numpy.ndarray
-            Correlation image.
-        peak_radius: int, optional
-            The l_inf radius of the area to extract around the peak. The default is 1.
-        cc_vs: numpy.ndarray, optional
-            The vertical coordinates of `cc`. The default is None.
-        cc_hs: numpy.ndarray, optional
-            The horizontal coordinates of `cc`. The default is None.
-
-        Returns
-        -------
-        f_vals: numpy.ndarray
-            The extracted function values.
-        fv: numpy.ndarray
-            The vertical coordinates of the extracted values.
-        fh: numpy.ndarray
-            The horizontal coordinates of the extracted values.
-        """
-        img_shape = np.array(cc.shape)
-        # get pixel having the maximum value of the correlation array
-        pix_max_corr = np.argmax(cc)
-        pv, ph = np.unravel_index(pix_max_corr, img_shape)
-
-        # select a n x n neighborhood for the sub-pixel fitting (with wrapping)
-        pv = np.arange(pv - peak_radius, pv + peak_radius + 1) % img_shape[-2]
-        ph = np.arange(ph - peak_radius, ph + peak_radius + 1) % img_shape[-1]
-
-        # extract the (v, h) pixel coordinates
-        fv = None if cc_vs is None else cc_vs[pv]
-        fh = None if cc_hs is None else cc_hs[ph]
-
-        # extract the correlation values
-        pv, ph = np.meshgrid(pv, ph, indexing="ij")
-        f_vals = cc[pv, ph]
-
-        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.
-        """
-        if len(cc.shape) == 1:
-            cc = cc[None, ...]
-        img_shape = np.array(cc.shape)
-        if not (len(img_shape) == 2):
-            raise ValueError(
-                "The input image should be either a 1 or 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[other_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)
-
-    def _determine_roi(self, img_shape, roi_yxhw):
-        if roi_yxhw is None:
-            # vertical and horizontal window sizes are reduced to a power of 2
-            # to accelerate fft if requested. Default is not.
-            roi_yxhw = previouspow2(img_shape)
-            if not self.truncate_vert_pow2:
-                roi_yxhw[0] = img_shape[0]
-            if not self.truncate_horz_pow2:
-                roi_yxhw[1] = img_shape[1]
-
-        roi_yxhw = np.array(roi_yxhw, dtype=np.intp)
-        if len(roi_yxhw) == 2:  # Convert centered 2-element roi into 4-element
-            roi_yxhw = np.concatenate(((img_shape - roi_yxhw) // 2, roi_yxhw))
-        return roi_yxhw
-
-    def _prepare_image(
-        self,
-        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.
-
-        Parameters
-        ----------
-        img: numpy.ndarray
-            image or stack of images
-        invalid_val: float
-            value to be used in replacement of nan and inf values
-        median_filt_shape: int or sequence of int
-            the width or the widths of the median window
-        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`
-
-        Returns
-        -------
-        numpy.array_like
-            The computed filter
-        """
-        img = np.squeeze(img)  # Removes singleton dimensions, but does a shallow copy
-        img = np.ascontiguousarray(img, dtype=self.data_type)
-
-        if roi_yxhw is not None:
-            img = img[
-                ...,
-                roi_yxhw[0] : roi_yxhw[0] + roi_yxhw[2],
-                roi_yxhw[1] : roi_yxhw[1] + roi_yxhw[3],
-            ]
-
-        img = img.copy()
-
-        img[np.isnan(img)] = invalid_val
-        img[np.isinf(img)] = invalid_val
-
-        if high_pass is not None or low_pass is not None:
-            img_filter = fourier_filters.get_bandpass_filter(
-                img.shape[-2:],
-                cutoff_lowpass=low_pass,
-                cutoff_highpass=high_pass,
-                use_rfft=True,
-                data_type=self.data_type,
-            )
-            # fft2 and iff2 use axes=(-2, -1) by default
-            img = local_ifftn(local_fftn(img, axes=(-2, -1)) * img_filter, axes=(-2, -1)).real
-
-        if median_filt_shape is not None:
-            img_shape = img.shape
-            # expanding filter shape with ones, to cover the stack of images
-            # but disabling inter-image filtering
-            median_filt_shape = np.concatenate(
-                (
-                    np.ones((len(img_shape) - len(median_filt_shape),), dtype=np.intp),
-                    median_filt_shape,
-                )
-            )
-            img = medfilt2d(img, kernel_size=median_filt_shape)
-        return img
-
-    def _transform_to_fft(
-        self, img_1: np.ndarray, img_2: np.ndarray, padding_mode, axes=(-2, -1), low_pass=None, high_pass=None
-    ):
-        do_circular_conv = padding_mode is None or padding_mode == "wrap"
-        img_shape = img_2.shape
-        if not do_circular_conv:
-            pad_size = np.ceil(np.array(img_shape) / 2).astype(np.intp)
-            pad_array = [(0,)] * len(img_shape)
-            for a in axes:
-                pad_array[a] = (pad_size[a],)
-
-            img_1 = np.pad(img_1, pad_array, mode=padding_mode)
-            img_2 = np.pad(img_2, pad_array, mode=padding_mode)
-        else:
-            pad_size = None
-        img_shape = img_2.shape
-
-        # compute fft's of the 2 images
-        img_fft_1 = local_fftn(img_1, axes=axes)
-        img_fft_2 = local_fftn(img_2, axes=axes)
-
-        if low_pass is not None or high_pass is not None:
-            filt = fourier_filters.get_bandpass_filter(
-                img_shape[-2:],
-                cutoff_lowpass=low_pass,
-                cutoff_highpass=high_pass,
-                use_rfft=True,
-                data_type=self.data_type,
-            )
-        else:
-            filt = None
-
-        return img_fft_1, img_fft_2, filt, pad_size
-
-    def _compute_correlation_fft(
-        self, img_1: np.ndarray, img_2: np.ndarray, padding_mode, axes=(-2, -1), low_pass=None, high_pass=None
-    ):
-        img_fft_1, img_fft_2, filt, pad_size = self._transform_to_fft(
-            img_1, img_2, padding_mode=padding_mode, axes=axes, low_pass=low_pass, high_pass=high_pass
-        )
-
-        img_prod = img_fft_1 * np.conjugate(img_fft_2)
-
-        if filt is not None:
-            img_prod *= filt
-
-        # inverse fft of the product to get cross_correlation of the 2 images
-        cc = np.real(local_ifftn(img_prod, axes=axes))
-
-        if pad_size is not None:
-            cc_shape = cc.shape
-            cc = np.fft.fftshift(cc, axes=axes)
-
-            slicing = [slice(None)] * len(cc_shape)
-            for a in axes:
-                slicing[a] = slice(pad_size[a], cc_shape[a] - pad_size[a])
-            cc = cc[tuple(slicing)]
-
-            cc = np.fft.ifftshift(cc, axes=axes)
-
-        return cc
-
-    def _add_plot_window(self, fig, ax=None):
-        self._plot_windows[fig.number] = {"figure": fig, "axes": ax}
-
-    def close_plot_window(self, n, errors="raise"):
-        """
-        Close a plot window. Applicable only if the class was instantiated with verbose=True.
-
-        Parameters
-        ----------
-        n: int
-            Figure number to close
-        errors: str, optional
-            What to do with errors. It can be either "raise", "log" or "ignore".
-        """
-        if not self.verbose:
-            return
-        if n not in self._plot_windows:
-            msg = "Cannot close plot window number %d: no such window" % n
-            if errors == "raise":
-                raise ValueError(msg)
-            elif errors == "log":
-                self.logger.error(msg)
-        fig_ax = self._plot_windows.pop(n)
-        plt.close(fig_ax["figure"].number)  # would also work with the object itself
-
-    def close_last_plot_windows(self, n=1):
-        """
-        Close the last "n" plot windows.
-        Applicable only if the class was instanciated with verbose=True.
-
-        Parameters
-        -----------
-        n: int, optional
-            Integer indicating how many plot windows should be closed.
-        """
-        figs_nums = sorted(self._plot_windows.keys(), reverse=True)
-        n = min(n, len(figs_nums))
-        for i in range(n):
-            self.close_plot_window(figs_nums[i], errors="ignore")