add estimate_sigma function to threshold features

pywt/_thresholding.py

@@ -11,7 +11,10 @@
 from __future__ import division, print_function, absolute_import
 import numpy as np
 
-__all__ = ['threshold', 'threshold_firm']
+from ._multidim import dwtn
+
+
+__all__ = ['threshold', 'threshold_firm', 'estimate_sigma']
 
 
 def soft(data, value, substitute=0):
@@ -245,3 +248,77 @@ def threshold_firm(data, value_low, value_high):
     if np.any(large_vals[0]):
         thresholded[large_vals] = data[large_vals]
     return thresholded
+
+
+def estimate_sigma(data):
+    """
+    Robust wavelet-based estimator of the (Gaussian) noise standard deviation.
+    Parameters
+    ----------
+    data : ndarray
+        The data used to estimate sigma.
+
+    Returns
+    -------
+    sigma : float
+        Estimated noise standard deviation.
+
+    Notes
+    -----
+    This function assumes the noise follows a Gaussian distribution. The
+    estimation algorithm is based on the median absolute deviation of the
+    wavelet detail coefficients as described in section 4.2 of [1]_.
+
+    References
+    ----------
+    .. [1] D. L. Donoho and I. M. Johnstone. "Ideal spatial adaptation
+       by wavelet shrinkage." Biometrika 81.3 (1994): 425-455.
+       DOI:10.1093/biomet/81.3.425
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> import pywt
+    >>> data = np.sin(np.linspace(0,10,1000))
+    >>> np.random.seed(42)
+    >>> noise = np.random.normal(0,1,1000)
+    >>> sigma = pywt.estimate_sigma(data + noise)
+    """
+
+    coeffs = dwtn(data, wavelet='db2')
+    detail_coeffs = coeffs['d' * data.ndim]
+    return _sigma_est_dwt(detail_coeffs, distribution='Gaussian')
+
+
+def _sigma_est_dwt(detail_coeffs, distribution='Gaussian'):
+    """Calculate the robust median estimator of the noise standard deviation.
+    Parameters
+    ----------
+    detail_coeffs : ndarray
+        The detail coefficients corresponding to the discrete wavelet
+        transform of an image.
+    distribution : str
+        The underlying noise distribution.
+    Returns
+    -------
+    sigma : float
+        The estimated noise standard deviation (see section 4.2 of [1]_).
+    References
+    ----------
+    .. [1] D. L. Donoho and I. M. Johnstone. "Ideal spatial adaptation
+       by wavelet shrinkage." Biometrika 81.3 (1994): 425-455.
+       DOI:10.1093/biomet/81.3.425
+    """
+    # Consider regions with detail coefficients exactly zero to be masked out
+    detail_coeffs = detail_coeffs[np.nonzero(detail_coeffs)]
+
+    if distribution.lower() == 'gaussian':
+        # 75th quantile of the underlying, symmetric noise distribution
+        # denom = scipy.stats.norm.ppf(0.75)
+        # magic number to fill in because no scipy
+        denom = 0.6744897501960817
+        sigma = np.median(np.abs(detail_coeffs)) / denom
+    else:
+        raise ValueError("Only Gaussian noise estimation is currently "
+                         "supported")
+    return sigma

pywt/tests/test_thresholding.py

@@ -170,5 +170,14 @@ def test_threshold_firm():
         assert_(np.all(mt_abs_firm > np.abs(d_soft[mt])))
 
 
+def test_estimate_sigma():
+    sigma_test_data = np.sin(np.linspace(0,10,1000))
+    np.random.seed(42)
+    sigma_test_noise = np.random.normal(0,1,1000)
+    # TODO: get rid of magic number for test...
+    # can test by creating a list of noise arrays to pass in with varying sigmas
+    assert_allclose(pywt.estimate_sigma(sigma_test_data + sigma_test_noise),
+                 0.9736668419858439)
+
 if __name__ == '__main__':
     run_module_suite()