func: adding squeeze functionallity and extrapolation indices

Author: Luiskitsu

src/diffpy/morph/morphs/morphsqueeze.py

@@ -1,3 +1,7 @@
+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
 
 
@@ -28,10 +32,13 @@ def morph(self, x_morph, y_morph, x_target, y_target):
 
         Returns
         -------
-            A tuple (x_morph_out, y_morph_out, x_target_out, y_target_out)
-            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.
+            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
         -------
@@ -54,4 +61,13 @@ def morph(self, x_morph, y_morph, x_target, y_target):
         """
         Morph.morph(self, x_morph, y_morph, x_target, y_target)
 
-        return self.xyallout
+        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)

tests/test_morphsqueeze.py

@@ -7,19 +7,19 @@
 squeeze_coeffs_list = [
     # The order of coefficients is [a0, a1, a2, ..., an]
     # Negative cubic squeeze coefficients
-    [-0.2, -0.01, -0.001, -0.001],
+    [-0.2, -0.01, -0.001, -0.0001],
     # Positive cubic squeeze coefficients
-    [0.2, 0.01, 0.001, 0.001],
+    [0.2, 0.01, 0.001, 0.0001],
     # Positive and negative cubic squeeze coefficients
-    [0.2, -0.01, 0.002, -0.001],
+    [0.2, -0.01, 0.002, -0.0001],
     # 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],
+    [0, 0.03, 0, -0.0001],
     # Testing zeros, expect no squeezing
     [0, 0, 0, 0, 0, 0],
 ]
@@ -51,10 +51,28 @@ def test_morphsqueeze(x_morph, x_target, squeeze_coeffs):
     y_morph_expected = np.sin(x_morph)
     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)
+    (
+        x_morph_actual,
+        y_morph_actual,
+        x_target_actual,
+        y_target_actual,
+        low_extrap_idx,
+        high_extrap_idx,
+    ) = morph(x_morph, y_morph, x_target, y_target)
+    if low_extrap_idx is None and high_extrap_idx is None:
+        assert np.allclose(y_morph_actual, y_morph_expected, atol=1e-6)
+    else:
+        interp_start = low_extrap_idx + 1 if low_extrap_idx is not None else 0
+        interp_end = (
+            high_extrap_idx
+            if high_extrap_idx is not None
+            else len(y_morph_actual)
+        )
+        assert np.allclose(
+            y_morph_actual[interp_start:interp_end],
+            y_morph_expected[interp_start:interp_end],
+            atol=1e-6,
+        )
     assert np.allclose(x_morph_actual, x_morph_expected)
     assert np.allclose(x_target_actual, x_target)
     assert np.allclose(y_target_actual, y_target)