extended savgoldiff to also be multidimensional

pynumdiff/polynomial_fit.py

@@ -100,10 +100,10 @@ def _polydiff(x, dt, degree, weights=None):
     return utility.slide_function(_polydiff, x, dt, kernel, degree, stride=step_size, pass_weights=True)
 
 
-def savgoldiff(x, dt, params=None, options=None, degree=None, window_size=None, smoothing_win=None):
+def savgoldiff(x, dt, params=None, options=None, degree=None, window_size=None, smoothing_win=None, axis=0):
     """Use the Savitzky-Golay to smooth the data and calculate the first derivative. It uses
     scipy.signal.savgol_filter. The Savitzky-Golay is very similar to the sliding polynomial fit,
-    but slightly noisier, and much faster.
+    but slightly noisier and much faster.
 
     :param np.array[float] x: data to differentiate
     :param float dt: step size
@@ -113,6 +113,7 @@ def savgoldiff(x, dt, params=None, options=None, degree=None, window_size=None,
     :param int window_size: size of the sliding window, must be odd (if not, 1 is added)
     :param int smoothing_win: size of the window used for gaussian smoothing, a good default is
         window_size, but smaller for high frequnecy data
+    :param int axis: data dimension along which differentiation is performed
 
     :return: - **x_hat** (np.array) -- estimated (smoothed) x
              - **dxdt_hat** (np.array) -- estimated derivative of x
@@ -130,12 +131,12 @@ def savgoldiff(x, dt, params=None, options=None, degree=None, window_size=None,
         warn("Kernel window size should be odd. Added 1 to length.")
     smoothing_win = min(smoothing_win, len(x)-1)
 
-    dxdt_hat = scipy.signal.savgol_filter(x, window_size, degree, deriv=1)/dt
+    dxdt_hat = scipy.signal.savgol_filter(x, window_size, degree, deriv=1, axis=axis)/dt
 
     kernel = utility.gaussian_kernel(smoothing_win)
-    dxdt_hat = utility.convolutional_smoother(dxdt_hat, kernel)
+    dxdt_hat = utility.convolutional_smoother(dxdt_hat, kernel, axis=axis)
 
-    x_hat = utility.integrate_dxdt_hat(dxdt_hat, dt)
-    x_hat += utility.estimate_integration_constant(x, x_hat)
+    x_hat = utility.integrate_dxdt_hat(dxdt_hat, dt, axis=axis)
+    x_hat += utility.estimate_integration_constant(x, x_hat, axis=axis)
 
     return x_hat, dxdt_hat

pynumdiff/tests/test_diff_methods.py

@@ -308,6 +308,7 @@ def test_diff_method(diff_method_and_params, test_func_and_deriv, request): # re
     (kerneldiff, {'kernel': 'gaussian', 'window_size': 5}),
     (butterdiff, {'filter_order': 3, 'cutoff_freq': 1 - 1e-6}),
     (finitediff, {}),
+    (savgoldiff, {'degree': 3, 'window_size': 11, 'smoothing_win': 3})
 ]
 
 # Similar to the error_bounds table, index by method first. But then we test against only one 2D function,
@@ -317,7 +318,8 @@ def test_diff_method(diff_method_and_params, test_func_and_deriv, request): # re
 multidim_error_bounds = {
     kerneldiff: [(2, 1), (3, 2)],
     butterdiff: [(0, -1), (1, -1)],
-    finitediff: [(0, -1), (1, -1)]
+    finitediff: [(0, -1), (1, -1)],
+    savgoldiff: [(0, -1), (1, 1)]
 }
 
 @mark.parametrize("multidim_method_and_params", multidim_methods_and_params)