Add axis parameter to splinediff for multidimensional support (#76)

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: pavelkomarov

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,22 +32,30 @@ 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]
 
+    x = np.moveaxis(np.asarray(x), axis, 0)
+    n = x.shape[0]
+
     if np.isscalar(dt_or_t):
-        t = np.arange(len(x))*dt_or_t
+        t = np.arange(n) * 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 n != 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
-
-    dspline = spline.derivative()
-    dxdt_hat = dspline(t)
-
-    return x_hat, dxdt_hat
+    x_hat = np.empty_like(x)
+    dxdt_hat = np.empty_like(x)
+
+    for idx in np.ndindex(x.shape[1:]):
+        sl = (slice(None),) + idx
+        xi = x[sl]
+        for _ in range(num_iterations):
+            obs = ~np.isnan(xi) # UnivariateSpline can't handle NaN, so fit only on observed points
+            spline = scipy.interpolate.UnivariateSpline(t[obs], xi[obs], k=degree, s=s)
+            xi = spline(t) # evaluate at all t, filling in NaN positions by interpolation
+        dspline = spline.derivative()
+        x_hat[sl] = xi
+        dxdt_hat[sl] = dspline(t)
+
+    return np.moveaxis(x_hat, 0, axis), np.moveaxis(dxdt_hat, 0, axis)
 
 
 def polydiff(x, dt_or_t, params=None, options=None, degree=None, window_size=None, step_size=1,

pynumdiff/tests/test_diff_methods.py

@@ -326,7 +326,8 @@ def test_diff_method(diff_method_and_params, test_func_and_deriv, request): # re
     (butterdiff, {'filter_order': 3, 'cutoff_freq': 1 - 1e-6}),
     (finitediff, {}),
     (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': 5, 's': 2}),
 ]
 
 # Similar to the error_bounds table, index by method first. But then we test against only one 2D function,
@@ -338,7 +339,8 @@ def test_diff_method(diff_method_and_params, test_func_and_deriv, request): # re
     butterdiff: [(0, -1), (1, -1)],
     finitediff: [(0, -1), (1, -1)],
     savgoldiff: [(0, -1), (1, 1)],
-    rtsdiff: [(1, -1), (1, 0)]
+    rtsdiff: [(1, -1), (1, 0)],
+    splinediff: [(3, 1), (3, 2)],
 }
 
 @mark.parametrize("multidim_method_and_params", multidim_methods_and_params)