Add waveletdiff: wavelet-based denoising differentiator

- Implements Donoho & Johnstone (1994) universal threshold estimator
- Vectorised thresholding over multi-column inputs
- Supports variable step size via np.gradient
- Supports arbitrary axis via ascontiguousarray + moveaxis
- Adds pywt dependency

Author: mariaprot

pynumdiff/__init__.py

@@ -15,6 +15,6 @@
 from .finite_difference import finitediff, first_order, second_order, fourth_order
 from .smooth_finite_difference import kerneldiff, meandiff, mediandiff, gaussiandiff, friedrichsdiff, butterdiff
 from .polynomial_fit import splinediff, polydiff, savgoldiff
-from .basis_fit import spectraldiff, rbfdiff
+from .basis_fit import spectraldiff, rbfdiff, waveletdiff
 from .total_variation_regularization import iterative_velocity
 from .kalman_smooth import kalman_filter, rts_smooth, rtsdiff, constant_velocity, constant_acceleration, constant_jerk

pynumdiff/basis_fit.py

@@ -2,6 +2,7 @@
 from warnings import warn
 import numpy as np
 from scipy import sparse
+import pywt
 
 from pynumdiff.utils import utility
 
@@ -137,3 +138,125 @@ def rbfdiff(x, dt_or_t, sigma=1, lmbd=0.01):
     alpha = sparse.linalg.spsolve(rbf_regularized, x) # solve sparse system targeting the noisy data, O(N sigma^2)
 
     return rbf @ alpha, drbfdt @ alpha # find samples of reconstructions using the smooth bases
+
+
+def waveletdiff(x, dt_or_t, wavelet='db4', level=None, threshold=1.0, axis=0, mode='periodization'):
+    """Smooth and differentiate noisy data via discrete wavelet denoising.
+
+    Decomposes x into wavelet detail and approximation coefficients, soft-thresholds
+    the detail coefficients to remove noise using the Donoho & Johnstone (1994)
+    universal threshold estimator, reconstructs a smoothed signal, then
+    differentiates with finite differences via np.gradient.
+
+    :param np.array x: data to differentiate
+    :param float or array dt_or_t: scalar dt or array of sample times. If an
+        array is provided it is passed directly to np.gradient, giving correct
+        results for non-uniformly sampled data.
+    :param str wavelet: PyWavelets wavelet name, e.g. 'db4', 'sym4', 'coif2'.
+        'db4' is a solid general-purpose default. Biorthogonal wavelets such as
+        'bior2.2' or 'bior4.4' are symmetric and designed for smooth reconstruction
+        but may need a lower threshold value.
+    :param int level: decomposition depth. None (default) resolves to
+        min(pywt.dwt_max_level(N, wavelet), 5) to avoid over-decomposing short
+        signals. Increase for heavily oversampled data.
+    :param float threshold: soft-thresholding scale factor in [0, inf).
+        Multiplies the universal threshold sigma * sqrt(2 * log(N)).
+        threshold=1.0 is the classical Donoho & Johnstone universal threshold
+        and is the recommended starting point. Values < 1.0 give less smoothing;
+        values > 1.0 give more aggressive smoothing. This parameter maps onto
+        tvgamma in the pynumdiff.optimize framework.
+    :param int axis: axis along which to differentiate (default 0).
+    :param str mode: PyWavelets signal extension mode passed to wavedec/waverec.
+        'periodization' (default) keeps coefficient arrays exactly length N and
+        is the most numerically stable choice for differentiation. 'reflect' is
+        a good alternative for clearly non-periodic signals.
+        See pywt.Modes.modes for all options.
+    :return: - **x_hat** (np.array) -- estimated (smoothed) x
+             - **dxdt_hat** (np.array) -- estimated derivative of x
+    """
+    N = x.shape[axis]
+
+    # Axis normalisation — bring target axis to front.
+    # Skip moveaxis when axis is already 0 to avoid an unnecessary allocation.
+    # When we do move, call ascontiguousarray immediately so the subsequent
+    # reshape is guaranteed zero-copy.
+    if axis == 0:
+        x_work = x if x.flags['C_CONTIGUOUS'] else np.ascontiguousarray(x)
+    else:
+        x_work = np.ascontiguousarray(np.moveaxis(x, axis, 0))
+
+    shape = x_work.shape
+    x_flat = x_work.reshape(N, -1)  # (N, M) contiguous, no hidden copy
+    M = x_flat.shape[1]
+
+    if np.isscalar(dt_or_t):
+        grad_arg = dt_or_t
+    else:
+        if len(dt_or_t) != N:
+            raise ValueError(
+                "`dt_or_t` array must have the same length as x along `axis`."
+            )
+        grad_arg = dt_or_t  # np.gradient accepts a full coordinate array
+
+    # Conservative level default avoids over-decomposing short signals
+    # (pywt default uses the maximum possible level).
+    if level is None:
+        max_level = pywt.dwt_max_level(N, wavelet)
+        level = min(max_level, 5)
+
+    # Decompose all columns and stack coefficients into 2-D arrays of shape
+    # (coeff_len_i, M). Probing column 0 first lets us pre-allocate correctly;
+    # the probe result is reused for col 0 so we pay N+1 wavedec calls total.
+    _probe = pywt.wavedec(x_flat[:, 0], wavelet, level=level, mode=mode)
+    coeff_lengths = [len(c) for c in _probe]
+    n_levels = len(_probe)
+
+    coeffs_all = [
+        np.empty((coeff_lengths[i], M), dtype=x_flat.dtype)
+        for i in range(n_levels)
+    ]
+    for i, c in enumerate(_probe):
+        coeffs_all[i][:, 0] = c
+
+    for col in range(1, M):
+        for i, c in enumerate(
+            pywt.wavedec(x_flat[:, col], wavelet, level=level, mode=mode)
+        ):
+            coeffs_all[i][:, col] = c
+
+    # Vectorised noise estimation and soft-thresholding over all columns at once.
+    # sigma: robust MAD estimator from finest detail level, shape (M,).
+    # thresh: per-column universal threshold, shape (M,).
+    # Approximation coefficients (index 0) are left untouched; only detail
+    # levels (indices 1..n_levels-1) are thresholded.
+    sigma = np.median(np.abs(coeffs_all[-1]), axis=0) / 0.6745
+    np.maximum(sigma, 1e-10, out=sigma)  # floor avoids zero threshold on clean signals
+
+    thresh = threshold * sigma * np.sqrt(2 * np.log(N))  # shape (M,)
+
+    coeffs_denoised = [coeffs_all[0]] + [
+        pywt.threshold(c, thresh[np.newaxis, :], mode='soft')
+        for c in coeffs_all[1:]
+    ]
+
+    # Reconstruct and differentiate — pywt.waverec is 1-D only so a column
+    # loop remains, but all Python-level arithmetic has been moved out above.
+    x_hat_flat    = np.empty_like(x_flat)
+    dxdt_hat_flat = np.empty_like(x_flat)
+
+    for col in range(M):
+        col_coeffs = [coeffs_denoised[i][:, col] for i in range(n_levels)]
+        x_hat_col = pywt.waverec(col_coeffs, wavelet, mode=mode)[:N]
+        x_hat_flat[:, col]    = x_hat_col
+        dxdt_hat_flat[:, col] = np.gradient(x_hat_col, grad_arg)
+
+    # Restore original shape and axis order.
+    # moveaxis on the way out is only needed when we moved on the way in.
+    x_hat    = x_hat_flat.reshape(shape)
+    dxdt_hat = dxdt_hat_flat.reshape(shape)
+
+    if axis != 0:
+        x_hat    = np.moveaxis(x_hat,    0, axis)
+        dxdt_hat = np.moveaxis(dxdt_hat, 0, axis)
+
+    return x_hat, dxdt_hat

pynumdiff/tests/test_diff_methods.py

@@ -5,7 +5,7 @@
 from ..smooth_finite_difference import kerneldiff, mediandiff, meandiff, gaussiandiff, friedrichsdiff, butterdiff
 from ..finite_difference import finitediff, first_order, second_order, fourth_order
 from ..polynomial_fit import polydiff, savgoldiff, splinediff
-from ..basis_fit import spectraldiff, rbfdiff
+from ..basis_fit import spectraldiff, rbfdiff, waveletdiff
 from ..total_variation_regularization import velocity, acceleration, jerk, iterative_velocity, smooth_acceleration
 from ..kalman_smooth import rtsdiff, constant_velocity, constant_acceleration, constant_jerk, robustdiff
 from ..linear_model import lineardiff
@@ -47,6 +47,8 @@ def spline_irreg_step(*args, **kwargs): return splinediff(*args, **kwargs)
     (spline_irreg_step, {'degree':5, 's':2}),
     (spectraldiff, {'high_freq_cutoff':0.2}), (spectraldiff, [0.2]),
     (rbfdiff, {'sigma':0.5, 'lmbd':0.001}),
+    (waveletdiff, {'wavelet':'db4', 'threshold':1.0}),
+    (waveletdiff, {'wavelet':'db4', 'threshold':1.0}),
     (constant_velocity, {'r':1e-2, 'q':1e3}), (constant_velocity, [1e-2, 1e3]),
     (constant_acceleration, {'r':1e-3, 'q':1e4}), (constant_acceleration, [1e-3, 1e4]),
     (constant_jerk, {'r':1e-4, 'q':1e5}), (constant_jerk, [1e-4, 1e5]),
@@ -162,6 +164,12 @@ def spline_irreg_step(*args, **kwargs): return splinediff(*args, **kwargs)
               [(-2, -2), (0, 0), (0, -1), (0, 0)],
               [(0, 0), (2, 2), (0, 0), (2, 2)],
               [(1, 1), (3, 3), (1, 1), (3, 3)]],
+    waveletdiff: [[(-14, -15), (-14, -14), (-1, -1), (0, 0)],
+                  [(-9, -9), (-8, -8), (0, 0), (1, 1)],
+                  [(-9, -9), (0, 0), (0, 0), (1, 1)],
+                  [(-1, -1), (0, 0), (0, 0), (1, 1)],
+                  [(1, 0), (2, 2), (1, 1), (2, 2)],
+                  [(0, 0), (3, 3), (1, 0), (3, 3)]], 
     velocity: [[(-25, -25), (-18, -19), (0, -1), (1, 0)],
                [(-12, -12), (-11, -12), (-1, -1), (-1, -2)],
                [(0, -1), (1, 0), (0, -1), (1, 0)],
@@ -308,7 +316,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}),
+    (waveletdiff, {'wavelet': 'db4', 'threshold': 1.0}),
 ]
 
 # Similar to the error_bounds table, index by method first. But then we test against only one 2D function,
@@ -319,7 +328,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)],
+    waveletdiff: [(1, 0), (2, 1)],
 }
 
 @mark.parametrize("multidim_method_and_params", multidim_methods_and_params)
@@ -372,4 +382,4 @@ def test_multidimensionality(multidim_method_and_params, request):
         ax2.plot_wireframe(T1, T2, computed_d2)
         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)
+        fig.suptitle(f'{diff_method.__name__}', fontsize=16)