Merge pull request #194 from florisvb/polydiff-multidim

Add axis parameter to polydiff (#76)

Author: pavelkomarov

pynumdiff/polynomial_fit.py

@@ -50,7 +50,7 @@ def splinediff(x, dt_or_t, params=None, options=None, degree=3, s=None, num_iter
 
 
 def polydiff(x, dt_or_t, params=None, options=None, degree=None, window_size=None, step_size=1,
-    kernel='friedrichs'):
+    kernel='friedrichs', axis=0):
     """Fit polynomials to the data, and differentiate the polynomials.
 
     :param np.array[float] x: data to differentiate. May contain NaN values (missing data); NaNs are excluded from
@@ -64,6 +64,7 @@ def polydiff(x, dt_or_t, params=None, options=None, degree=None, window_size=Non
     :param int window_size: size of the sliding window, if not given no sliding
     :param int step_size: step size for sliding
     :param str kernel: name of kernel to use for weighting and smoothing windows ('gaussian' or 'friedrichs')
+    :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
@@ -80,11 +81,13 @@ def polydiff(x, dt_or_t, params=None, options=None, degree=None, window_size=Non
     elif degree is None or window_size is None:
         raise ValueError("`degree` and `window_size` must be given.")
 
-    if window_size < degree*3:
-        window_size = degree*3+1
-    if window_size % 2 == 0:
-        window_size += 1
-        warn("Kernel window size should be odd. Added 1 to length.")
+    if window_size:
+        if window_size < degree*3:
+            window_size = degree*3+1
+        if window_size % 2 == 0:
+            window_size += 1
+            warn("Kernel window size should be odd. Added 1 to length.")
+        kernel = {'gaussian':utility.gaussian_kernel, 'friedrichs':utility.friedrichs_kernel}[kernel](window_size)
 
     def _polydiff(x, dt_or_t, degree, weights=None):
         t = dt_or_t if not np.isscalar(dt_or_t) else np.arange(len(x)) * dt_or_t # sample locations
@@ -99,10 +102,13 @@ def _polydiff(x, dt_or_t, degree, weights=None):
 
         return x_hat, dxdt_hat
 
-    if not window_size: return _polydiff(x, dt_or_t, degree)
-
-    kernel = {'gaussian':utility.gaussian_kernel, 'friedrichs':utility.friedrichs_kernel}[kernel](window_size)
-    return utility.slide_function(_polydiff, x, dt_or_t, kernel, degree, stride=step_size, pass_weights=True)
+    x_hat = np.empty_like(x); dxdt_hat = np.empty_like(x)
+    for vec_idx in np.ndindex(x.shape[:axis] + x.shape[axis+1:]):
+        s = vec_idx[:axis] + (slice(None),) + vec_idx[axis:]
+        x_hat[s], dxdt_hat[s] = _polydiff(x[s], dt_or_t, degree) if not window_size else \
+            utility.slide_function(_polydiff, x[s], dt_or_t, kernel, degree, stride=step_size, pass_weights=True)
+    
+    return x_hat, dxdt_hat
 
 
 def savgoldiff(x, dt, params=None, options=None, degree=None, window_size=None, smoothing_win=None, axis=0):

pynumdiff/tests/test_diff_methods.py

@@ -325,9 +325,10 @@ 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, {}),
+    (polydiff, {'degree': 2, 'window_size': 5}),
     (savgoldiff, {'degree': 3, 'window_size': 11, 'smoothing_win': 3}),
     (rtsdiff, {'order':2, 'log_qr_ratio':7, 'forwardbackward':True}),
-    (robustdiff, {'order':2, 'log_q':7, 'log_r':2}),
+    (robustdiff, {'order':2, 'log_q':7, 'log_r':2})
 ]
 
 # Similar to the error_bounds table, index by method first. But then we test against only one 2D function,
@@ -338,6 +339,7 @@ def test_diff_method(diff_method_and_params, test_func_and_deriv, request): # re
     kerneldiff: [(2, 1), (3, 2)],
     butterdiff: [(0, -1), (1, -1)],
     finitediff: [(0, -1), (1, -1)],
+    polydiff: [(1, -1), (1, 0)],
     savgoldiff: [(0, -1), (1, 1)],
     rtsdiff: [(1, -1), (1, 0)],
     robustdiff: [(-2, -3), (0, -1)]