cleaning more

Author: pavelkomarov

notebooks/7_circular_domain.ipynb

@@ -80,9 +80,8 @@
    ],
    "source": [
     "x_hat, dxdt_hat = pynumdiff.kalman_smooth.rtsdiff(x, dt, order=1, log_qr_ratio=5, axis=0, circular=False)\n",
-    "x_hat_wrapped = (x_hat + np.pi) % (2*np.pi) - np.pi\n",
     "\n",
-    "evaluate.plot(x, dt, x_hat_wrapped, dxdt_hat, x_truth, dxdt_truth);"
+    "evaluate.plot(x, dt, x_hat, dxdt_hat, x_truth, dxdt_truth);"
    ]
   },
   {
@@ -112,9 +111,8 @@
    ],
    "source": [
     "x_hat, dxdt_hat = pynumdiff.kalman_smooth.rtsdiff(x, dt, order=1, log_qr_ratio=3, circular=True)\n",
-    "x_hat_wrapped = (x_hat + np.pi) % (2*np.pi) - np.pi\n",
     "\n",
-    "evaluate.plot(x, dt, x_hat_wrapped, dxdt_hat, x_truth, dxdt_truth);"
+    "evaluate.plot(x, dt, x_hat, dxdt_hat, x_truth, dxdt_truth);"
    ]
   }
  ],

pynumdiff/kalman_smooth.py

@@ -24,9 +24,9 @@ def kalman_filter(y, xhat0, P0, A, Q, C, R, B=None, u=None, save_P=True, innovat
     :param np.array u: optional control inputs, stacked in the direction of axis 0
     :param bool save_P: whether to save history of error covariance and a priori state estimates, used with rts
         smoothing but nonstandard to compute for ordinary filtering
-    :param callable innovation_fn: optional function :code:`(y_n, pred)` returning the innovation :code:`y_n - pred`,
-        where :code:`pred = C @ xhat_`. When :code:`None` (default), standard subtraction is used. Use this to handle
-        circular quantities, e.g. :code:`lambda y, pred: (y - pred + np.pi) % (2*np.pi) - np.pi` for angular measurements in radians.
+    :param callable innovation_fn: optional function taking measurements and predicted measurements and returning the innovation.
+        When :code:`None`, traditional subtraction is used. This is exposed to handle cases like wrapped domains, where alternative
+        displacement measures may be more appropriate. See e.g. the function passed by :code:`rtsdiff` with :code:`circular=True`.
 
     :return: - **xhat_pre** (np.array) -- a priori estimates of xhat, with axis=0 the batch dimension, so xhat[n] gets the nth step
              - **xhat_post** (np.array) -- a posteriori estimates of xhat
@@ -60,7 +60,7 @@ def kalman_filter(y, xhat0, P0, A, Q, C, R, B=None, u=None, save_P=True, innovat
         P = P_.copy()
         if not np.isnan(y[n]): # handle missing data
             K = P_ @ C.T @ np.linalg.inv(C @ P_ @ C.T + R)
-            innovation = innovation_fn(y[n], C @ xhat_) if innovation_fn is not None else y[n] - C @ xhat_
+            innovation = y[n] - C @ xhat_ if innovation_fn is None else innovation_fn(y[n], C @ xhat_)
             xhat += K @ innovation
             P -= K @ C @ P_
         # the [n]th index of pre variables holds _{n|n-1} info; the [n]th index of post variables holds _{n|n} info
@@ -116,7 +116,8 @@ def rtsdiff(x, dt_or_t, order, log_qr_ratio, forwardbackward=False, axis=0, circ
     :param bool circular: if :code:`True`, treat the measured quantity as a circular variable in radians, wrapping the
         innovation to :math:`[-\\pi, \\pi]`. The input :code:`x` must be in radians; convert degrees with :code:`np.deg2rad`.
 
-    :return: - **x_hat** (np.array) -- estimated (smoothed) x, same shape as input :code:`x`
+    :return: - **x_hat** (np.array) -- estimated (smoothed) x, same shape as input :code:`x`.
+             When :code:`circular=True`, wrapped to :math:`[-\\pi, \\pi]`.
              - **dxdt_hat** (np.array) -- estimated derivative of x, same shape as input :code:`x`
     """
     N = x.shape[axis]
@@ -153,22 +154,24 @@ def rtsdiff(x, dt_or_t, order, log_qr_ratio, forwardbackward=False, axis=0, circ
 
     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
+        xhat0 = np.zeros(order+1)
+        if not np.isnan(x[s][0]): xhat0[0] = x[s][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, innovation_fn=innovation_fn)
         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 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
+            if not np.isnan(x[s][-1]): xhat0[0] = x[s][-1]
             xhat_pre, xhat_post, P_pre, P_post = kalman_filter(x[s][::-1], xhat0, P0, A_d_bwd,
                 Q_d if Q_d.ndim == 2 else Q_d[::-1], C, R, innovation_fn=innovation_fn) # Use same Q matrices as before, because noise should still grow in reverse time
             xhat_smooth = rts_smooth(A_d_bwd, xhat_pre, xhat_post, P_pre, P_post, compute_P_smooth=False)
 
             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)
 
+    if circular: x_hat = (x_hat + np.pi) % (2*np.pi) - np.pi # wrap output to match the input domain
     return x_hat, dxdt_hat
 
 

pynumdiff/tests/test_diff_methods.py

@@ -405,22 +405,34 @@ def test_multidimensionality(multidim_method_and_params, request):
         legend = ax3.legend(bbox_to_anchor=(0.7, 0.8)); legend.legend_handles[0].set_facecolor(pyplot.cm.viridis(0.6))
         fig.suptitle(f'{diff_method.__name__}', fontsize=16)
 
-def test_circular_rtsdiff():
+def test_circular_rtsdiff(request):
     """Ensure rtsdiff with circular=True correctly differentiates a wrapping angle signal in radians"""
-    np.random.seed(42)
-    N = 200
-    dt_circ = 0.05
-    t_circ = np.arange(N) * dt_circ
-    true_dtheta = 2.0 # constant angular velocity in rad/s
-    theta_true = true_dtheta * t_circ # linearly increasing angle, crosses 2*pi boundaries
-    theta_noisy = np.angle(np.exp(1j * (theta_true + 0.1 * np.random.randn(N)))) # wrap to [-pi, pi] and add noise
+    dtheta = 5 # constant angular velocity in rad/s
+    theta = dtheta * t # linearly increasing angle, crosses 2*pi boundaries
+    theta_noisy = np.angle(np.exp(1j * (theta + noise))) # add noise and wrap to [-pi, pi]
 
-    _, dxdt_hat = rtsdiff(theta_noisy, dt_circ, order=1, log_qr_ratio=1, forwardbackward=True, circular=True)
+    theta_hat_naive, dxdt_hat_naive = rtsdiff(theta_noisy, dt, order=1, log_qr_ratio=1, circular=False)
+    theta_hat, dxdt_hat = rtsdiff(theta_noisy, dt, order=1, log_qr_ratio=1, circular=True)
+    
+    naive_rmse = np.sqrt(np.mean((dxdt_hat_naive - dtheta)**2))
+    wrapped_rmse = np.sqrt(np.mean((dxdt_hat - dtheta)**2))
+    assert wrapped_rmse < naive_rmse
 
-    # The interior of the signal (away from endpoints) should recover the true angular velocity well
-    interior = slice(10, N-10)
-    rmse = np.sqrt(np.mean((dxdt_hat[interior] - true_dtheta)**2))
-    assert rmse < 0.5, f"RMSE of angular velocity estimate too large: {rmse:.3f} rad/s"
+    if request.config.getoption("--plot"):
+        from matplotlib import pyplot
+        fig, (ax1, ax2) = pyplot.subplots(2, 1, figsize=(10, 6), sharex=True)
+        ax1.plot(t, theta_noisy, 'k+', label=r'$\theta$ noisy (wrapped)')
+        ax1.plot(t, theta_hat_naive, 'C1--', label=r'$\hat{\theta}$ with circular=False')
+        ax1.plot(t, theta_hat, 'C0', label=r'$\hat{\theta}$ with circular=True')
+        ax1.set_ylabel(r'$\theta$ (rad)')
+        ax1.legend()
+        ax2.axhline(dtheta, color='C2', xmin=0.045, xmax=0.955, label=r'true $\dot{\theta}$')
+        ax2.plot(t, dxdt_hat_naive, 'C1--', label=r'$\hat{\dot{\theta}}$ circular=False')
+        ax2.plot(t, dxdt_hat, 'C0', label=r'$\hat{\dot{\theta}}$ circular=True')
+        ax2.set_ylabel(r'$\dot{\theta}$ (rad/time)')
+        ax2.set_xlabel('t')
+        ax2.legend()
+        fig.suptitle('rtsdiff with circular domain', fontsize=16)
 
 
 # List of methods that can handle missing values

pynumdiff/utils/utility.py

@@ -7,7 +7,6 @@
 from scipy.ndimage import convolve1d
 
 
-
 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. Here :code:`huber`