diff --git a/pynumdiff/kalman_smooth.py b/pynumdiff/kalman_smooth.py index a34da12..00219cc 100644 --- a/pynumdiff/kalman_smooth.py +++ b/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`.") diff --git a/pynumdiff/polynomial_fit.py b/pynumdiff/polynomial_fit.py index 7feb843..e720d26 100644 --- a/pynumdiff/polynomial_fit.py +++ b/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 " diff --git a/pynumdiff/tests/test_negative_axes.py b/pynumdiff/tests/test_negative_axes.py new file mode 100644 index 0000000..3e6cca2 --- /dev/null +++ b/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) diff --git a/pynumdiff/total_variation_regularization.py b/pynumdiff/total_variation_regularization.py index 23bed4d..9801ce1 100644 --- a/pynumdiff/total_variation_regularization.py +++ b/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.")