numpy · jsilter · Sep 9, 2011 · Sep 9, 2011 · Sep 9, 2011
diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py
@@ -2,7 +2,8 @@
 __all__ = ['select', 'piecewise', 'trim_zeros', 'copy', 'iterable',
         'percentile', 'diff', 'gradient', 'angle', 'unwrap', 'sort_complex',
         'disp', 'extract', 'place', 'nansum', 'nanmax', 'nanargmax',
-        'nanargmin', 'nanmin', 'vectorize', 'asarray_chkfinite', 'average',
+        'nanargmin', 'nanmin', 'argrelmax', 'argrelmin', 'vectorize',
+        'asarray_chkfinite', 'average',
         'histogram', 'histogramdd', 'bincount', 'digitize', 'cov', 'corrcoef',
         'msort', 'median', 'sinc', 'hamming', 'hanning', 'bartlett',
         'blackman', 'kaiser', 'trapz', 'i0', 'add_newdoc', 'add_docstring',
@@ -1647,6 +1648,90 @@ def nanargmax(a, axis=None):
     """
     return _nanop(np.argmax, -np.inf, a, axis)
 
+def _argrelextrema(data,comparator,axis=0,order=1,mode='clip'):
+
+    int_order = int(order)
+    if( (int_order != order) or (order < 1) ):
+        raise ValueError('Order must be an integer >= 1')
+
+    datalen = data.shape[axis]
+    locs = np.arange(0,datalen)
+
+    results = ones(data.shape,dtype=bool)
+    main = data.take(locs,axis=axis,mode=mode)
+    for shift in xrange(1,int_order+1):
+        plus = data.take(locs + shift,axis=axis,mode=mode)
+        minus = data.take(locs - shift,axis=axis,mode=mode)
+        results &= comparator(main,plus)
+        results &= comparator(main,minus)
+        if(~results.any()):
+            return np.array([],)*2
+    arglocs = np.where(results)
+    return arglocs
+
+def argrelmin(data,axis=0,order=1,mode='clip'):
+    """
+    Calculate the relative minima of `data`.
+
+    See also
+    --------
+    argmax : Similar function.  Please refer to `numpy.argmax` for detailed
+        documentation.
+    """
+    return _argrelextrema(data,np.less,axis,order,mode)
+
+def argrelmax(data,axis=0,order=1,mode='clip'):
+    """
+    Calculate the relative maxima of `data`.
+
+    Parameters
+    -----------
+    data: array-like
+    axis: integer, optional
+        axis over which to select data. Default is 0.
+    order: integer, optional
+        Default is 1
+    mode: string, optional
+        How the edges of the vector are treated.
+        'wrap' (wrap around) or 'clip' (treat overflow
+        as the same as the last (or first) element).
+        Default 'clip'. See numpy.take
+
+    Returns
+    ----------
+    maxima: array-like
+        Indices of the maxima, same format as
+        numpy.where
+
+    See also
+    --------
+    argrelmin : Similar function
+
+    Examples
+    ---------
+    >>>> totlen = 30
+    >>>> nx = 15
+    >>>> xinds = np.arange(0,np.pi,np.pi/nx)
+    >>>> sininds = np.arange(8,8+nx)
+    >>>> a = np.ones(totlen)
+    >>>> a[sininds] = 0.9*np.sin(xinds)
+    >>>> print np.argrelmax(a)
+    >>>> (array([16]),)
+
+    This function is 1 on each side, with a curve in the center.
+    The curve in the center represents a relative maximum,
+    and that is what is returned.
+
+    Notes
+    --------
+    Relative maxima are calculated by finding locations where
+    data[n] > data[n+1:n+order+1:+1] &
+    data[n] > data[n-1:n-order-1:-1]. 
+
+    """
+
+    return _argrelextrema(data,np.greater,axis,order,mode)
+
 def disp(mesg, device=None, linefeed=True):
     """
     Display a message on a device.
@@ -1697,6 +1782,7 @@ def _get_nargs(obj):
 
     terr = re.compile(r'.*? takes (exactly|at least) (?P<exargs>(\d+)|(\w+))' +
             r' argument(s|) \((?P<gargs>(\d+)|(\w+)) given\)')
+
     def _convert_to_int(strval):
         try:
             result = int(strval)

diff --git a/numpy/lib/tests/test_function_base.py b/numpy/lib/tests/test_function_base.py
@@ -802,6 +802,51 @@ def test_dtype_order(self):
         a = numpy.lib.asarray_chkfinite(a, order='F', dtype=numpy.float64)
         assert_(a.dtype == numpy.float64)
 
+class TestRelExtFuncts(TestCase):
+    def _gen_gaussians(self,center_locs,sigmas,total_length):
+        num_peaks = len(sigmas)
+        xdata = np.arange(0,total_length).astype(float)
+        out_data = np.zeros(total_length,dtype=float)
+        for ind,sigma in enumerate(sigmas):
+            tmp = (xdata - center_locs[ind])/sigma
+            out_data += np.exp(-(tmp**2))
+        return out_data
+
+    def _gen_gaussians_even(self,sigmas,total_length):
+        num_peaks = len(sigmas)
+        delta = total_length / num_peaks
+        center_locs = np.arange(1,total_length,delta)
+        out_data = self._gen_gaussians(center_locs,sigmas,total_length)
+        return out_data,center_locs
+
+    def test_argrelmax_highorder(self,order = 2):
+        sigmas = [1.0,2.0,10.0,5.0,15.0]
+        test_data,act_locs = self._gen_gaussians_even(sigmas,500)
+        test_data[act_locs + order] = test_data[act_locs]*0.99999
+        test_data[act_locs - order] = test_data[act_locs]*0.99999
+        rel_max_locs = argrelmax(test_data,order=2,mode='clip')[0]
+
+        assert_(len(rel_max_locs) == len(act_locs))
+        assert_((rel_max_locs == act_locs).all())
+
+    def test_argrelmax_2d_gaussians(self):
+        sigmas = [1.0,2.0,10.0]
+        test_data,act_locs = self._gen_gaussians_even(sigmas,100)
+        rot_factor = 20
+        rot_range = numpy.arange(0,len(test_data)) - rot_factor
+        test_data_2 = numpy.vstack([test_data,test_data[rot_range]])
+        rel_max_rows,rel_max_cols = argrelmax(test_data_2,axis=1,order=1.0)
+
+        for rw in xrange(0,test_data_2.shape[0]):
+            inds = (rel_max_rows == rw)
+            assert_(len(rel_max_cols[inds]) == len(act_locs))
+            assert_((act_locs == (rel_max_cols[inds]-rot_factor*rw)).all())
+
+    def test_argrelextrama_badorder(self):
+        data = np.arange(0,10)
+        comparator = np.greater
+        assert_raises(ValueError,argrelmax,data,comparator,order = 0)
+        assert_raises(ValueError,argrelmax,data,comparator,order = 0.5)
 
 class TestNaNFuncts(TestCase):
     def setUp(self):