alignment.py

import numpy as np
import math
from numpy.polynomial.polynomial import polyroots

from nabu.utils import previouspow2

try:
    from scipy.ndimage.filters import median_filter
    import scipy.special as spspe

    __have_scipy__ = True
except ImportError:
    from silx.math.medianfilter import medfilt2d as median_filter

    __have_scipy__ = False


def _get_lowpass_filter(img_shape, cutoff_pix, cutoff_trans_pix):
    """Computes a low pass filter with the erfc.
    :param img_shape: Shape of the image
    :type img_shape: tuple
    :param cutoff_pix: Position of the cut off in pixels
    :type cutoff_pix: float
    :param cutoff_trans_pix: Size of the cut off transition in pixels
    :type cutoff_trans_pix: float

    :return: The computes filter
    :rtype: `numpy.array_like`
    """
    coords = [np.fft.fftfreq(s, 1) for s in img_shape]
    coords = np.meshgrid(*coords, indexing="ij")

    Kcut = 0.5 / cutoff_pix
    Kcut_width = 0.5 / cutoff_trans_pix

    r = np.sqrt(np.sum(np.array(coords) ** 2, axis=0))
    myargs = (r - Kcut) / Kcut_width
    if __have_scipy__:
        res = spspe.erfc(myargs) / 2
    else:
        res = np.array(list(map(math.erfc, myargs))) / 2

    return res


class AlignmentBase(object):
    @staticmethod
    def _biquadratic_refinement(cc_vals):

        X, Y = np.meshgrid([-1, 0, 1], [-1, 0, 1])
        X = X.flatten()
        Y = Y.flatten()
        coords = np.array([np.ones(X.size), X, Y, X * Y, X * X, Y * Y])

        a = np.linalg.lstsq(coords.T, cc_vals.flatten(), rcond=None)[0]
        zero_derivative_coords = np.linalg.lstsq(
            [[2 * a[4], a[3]], [a[3], 2 * a[5]]], [-a[1], -a[2]], rcond=None
        )[0]
        zero_derivative_coords = np.clip(zero_derivative_coords, -1, 1)

        return zero_derivative_coords

    @staticmethod
    def _determine_roi(img_shape, roi_yxhw, do_truncate_horz_pow2):
        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:
                roi_yxhw[1] = img_shape[1]

        if len(roi_yxhw) == 2:  # Convert centered 2-element roi into 4-element
            roi_yxhw = np.array(roi_yxhw, dtype=np.int)
            roi_yxhw = np.concatenate(((img_shape - roi_yxhw) // 2, roi_yxhw))
        return roi_yxhw

    @staticmethod
    def _prepare_image(
        img,
        invalid_val=1e-5,
        roi_yxhw=None,
        median_filt_shape=None,
        cutoff_pix=None,
        cutoff_trans_pix=None,
    ):
        img = np.squeeze(img)  # Removes singleton dimensions, but does a shallow copy
        img = np.ascontiguousarray(img)

        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 cutoff_pix is not None and cutoff_trans_pix is not None:
            myfilter = _get_lowpass_filter(img.shape[-2:], cutoff_pix, cutoff_trans_pix)
            img_shape = img.shape
            if len(img_shape) == 2:
                img = np.fft.ifft2(np.fft.fft2(img) * myfilter).real
            elif len(img_shape) > 2:
                # if dealing with a stack of images, we have to do them one by one
                img = np.reshape(img, (-1,) + tuple(img_shape[-2:]))
                for ii in range(img.shape[0]):
                    img[ii, ...] = np.fft.ifft2(
                        np.fft.fft2(img[ii, ...]) * myfilter
                    ).real
                img = np.reshape(img, img_shape)

        if median_filt_shape is not None:
            img_shape = img.shape
            if __have_scipy__:
                # 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.int
                        ),
                        median_filt_shape,
                    )
                )
                img = median_filter(img, size=median_filt_shape)
            else:
                if len(img_shape) == 2:
                    img = median_filter(img, kernel_size=median_filt_shape)
                elif len(img_shape) > 2:
                    # if dealing with a stack of images, we have to do them one by one
                    img = np.reshape(img, (-1,) + tuple(img_shape[-2:]))
                    for ii in range(img.shape[0]):
                        img[ii, ...] = median_filter(
                            img[ii, ...], kernel_size=median_filt_shape
                        )
                    img = np.reshape(img, img_shape)

        return img

    @staticmethod
    def _compute_correlation_fft(img_1, img_2, padding_mode):
        do_circular_conv = padding_mode is None or padding_mode == "wrap"
        if not do_circular_conv:
            padding = np.ceil(np.array(img_2.img_shape) / 2).astype(np.int)
            img_1 = np.pad(img_1, ((padding[0],), (padding[1],)), mode=padding_mode)
            img_2 = np.pad(img_2, ((padding[0],), (padding[1],)), mode=padding_mode)

        # compute fft's of the 2 images
        img_fft_1 = np.fft.fft2(img_1)
        img_fft_2 = np.conjugate(np.fft.fft2(img_2))
        # inverse fft of the product to get cross_correlation of the 2 images
        cc = np.real(np.fft.ifft2(img_fft_1 * img_fft_2))

        if not do_circular_conv:
            cc = cc[padding[0] : -padding[0], padding[1] : -padding[1]]

        return cc

    def refine_max_position(self, xf, yf):
        """Computes the sub-pixel max position of the given function sampling.

        Parameters
        ----------
        xf: numpy.ndarray
            Coordinates of the sampled points
        yf: numpy.ndarray
            Function values of the sampled points

        Raises
        ------
        ValueError
            In case position and values do not have the same size.

        Returns
        -------
        float
            Estimated function max, according to the coordinates in xf.
        """
        npix = xf.size
        if not npix == yf.size:
            raise ValueError(
                "Coordinates and values should have the same length. Sizes of xf: %d, yf: %d"
                % (xf.size, yf.size)
            )

        matr = xf[:, None] ** (npix - np.arange(1, npix + 1))

        # kf is the coeffs of the polynom passing trough the npix pixels
        kf = np.linalg.inv(matr).dot(yf[:, None])

        # compute coefficients of the 1st order derivative of the polynom
        cf = ((npix - 1) - np.arange(npix - 1)) * kf[:-1, 0]
        # revert for polyroots
        cf = np.flip(cf)

        # maximum is the 0 of the 1st derivative of the polynom
        zeros_pol = polyroots(cf)

        # The pixel within the roots of the derivative of the polynom is
        # the one near the approximate center of rotation (CoR) px
        approx_max = xf[np.argmax(yf)]
        closest_pix_ind = np.argmin((zeros_pol - approx_max) ** 2)
        return np.real(zeros_pol[closest_pix_ind])


class CenterOfRotation(AlignmentBase):
    def __init__(self, poly_deg=2, horz_fft_width=False):
        """
        Center of Rotation (CoR) computation object.
        This class is used on radios.

        Parameters
        ----------
        poly_deg: int, optional
            Degree of polynom to interpolate the global value to a
            sub-pixel precision. The number of pixels involved will be

            >>> npix = poly_deg + 1

        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(poly_deg, horz_fft_width)

    def _init_parameters(self, poly_deg, horz_fft_width):
        self.truncate_horz_pow2 = horz_fft_width
        self.npix = poly_deg + 1

    @staticmethod
    def _check_img_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)),
                )
            )

    def find_shift(
        self,
        img_1: np.ndarray,
        img_2: np.ndarray,
        roi_yxhw=None,
        median_filt_shape=None,
        padding_mode=None,
        cutoff_pix=None,
        cutoff_trans_pix=None,
    ):
        """Find the Center of Rotation (CoR), given to 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 * s) * 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'

        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(poly_deg=2)
        ... 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_sizes(img_1, img_2)

        img_shape = img_2.shape
        roi_yxhw = self._determine_roi(img_shape, roi_yxhw, self.truncate_horz_pow2)

        img_1 = self._prepare_image(
            img_1,
            roi_yxhw=roi_yxhw,
            median_filt_shape=median_filt_shape,
            cutoff_pix=cutoff_pix,
            cutoff_trans_pix=cutoff_trans_pix,
        )
        img_2 = self._prepare_image(
            img_2,
            roi_yxhw=roi_yxhw,
            median_filt_shape=median_filt_shape,
            cutoff_pix=cutoff_pix,
            cutoff_trans_pix=cutoff_trans_pix,
        )

        cc = self._compute_correlation_fft(img_1, img_2, padding_mode)

        img_shape = cc.shape  # Re-define image shape after RoI

        # get pixel having the maximum value of the correlation array
        pix_max_corr = np.argmax(cc)
        NY, NX = cc.shape
        py, px = np.unravel_index(pix_max_corr, (NY, NX))

        payload_vals = np.zeros([3, 3], "f")

        ly = max(0, py - 1)
        lx = max(0, px - 1)
        hy = min(NY, py + 1)
        hx = min(NX, px + 1)
        payload_vals[:] = (cc[ly : hy + 1, lx : hx + 1]).min()

        payload_vals[
            (ly - (py - 1)) : (hy + 1 - (py - 1)), (lx - (px - 1)) : (hx + 1 - (px - 1))
        ] = cc[ly : hy + 1, lx : hx + 1]

        Dpy, Dpx = self._biquadratic_refinement(payload_vals)

        ppy = py + Dpy
        ppx = px + Dpx

        if ppy > (cc.shape[0] - 1) / 2:
            ppy = ppy - cc.shape[0]
        if ppx > (cc.shape[1] - 1) / 2:
            ppx = ppx - cc.shape[1]

        fitted_shift = ppx

        return fitted_shift / 2.0

    __call__ = find_shift


class DetectorTranslationAlongBeam(AlignmentBase):
    @staticmethod
    def _check_img_sizes(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)),
                )
            )

    def find_shift(
        self,
        img_stack: np.ndarray,
        img_pos: np.array,
        roi_yxhw=None,
        median_filt_shape=None,
        padding_mode=None,
        poly_deg=2,
    ):
        """Find the deviation of the translation axis of the area detector
        along the beam propagation direction.

        TODO: Add more information here! Including interpretation of the result

        Parameters
        ----------
        img_stack: numpy.ndarray
            A stack of images (usually 4) at different distances
        img_pos: numpy.ndarray
            Position of the images along the translation axis
        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'

        Returns
        -------
        tuple(float, float)
            Estimated (vertical, horizontal) increment per unit-distance of the shift.

        Examples
        --------
        TODO: Add examples here!
        """
        self._check_img_sizes(img_stack, img_pos)

        num_imgs = img_stack.shape[0]
        img_shape = img_stack.shape[-2:]
        roi_yxhw = self._determine_roi(img_shape, roi_yxhw, False)

        img_stack = self._prepare_image(
            img_stack, roi_yxhw=roi_yxhw, median_filt_shape=median_filt_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])
        # do correlations
        ccs = [
            self._compute_correlation_fft(
                img_stack[0, ...], img_stack[ii, ...], padding_mode
            )
            for ii in range(1, num_imgs)
        ]

        npix = poly_deg + 1
        poly_center = npix // 2

        shifts_v, shifts_h = np.empty((num_imgs - 1,)), np.empty((num_imgs - 1,))
        for ii, cc in enumerate(ccs):
            # 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)

            xh_pos = ((ph - poly_center + np.arange(npix)) % img_shape[1]).astype(
                dtype=np.int
            )
            shifts_h[ii] = self.refine_max_position(cc_hs[xh_pos], cc[pv, xh_pos])

            xv_pos = ((pv - poly_center + np.arange(npix)) % img_shape[0]).astype(
                dtype=np.int
            )
            shifts_v[ii] = self.refine_max_position(cc_vs[xv_pos], cc[xv_pos, ph])

        """
        TODO: finish implementing the following logic!

        pos -= repmat(pos(:,1),1,ni);

        ph = polyfit(motpos,pos(1,:),1);
        thz = ph(1)*pixsize/1000*180/pi;
        printf('Correction in horizontal angle: %5.2f urad -> mvr thz %5.4f\n',thz*pi/180*1e6,thz)
        figure(1);
        plot(motpos,pos(1,:))
        title('Horizontal displacement')

        pv = polyfit(motpos,pos(2,:),1);
        thy = pv(1)*pixsize/1000*180/pi;
        printf('Correction in vertical angle: %5.2f urad -> mvr thy %5.4f\n',thy*pi/180*1e6,thy)
        figure(2);
        plot(motpos,pos(2,:))
        title('Vertical displacement')
        """