test: add test for warning extrapolation in morphsqueeze

Author: ycexiao

src/diffpy/morph/morphs/morphsqueeze.py

@@ -10,6 +10,13 @@
 from diffpy.morph.morphs.morph import LABEL_GR, LABEL_RA, Morph
 
 
+def custom_formatwarning(msg, *args, **kwargs):
+    return f"{msg}\n"
+
+
+warnings.formatwarning = custom_formatwarning
+
+
 class MorphSqueeze(Morph):
     """Squeeze the morph function.
 
@@ -88,14 +95,19 @@ def morph(self, x_morph, y_morph, x_target, y_target):
         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
 
-        low_extrap_x = x_squeezed[0] - self.x_morph_in[0]
-        high_extrap_x = self.x_morph_in[-1] - x_squeezed[-1]
-        if low_extrap_x > 0 or high_extrap_x > 0:
-            wmsg = "Extrapolating the morphed function: "
-            if low_extrap_x > 0:
-                wmsg += f"extrapolating length in the lowe r {low_extrap_x} "
-            if high_extrap_x > 0:
-                wmsg += f"extrapolating length in the high r {high_extrap_x} "
+        begin_end_sqeeze = min(x_squeezed), max(x_squeezed)
+        begin_end_in = min(self.x_morph_in), max(self.x_morph_in)
+        if not (
+            begin_end_sqeeze[0] <= begin_end_in[0]
+            and begin_end_in[-1] >= begin_end_in[-1]
+        ):
+            wmsg = (
+                "\nExtrapolating the morphed function via CubicSpline:\n"
+                f"Obtaining grid points between {begin_end_in[0]} and "
+                f"{begin_end_in[1]}.\n"
+                f"Points below {begin_end_sqeeze[0]} and "
+                f"above {begin_end_sqeeze[1]} will be extrapolated."
+            )
             warnings.warn(
                 wmsg,
                 UserWarning,

tests/test_morphsqueeze.py

@@ -85,3 +85,30 @@ def test_morphsqueeze(x_morph, x_target, squeeze_coeffs):
     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)
+
+
+def test_morphsqueeze_extrapolate():
+    x_morph = np.linspace(0, 10, 101)
+    y_morph = np.sin(x_morph)
+    x_target = x_morph
+    y_target = y_morph
+    squeeze_coeff = {"a0": 0.01, "a1": -0.0005, "a2": -0.0005, "a3": -1e-6}
+    morph = MorphSqueeze()
+    morph.squeeze = squeeze_coeff
+    coeffs = [squeeze_coeff[f"a{i}"] for i in range(len(squeeze_coeff))]
+    squeeze_polynomial = Polynomial(coeffs)
+    x_squeezed = x_morph + squeeze_polynomial(x_morph)
+    with pytest.warns() as w:
+        x_morph_actual, y_morph_actual, x_target_actual, y_target_actual = (
+            morph(x_morph, y_morph, x_target, y_target)
+        )
+        assert len(w) == 1
+        assert w[0].category is UserWarning
+        actual_wmsg = str(w[0].message)
+    expected_wmsg = (
+        "\nExtrapolating the morphed function via CubicSpline:\n"
+        f"Obtaining grid points between {x_morph[0]} and {x_morph[-1]}.\n"
+        f"Points below {x_squeezed[0]} and "
+        f"above {x_squeezed[-1]} will be extrapolated."
+    )
+    assert actual_wmsg == expected_wmsg