Fixes the bug where the smoother drops the pandas series index:

* restore the index after smoothing
* test to match

Author: dshemetov

_delphi_utils_python/delphi_utils/smooth.py

@@ -172,6 +172,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
@@ -197,7 +198,10 @@ 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):

_delphi_utils_python/tests/test_smooth.py

@@ -2,7 +2,6 @@
 Tests for the smoothing utility.
 Authors: Dmitry Shemetov, Addison Hu, Maria Jahja
 """
-from numpy.lib.polynomial import poly
 import pytest
 
 import numpy as np
@@ -254,3 +253,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)