duplicate idl median function

pydl/median.py

@@ -0,0 +1,75 @@
+"""This module corresponds to the math directory in idlutils.
+"""
+import numpy as np
+
+
+def median(array, width=None, axis=None, even=False):
+    """Wrap medfilt so that the results more closely resemble IDL MEDIAN().
+
+    Parameters
+    ----------
+    array : array-like
+        Compute the median of this array.
+    width : :class:`int`, optional
+        Size of the neighborhood in which to compute the median (*i.e.*,
+        perform median filtering).  If omitted, the median of the whole
+        array is returned.
+    axis : :class:`int`, optional
+        Compute the median over this axis for a multi-dimensional array.  If
+        ommitted, the median over the entire array will be returned.  If
+        set, this function will behave as though `even` is ``True``.
+    even : :class:`bool`, optional
+        If set to ``True``, the median of arrays with an even number of elements
+        will be the average of the middle two values.
+
+    Returns
+    -------
+    median : array-like
+        The median of the array.
+
+    Raises
+    ------
+    ValueError
+        If the input `array` is not 1 or 2 dimensional.
+
+    Notes
+    -----
+    * For arrays with an even number of elements, the :func:`numpy.median`
+      function behaves like ``MEDIAN(array, /EVEN)``, so the absence of
+      the `even` keyword has to turn *off* that behavior.
+    * For median filtering, this uses :func:`~scipy.signal.medfilt` under the
+      hood, but patches up the values on the array boundaries to match the
+      return values of the IDL MEDIAN() function.
+    """
+    from scipy.signal import medfilt, medfilt2d
+    if width is None:
+        if axis is None:
+            f = array.flatten()
+            if f.size % 2 == 1 or even:
+                return np.median(array)
+            else:
+                i = f.argsort()
+                return f[i[f.size/2]]
+        else:
+            return np.median(array, axis=axis)
+    else:
+        if array.ndim == 1:
+            medarray = medfilt(array, min(width, array.size))
+            istart = int((width - 1)/2)
+            iend = array.size - int((width + 1)/2)
+            i = np.arange(array.size)
+            w = (i < istart) | (i > iend)
+            medarray[w] = array[w]
+            return medarray
+        elif array.ndim == 2:
+            medarray = medfilt2d(array, min(width, array.size))
+            istart = int((width-1)/2)
+            iend = (array.shape[0] - int((width+1)/2), array.shape[1] - int((width+1)/2))
+            i = np.arange(array.shape[0])
+            j = np.arange(array.shape[1])
+            w = ((i < istart) | (i > iend[0]), (j < istart) | (j > iend[1]))
+            medarray[w[0], :] = array[w[0], :]
+            medarray[:, w[1]] = array[:, w[1]]
+            return medarray
+        else:
+            raise ValueError("Invalid number of dimensions for input array!")

pydl/pydlutils/math.py

@@ -108,33 +108,6 @@ def var(self):
         return np.diag(self.covar)
 
 
-def _fix_medfilt(array, width):
-    """Wrap medfilt so that the results more closely resemble IDL MEDIAN().
-    """
-    from scipy.signal import medfilt
-    medarray = medfilt(array, min(width, array.size))
-    istart = int((width - 1)/2)
-    iend = array.size - int((width + 1)/2)
-    i = np.arange(array.size)
-    w = (i < istart) | (i > iend)
-    medarray[w] = array[w]
-    return medarray
-
-
-def _fix_medfilt2d(array, width):
-    """Wrap medfilt2d so that the results more closely resemble IDL MEDIAN().
-    """
-    from scipy.signal import medfilt2d
-    medarray = medfilt2d(array, min(width, array.size))
-    istart = int((width-1)/2)
-    iend = (array.shape[0] - int((width+1)/2), array.shape[1] - int((width+1)/2))
-    i = np.arange(array.shape[0])
-    j = np.arange(array.shape[1])
-    w = ((i < istart) | (i > iend[0]), (j < istart) | (j > iend[1]))
-    medarray[w[0], w[1]] = array[w[0], w[1]]
-    return medarray
-
-
 def djs_median(array, dimension=None, width=None, boundary='none'):
     """Compute the median of an array.
 
@@ -165,6 +138,7 @@ def djs_median(array, dimension=None, width=None, boundary='none'):
         then the result simply ``numpy.median(array,dimension)``.
         If `width` is set, the result has the same shape as the input array.
     """
+    from ..median import median
     if dimension is None and width is None:
         return np.median(array)
     elif width is None:
@@ -174,9 +148,9 @@ def djs_median(array, dimension=None, width=None, boundary='none'):
             return array
         if boundary == 'none':
             if array.ndim == 1:
-                return _fix_medfilt(array, width)
+                return median(array, width)
             elif array.ndim == 2:
-                return _fix_medfilt2d(array, width)
+                return median(array, width)
             else:
                 raise ValueError('Unsupported number of dimensions with ' +
                                 'this boundary condition.')
@@ -188,7 +162,7 @@ def djs_median(array, dimension=None, width=None, boundary='none'):
                 bigarr[0:padsize] = array[0:padsize][::-1]
                 bigarr[padsize+array.shape[0]:padsize*2+array.shape[0]] = (
                         array[array.shape[0]-padsize:array.shape[0]][::-1])
-                f = _fix_medfilt(bigarr, width)
+                f = median(bigarr, width)
                 medarray = f[padsize:padsize+array.shape[0]]
                 return medarray
             elif array.ndim == 2:
@@ -210,7 +184,7 @@ def djs_median(array, dimension=None, width=None, boundary='none'):
                 bigarr[bigarr.shape[0]-padsize:bigarr.shape[0], 0:padsize] = array[array.shape[0]-padsize:array.shape[0], 0:padsize][::-1, ::-1]
                 # Copy into bottom right
                 bigarr[bigarr.shape[0]-padsize:bigarr.shape[0], bigarr.shape[1]-padsize:bigarr.shape[1]] = array[array.shape[0]-padsize:array.shape[0], array.shape[1]-padsize:array.shape[1]][::-1, ::-1]
-                f = _fix_medfilt2d(bigarr, min(width, array.size))
+                f = median(bigarr, min(width, array.size))
                 medarray = f[padsize:array.shape[0]+padsize, padsize:array.shape[1]+padsize]
                 return medarray
             else:

pydl/tests/test_pydl.py

@@ -9,6 +9,7 @@
 from astropy.tests.helper import raises
 from os.path import basename, dirname, join
 from ..file_lines import file_lines
+from ..median import median
 from ..pcomp import pcomp
 from ..smooth import smooth
 from ..uniq import uniq
@@ -54,6 +55,22 @@ def test_file_lines(self):
         n = file_lines(join(self.data_dir, 'this-file-is-empty.txt'))
         assert n == 0
 
+    def test_median(self):
+        odd_data = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13], dtype=np.float32)
+        even_data = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], dtype=np.float32)
+        assert median(odd_data) == 7
+        assert median(odd_data, even=True) == 7
+        assert median(even_data) == 7
+        assert median(even_data, even=True) == 6.5
+        assert (median(odd_data, 3) == odd_data).all()
+        with raises(ValueError):
+            foo = median(np.ones((9,9,9)),3)
+        odd_data2 = np.vstack((odd_data, odd_data, odd_data, odd_data, odd_data))
+        assert (median(odd_data2, 3) == odd_data2).all()
+        assert (median(odd_data2, axis=0) == odd_data).all()
+        assert (median(odd_data2, axis=1) ==
+            7*np.ones((odd_data2.shape[0],), dtype=odd_data2.dtype)).all()
+
     def test_pcomp(self):
         test_data_file = join(self.data_dir, 'pcomp_data.txt')
         test_data = np.loadtxt(test_data_file, dtype='d', delimiter=',')