morphsqueeze.py

import numpy as np
from numpy.polynomial import Polynomial
from scipy.interpolate import CubicSpline

from diffpy.morph.morphs.morph import LABEL_GR, LABEL_RA, Morph


class MorphSqueeze(Morph):
    """Apply a polynomial to squeeze the morph function. The morphed
    data is returned on the same grid as the unmorphed data."""

    # Define input output types
    summary = "Squeeze morph by polynomial shift"
    xinlabel = LABEL_RA
    yinlabel = LABEL_GR
    xoutlabel = LABEL_RA
    youtlabel = LABEL_GR
    parnames = ["squeeze"]

    def morph(self, x_morph, y_morph, x_target, y_target):
        """Squeeze the morph function.

        This applies a polynomial to squeeze the morph non-linearly.

        Configuration Variables
        -----------------------
        squeeze : list
            The polynomial coefficients [a0, a1, ..., an] for the squeeze
            function where the polynomial would be of the form
            a0 + a1*x + a2*x^2 and so on.  The order of the polynomial is
            determined by the length of the list.

        Returns
        -------
            A tuple (x_morph_out, y_morph_out, x_target_out, y_target_out,
            min_index, max_index) where the target values remain the same and
            the morph data is shifted according to the squeeze. The morphed
            data is returned on the same grid as the unmorphed data.
            The min_index and max_index are the last index before the
            interpolated region and the first index after the interpolated
            region, respectively. If there is no extrapolation it returns None.

        Example
        -------
        Import the squeeze morph function:
        >>> from diffpy.morph.morphs.morphsqueeze import MorphSqueeze
        Provide initial guess for squeezing coefficients:
        >>> squeeze_coeff = [0.1, -0.01, 0.005]
        Run the squeeze morph given input morph array (x_morph, y_morph)
        and target array (x_target, y_target):
        >>> morph = MorphSqueeze()
        >>> morph.squeeze = squeeze_coeff
        >>> x_morph_out, y_morph_out, x_target_out, y_target_out = morph(
        ...     x_morph, y_morph, x_target, y_target)
        To access parameters from the morph instance:
        >>> x_morph_in = morph.x_morph_in
        >>> y_morph_in = morph.y_morph_in
        >>> x_target_in = morph.x_target_in
        >>> y_target_in = morph.y_target_in
        >>> squeeze_coeff_out = morph.squeeze
        """
        Morph.morph(self, x_morph, y_morph, x_target, y_target)

        squeeze_polynomial = Polynomial(self.squeeze)
        x_squeezed = self.x_morph_in + squeeze_polynomial(self.x_morph_in)
        self.y_morph_out = CubicSpline(x_squeezed, self.y_morph_in)(
            self.x_morph_in
        )
        left_extrap = np.where(self.x_morph_in < x_squeezed[0])[0]
        right_extrap = np.where(self.x_morph_in > x_squeezed[-1])[0]
        min_index = left_extrap[-1] if left_extrap.size else None
        max_index = right_extrap[0] if right_extrap.size else None
        return self.xyallout + (min_index, max_index)