test: adding a test to unsqueeze squeezed data

news/test_squeeze.rst

@@ -0,0 +1,23 @@
+**Added:**
+
+* Polynomial squeeze of x-axis of morphed data
+
+**Changed:**
+
+* <news item>
+
+**Deprecated:**
+
+* <news item>
+
+**Removed:**
+
+* <news item>
+
+**Fixed:**
+
+* <news item>
+
+**Security:**
+
+* <news item>

src/diffpy/morph/morphs/morphsqueeze.py

@@ -0,0 +1,43 @@
+import numpy as np
+from numpy.polynomial import Polynomial
+from scipy.interpolate import interp1d
+
+from diffpy.morph.morphs.morph import LABEL_GR, LABEL_RA, Morph
+
+
+class MorphSqueeze(Morph):
+    """Squeeze the morph function.
+
+    This applies a polynomial to squeeze the morph non-linearly.
+
+    Configuration Variables
+    -----------------------
+    squeeze
+        list or array-like
+        Polynomial coefficients [a0, a1, ..., an] for the squeeze function.
+    """
+
+    # 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):
+        Morph.morph(self, x_morph, y_morph, x_target, y_target)
+        if self.squeeze is None or np.allclose(self.squeeze, 0):
+            self.x_morph_out = self.x_morph_in
+            self.y_morph_out = self.y_morph_in
+            return self.xyallout
+
+        squeeze_polynomial = Polynomial(self.squeeze)
+        x_squeezed = self.x_morph_in + squeeze_polynomial(self.x_morph_in)
+
+        self.y_morph_out = interp1d(
+            x_squeezed, self.y_morph_in, kind="cubic", bounds_error=False
+        )(self.x_morph_in)
+
+        self.x_morph_out = self.x_morph_in
+        return self.xyallout

tests/test_morphsqueeze.py

@@ -0,0 +1,60 @@
+import numpy as np
+import pytest
+from numpy.polynomial import Polynomial
+
+from diffpy.morph.morphs.morphsqueeze import MorphSqueeze
+
+
+@pytest.mark.parametrize(
+    "squeeze_coeffs",
+    [
+        # The order of coefficients is [a0, a1, a2, ..., an]
+        # Negative cubic squeeze coefficients
+        [-0.2, -0.01, -0.001, -0.001],
+        # Positive cubic squeeze coefficients
+        [0.2, 0.01, 0.001, 0.001],
+        # Positive and negative cubic squeeze coefficients
+        [0.2, -0.01, 0.002, -0.001],
+        # Quadratic squeeze coefficients
+        [-0.2, 0.005, -0.007],
+        # Linear squeeze coefficients
+        [0.1, 0.3],
+        # 4th order squeeze coefficients
+        [0.2, -0.01, 0.001, -0.001, 0.0004],
+        # Zeros and non-zeros, expect 0 + a1x + 0 + a3x**3
+        [0, 0.03, 0, -0.001],
+        # Testing zeros, expect no squeezing
+        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+    ],
+)
+def test_morphsqueeze(squeeze_coeffs):
+    x_target = np.linspace(0, 10, 1001)
+    y_target = np.sin(x_target)
+
+    x_make = np.linspace(-3, 13, 1601)
+    lower_idx = np.where(x_make == 0.0)[0][0]
+    upper_idx = np.where(x_make == 10.0)[0][0]
+
+    squeeze_polynomial = Polynomial(squeeze_coeffs)
+    x_squeezed = x_make + squeeze_polynomial(x_make)
+
+    x_morph = x_make.copy()
+    y_morph = np.sin(x_squeezed)
+
+    morph = MorphSqueeze()
+    morph.squeeze = squeeze_coeffs
+
+    x_actual, y_actual, x_expected, y_expected = morph(
+        x_morph, y_morph, x_target, y_target
+    )
+    y_actual = y_actual[lower_idx : upper_idx + 1]
+    assert np.allclose(y_actual, y_expected)
+
+    # Plotting code used for figures in PR comments
+    # https://github.com/diffpy/diffpy.morph/pull/180
+    # plt.figure()
+    # plt.scatter(x_expected, y_expected, color='black', label='Expected')
+    # plt.plot(x_morph, y_morph, color='purple', label='morph')
+    # plt.plot(x_actual, y_actual, '--', color='gold', label='Actual')
+    # plt.legend()
+    # plt.show()