Merge branch 'master' of github.com:florisvb/PyNumDiff into spectraldiff-multidim

Author: pavelkomarov

pynumdiff/kalman_smooth.py

@@ -37,9 +37,9 @@ def kalman_filter(y, xhat0, P0, A, Q, C, R, B=None, u=None, save_P=True):
     m = xhat0.shape[0] # dimension of the state
     xhat_post = np.empty((N,m))
     if save_P:
-        xhat_pre = np.empty((N,m)) # _pre = a priori predictions based on only past information
+        xhat_pre = np.empty_like(xhat_post) # _pre = a priori predictions based on only past information
         P_pre = np.empty((N,m,m)) # _post = a posteriori combinations of all information available at a step
-        P_post = np.empty((N,m,m))
+        P_post = np.empty_like(P_pre)
 
     control = isinstance(B, np.ndarray) and isinstance(u, np.ndarray) # whether there is a control input
     if A.ndim == 2: An, Qn, Bn = A, Q, B # single matrices, assign once outside the loop
@@ -113,8 +113,8 @@ def rtsdiff(x, dt_or_t, order, log_qr_ratio, forwardbackward, axis=0):
     :return: - **x_hat** (np.array) -- estimated (smoothed) x, same shape as input :code:`x`
              - **dxdt_hat** (np.array) -- estimated derivative of x, same shape as input :code:`x`
     """
-    x = np.moveaxis(np.asarray(x), axis, 0) # move axis of differentiation to standard position
-    if not np.isscalar(dt_or_t) and x.shape[0] != len(dt_or_t):
+    N = x.shape[axis]
+    if not np.isscalar(dt_or_t) and N != len(dt_or_t):
         raise ValueError("If `dt_or_t` is given as array-like, must have same length as x along `axis`.")
 
     q = 10**int(log_qr_ratio/2) # even-ish split of the powers across 0
@@ -132,25 +132,25 @@ def rtsdiff(x, dt_or_t, order, log_qr_ratio, forwardbackward, axis=0):
         Q_d = eM[:order+1, order+1:] @ A_d.T
         if forwardbackward: A_d_bwd = np.linalg.inv(A_d)
     else:
-        A_d = np.empty((len(x)-1, order+1, order+1))
+        A_d = np.empty((N-1, order+1, order+1))
         Q_d = np.empty_like(A_d)
-        for i, dt in enumerate(np.diff(dt_or_t)):
+        for n,dt in enumerate(np.diff(dt_or_t)):
             eM = expm(M * dt)
-            A_d[i] = eM[:order+1, :order+1]
-            Q_d[i] = eM[:order+1, order+1:] @ A_d[i].T
+            A_d[n] = eM[:order+1, :order+1]
+            Q_d[n] = eM[:order+1, order+1:] @ A_d[n].T
         if forwardbackward: A_d_bwd = np.linalg.inv(A_d[::-1]) # properly broadcasts, taking inv of each stacked 2D array
 
     x_hat = np.empty_like(x); dxdt_hat = np.empty_like(x)
-    if forwardbackward: w = np.linspace(0, 1, len(x)) # weights used to combine forward and backward results
+    if forwardbackward: w = np.linspace(0, 1, N) # weights used to combine forward and backward results
 
-    for vec_idx in np.ndindex(x.shape[1:]):
-        s = (slice(None),) + vec_idx # for indexing the the vector we wish to differentiate
+    for vec_idx in np.ndindex(x.shape[:axis] + x.shape[axis+1:]): # works properly for 1D case too
+        s = vec_idx[:axis] + (slice(None),) + vec_idx[axis:] # for indexing the vector we wish to differentiate
         xhat0 = np.zeros(order+1); xhat0[0] = x[s][0] if not np.isnan(x[s][0]) else 0 # The first estimate is the first seen state. See #110
 
         xhat_pre, xhat_post, P_pre, P_post = kalman_filter(x[s], xhat0, P0, A_d, Q_d, C, R)
         xhat_smooth = rts_smooth(A_d, xhat_pre, xhat_post, P_pre, P_post, compute_P_smooth=False)
-        x_hat[s] = xhat_smooth[:, 0] # first dimension is time, so slice first element at all times
-        dxdt_hat[s] = xhat_smooth[:, 1]
+        x_hat[s] = xhat_smooth[:,0] # first dimension is time, so slice first and second states at all times
+        dxdt_hat[s] = xhat_smooth[:,1]
 
         if forwardbackward:
             xhat0[0] = x[s][-1] if not np.isnan(x[s][-1]) else 0
@@ -161,7 +161,7 @@ def rtsdiff(x, dt_or_t, order, log_qr_ratio, forwardbackward, axis=0):
             x_hat[s] = x_hat[s] * w + xhat_smooth[:, 0][::-1] * (1-w)
             dxdt_hat[s] = dxdt_hat[s] * w + xhat_smooth[:, 1][::-1] * (1-w)
 
-    return np.moveaxis(x_hat, 0, axis), np.moveaxis(dxdt_hat, 0, axis)
+    return x_hat, dxdt_hat
 
 
 def constant_velocity(x, dt, params=None, options=None, r=None, q=None, forwardbackward=True):
@@ -254,7 +254,7 @@ def constant_jerk(x, dt, params=None, options=None, r=None, q=None, forwardbackw
     return rtsdiff(x, dt, 3, np.log10(q/r), forwardbackward)
 
 
-def robustdiff(x, dt_or_t, order, log_q, log_r, proc_huberM=6, meas_huberM=0):
+def robustdiff(x, dt_or_t, order, log_q, log_r, proc_huberM=6, meas_huberM=0, axis=0):
     """Perform outlier-robust differentiation by solving the Maximum A Priori optimization problem:
     :math:`\\text{argmin}_{\\{x_n\\}} \\sum_{n=0}^{N-1} V(R^{-1/2}(y_n - C x_n)) + \\sum_{n=1}^{N-1} J(Q_{n-1}^{-1/2}(x_n - A_{n-1} x_{n-1}))`,
     where :math:`A,Q,C,R` come from an assumed constant derivative model and :math:`V,J` are the :math:`\\ell_1` norm or Huber
@@ -287,38 +287,43 @@ def robustdiff(x, dt_or_t, order, log_q, log_r, proc_huberM=6, meas_huberM=0):
     :param float log_r: base 10 logarithm of measurement noise variance, so :code:`r = 10**log_r`
     :param float proc_huberM: quadratic-to-linear transition point for process loss
     :param float meas_huberM: quadratic-to-linear transition point for measurement loss
+    :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
+    :return: - **x_hat** (np.array) -- estimated (smoothed) x, same shape as input :code:`x`
+             - **dxdt_hat** (np.array) -- estimated derivative of x, same shape as input :code:`x`
     """
-    equispaced = np.isscalar(dt_or_t)
-    if not equispaced and len(x) != len(dt_or_t):
-        raise ValueError("If `dt_or_t` is given as array-like, must have same length as `x`.")
+    N = x.shape[axis]
+    if not np.isscalar(dt_or_t) and N != len(dt_or_t):
+        raise ValueError("If `dt_or_t` is given as array-like, must have same length as `x` along `axis`.")
 
     A_c = np.diag(np.ones(order), 1) # continuous-time A just has 1s on the first diagonal (where 0th is main diagonal)
     Q_c = np.zeros(A_c.shape); Q_c[-1,-1] = 10**log_q # continuous-time uncertainty around the last derivative
     C = np.zeros((1, order+1)); C[0,0] = 1 # we measure only y = noisy x
     R = np.array([[10**log_r]]) # 1 observed state, so this is 1x1
     M = np.block([[A_c, Q_c], [np.zeros(A_c.shape), -A_c.T]])  # exponentiate per step
 
-    if equispaced:
-        # convert to discrete time using matrix exponential
-        eM = expm(M * dt_or_t) # Note this could handle variable dt, similar to rtsdiff
+    if np.isscalar(dt_or_t): # convert to discrete time using matrix exponential
+        eM = expm(M * dt_or_t)
         A_d = eM[:order+1, :order+1]
         Q_d = eM[:order+1, order+1:] @ A_d.T
-        if np.linalg.cond(Q_d) > 1e12: Q_d += np.eye(order + 1)*1e-12 # for numerical stability with convex solver. Doesn't change answers appreciably (or at all).
+        if np.linalg.cond(Q_d) > 1e12: Q_d += np.eye(order+1)*1e-12 # for numerical stability with convex solver. Doesn't change answers appreciably (or at all).
     else: # support variable step size for this function
-        A_d = np.empty((len(x)-1, order+1, order+1)) # stack all the evolution matrices
-        Q_d = np.empty((len(x)-1, order+1, order+1))
-
-        for n,dt in enumerate(np.diff(dt_or_t)): # for each variable time step
+        A_d = np.empty((N-1, order+1, order+1))
+        Q_d = np.empty_like(A_d)
+        for n,dt in enumerate(np.diff(dt_or_t)):
             eM = expm(M * dt)
             A_d[n] = eM[:order+1, :order+1] # extract discrete time A matrix
             Q_d[n] = eM[:order+1, order+1:] @ A_d[n].T # extract discrete time Q matrix
-            if np.linalg.cond(Q_d[n]) > 1e12: Q_d[n] += np.eye(order + 1)*1e-12
+            if np.linalg.cond(Q_d[n]) > 1e12: Q_d[n] += np.eye(order+1)*1e-12
+
+    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:]): # works properly for 1D case too
+        s = vec_idx[:axis] + (slice(None),) + vec_idx[axis:]
+        x_states = convex_smooth(x[s], A_d, Q_d, C, R, proc_huberM=proc_huberM, meas_huberM=meas_huberM) # outsource solution of the convex optimization problem
+        x_hat[s] = x_states[:,0]; dxdt_hat[s] = x_states[:,1]
 
-    x_states = convex_smooth(x, A_d, Q_d, C, R, proc_huberM=proc_huberM, meas_huberM=meas_huberM) # outsource solution of the convex optimization problem
-    return x_states[:,0], x_states[:,1]
+    return x_hat, dxdt_hat
 
 
 def convex_smooth(y, A, Q, C, R, B=None, u=None, proc_huberM=6, meas_huberM=0):

pynumdiff/polynomial_fit.py

@@ -6,7 +6,7 @@
 from pynumdiff.utils import utility
 
 
-def splinediff(x, dt_or_t, params=None, options=None, degree=3, s=None, num_iterations=1):
+def splinediff(x, dt_or_t, params=None, options=None, degree=3, s=None, num_iterations=1, axis=0):
     """Find smoothed data and derivative estimates by fitting a smoothing spline to the data with
     scipy.interpolate.UnivariateSpline. Variable step size is supported with equal ease as uniform step size.
 
@@ -20,6 +20,7 @@ def splinediff(x, dt_or_t, params=None, options=None, degree=3, s=None, num_iter
     :param float s: positive smoothing factor used to choose the number of knots. Number of knots will be increased
         until the smoothing condition is satisfied: :math:`\\sum_t (x[t] - \\text{spline}[t])^2 \\leq s`
     :param int num_iterations: how many times to apply smoothing
+    :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
@@ -31,26 +32,29 @@ def splinediff(x, dt_or_t, params=None, options=None, degree=3, s=None, num_iter
         if options is not None:
             if 'iterate' in options and options['iterate']: num_iterations = params[2]
 
+    if num_iterations < 1: raise ValueError("`num_iterations` should be >=1")
     if np.isscalar(dt_or_t):
-        t = np.arange(len(x))*dt_or_t
+        t = np.arange(x.shape[axis]) * dt_or_t
     else: # support variable step size for this function
-        if len(x) != len(dt_or_t): raise ValueError("If `dt_or_t` is given as array-like, must have same length as `x`.")
+        if x.shape[axis] != len(dt_or_t): raise ValueError("If `dt_or_t` is given as array-like, must have same length as `x`.")
         t = dt_or_t
 
-    x_hat = x
-    for _ in range(num_iterations):
-        obs = ~np.isnan(x_hat) # UnivariateSpline can't handle NaN, so fit only on observed points
-        spline = scipy.interpolate.UnivariateSpline(t[obs], x_hat[obs], k=degree, s=s)
-        x_hat = spline(t) # evaluate at all t, filling in NaN positions by interpolation
+    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:]):
+        i = vec_idx[:axis] + (slice(None),) + vec_idx[axis:] # use i instead of s, becase s is already used as smoothness param
 
-    dspline = spline.derivative()
-    dxdt_hat = dspline(t)
+        obs = ~np.isnan(x[i]) # make_splrep can't handle NaN, so use only observed points for first fit
+        x_hat[i] = scipy.interpolate.make_splrep(t[obs], x[i][obs], k=degree, s=s)(t) # interpolate at all t
+        for _ in range(num_iterations-1):
+            x_hat[i] = scipy.interpolate.make_splrep(t, x_hat[i], k=degree, s=s)(t)
+        dxdt_hat[i] = spline.derivative()(t) # evaluate derivative at sample points
 
     return x_hat, dxdt_hat
 
 
 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 +68,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 +85,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 +106,14 @@ 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)
+    x_hat = np.empty_like(x); dxdt_hat = np.empty_like(x)
 
-    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)
+    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

@@ -149,7 +149,7 @@ def polydiff_irreg_step(*args, **kwargs): return polydiff(*args, **kwargs)
                  [(0, -1), (0, 0), (0, 0), (1, 0)],
                  [(1, 1), (2, 2), (1, 1), (2, 2)],
                  [(1, 1), (3, 3), (1, 1), (3, 3)]],
-    splinediff: [[(-14, -15), (-14, -15), (-1, -1), (0, 0)],
+    splinediff: [[(-14, -14), (-14, -14), (-1, -1), (0, 0)],
                  [(-14, -14), (-13, -14), (-1, -1), (0, 0)],
                  [(-14, -14), (-13, -13), (-1, -1), (0, 0)],
                  [(0, 0), (1, 1), (0, 0), (1, 1)],
@@ -325,8 +325,11 @@ 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})
+    (rtsdiff, {'order':2, 'log_qr_ratio':7, 'forwardbackward':True}),
+    (splinediff, {'degree': 9, 's': 1e-6}),
+    (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,
@@ -337,8 +340,11 @@ 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)]
+    rtsdiff: [(1, -1), (1, 0)],
+    splinediff: [(0, -1), (1, 0)],
+    robustdiff: [(-2, -3), (0, -1)]
 }
 
 @mark.parametrize("multidim_method_and_params", multidim_methods_and_params)