Normalize negative axes before iterating signal vectors

pynumdiff/kalman_smooth.py

@@ -121,6 +121,7 @@ def rtsdiff(x, dt_or_t, order, log_qr_ratio, forwardbackward=False, axis=0, circ
                When :code:`circular=True`, wrapped to :math:`[-\\pi, \\pi]`.
              - **dxdt_hat** (np.array) -- estimated derivative of x, same shape as input :code:`x`
     """
+    axis = range(x.ndim)[axis] # normalize negative axes before constructing per-vector slices
     N = x.shape[axis]
     if not np.isscalar(dt_or_t):
         if N != len(dt_or_t): raise ValueError("If `dt_or_t` is given as array-like, must have same length as x along `axis`.")
@@ -223,6 +224,7 @@ def robustdiff(x, dt_or_t, order, log_q, log_r, proc_huberM=6, meas_huberM=0, ax
     :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`
     """
+    axis = range(x.ndim)[axis] # normalize negative axes before constructing per-vector slices
     N = x.shape[axis]
     if not np.isscalar(dt_or_t):
         if N != len(dt_or_t): raise ValueError("If `dt_or_t` is given as array-like, must have same length as `x` along `axis`.")

pynumdiff/polynomial_fit.py

@@ -24,6 +24,7 @@ def splinediff(x, dt_or_t, degree=3, s=1, axis=0):
     :return: - **x_hat** (np.array) -- estimated (smoothed) x
              - **dxdt_hat** (np.array) -- estimated derivative of x
     """
+    axis = range(x.ndim)[axis] # normalize negative axes before constructing per-vector slices
     if np.isscalar(dt_or_t):
         t = np.arange(x.shape[axis]) * dt_or_t
     else: # support variable step size for this function
@@ -68,6 +69,7 @@ def polydiff(x, dt_or_t, degree, window_size=None, stride=1, kernel='friedrichs'
     :return: - **x_hat** (np.array) -- estimated (smoothed) x
              - **dxdt_hat** (np.array) -- estimated derivative of x
     """
+    axis = range(x.ndim)[axis] # normalize negative axes before constructing per-vector slices
     if not np.isscalar(dt_or_t): # check once here rather than per window, since `slide_function` hands slices to `_polydiff`
         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`.")
         if np.any(np.diff(dt_or_t) <= 0): raise ValueError("`dt_or_t` must be strictly increasing. Out-of-order or repeated "

pynumdiff/tests/test_negative_axes.py

@@ -0,0 +1,29 @@
+"""Equivalent positive and negative axes must select the same signal vectors."""
+import numpy as np
+import pytest
+
+from pynumdiff import splinediff, polydiff, rtsdiff, robustdiff, tvrdiff
+
+
+@pytest.mark.parametrize('method,kwargs', [
+    (splinediff, {'s': 0}),
+    (polydiff, {'degree': 2, 'window_size': 7}),
+    (rtsdiff, {'order': 2, 'log_qr_ratio': 3}),
+    (robustdiff, {'order': 2, 'log_q': 3, 'log_r': 0}),
+    (tvrdiff, {'order': 2, 'gamma': .1}),
+])
+@pytest.mark.parametrize('axis', [0, 1, 2])
+def test_negative_signal_axes(method, kwargs, axis):
+    """Check the same signals along both signs of each valid dimension."""
+    t = np.linspace(0, 1, 15)
+    x = np.stack([t**2, 2*t + 1, np.sin(t), np.cos(t)]).reshape(2, 2, -1)
+    x = np.moveaxis(x, -1, axis)
+    original = x.copy()
+    expected = method(x, .1, axis=axis, **kwargs)
+    actual = method(x, .1, axis=axis-x.ndim, **kwargs)
+    for value, reference in zip(actual, expected):
+        np.testing.assert_array_equal(value, reference)
+    np.testing.assert_array_equal(x, original)
+    for invalid in [-x.ndim-1, x.ndim]:
+        with pytest.raises(IndexError):
+            method(x, .1, axis=invalid, **kwargs)

pynumdiff/total_variation_regularization.py

@@ -53,6 +53,7 @@ def tvrdiff(x, dt, order, gamma, huberM=float('inf'), axis=0):
     :return: - **x_hat** (np.array) -- estimated (smoothed) x
              - **dxdt_hat** (np.array) -- estimated derivative of x
     """
+    axis = range(x.ndim)[axis] # normalize negative axes before constructing per-vector slices
     if np.any(np.isnan(x)): raise ValueError("`x` may not contain NaN. CVXPY cannot form a problem with missing data.")
     if not np.isscalar(dt): raise ValueError("`dt` must be a scalar. The convex problem setup integrates with a cumulative "
         "sum and penalizes variation between consecutive samples, both of which assume uniform steps.")