Test/funct: trim target data to overlapping region and added more test

Author: Luiskitsu

src/diffpy/morph/morphs/morphsqueeze.py

@@ -1,3 +1,4 @@
+import numpy as np
 from numpy.polynomial import Polynomial
 from scipy.interpolate import CubicSpline
 
@@ -7,7 +8,12 @@
 class MorphSqueeze(Morph):
     """Squeeze the morph function.
 
-    This applies a polynomial to squeeze the morph non-linearly.
+    This applies a polynomial to squeeze the morph non-linearly. The resulting
+    squeezed morph is interpolated to the (trimmed) target grid.
+    Only the overlapping region between the squeezed morph and the target
+    grid is used. The target is trimmed (or not) accordingly, and the final
+    outputs (morph and target) are returned on the same grid, defined by this
+    trimmed target range.
 
     Configuration Variables
     -----------------------
@@ -33,11 +39,14 @@ def 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)
+        x_min = max(float(self.x_target_in[0]), float(x_squeezed[0]))
+        x_max = min(float(self.x_target_in[-1]), float(x_squeezed[-1]))
+        min_index = np.where(self.x_target_in >= x_min)[0][0]
+        max_index = np.where(self.x_target_in <= x_max)[0][-1]
+        self.x_target_out = self.x_target_in[min_index : max_index + 1]
+        self.y_target_out = self.y_target_in[min_index : max_index + 1]
         self.y_morph_out = CubicSpline(x_squeezed, self.y_morph_in)(
-            self.x_morph_in
+            self.x_target_out
         )
-        self.y_morph_out = CubicSpline(self.x_morph_in, self.y_morph_out)(
-            self.x_target_in
-        )
-        self.x_morph_out = self.x_target_in
+        self.x_morph_out = self.x_target_out
         return self.xyallout

tests/test_morphsqueeze.py

@@ -27,21 +27,87 @@
         [0, 0, 0, 0, 0, 0],
     ],
 )
-def test_morphsqueeze(squeeze_coeffs):
-    x_target_expected = np.linspace(0, 10, 101)
-    y_target_expected = np.sin(x_target_expected)
-    # Different grid for morph data to test inputs with different grids
-    # Morph grid must be finer than the target to avoid interpolation issues
-    x_morph = np.linspace(-3, 13, 301)
+def test_morphsqueeze_target_extends_beyond_morph(squeeze_coeffs):
+    # Target data extends beyond morph and different grids
+    x_target = np.linspace(-3, 25, 401)
+    y_target = np.sin(x_target)
+    x_morph = np.linspace(0, 10, 301)
     squeeze_polynomial = Polynomial(squeeze_coeffs)
     x_squeezed = x_morph + squeeze_polynomial(x_morph)
     y_morph = np.sin(x_squeezed)
+    # Trim target data to the region overlapping with the squeezed morph
+    x_min = max(float(x_target[0]), float(x_squeezed[0]))
+    x_max = min(float(x_target[-1]), float(x_squeezed[-1]))
+    min_index = np.where(x_target >= x_min)[0][0]
+    max_index = np.where(x_target <= x_max)[0][-1]
+    x_morph_expected = x_target[min_index : max_index + 1]
+    y_morph_expected = y_target[min_index : max_index + 1]
     morph = MorphSqueeze()
     morph.squeeze = squeeze_coeffs
     x_morph_actual, y_morph_actual, x_target_actual, y_target_actual = morph(
-        x_morph, y_morph, x_target_expected, y_target_expected
+        x_morph, y_morph, x_target, y_target
     )
-    assert np.allclose(y_morph_actual, y_target_expected)
-    assert np.allclose(x_morph_actual, x_target_expected)
-    assert np.allclose(x_target_actual, x_target_expected)
-    assert np.allclose(y_target_actual, y_target_expected)
+    assert np.allclose(y_morph_actual, y_morph_expected)
+    assert np.allclose(x_morph_actual, x_morph_expected)
+    assert np.allclose(x_target_actual, x_morph_expected)
+    assert np.allclose(y_target_actual, y_morph_expected)
+
+    import matplotlib.pyplot as plt
+
+    plt.figure()
+    plt.plot(x_target, y_target, color="gray", label="target")
+    plt.plot(x_morph, y_morph, color="black", label="morph")
+    plt.scatter(x_morph_actual, y_morph_actual, color="gold", label="actual")
+    plt.plot(
+        x_morph_expected, y_morph_expected, color="purple", label="expected"
+    )
+    plt.legend()
+    plt.show()
+
+
+@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, the full polynomial is applied
+        [0, 0.03, 0, -0.001],
+        # Testing zeros, expect no squeezing
+        [0, 0, 0, 0, 0, 0],
+    ],
+)
+def test_morphsqueeze_morph_extends_beyond_target(squeeze_coeffs):
+    # Different grid for morph and target data to test different grids
+    x_target = np.linspace(0, 10, 101)
+    y_target = np.sin(x_target)
+    x_morph = np.linspace(-3, 15, 301)
+    squeeze_polynomial = Polynomial(squeeze_coeffs)
+    x_squeezed = x_morph + squeeze_polynomial(x_morph)
+    y_morph = np.sin(x_squeezed)
+    # Trim target data to the region overlapping with the squeezed morph
+    x_min = max(float(x_target[0]), float(x_squeezed[0]))
+    x_max = min(float(x_target[-1]), float(x_squeezed[-1]))
+    min_index = np.where(x_target >= x_min)[0][0]
+    max_index = np.where(x_target <= x_max)[0][-1]
+    x_morph_expected = x_target[min_index : max_index + 1]
+    y_morph_expected = y_target[min_index : max_index + 1]
+    morph = MorphSqueeze()
+    morph.squeeze = squeeze_coeffs
+    x_morph_actual, y_morph_actual, x_target_actual, y_target_actual = morph(
+        x_morph, y_morph, x_target, y_target
+    )
+    assert np.allclose(y_morph_actual, y_morph_expected)
+    assert np.allclose(x_morph_actual, x_morph_expected)
+    assert np.allclose(x_target_actual, x_morph_expected)
+    assert np.allclose(y_target_actual, y_morph_expected)