func: not returning extrapolation indices

Author: Luiskitsu

src/diffpy/morph/morphs/morphsqueeze.py

@@ -16,6 +16,10 @@ class MorphSqueeze(Morph):
     xoutlabel = LABEL_RA
     youtlabel = LABEL_GR
     parnames = ["squeeze"]
+    # extrap_index_low: last index before interpolation region
+    # extrap_index_high: first index after interpolation region
+    extrap_index_low = None
+    extrap_index_high = None
 
     def morph(self, x_morph, y_morph, x_target, y_target):
         """Squeeze the morph function.
@@ -32,13 +36,10 @@ 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,
-            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.
+            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.
 
         Example
         -------
@@ -66,8 +67,8 @@ def morph(self, x_morph, y_morph, x_target, y_target):
         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)
+        low_extrap = np.where(self.x_morph_in < x_squeezed[0])[0]
+        high_extrap = np.where(self.x_morph_in > x_squeezed[-1])[0]
+        self.extrap_index_low = low_extrap[-1] if low_extrap.size else None
+        self.extrap_index_high = high_extrap[0] if high_extrap.size else None
+        return self.xyallout

tests/test_morphsqueeze.py

@@ -47,32 +47,40 @@ def test_morphsqueeze(x_morph, x_target, squeeze_coeffs):
     squeeze_polynomial = Polynomial(squeeze_coeffs)
     x_squeezed = x_morph + squeeze_polynomial(x_morph)
     y_morph = np.sin(x_squeezed)
+    low_extrap = np.where(x_morph < x_squeezed[0])[0]
+    high_extrap = np.where(x_morph > x_squeezed[-1])[0]
+    extrap_index_low_expected = low_extrap[-1] if low_extrap.size else None
+    extrap_index_high_expected = high_extrap[0] if high_extrap.size else None
     x_morph_expected = x_morph
     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,
-        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,
-        )
+    x_morph_actual, y_morph_actual, x_target_actual, y_target_actual = morph(
+        x_morph, y_morph, x_target, y_target
+    )
+    extrap_index_low = morph.extrap_index_low
+    extrap_index_high = morph.extrap_index_high
+    if extrap_index_low is None:
+        extrap_index_low = 0
+    elif extrap_index_high is None:
+        extrap_index_high = -1
+    assert np.allclose(
+        y_morph_actual[extrap_index_low + 1 : extrap_index_high],
+        y_morph_expected[extrap_index_low + 1 : extrap_index_high],
+        atol=1e-6,
+    )
+    assert np.allclose(
+        y_morph_actual[:extrap_index_low],
+        y_morph_expected[:extrap_index_low],
+        atol=1e-3,
+    )
+    assert np.allclose(
+        y_morph_actual[extrap_index_high:],
+        y_morph_expected[extrap_index_high:],
+        atol=1e-3,
+    )
+    assert morph.extrap_index_low == extrap_index_low_expected
+    assert morph.extrap_index_high == extrap_index_high_expected
     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)