Merge pull request #490 from cmu-delphi/fix_smoother_imputing_polyorder

Fix smoother imputing polynomial order

Author: krivard

_delphi_utils_python/delphi_utils/smooth.py

@@ -150,7 +150,7 @@ def __init__(
         else:
             self.coeffs = None
 
-    def smooth(self, signal: Union[np.ndarray, pd.Series]) -> Union[np.ndarray, pd.Series]:
+    def smooth(self, signal: Union[np.ndarray, pd.Series], impute_order=2) -> Union[np.ndarray, pd.Series]:
         """Apply a smoother to a signal.
 
         The major workhorse smoothing function. Imputes the nans and then applies
@@ -160,6 +160,9 @@ def smooth(self, signal: Union[np.ndarray, pd.Series]) -> Union[np.ndarray, pd.S
         ----------
         signal: np.ndarray or pd.Series
             A 1D signal to be smoothed.
+        impute_order: int
+            The polynomial order of the fit used for imputation. By default, this is set to
+            2.
 
         Returns
         ----------
@@ -172,6 +175,7 @@ def smooth(self, signal: Union[np.ndarray, pd.Series]) -> Union[np.ndarray, pd.S
             return signal
 
         is_pandas_series = isinstance(signal, pd.Series)
+        pandas_index = signal.index if is_pandas_series else None
         signal = signal.to_numpy() if is_pandas_series else signal
 
         # Find where the first non-nan value is located and truncate the initial nans
@@ -183,7 +187,7 @@ def smooth(self, signal: Union[np.ndarray, pd.Series]) -> Union[np.ndarray, pd.S
             signal_smoothed = signal.copy()
         else:
             # Impute
-            signal = self.impute(signal)
+            signal = self.impute(signal, impute_order=impute_order)
 
             # Smooth
             if self.smoother_name == "savgol":
@@ -197,10 +201,13 @@ def smooth(self, signal: Union[np.ndarray, pd.Series]) -> Union[np.ndarray, pd.S
 
         # Append the nans back, since we want to preserve length
         signal_smoothed = np.hstack([np.nan*np.ones(ix), signal_smoothed])
-        signal_smoothed = signal_smoothed if not is_pandas_series else pd.Series(signal_smoothed)
+        # Convert back to pandas if necessary
+        if is_pandas_series:
+            signal_smoothed = pd.Series(signal_smoothed)
+            signal_smoothed.index = pandas_index
         return signal_smoothed
 
-    def impute(self, signal):
+    def impute(self, signal, impute_order=2):
         """Impute the nan values in the signal.
 
         See the class docstring for an explanation of the impute methods.
@@ -209,6 +216,8 @@ def impute(self, signal):
         ----------
         signal: np.ndarray
             1D signal to be imputed.
+        impute_order: int
+            The polynomial order of the fit used for imputation.
 
         Returns
         -------
@@ -220,7 +229,7 @@ def impute(self, signal):
             # To preserve input-output array lengths, this util will not drop NaNs for you.
             if np.isnan(signal[0]):
                 raise ValueError("The signal should not begin with a nan value.")
-            imputed_signal = self.savgol_impute(signal)
+            imputed_signal = self.savgol_impute(signal, impute_order)
         elif self.impute_method == "zeros":
             imputed_signal = np.nan_to_num(signal)
         elif self.impute_method is None:
@@ -424,10 +433,10 @@ def savgol_smoother(self, signal):
         elif self.boundary_method == "nan":
             return signal_smoothed
 
-    def savgol_impute(self, signal):
+    def savgol_impute(self, signal, impute_order):
         """Impute the nan values in signal using savgol.
 
-        This method fills the nan values in the signal with a quadratic polynomial fit
+        This method fills the nan values in the signal with polynomial interpolation
         on a rolling window of the immediate past up to window_length data points.
 
         A number of boundary cases are handled involving nan filling close to the boundary.
@@ -439,34 +448,36 @@ def savgol_impute(self, signal):
         ----------
         signal: np.ndarray
             A 1D signal to be imputed.
+        impute_order: int
+            The polynomial order of the fit used for imputation.
 
         Returns
         ----------
         signal_imputed: np.ndarray
             An imputed 1D signal.
         """
+        if impute_order > self.window_length:
+            raise ValueError("Impute order must be smaller than window length.")
+
         signal_imputed = np.copy(signal)
         for ix in np.where(np.isnan(signal_imputed))[0]:
             # Boundary cases
             if ix < self.window_length:
                 # At the boundary, a single value should just be extended
                 if ix == 1:
                     signal_imputed[ix] = signal_imputed[ix - 1]
-                # Reduce the polynomial degree if needed
-                elif ix == 2:
-                    signal_imputed[ix] = self.savgol_predict(
-                        signal_imputed[:ix], 1, -1
-                    )
-                # Otherwise, use savgol fitting on the largest window prior
+                # Otherwise, use savgol fitting on the largest window prior,
+                # reduce the polynomial degree if needed (can't fit if the
+                # imputation order is larger than the available data)
                 else:
                     signal_imputed[ix] = self.savgol_predict(
-                        signal_imputed[:ix], self.poly_fit_degree, -1
+                        signal_imputed[:ix], min(ix-1, impute_order), -1
                     )
             # Away from the boundary, use savgol fitting on a fixed window
             else:
                 signal_imputed[ix] = self.savgol_predict(
                     signal_imputed[ix - self.window_length : ix],
-                    self.poly_fit_degree,
+                    impute_order,
                     -1,
                 )
         return signal_imputed

_delphi_utils_python/tests/test_smooth.py

@@ -158,6 +158,12 @@ def test_causal_savgol_smoother(self):
             smoothed_signal = smoother.smooth(signal)
             assert np.allclose(smoothed_signal, signal, equal_nan=True)
 
+        # test window_length > len(signal) and boundary_method="identity"
+        signal = np.arange(20)
+        smoother = Smoother(boundary_method="identity", window_length=30)
+        smoothed_signal = smoother.smooth(signal)
+        assert np.allclose(signal, smoothed_signal)
+
     def test_impute(self):
         # test front nan error
         with pytest.raises(ValueError):
@@ -178,7 +184,7 @@ def test_impute(self):
         signal = np.array([i if i % 3 else np.nan for i in range(1, 40)])
         # test that the non-nan values are unchanged
         not_nans_ixs = np.bitwise_xor(np.isnan(signal, where=True), np.full(len(signal), True))
-        smoothed_signal = Smoother().savgol_impute(signal)
+        smoothed_signal = Smoother().impute(signal)
         assert np.allclose(signal[not_nans_ixs], smoothed_signal[not_nans_ixs])
         # test that the imputer is close to the true line
         assert np.allclose(range(1, 40), smoothed_signal, atol=0.5)
@@ -187,47 +193,41 @@ def test_impute(self):
         signal = np.hstack([np.arange(10), [np.nan], np.arange(10)])
         window_length = 10
         smoother = Smoother(
-            smoother_name="savgol", window_length=window_length, poly_fit_degree=1
+            window_length=window_length, poly_fit_degree=1
         )
-        imputed_signal = smoother.savgol_impute(signal)
+        imputed_signal = smoother.impute(signal)
         assert np.allclose(imputed_signal, np.hstack([np.arange(11), np.arange(10)]))
         smoother = Smoother(
-            smoother_name="savgol", window_length=window_length, poly_fit_degree=2
+            window_length=window_length, poly_fit_degree=2
         )
-        imputed_signal = smoother.savgol_impute(signal)
+        imputed_signal = smoother.impute(signal)
         assert np.allclose(imputed_signal, np.hstack([np.arange(11), np.arange(10)]))
 
         # if there are nans on the boundary, should dynamically change window
         signal = np.hstack(
             [np.arange(5), [np.nan], np.arange(20), [np.nan], np.arange(5)]
         )
         smoother = Smoother(
-            smoother_name="savgol", window_length=window_length, poly_fit_degree=2
+            window_length=window_length, poly_fit_degree=2
         )
-        imputed_signal = smoother.savgol_impute(signal)
+        imputed_signal = smoother.impute(signal)
         assert np.allclose(
             imputed_signal, np.hstack([np.arange(6), np.arange(21), np.arange(5)]),
         )
 
         # if the array begins with np.nan, we should tell the user to peel it off before sending
         signal = np.hstack([[np.nan], np.arange(20), [np.nan], np.arange(5)])
         smoother = Smoother(
-            smoother_name="savgol", window_length=window_length, poly_fit_degree=2
+            window_length=window_length, poly_fit_degree=2
         )
         with pytest.raises(ValueError):
-            imputed_signal = smoother.savgol_impute(signal)
-
-        # test window_length > len(signal) and boundary_method="identity"
-        signal = np.arange(20)
-        smoother = Smoother(smoother_name="savgol", boundary_method="identity", window_length=30)
-        smoothed_signal = smoother.smooth(signal)
-        assert np.allclose(signal, smoothed_signal)
+            imputed_signal = smoother.impute(signal)
 
         # test the boundary methods
         signal = np.arange(20)
-        smoother = Smoother(smoother_name="savgol", poly_fit_degree=0,
+        smoother = Smoother(poly_fit_degree=0,
                             boundary_method="identity", window_length=10)
-        smoothed_signal = smoother.savgol_impute(signal)
+        smoothed_signal = smoother.impute(signal)
         assert np.allclose(smoothed_signal, signal)
 
         # test that we don't hit a matrix inversion error when there are
@@ -238,6 +238,13 @@ def test_impute(self):
         smoothed_signal = smoother.impute(signal)
         assert np.allclose(smoothed_signal, np.hstack([[1], np.ones(12), np.arange(5)]))
 
+        # test the impute_order argument
+        signal = np.hstack([[1, np.nan, np.nan, 2], np.arange(5)])
+        smoother = Smoother()
+        smoothed_signal = smoother.impute(signal, impute_order=1)
+        assert np.allclose(smoothed_signal, np.hstack([[1, 1, 1, 2], np.arange(5)]))
+
+
     def test_pandas_series_input(self):
         # The savgol method should match the linear regression method on the first
         # window_length-many values of the signal, if the savgol_weighting is set to true,
@@ -281,3 +288,11 @@ def test_pandas_series_input(self):
         assert np.allclose(
             signal[window_length - 1 :], smoothed_signal[window_length - 1 :]
         )
+
+        # Test that the index of the series gets preserved
+        signal = pd.Series(np.ones(30), index=np.arange(50, 80))
+        smoother = Smoother(smoother_name="moving_average", window_length=10)
+        smoothed_signal = signal.transform(smoother.smooth)
+        ix1 = signal.index
+        ix2 = smoothed_signal.index
+        assert ix1.equals(ix2)