13
13
import warnings
14
14
15
15
import numpy as np
16
+ import pandas as pd
16
17
17
18
18
19
class Smoother :
19
20
"""
20
21
This is the smoothing utility class. This class holds the parameter settings for its smoother
21
22
methods and provides reasonable defaults. Basic usage can be found in the examples below.
22
23
23
- The smoother function takes numpy arrays as input, expecting the values to come from a
24
- regularly-spaced time grid. NANs are ok, as long as the array does not begin with a NAN. The
25
- rest of the NANs will be handled via imputation by default, though this can be turned off.
24
+ The smoother function takes numpy arrays or pandas Series as input, expecting the values to be
25
+ on a regularly-spaced time grid. NANs are ok, as long as the array does not begin with a NAN.
26
+ The rest of the NANs will be handled via imputation by default, though this can be turned off.
26
27
27
28
Parameters
28
29
----------
@@ -64,8 +65,9 @@ class Smoother:
64
65
65
66
Methods
66
67
----------
67
- smooth: np.ndarray
68
- Takes a 1D signal and returns a smoothed version. Both arrays have the same length.
68
+ smooth: np.ndarray or pd.Series
69
+ Takes a 1D signal and returns a smoothed version. The input and the output have the same length
70
+ and type.
69
71
70
72
Example Usage
71
73
-------------
@@ -75,7 +77,7 @@ class Smoother:
75
77
76
78
Example 2. Smooth a dataframe column.
77
79
>>> smoother = Smoother(smoother_name='savgol')
78
- >>> df[col] = pd.Series(smoother.smooth( df[col].to_numpy()) )
80
+ >>> df[col] = df[col].transform(smoother.smooth )
79
81
80
82
Example 3. Apply a rolling weighted average smoother, with 95% weight on the recent 2 weeks and
81
83
a sharp cutoff after 4 weeks.
@@ -133,17 +135,21 @@ def smooth(self, signal):
133
135
134
136
Parameters
135
137
----------
136
- signal: np.ndarray
138
+ signal: np.ndarray or pd.Series
137
139
A 1D signal to be smoothed.
138
140
139
- signal_smoothed: np.ndarray
140
- A smoothed 1D signal.
141
+ signal_smoothed: np.ndarray or pd.Series
142
+ A smoothed 1D signal. Returns an array of the same type and length as
143
+ the input.
141
144
"""
142
145
if len (signal ) < self .window_length :
143
146
raise ValueError (
144
147
"The window_length must be smaller than the length of the signal."
145
148
)
146
149
150
+ is_pandas_series = isinstance (signal , pd .Series )
151
+ signal = signal .to_numpy () if is_pandas_series else signal
152
+
147
153
signal = self .impute (signal )
148
154
149
155
if self .smoother_name == "savgol" :
@@ -155,7 +161,7 @@ def smooth(self, signal):
155
161
elif self .smoother_name == "identity" :
156
162
signal_smoothed = signal
157
163
158
- return signal_smoothed
164
+ return signal_smoothed if not is_pandas_series else pd . Series ( signal_smoothed )
159
165
160
166
def impute (self , signal ):
161
167
"""
@@ -281,7 +287,7 @@ def savgol_coeffs(self, nl, nr):
281
287
"""
282
288
Solves for the Savitzky-Golay coefficients. The coefficients c_i
283
289
give a filter so that
284
- y = \ sum_{i=-{n_l}}^{n_r} c_i x_i
290
+ y = sum_{i=-{n_l}}^{n_r} c_i x_i
285
291
is the value at 0 (thus the constant term) of the polynomial fit
286
292
through the points {x_i}. The coefficients are c_i are calculated as
287
293
c_i = ((A.T @ A)^(-1) @ (A.T @ e_i))_0
0 commit comments