Update smoother to gracefully handle nans:

* entire array of nans is handled
* left-padded nans are now ignored
* a few other edge cases
* add tests to match

Author: dshemetov

_delphi_utils_python/delphi_utils/smooth.py

@@ -138,6 +138,8 @@ def __init__(
             raise ValueError("Invalid impute_method given.")
         if self.boundary_method not in valid_boundary_methods:
             raise ValueError("Invalid boundary_method given.")
+        if self.window_length <= 1:
+            raise ValueError("Window length is too short.")
 
         if smoother_name == "savgol":
             # The polynomial fitting is done on a past window of size window_length
@@ -165,20 +167,36 @@ def smooth(self, signal: Union[np.ndarray, pd.Series]) -> Union[np.ndarray, pd.S
             A smoothed 1D signal. Returns an array of the same type and length as
             the input.
         """
+        # If all nans, pass through
+        if np.all(np.isnan(signal)):
+            return signal
+
         is_pandas_series = isinstance(signal, pd.Series)
         signal = signal.to_numpy() if is_pandas_series else signal
 
-        signal = self.impute(signal)
+        # Find where the first non-nan value is located and truncate the initial nans
+        ix = np.where(~np.isnan(signal))[0][0]
+        signal = signal[ix:]
 
-        if self.smoother_name == "savgol":
-            signal_smoothed = self.savgol_smoother(signal)
-        elif self.smoother_name == "left_gauss_linear":
-            signal_smoothed = self.left_gauss_linear_smoother(signal)
-        elif self.smoother_name == "moving_average":
-            signal_smoothed = self.moving_average_smoother(signal)
-        else:
+        # Don't smooth in certain edge cases
+        if len(signal) < self.poly_fit_degree or len(signal) == 1:
             signal_smoothed = signal.copy()
-
+        else:
+            # Impute
+            signal = self.impute(signal)
+
+            # Smooth
+            if self.smoother_name == "savgol":
+                signal_smoothed = self.savgol_smoother(signal)
+            elif self.smoother_name == "left_gauss_linear":
+                signal_smoothed = self.left_gauss_linear_smoother(signal)
+            elif self.smoother_name == "moving_average":
+                signal_smoothed = self.moving_average_smoother(signal)
+            elif self.smoother_name == "identity":
+                signal_smoothed = signal
+
+        # 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)
         return signal_smoothed
 
@@ -282,7 +300,7 @@ def left_gauss_linear_smoother(self, signal):
 
     def savgol_predict(self, signal, poly_fit_degree, nr):
         """Predict a single value using the savgol method.
-
+        
         Fits a polynomial through the values given by the signal and returns the value
         of the polynomial at the right-most signal-value. More precisely, for a signal of length
         n, fits a poly_fit_degree polynomial through the points signal[-n+1+nr], signal[-n+2+nr],
@@ -311,7 +329,8 @@ def savgol_predict(self, signal, poly_fit_degree, nr):
     def savgol_coeffs(self, nl, nr, poly_fit_degree):
         """Solve for the Savitzky-Golay coefficients.
 
-        The coefficients c_i give a filter so that
+        Solves for the Savitzky-Golay coefficients. The coefficients c_i
+        give a filter so that
             y = sum_{i=-{n_l}}^{n_r} c_i x_i
         is the value at 0 (thus the constant term) of the polynomial fit
         through the points {x_i}. The coefficients are c_i are calculated as
@@ -385,7 +404,7 @@ def savgol_smoother(self, signal):
         # - identity keeps the original signal (doesn't smooth)
         # - nan writes nans
         if self.boundary_method == "shortened_window":
-            for ix in range(len(self.coeffs)):
+            for ix in range(min(len(self.coeffs), len(signal))):
                 if ix == 0:
                     signal_smoothed[ix] = signal[ix]
                 else:

_delphi_utils_python/tests/test_smooth.py

@@ -2,6 +2,7 @@
 Tests for the smoothing utility.
 Authors: Dmitry Shemetov, Addison Hu, Maria Jahja
 """
+from numpy.lib.polynomial import poly
 import pytest
 
 import numpy as np
@@ -109,9 +110,39 @@ def test_causal_savgol_smoother(self):
             smoother_name="savgol", window_length=window_length, poly_fit_degree=1,
         )
         smoothed_signal2 = smoother.smooth(signal)
-
         assert np.allclose(smoothed_signal1, smoothed_signal2)
 
+        # Test the all nans case
+        signal = np.nan * np.ones(10)
+        smoother = Smoother(window_length=9)
+        smoothed_signal = smoother.smooth(signal)
+        assert np.all(np.isnan(smoothed_signal))
+
+        # Test the case where the signal is length 1
+        signal = np.ones(1)
+        smoother = Smoother()
+        smoothed_signal = smoother.smooth(signal)
+        assert np.allclose(smoothed_signal, signal)
+
+        # Test the case where the signal length is less than polynomial_fit_degree
+        signal = np.ones(2)
+        smoother = Smoother(poly_fit_degree=3)
+        smoothed_signal = smoother.smooth(signal)
+        assert np.allclose(smoothed_signal, signal)
+
+        # Test an edge fitting case
+        signal = np.array([np.nan, 1, np.nan])
+        smoother = Smoother(poly_fit_degree=1, window_length=2)
+        smoothed_signal = smoother.smooth(signal)
+        assert np.allclose(smoothed_signal, np.array([np.nan, 1, 1]), equal_nan=True)
+
+        # Test a range of cases where the signal size following a sequence of nans is returned
+        for i in range(10):
+            signal = np.hstack([[np.nan, np.nan, np.nan], np.ones(i)])
+            smoother = Smoother(poly_fit_degree=0, window_length=5)
+            smoothed_signal = smoother.smooth(signal)
+            assert np.allclose(smoothed_signal, signal, equal_nan=True)
+
     def test_impute(self):
         # test the nan imputer
         signal = np.array([i if i % 3 else np.nan for i in range(1, 40)])