Merge pull request #193 from florisvb/fix/191-splinediff-nan

Extend splinediff to handle missing data (NaN), add missing-data tests

Author: pavelkomarov

pynumdiff/polynomial_fit.py

@@ -10,7 +10,8 @@ def splinediff(x, dt_or_t, params=None, options=None, degree=3, s=None, num_iter
     """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.
 
-    :param np.array[float] x: data to differentiate
+    :param np.array[float] x: data to differentiate. May contain NaN values (missing data); NaNs are excluded from
+        fitting and imputed by spline interpolation.
     :param float or array[float] dt_or_t: This function supports variable step size. This parameter is either the constant
         :math:`\\Delta t` if given as a single float, or data locations if given as an array of same length as :code:`x`.
     :param list params: (**deprecated**, prefer :code:`degree`, :code:`cutoff_freq`, and :code:`num_iterations`)
@@ -38,8 +39,9 @@ def splinediff(x, dt_or_t, params=None, options=None, degree=3, s=None, num_iter
 
     x_hat = x
     for _ in range(num_iterations):
-        spline = scipy.interpolate.UnivariateSpline(t, x_hat, k=degree, s=s)
-        x_hat = spline(t)
+        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)

pynumdiff/tests/test_diff_methods.py

@@ -325,7 +325,8 @@ def test_diff_method(diff_method_and_params, test_func_and_deriv, request): # re
     (kerneldiff, {'kernel': 'gaussian', 'window_size': 5}),
     (butterdiff, {'filter_order': 3, 'cutoff_freq': 1 - 1e-6}),
     (finitediff, {}),
-    (savgoldiff, {'degree': 3, 'window_size': 11, 'smoothing_win': 3})
+    (savgoldiff, {'degree': 3, 'window_size': 11, 'smoothing_win': 3}),
+    (rtsdiff, {'order':2, 'log_qr_ratio':7, 'forwardbackward':True})
 ]
 
 # Similar to the error_bounds table, index by method first. But then we test against only one 2D function,
@@ -336,7 +337,8 @@ def test_diff_method(diff_method_and_params, test_func_and_deriv, request): # re
     kerneldiff: [(2, 1), (3, 2)],
     butterdiff: [(0, -1), (1, -1)],
     finitediff: [(0, -1), (1, -1)],
-    savgoldiff: [(0, -1), (1, 1)]
+    savgoldiff: [(0, -1), (1, 1)],
+    rtsdiff: [(1, -1), (1, 0)]
 }
 
 @mark.parametrize("multidim_method_and_params", multidim_methods_and_params)
@@ -390,3 +392,23 @@ def test_multidimensionality(multidim_method_and_params, request):
         ax3.plot_wireframe(T1, T2, computed_laplacian, label='computed')
         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)
+
+# List of methods that can handle missing values
+nan_methods_and_params = [
+    (splinediff, {'degree': 5, 's': 2}),
+    (polydiff, {'degree': 2, 'window_size': 3}),
+    (rtsdiff, {'order': 2, 'log_qr_ratio': 7, 'forwardbackward': True}),
+    (robustdiff, {'order': 3, 'log_q': 7, 'log_r': 2}),
+]
+
+@mark.parametrize("diff_method_and_params", nan_methods_and_params)
+def test_missing_data(diff_method_and_params):
+    """Ensure methods that support missing data return finite outputs when NaN values are present"""
+    diff_method, params = diff_method_and_params
+
+    x_nan = np.sin(t)
+    x_nan[[5, 10, 15, 20]] = np.nan # introduce missing data at several locations
+    x_hat, dxdt_hat = diff_method(x_nan, dt, **params)
+    
+    assert np.all(np.isfinite(x_hat))
+    assert np.all(np.isfinite(dxdt_hat))