feat: add set_extrapolation_info function in morph.Morph (#255)

* feat: add `checkExtrapolation` function in `morph.Morph`

* refactor: rename the function to `set_extrapolation_info`

Author: ycexiao

news/extrap-warnings.rst

@@ -0,0 +1,23 @@
+**Added:**
+
+* Enable ``diffpy.morph`` to detect extrapolation.
+
+**Changed:**
+
+* <news item>
+
+**Deprecated:**
+
+* <news item>
+
+**Removed:**
+
+* <news item>
+
+**Fixed:**
+
+* <news item>
+
+**Security:**
+
+* <news item>

src/diffpy/morph/morph_io.py

@@ -410,29 +410,35 @@ def tabulate_results(multiple_morph_results):
 
 def handle_warnings(squeeze_morph):
     if squeeze_morph is not None:
-        eil = squeeze_morph.extrap_index_low
-        eih = squeeze_morph.extrap_index_high
-
-        if eil is not None or eih is not None:
-            if eih is None:
-                wmsg = (
-                    "Warning: points with grid value below "
-                    f"{squeeze_morph.squeeze_cutoff_low} "
-                    f"will be extrapolated."
-                )
-            elif eil is None:
-                wmsg = (
-                    "Warning: points with grid value above "
-                    f"{squeeze_morph.squeeze_cutoff_high} "
-                    f"will be extrapolated."
-                )
-            else:
-                wmsg = (
-                    "Warning: points with grid value below "
-                    f"{squeeze_morph.squeeze_cutoff_low} and above "
-                    f"{squeeze_morph.squeeze_cutoff_high} "
-                    f"will be extrapolated."
-                )
+        extrapolation_info = squeeze_morph.extrapolation_info
+        is_extrap_low = extrapolation_info["is_extrap_low"]
+        is_extrap_high = extrapolation_info["is_extrap_high"]
+        cutoff_low = extrapolation_info["cutoff_low"]
+        cutoff_high = extrapolation_info["cutoff_high"]
+
+        if is_extrap_low and is_extrap_high:
+            wmsg = (
+                "Warning: points with grid value below "
+                f"{cutoff_low} and above "
+                f"{cutoff_high} "
+                f"are extrapolated."
+            )
+        elif is_extrap_low:
+            wmsg = (
+                "Warning: points with grid value below "
+                f"{cutoff_low} "
+                f"are extrapolated."
+            )
+        elif is_extrap_high:
+            wmsg = (
+                "Warning: points with grid value above "
+                f"{cutoff_high} "
+                f"are extrapolated."
+            )
+        else:
+            wmsg = None
+
+        if wmsg:
             warnings.warn(
                 wmsg,
                 UserWarning,

src/diffpy/morph/morphapp.py

@@ -610,10 +610,12 @@ def single_morph(
         config["smear"] = smear_in
     # Shift
     # Only enable hshift is squeeze is not enabled
+    shift_morph = None
     if (
         opts.hshift is not None and squeeze_poly_deg < 0
     ) or opts.vshift is not None:
-        chain.append(morphs.MorphShift())
+        shift_morph = morphs.MorphShift()
+        chain.append(shift_morph)
     if opts.hshift is not None and squeeze_poly_deg < 0:
         hshift_in = opts.hshift
         config["hshift"] = hshift_in
@@ -700,6 +702,7 @@ def single_morph(
 
     # THROW ANY WARNINGS HERE
     io.handle_warnings(squeeze_morph)
+    io.handle_warnings(shift_morph)
 
     # Get Rw for the morph range
     rw = tools.getRw(chain)

src/diffpy/morph/morphs/morph.py

@@ -12,9 +12,8 @@
 # See LICENSE.txt for license information.
 #
 ##############################################################################
-"""Morph -- base class for defining a morph.
-"""
-
+"""Morph -- base class for defining a morph."""
+import numpy
 
 LABEL_RA = "r (A)"  # r-grid
 LABEL_GR = "G (1/A^2)"  # PDF G(r)
@@ -246,6 +245,36 @@ def plotOutputs(self, xylabels=True, **plotargs):
             ylabel(self.youtlabel)
         return rv
 
+    def set_extrapolation_info(self, x_true, x_extrapolate):
+        """Set extrapolation information of the concerned morphing
+        process.
+
+        Parameters
+        ----------
+        x_true : array
+            original x values
+        x_extrapolate : array
+            x values after a morphing process
+        """
+
+        cutoff_low = min(x_true)
+        extrap_low_x = numpy.where(x_extrapolate < cutoff_low)[0]
+        is_extrap_low = False if len(extrap_low_x) == 0 else True
+        cutoff_high = max(x_true)
+        extrap_high_x = numpy.where(x_extrapolate > cutoff_high)[0]
+        is_extrap_high = False if len(extrap_high_x) == 0 else True
+        extrap_index_low = extrap_low_x[-1] if is_extrap_low else 0
+        extrap_index_high = extrap_high_x[0] if is_extrap_high else -1
+        extrapolation_info = {
+            "is_extrap_low": is_extrap_low,
+            "cutoff_low": cutoff_low,
+            "extrap_index_low": extrap_index_low,
+            "is_extrap_high": is_extrap_high,
+            "cutoff_high": cutoff_high,
+            "extrap_index_high": extrap_index_high,
+        }
+        self.extrapolation_info = extrapolation_info
+
     def __getattr__(self, name):
         """Obtain the value from self.config, when normal lookup fails.
 

src/diffpy/morph/morphs/morphshift.py

@@ -57,6 +57,7 @@ def morph(self, x_morph, y_morph, x_target, y_target):
         r = self.x_morph_in - hshift
         self.y_morph_out = numpy.interp(r, self.x_morph_in, self.y_morph_in)
         self.y_morph_out += vshift
+        self.set_extrapolation_info(self.x_morph_in, r)
         return self.xyallout
 
 

src/diffpy/morph/morphs/morphsqueeze.py

@@ -1,7 +1,6 @@
 """Class MorphSqueeze -- Apply a polynomial to squeeze the morph
 function."""
 
-import numpy as np
 from numpy.polynomial import Polynomial
 from scipy.interpolate import CubicSpline
 
@@ -83,14 +82,9 @@ def morph(self, x_morph, y_morph, x_target, y_target):
         coeffs = [self.squeeze[f"a{i}"] for i in range(len(self.squeeze))]
         squeeze_polynomial = Polynomial(coeffs)
         x_squeezed = self.x_morph_in + squeeze_polynomial(self.x_morph_in)
-        self.squeeze_cutoff_low = min(x_squeezed)
-        self.squeeze_cutoff_high = max(x_squeezed)
         self.y_morph_out = CubicSpline(x_squeezed, self.y_morph_in)(
             self.x_morph_in
         )
-        low_extrap = np.where(self.x_morph_in < self.squeeze_cutoff_low)[0]
-        high_extrap = np.where(self.x_morph_in > self.squeeze_cutoff_high)[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
+        self.set_extrapolation_info(x_squeezed, self.x_morph_in)
 
         return self.xyallout

tests/test_morphsqueeze.py

@@ -46,47 +46,56 @@
 @pytest.mark.parametrize("squeeze_coeffs", squeeze_coeffs_dic)
 def test_morphsqueeze(x_morph, x_target, squeeze_coeffs):
     y_target = np.sin(x_target)
+    y_morph = np.sin(x_morph)
+    # expected output
+    y_morph_expected = y_morph
+    x_morph_expected = x_morph
+    x_target_expected = x_target
+    y_target_expected = y_target
+    # actual output
     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)
     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 = 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
+
+    extrap_low = np.where(x_morph < min(x_squeezed))[0]
+    extrap_high = np.where(x_morph > max(x_squeezed))[0]
+    extrap_index_low_expected = extrap_low[-1] if extrap_low.size else 0
+    extrap_index_high_expected = extrap_high[0] if extrap_high.size else -1
+
+    extrapolation_info = morph.extrapolation_info
+    extrap_index_low_actual = extrapolation_info["extrap_index_low"]
+    extrap_index_high_actual = extrapolation_info["extrap_index_high"]
+
     assert np.allclose(
-        y_morph_actual[extrap_index_low + 1 : extrap_index_high],
-        y_morph_expected[extrap_index_low + 1 : extrap_index_high],
+        y_morph_actual[
+            extrap_index_low_expected + 1 : extrap_index_high_expected
+        ],
+        y_morph_expected[
+            extrap_index_low_expected + 1 : extrap_index_high_expected
+        ],
         atol=1e-6,
     )
     assert np.allclose(
-        y_morph_actual[:extrap_index_low],
-        y_morph_expected[:extrap_index_low],
+        y_morph_actual[:extrap_index_low_expected],
+        y_morph_expected[:extrap_index_low_expected],
         atol=1e-3,
     )
     assert np.allclose(
-        y_morph_actual[extrap_index_high:],
-        y_morph_expected[extrap_index_high:],
+        y_morph_actual[extrap_index_high_expected:],
+        y_morph_expected[extrap_index_high_expected:],
         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)
+    assert np.allclose(x_target_actual, x_target_expected)
+    assert np.allclose(y_target_actual, y_target_expected)
+    assert extrap_index_low_actual == extrap_index_low_expected
+    assert extrap_index_high_actual == extrap_index_high_expected
 
 
 @pytest.mark.parametrize(
@@ -97,23 +106,23 @@ def test_morphsqueeze(x_morph, x_target, squeeze_coeffs):
             {"a0": 0.01},
             lambda x: (
                 "Warning: points with grid value below "
-                f"{x[0]} will be extrapolated."
+                f"{x[0]} are extrapolated."
             ),
         ),
         # extrapolate above
         (
             {"a0": -0.01},
             lambda x: (
                 "Warning: points with grid value above "
-                f"{x[1]} will be extrapolated."
+                f"{x[1]} are 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 "
+                f"{x[0]} and above {x[1]} are "
                 "extrapolated."
             ),
         ),