Merge pull request #192 from Ashton-Graves/robustdiff-variable-step

Robustdiff variable step

Author: pavelkomarov

notebooks/1_basic_tutorial.ipynb

@@ -259,7 +259,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, 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):
     """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^{-1/2}(x_n - A 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
@@ -295,16 +295,31 @@ def robustdiff(x, dt, order, log_q, log_r, proc_huberM=6, meas_huberM=0):
     :return: - **x_hat** (np.array) -- estimated (smoothed) x
              - **dxdt_hat** (np.array) -- estimated derivative of 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`.")
+
     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
 
-    # convert to discrete time using matrix exponential
-    eM = expm(np.block([[A_c, Q_c], [np.zeros(A_c.shape), -A_c.T]]) * dt) # Note this could handle variable dt, similar to rtsdiff
-    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 equispaced:
+        # convert to discrete time using matrix exponential
+        eM = expm(M * dt_or_t) # Note this could handle variable dt, similar to rtsdiff
+        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).
+    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 i, dt in enumerate(np.diff(dt_or_t)): # for each variable time step
+            eM = expm(M * dt)
+            A_d[i] = eM[:order+1, :order+1] # extract discrete time A matrix
+            Q_d[i] = eM[:order+1, order+1:] @ A_d[i].T # extract discrete time Q matrix
+            if np.linalg.cond(Q_d[i]) > 1e12: Q_d[i] += np.eye(order + 1)*1e-12
 
     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]
@@ -327,12 +342,20 @@ def convex_smooth(y, A, Q, C, R, B=None, u=None, proc_huberM=6, meas_huberM=0):
     :return: (np.array) -- state estimates (state_dim x N)
     """
     N = len(y)
-    x_states = cvxpy.Variable((A.shape[0], N)) # each column is [position, velocity, acceleration, ...] at step n
+    state_dim = A.shape[-1]
+    x_states = cvxpy.Variable((state_dim, N)) # each column is [position, velocity, acceleration, ...] at step n
     control = isinstance(B, np.ndarray) and isinstance(u, np.ndarray) # whether there is a control input
 
-    # It is extremely important to run time that CVXPY expressions be in vectorized form
-    proc_resids = np.linalg.inv(sqrtm(Q)) @ (x_states[:,1:] - A @ x_states[:,:-1] - (0 if not control else B @ u[1:].T)) # all Q^(-1/2)(x_n - (A x_{n-1} + B u_n))
-    meas_resids = np.linalg.inv(sqrtm(R)) @ (y.reshape(C.shape[0],-1) - C @ x_states) # all R^(-1/2)(y_n - C x_n)
+    if A.ndim == 3: # It is extremely important to runtime that CVXPY expressions be in vectorized form
+        Ax = cvxpy.einsum('nij,jn->in', A, x_states[:, :-1]) # multipy each A matrix by the corresponding x_states at that time step
+        Q_inv_sqrts = np.array([np.linalg.inv(sqrtm(Q[n])) for n in range(N-1)]) # precompute Q^(-1/2) for each time step
+        proc_resids = cvxpy.einsum('nij,jn->in', Q_inv_sqrts, x_states[:,1:] - Ax - (0 if not control else B @ u[1:].T))
+    else: # all Q^(-1/2)(x_n - (A x_{n-1} + B u_n))
+        proc_resids = np.linalg.inv(sqrtm(Q)) @ (x_states[:,1:] - A @ x_states[:,:-1] - (0 if not control else B @ u[1:].T))
+    
+    obs = ~np.isnan(y) # boolean mask of non-NaN observations
+    meas_resids = np.linalg.inv(sqrtm(R)) @ (y[obs].reshape(C.shape[0],-1) - C @ x_states[:,obs]) # all R^(-1/2)(y_n - C x_n)
+
     # Process terms: sum of J(proc_resids)
     objective = 0.5*cvxpy.sum_squares(proc_resids) if proc_huberM == float('inf') \
                 else np.sqrt(2)*cvxpy.sum(cvxpy.abs(proc_resids)) if proc_huberM < 1e-3 \
@@ -345,8 +368,8 @@ def convex_smooth(y, A, Q, C, R, B=None, u=None, proc_huberM=6, meas_huberM=0):
     # function https://www.cvxpy.org/api_reference/cvxpy.atoms.elementwise.html#huber, so correct with a factor of 0.5.
 
     problem = cvxpy.Problem(cvxpy.Minimize(objective))
-    try: problem.solve(solver=cvxpy.CLARABEL)
+    try: problem.solve(solver=cvxpy.CLARABEL, canon_backend=cvxpy.SCIPY_CANON_BACKEND)
     except cvxpy.error.SolverError: pass # Could try another solver here, like SCS, but slows things down
 
-    if x_states.value is None: return np.full((N, A.shape[0]), np.nan) # There can be solver failure, even without error
+    if x_states.value is None: return np.full((N, state_dim), np.nan) # There can be solver failure, even without error
     return x_states.value.T

pynumdiff/kalman_smooth.py

@@ -13,8 +13,9 @@
 def iterated_second_order(*args, **kwargs): return second_order(*args, **kwargs)
 def iterated_fourth_order(*args, **kwargs): return fourth_order(*args, **kwargs)
 def spline_irreg_step(*args, **kwargs): return splinediff(*args, **kwargs)
+def robust_irreg_step(*args, **kwargs): return robustdiff(*args, **kwargs)
 def polydiff_irreg_step(*args, **kwargs): return polydiff(*args, **kwargs)
-irreg_list = [spline_irreg_step, polydiff_irreg_step, rbfdiff, rtsdiff] # methods to test with irregular time steps
+irreg_list = [spline_irreg_step, polydiff_irreg_step, rbfdiff, rtsdiff, robust_irreg_step] # methods to test with irregular time steps
 
 dt = 0.1
 t = np.linspace(0, 3, 31) # sample locations, including the endpoint
@@ -55,6 +56,7 @@ def polydiff_irreg_step(*args, **kwargs): return polydiff(*args, **kwargs)
     (constant_jerk, {'r':1e-4, 'q':1e5}), (constant_jerk, [1e-4, 1e5]),
     (rtsdiff, {'order':2, 'log_qr_ratio':7, 'forwardbackward':True}),
     (robustdiff, {'order':3, 'log_q':7, 'log_r':2}),
+    (robust_irreg_step, {'order':3, 'log_q':7, 'log_r':2}),
     (velocity, {'gamma':0.5}), (velocity, [0.5]),
     (acceleration, {'gamma':1}), (acceleration, [1]),
     (jerk, {'gamma':10}), (jerk, [10]),
@@ -231,6 +233,12 @@ def polydiff_irreg_step(*args, **kwargs): return polydiff(*args, **kwargs)
                  [(-7, -7), (-2, -2), (0, -1), (1, 1)],
                  [(0, 0), (2, 2), (0, 0), (2, 2)],
                  [(1, 1), (3, 3), (1, 1), (3, 3)]],
+    robust_irreg_step: [[(-15, -15), (-13, -14), (0, -1), (1, 1)],
+                        [(-14, -14), (-13, -13), (0, -1), (1, 1)],
+                        [(-14, -14), (-13, -13), (0, -1), (1, 1)],
+                        [(-8, -8), (-2, -2), (0, -1), (1, 1)],
+                        [(0, 0), (2, 2), (0, 0), (2, 2)],
+                        [(1, 1), (3, 3), (1, 1), (3, 3)]],
     lineardiff: [[(-3, -4), (-3, -3), (0, -1), (1, 0)],
                  [(-1, -2), (0, 0), (0, -1), (1, 0)],
                  [(-1, -1), (0, 0), (0, -1), (1, 1)],

pynumdiff/tests/test_diff_methods.py

@@ -2,6 +2,7 @@
 import numpy as np
 import matplotlib.pyplot as plt
 from scipy import stats
+from scipy.special import huber
 
 from pynumdiff.utils import utility
 
@@ -95,7 +96,7 @@ def robust_rme(u, v, padding=0, M=6):
     s = slice(padding, len(u)-padding) # slice out data we want to measure
 
     sigma = stats.median_abs_deviation(u[s] - v[s], scale='normal') # M is in units of this robust scatter metric
-    return np.sqrt(2*np.mean(utility.huber(u[s] - v[s], M*sigma)))
+    return np.sqrt(2*np.mean(huber(M*sigma, u[s] - v[s])))
 
 
 def rmse(u, v, padding=0):

pynumdiff/utils/evaluate.py

@@ -2,24 +2,15 @@
 import numpy as np
 from scipy.integrate import cumulative_trapezoid
 from scipy.optimize import minimize
+from scipy.special import huber
 from scipy.stats import median_abs_deviation, norm
 from scipy.ndimage import convolve1d
 
 
-def huber(x, M):
-    """Huber loss function, for outlier-robust applications,
-    `see here <https://www.cvxpy.org/api_reference/cvxpy.atoms.elementwise.html#huber>`_
-
-    :param np.array[float] x: data points on which to evaluate the Huber function pointwise
-    :param float M: where the loss turns from quadratic to linear
-    :return: (np.array[float]) -- pointwise evaluations of the Huber function
-    """
-    absx = np.abs(x)
-    return np.where(absx <= M, 0.5*x**2, M*(absx - 0.5*M))
-
 def huber_const(M):
     """Scale that makes :code:`sum(huber())` interpolate :math:`\\sqrt{2}\\|\\cdot\\|_1` and :math:`\\frac{1}{2}\\|\\cdot\\|_2^2`,
-    from https://jmlr.org/papers/volume14/aravkin13a/aravkin13a.pdf, with correction for missing sqrt
+    from https://jmlr.org/papers/volume14/aravkin13a/aravkin13a.pdf, with correction for missing sqrt. Here :code:`huber`
+    refers to `scipy.special.huber <https://docs.scipy.org/doc/scipy/reference/generated/scipy.special.huber.html>`_.
 
     :param float M: Huber parameter, where the function turns from quadratic to linear
     :return: (float) -- appropriate scale factor to normalize the Huber function
@@ -65,7 +56,7 @@ def estimate_integration_constant(x, x_hat, M=6, axis=0):
     elif M < 1e-3: # small M looks like l1 loss, and Huber gets too flat to work well
         return np.median(x - x_hat, axis=axis).reshape(s) # Solves the l1 distance minimization, argmin_c ||x_hat + c - x||_1
     else:
-        return minimize(lambda c: np.sum(huber(x_hat + c.reshape(s) - x, M*sigma)), # fn to minimize in 1st argument
+        return minimize(lambda c: np.sum(huber(M*sigma, x_hat + c.reshape(s) - x)), # fn to minimize in 1st argument
             np.zeros(np.prod(s)), method='SLSQP').x.reshape(s) # initial guess is zeros; vector result must be reshaped