refactor: rename the function to set_extrapolation_info

Author: ycexiao

src/diffpy/morph/morph_io.py

@@ -416,26 +416,29 @@ def handle_warnings(squeeze_morph):
         cutoff_low = extrapolation_info["cutoff_low"]
         cutoff_high = extrapolation_info["cutoff_high"]
 
-        if is_extrap_low or is_extrap_high:
-            if not is_extrap_high:
-                wmsg = (
-                    "Warning: points with grid value below "
-                    f"{cutoff_low} "
-                    f"will be extrapolated."
-                )
-            elif not is_extrap_low:
-                wmsg = (
-                    "Warning: points with grid value above "
-                    f"{cutoff_high} "
-                    f"will be extrapolated."
-                )
-            else:
-                wmsg = (
-                    "Warning: points with grid value below "
-                    f"{cutoff_low} and above "
-                    f"{cutoff_high} "
-                    f"will be extrapolated."
-                )
+        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/morphs/morph.py

@@ -13,7 +13,7 @@
 #
 ##############################################################################
 """Morph -- base class for defining a morph."""
-
+import numpy
 
 LABEL_RA = "r (A)"  # r-grid
 LABEL_GR = "G (1/A^2)"  # PDF G(r)
@@ -245,22 +245,35 @@ def plotOutputs(self, xylabels=True, **plotargs):
             ylabel(self.youtlabel)
         return rv
 
-    def checkExtrapolation(self, x_true, x_extrapolate):
-        import numpy
+    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)
-        low_extrap = numpy.where(x_extrapolate < cutoff_low)[0]
-        high_extrap = numpy.where(x_extrapolate > cutoff_high)[0]
-        is_extrap_low = False if len(low_extrap) == 0 else True
-        is_extrap_high = False if len(high_extrap) == 0 else 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,
         }
-        return extrapolation_info
+        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,7 +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.extrapolation_info = self.checkExtrapolation(self.x_morph_in, r)
+        self.set_extrapolation_info(self.x_morph_in, r)
         return self.xyallout
 
 

src/diffpy/morph/morphs/morphsqueeze.py

@@ -85,8 +85,6 @@ 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
         )
-        self.extrapolation_info = self.checkExtrapolation(
-            x_squeezed, self.x_morph_in
-        )
+        self.set_extrapolation_info(x_squeezed, self.x_morph_in)
 
         return self.xyallout

tests/test_morphsqueeze.py

@@ -48,47 +48,54 @@ def test_morphsqueeze(x_morph, x_target, squeeze_coeffs):
     y_target = np.sin(x_target)
     y_morph = np.sin(x_morph)
     # expected 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_expected = y_morph
     x_morph_expected = x_morph
     x_target_expected = x_target
     y_target_expected = y_target
     # actual output
-    morph = MorphSqueeze()
+    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)
+    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_low = np.where(x_morph < min(x_squeezed))[0]
     extrap_high = np.where(x_morph > max(x_squeezed))[0]
-    extrap_index_low = extrap_low[-1] if extrap_low.size else None
-    extrap_index_high = extrap_high[0] if extrap_high.size else None
-    if extrap_index_low is None:
-        extrap_index_low = 0
-    elif extrap_index_high is None:
-        extrap_index_high = -1
+    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 np.allclose(x_morph_actual, x_morph_expected)
     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(
@@ -99,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."
             ),
         ),