update interpolate util module to add constant extrapolation option

Author: ValentinGebhart

climada/util/interpolation.py

@@ -37,7 +37,7 @@ def interpolate_ev(
         logy = False,
         x_threshold = None,
         y_threshold = None,
-        extrapolation = False,
+        extrapolation = None,
         y_asymptotic = np.nan
     ):
     """
@@ -60,14 +60,16 @@ def interpolate_ev(
             Lower threshold to filter x_train. Defaults to None.
         y_threshold : float, optional
             Lower threshold to filter y_train. Defaults to None.
-        extrapolation : bool, optional
-            If set to True, values will be extrapolated. If set to False, x_test values
-            smaller than x_train will be assigned y_train[0] (x_train must be sorted in
-            ascending order), and x_test values larger than x_train will be assigned
-            y_asymptotic. Defaults to False
+        extrapolation : str, optional
+            If set to 'extrapolate', values will be extrapolated. If set to 'extrapolate_constant',
+            x_test values smaller than x_train will be assigned y_train[0] (x_train must be sorted
+            in ascending order), and x_test values larger than x_train will be assigned
+            y_asymptotic. If set to None, x_test values outside of the range of x_train will be
+            assigned np.nan. Defaults to None.
         y_asymptotic : float, optional
-            Return value and if extrapolation is True or x_train.size < 2, for x_test
-            values larger than x_train. Defaults to np.nan.
+            Has no effect if extrapolation is None. Else, provides return value and if
+            for x_test values larger than x_train, for x_train.size < 2 or if extrapolation is set
+            to 'extrapolate_constant'. Defaults to np.nan.
 
     Returns
     -------
@@ -80,20 +82,27 @@ def interpolate_ev(
         x_test, x_train, y_train, logx, logy, x_threshold, y_threshold
     )
 
-      # handle case of small training data sizes
+    # handle case of small training data sizes
     if x_train.size < 2:
+        if not extrapolation:
+            return np.full_like(x_test, np.nan)
+        else:
+            LOGGER.warning('Data is being extrapolated.')
+            return _interpolate_small_input(x_test, x_train, y_train, logy, y_asymptotic)
+
+    # warn if values are being extrapolated
+    if extrapolation and (np.min(x_test) < np.min(x_train) or np.max(x_test) > np.max(x_train)):
         LOGGER.warning('Data is being extrapolated.')
-        return _interpolate_small_input(x_test, x_train, y_train, logy, y_asymptotic)
 
     # calculate fill values
-    if extrapolation:
+    if extrapolation == 'extrapolate':
         fill_value = 'extrapolate'
-        if np.min(x_test) < np.min(x_train) or np.max(x_test) > np.max(x_train):
-            LOGGER.warning('Data is being extrapolated.')
-    else:
+    elif extrapolation == 'extrapolate_constant':
         if not all(sorted(x_train) == x_train):
             raise ValueError('x_train array must be sorted in ascending order.')
         fill_value = (y_train[0], np.log10(y_asymptotic) if logy else y_asymptotic) 
+    else:
+        fill_value = np.nan
 
     interpolation = interpolate.interp1d(
         x_train, y_train, fill_value=fill_value, bounds_error=False)
@@ -142,9 +151,16 @@ def stepfunction_ev(
         x_test, x_train, y_train, None, None, x_threshold, y_threshold
     )
 
+    
+
     # handle case of small training data sizes
     if x_train.size < 2:
+        LOGGER.warning('Data is being extrapolated.')
         return _interpolate_small_input(x_test, x_train, y_train, None, y_asymptotic)
+    
+    # warn if values are being extrapolated
+    if (np.min(x_test) < np.min(x_train) or np.max(x_test) > np.max(x_train)):
+        LOGGER.warning('Data is being extrapolated.')
 
     # find indices of x_test if sorted into x_train
     if not all(sorted(x_train) == x_train):

climada/util/test/test_interpolation.py

@@ -35,10 +35,15 @@ def test_interpolate_ev_linear_interp(self):
         x_test = np.array([0., 3., 4., 6.])
         np.testing.assert_allclose(
             interpolate_ev(x_test, x_train, y_train),
+            np.array([np.nan, 4., 3., np.nan])
+        )
+        np.testing.assert_allclose(
+            interpolate_ev(x_test, x_train, y_train, extrapolation='extrapolate_constant'),
             np.array([8., 4., 3., np.nan])
         )
         np.testing.assert_allclose(
-            interpolate_ev(x_test, x_train, y_train, y_asymptotic = 0),
+            interpolate_ev(x_test, x_train, y_train,
+                        extrapolation='extrapolate_constant', y_asymptotic = 0),
             np.array([8., 4., 3., 0.])
         )
 
@@ -48,15 +53,17 @@ def test_interpolate_ev_threshold_parameters(self):
         y_train = np.array([4., 1., 4.])
         x_test = np.array([-1., 3., 4.])
         np.testing.assert_allclose(
-            interpolate_ev(x_test, x_train, y_train),
+            interpolate_ev(x_test, x_train, y_train, extrapolation='extrapolate_constant'),
             np.array([4., 1., 2.])
         )
         np.testing.assert_allclose(
-            interpolate_ev(x_test, x_train, y_train, x_threshold=1.),
+            interpolate_ev(x_test, x_train, y_train, x_threshold=1.,
+                           extrapolation='extrapolate_constant'),
             np.array([1., 1., 2.])
         )
         np.testing.assert_allclose(
-            interpolate_ev(x_test, x_train, y_train, y_threshold=2.),
+            interpolate_ev(x_test, x_train, y_train, y_threshold=2.,
+                           extrapolation='extrapolate_constant'),
             np.array([4., 4., 4.])
         )
 
@@ -66,25 +73,27 @@ def test_interpolate_ev_scale_parameters(self):
         y_train = np.array([1., 3.])
         x_test = np.array([1e0, 1e2])
         np.testing.assert_allclose(
-            interpolate_ev(x_test, x_train, y_train, logx=True, extrapolation=True),
+            interpolate_ev(x_test, x_train, y_train, logx=True, extrapolation='extrapolate'),
             np.array([0., 2.])
         )
         np.testing.assert_allclose(
-            interpolate_ev(x_test, x_train, y_train, logx=True),
+            interpolate_ev(x_test, x_train, y_train, logx=True,
+                           extrapolation='extrapolate_constant'),
             np.array([1., 2.])
         )
         x_train = np.array([1., 3.])
         y_train = np.array([1e1, 1e3])
         x_test = np.array([0., 2.])
         np.testing.assert_allclose(
-            interpolate_ev(x_test, x_train, y_train, logy=True, extrapolation=True),
+            interpolate_ev(x_test, x_train, y_train, logy=True, extrapolation='extrapolate'),
             np.array([1e0, 1e2])
         )
         x_train = np.array([1e1, 1e3])
         y_train = np.array([1e1, 1e5])
         x_test = np.array([1e0, 1e2])
         np.testing.assert_allclose(
-            interpolate_ev(x_test, x_train, y_train, logx=True, logy=True, extrapolation=True),
+            interpolate_ev(x_test, x_train, y_train, logx=True, logy=True,
+                           extrapolation='extrapolate'),
             np.array([1e-1, 1e3])
         )
 
@@ -95,7 +104,7 @@ def test_interpolate_ev_degenerate_input(self):
         y_train = np.zeros(3)
         np.testing.assert_allclose(
             interpolate_ev(x_test, x_train, y_train),
-            np.array([0., 0., 0.])
+            np.array([np.nan, 0., 0.])
         )
 
     def test_interpolate_ev_small_input(self):
@@ -104,13 +113,18 @@ def test_interpolate_ev_small_input(self):
         y_train = np.array([2.])
         x_test = np.array([0., 1., 2.])
         np.testing.assert_allclose(
-            interpolate_ev(x_test, x_train, y_train),
+            interpolate_ev(x_test, x_train, y_train, extrapolation='extrapolate'),
             np.array([2., 2., np.nan])
         )
         np.testing.assert_allclose(
-            interpolate_ev(x_test, x_train, y_train, y_asymptotic=0),
+            interpolate_ev(x_test, x_train, y_train, extrapolation='extrapolate', y_asymptotic=0),
             np.array([2., 2., 0.])
         )
+        np.testing.assert_allclose(
+            interpolate_ev(x_test, x_train, y_train),
+            np.full(3, np.nan)
+        )
+
         x_train = np.array([])
         y_train = np.array([])
         x_test = np.array([0., 1., 2.])
@@ -119,7 +133,8 @@ def test_interpolate_ev_small_input(self):
             np.full(3, np.nan)
         )
         np.testing.assert_allclose(
-            interpolate_ev(x_test, x_train, y_train, y_asymptotic=0),
+            interpolate_ev(x_test, x_train, y_train,
+                           extrapolation='extrapolate_constant', y_asymptotic=0),
             np.zeros(3)
         )