feat: add warning for extrapolation in morphsqueeze (#250)

* feat: add warning for extrapolation in `morphsqueeze`

* [pre-commit.ci] auto fixes from pre-commit hooks

* chore: add news

* test: add test for warning extrapolation in `morphsqueeze`

* chore: update warning message for extrapolation in `morphsqueeze.py`

* test: add CLI test for extrapolation warning

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: ycexiao

news/extrapolate-warning.rst

@@ -0,0 +1,23 @@
+**Added:**
+
+* No news added: Add warning for extrapolation in morphsqueeze.py.
+
+**Changed:**
+
+* <news item>
+
+**Deprecated:**
+
+* <news item>
+
+**Removed:**
+
+* <news item>
+
+**Fixed:**
+
+* <news item>
+
+**Security:**
+
+* <news item>

src/diffpy/morph/morphs/morphsqueeze.py

@@ -1,13 +1,22 @@
 """Class MorphSqueeze -- Apply a polynomial to squeeze the morph
 function."""
 
+import warnings
+
 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
 
 
+def custom_formatwarning(msg, *args, **kwargs):
+    return f"{msg}\n"
+
+
+warnings.formatwarning = custom_formatwarning
+
+
 class MorphSqueeze(Morph):
     """Squeeze the morph function.
 
@@ -85,4 +94,27 @@ def morph(self, x_morph, y_morph, x_target, y_target):
         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
+        below_extrap = min(x_morph) < min(x_squeezed)
+        above_extrap = max(x_morph) > max(x_squeezed)
+        if below_extrap or above_extrap:
+            if not above_extrap:
+                wmsg = (
+                    "Warning: points with grid value below "
+                    f"{min(x_squeezed)} will be extrapolated."
+                )
+            elif not below_extrap:
+                wmsg = (
+                    "Warning: points with grid value above "
+                    f"{max(x_squeezed)} will be extrapolated."
+                )
+            else:
+                wmsg = (
+                    "Warning: points with grid value below "
+                    f"{min(x_squeezed)} and above {max(x_squeezed)} will be "
+                    "extrapolated."
+                )
+            warnings.warn(
+                wmsg,
+                UserWarning,
+            )
         return self.xyallout

tests/test_morphsqueeze.py

@@ -1,3 +1,5 @@
+import subprocess
+
 import numpy as np
 import pytest
 from numpy.polynomial import Polynomial
@@ -85,3 +87,83 @@ 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)
+
+
+@pytest.mark.parametrize(
+    "squeeze_coeffs, wmsg_gen",
+    [
+        # extrapolate below
+        (
+            {"a0": 0.01},
+            lambda x: (
+                "Warning: points with grid value below "
+                f"{x[0]} will be extrapolated."
+            ),
+        ),
+        # extrapolate above
+        (
+            {"a0": -0.01},
+            lambda x: (
+                "Warning: points with grid value above "
+                f"{x[1]} will be extrapolated."
+            ),
+        ),
+        # extrapolate below and above
+        (
+            {"a0": 0.01, "a1": -0.002},
+            lambda x: (
+                "Warning: points with grid value below "
+                f"{x[0]} and above {x[1]} will be "
+                "extrapolated."
+            ),
+        ),
+    ],
+)
+def test_morphsqueeze_extrapolate(user_filesystem, squeeze_coeffs, wmsg_gen):
+    x_morph = np.linspace(0, 10, 101)
+    y_morph = np.sin(x_morph)
+    x_target = x_morph
+    y_target = y_morph
+    morph = MorphSqueeze()
+    morph.squeeze = squeeze_coeffs
+    coeffs = [squeeze_coeffs[f"a{i}"] for i in range(len(squeeze_coeffs))]
+    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 = wmsg_gen([min(x_squeezed), max(x_squeezed)])
+    assert actual_wmsg == expected_wmsg
+
+    # CLI test
+    morph_file, target_file = create_morph_data_file(
+        user_filesystem / "cwd_dir", x_morph, y_morph, x_target, y_target
+    )
+    run_cmd = ["diffpy.morph"]
+    run_cmd.extend(["--squeeze=" + ",".join(map(str, coeffs))])
+    run_cmd.extend([str(morph_file), str(target_file)])
+    run_cmd.append("-n")
+    result = subprocess.run(run_cmd, capture_output=True, text=True)
+    assert expected_wmsg in result.stderr
+
+
+def create_morph_data_file(
+    data_dir_path, x_morph, y_morph, x_target, y_target
+):
+    morph_file = data_dir_path / "morph_data"
+    morph_data_text = [
+        str(x_morph[i]) + " " + str(y_morph[i]) for i in range(len(x_morph))
+    ]
+    morph_data_text = "\n".join(morph_data_text)
+    morph_file.write_text(morph_data_text)
+    target_file = data_dir_path / "target_data"
+    target_data_text = [
+        str(x_target[i]) + " " + str(y_target[i]) for i in range(len(x_target))
+    ]
+    target_data_text = "\n".join(target_data_text)
+    target_file.write_text(target_data_text)
+    return morph_file, target_file