ENH: Add rfftfreq() for numpy's rfft(), which behaves differently from scipy's rfft()/rfftfreq().

1 files changed, 49 insertions, 1 deletions

diff --git a/numpy/fft/helper.py b/numpy/fft/helper.py
index f6c570445..903ea257e 100644
--- a/numpy/fft/helper.py
+++ b/numpy/fft/helper.py
@@ -3,7 +3,7 @@ Discrete Fourier Transforms - helper.py
 """
 # Created by Pearu Peterson, September 2002
 
-__all__ = ['fftshift','ifftshift','fftfreq']
+__all__ = ['fftshift', 'ifftshift', 'fftfreq', 'rfftfreq']
 
 from numpy.core import asarray, concatenate, arange, take, \
     integer, empty
@@ -160,3 +160,51 @@ def fftfreq(n,d=1.0):
     results[N:] = p2
     return results * val
     #return hstack((arange(0,(n-1)/2 + 1), arange(-(n/2),0))) / (n*d)
+
+
+def rfftfreq(n, d=1.0):
+    """
+    Return the Discrete Fourier Transform sample frequencies 
+    (for usage with rfft, irfft).
+
+    The returned float array contains the frequency bins in
+    cycles/unit (with zero at the start) given a window length `n` and a
+    sample spacing `d`::
+
+      f = [0, 1, ...,     n/2-1,     n/2] / (d*n)   if n is even
+      f = [0, 1, ..., (n-1)/2-1, (n-1)/2] / (d*n)   if n is odd
+
+    Unlike `fftfreq` (but like `scipy.fftpack.rfftfreq`)
+    the Nyquist frequency component is considered to be positive.
+
+    Parameters
+    ----------
+    n : int
+        Window length.
+    d : scalar, optional
+        Sample spacing. Default is 1.
+
+    Returns
+    -------
+    out : ndarray
+        The array of length `n`, containing the sample frequencies.
+
+    Examples
+    --------
+    >>> signal = np.array([-2, 8, 6, 4, 1, 0, 3, 5, -3, 4], dtype=float)
+    >>> fourier = np.fft.rfft(signal)
+    >>> n = signal.size
+    >>> sample_rate = 100
+    >>> freq = np.fft.fftfreq(n, d=1./sample_rate)
+    >>> freq
+    array([  0.,  10.,  20.,  30.,  40., -50., -40., -30., -20., -10.])
+    >>> freq = np.fft.rfftfreq(n, d=1./sample_rate)
+    >>> freq
+    array([  0.,  10.,  20.,  30.,  40.,  50.])
+
+    """
+    assert isinstance(n,types.IntType) or isinstance(n, integer)
+    val = 1.0/(n*d)
+    N = n//2 + 1
+    results = arange(0, N, dtype=int)
+    return results * val