codekingpro/portable-devtools

1"""
2Discrete Fourier Transforms - _helper.py
3
4"""
5from numpy._core import arange, asarray, empty, integer, roll
6from numpy._core.overrides import array_function_dispatch, set_module
7
8# Created by Pearu Peterson, September 2002
9
10__all__ = ['fftshift', 'ifftshift', 'fftfreq', 'rfftfreq']
11
12integer_types = (int, integer)
13
14
15def _fftshift_dispatcher(x, axes=None):
16    return (x,)
17
18
19@array_function_dispatch(_fftshift_dispatcher, module='numpy.fft')
20def fftshift(x, axes=None):
21    """
22    Shift the zero-frequency component to the center of the spectrum.
23
24    This function swaps half-spaces for all axes listed (defaults to all).
25    Note that ``y[0]`` is the Nyquist component only if ``len(x)`` is even.
26
27    Parameters
28    ----------
29    x : array_like
30        Input array.
31    axes : int or shape tuple, optional
32        Axes over which to shift.  Default is None, which shifts all axes.
33
34    Returns
35    -------
36    y : ndarray
37        The shifted array.
38
39    See Also
40    --------
41    ifftshift : The inverse of `fftshift`.
42
43    Examples
44    --------
45    >>> import numpy as np
46    >>> freqs = np.fft.fftfreq(10, 0.1)
47    >>> freqs
48    array([ 0.,  1.,  2., ..., -3., -2., -1.])
49    >>> np.fft.fftshift(freqs)
50    array([-5., -4., -3., -2., -1.,  0.,  1.,  2.,  3.,  4.])
51
52    Shift the zero-frequency component only along the second axis:
53
54    >>> freqs = np.fft.fftfreq(9, d=1./9).reshape(3, 3)
55    >>> freqs
56    array([[ 0.,  1.,  2.],
57           [ 3.,  4., -4.],
58           [-3., -2., -1.]])
59    >>> np.fft.fftshift(freqs, axes=(1,))
60    array([[ 2.,  0.,  1.],
61           [-4.,  3.,  4.],
62           [-1., -3., -2.]])
63
64    """
65    x = asarray(x)
66    if axes is None:
67        axes = tuple(range(x.ndim))
68        shift = [dim // 2 for dim in x.shape]
69    elif isinstance(axes, integer_types):
70        shift = x.shape[axes] // 2
71    else:
72        shift = [x.shape[ax] // 2 for ax in axes]
73
74    return roll(x, shift, axes)
75
76
77@array_function_dispatch(_fftshift_dispatcher, module='numpy.fft')
78def ifftshift(x, axes=None):
79    """
80    The inverse of `fftshift`. Although identical for even-length `x`, the
81    functions differ by one sample for odd-length `x`.
82
83    Parameters
84    ----------
85    x : array_like
86        Input array.
87    axes : int or shape tuple, optional
88        Axes over which to calculate.  Defaults to None, which shifts all axes.
89
90    Returns
91    -------
92    y : ndarray
93        The shifted array.
94
95    See Also
96    --------
97    fftshift : Shift zero-frequency component to the center of the spectrum.
98
99    Examples
100    --------
101    >>> import numpy as np
102    >>> freqs = np.fft.fftfreq(9, d=1./9).reshape(3, 3)
103    >>> freqs
104    array([[ 0.,  1.,  2.],
105           [ 3.,  4., -4.],
106           [-3., -2., -1.]])
107    >>> np.fft.ifftshift(np.fft.fftshift(freqs))
108    array([[ 0.,  1.,  2.],
109           [ 3.,  4., -4.],
110           [-3., -2., -1.]])
111
112    """
113    x = asarray(x)
114    if axes is None:
115        axes = tuple(range(x.ndim))
116        shift = [-(dim // 2) for dim in x.shape]
117    elif isinstance(axes, integer_types):
118        shift = -(x.shape[axes] // 2)
119    else:
120        shift = [-(x.shape[ax] // 2) for ax in axes]
121
122    return roll(x, shift, axes)
123
124
125@set_module('numpy.fft')
126def fftfreq(n, d=1.0, device=None):
127    """
128    Return the Discrete Fourier Transform sample frequencies.
129
130    The returned float array `f` contains the frequency bin centers in cycles
131    per unit of the sample spacing (with zero at the start).  For instance, if
132    the sample spacing is in seconds, then the frequency unit is cycles/second.
133
134    Given a window length `n` and a sample spacing `d`::
135
136      f = [0, 1, ...,   n/2-1,     -n/2, ..., -1] / (d*n)   if n is even
137      f = [0, 1, ..., (n-1)/2, -(n-1)/2, ..., -1] / (d*n)   if n is odd
138
139    Parameters
140    ----------
141    n : int
142        Window length.
143    d : scalar, optional
144        Sample spacing (inverse of the sampling rate). Defaults to 1.
145    device : str, optional
146        The device on which to place the created array. Default: ``None``.
147        For Array-API interoperability only, so must be ``"cpu"`` if passed.
148
149        .. versionadded:: 2.0.0
150
151    Returns
152    -------
153    f : ndarray
154        Array of length `n` containing the sample frequencies.
155
156    Examples
157    --------
158    >>> import numpy as np
159    >>> signal = np.array([-2, 8, 6, 4, 1, 0, 3, 5], dtype=float)
160    >>> fourier = np.fft.fft(signal)
161    >>> n = signal.size
162    >>> timestep = 0.1
163    >>> freq = np.fft.fftfreq(n, d=timestep)
164    >>> freq
165    array([ 0.  ,  1.25,  2.5 , ..., -3.75, -2.5 , -1.25])
166
167    """
168    if not isinstance(n, integer_types):
169        raise ValueError("n should be an integer")
170    val = 1.0 / (n * d)
171    results = empty(n, int, device=device)
172    N = (n - 1) // 2 + 1
173    p1 = arange(0, N, dtype=int, device=device)
174    results[:N] = p1
175    p2 = arange(-(n // 2), 0, dtype=int, device=device)
176    results[N:] = p2
177    return results * val
178
179
180@set_module('numpy.fft')
181def rfftfreq(n, d=1.0, device=None):
182    """
183    Return the Discrete Fourier Transform sample frequencies
184    (for usage with rfft, irfft).
185
186    The returned float array `f` contains the frequency bin centers in cycles
187    per unit of the sample spacing (with zero at the start).  For instance, if
188    the sample spacing is in seconds, then the frequency unit is cycles/second.
189
190    Given a window length `n` and a sample spacing `d`::
191
192      f = [0, 1, ...,     n/2-1,     n/2] / (d*n)   if n is even
193      f = [0, 1, ..., (n-1)/2-1, (n-1)/2] / (d*n)   if n is odd
194
195    Unlike `fftfreq` (but like `scipy.fftpack.rfftfreq`)
196    the Nyquist frequency component is considered to be positive.
197
198    Parameters
199    ----------
200    n : int
201        Window length.
202    d : scalar, optional
203        Sample spacing (inverse of the sampling rate). Defaults to 1.
204    device : str, optional
205        The device on which to place the created array. Default: ``None``.
206        For Array-API interoperability only, so must be ``"cpu"`` if passed.
207
208        .. versionadded:: 2.0.0
209
210    Returns
211    -------
212    f : ndarray
213        Array of length ``n//2 + 1`` containing the sample frequencies.
214
215    Examples
216    --------
217    >>> import numpy as np
218    >>> signal = np.array([-2, 8, 6, 4, 1, 0, 3, 5, -3, 4], dtype=float)
219    >>> fourier = np.fft.rfft(signal)
220    >>> n = signal.size
221    >>> sample_rate = 100
222    >>> freq = np.fft.fftfreq(n, d=1./sample_rate)
223    >>> freq
224    array([  0.,  10.,  20., ..., -30., -20., -10.])
225    >>> freq = np.fft.rfftfreq(n, d=1./sample_rate)
226    >>> freq
227    array([  0.,  10.,  20.,  30.,  40.,  50.])
228
229    """
230    if not isinstance(n, integer_types):
231        raise ValueError("n should be an integer")
232    val = 1.0 / (n * d)
233    N = n // 2 + 1
234    results = arange(0, N, dtype=int, device=device)
235    return results * val
236