1"""

2 Discrete Fourier Transforms - helper.py

3

4"""

5 from numpy.core import integer, empty, arange, asarray, roll

6 from 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

12 integer_types = (int, integer)

13

14

15 def _fftshift_dispatcher(x, axes=None):

16 return (x,)

17

18

19@array_function_dispatch(_fftshift_dispatcher, module='numpy.fft')

20 def 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 >>> freqs = np.fft.fftfreq(10, 0.1)

46 >>> freqs

47 array([ 0., 1., 2., ..., -3., -2., -1.])

48 >>> np.fft.fftshift(freqs)

49 array([-5., -4., -3., -2., -1., 0., 1., 2., 3., 4.])

50

51 Shift the zero-frequency component only along the second axis:

52

53 >>> freqs = np.fft.fftfreq(9, d=1./9).reshape(3, 3)

54 >>> freqs

55 array([[ 0., 1., 2.],

56 [ 3., 4., -4.],

57 [-3., -2., -1.]])

58 >>> np.fft.fftshift(freqs, axes=(1,))

59 array([[ 2., 0., 1.],

60 [-4., 3., 4.],

61 [-1., -3., -2.]])

62

63 """

64 x = asarray(x)

65 if axes is None:

66 axes = tuple(range(x.ndim))

67 shift = [dim // 2 for dim in x.shape]

68 elif isinstance(axes, integer_types):

69 shift = x.shape[axes] // 2

70 else:

71 shift = [x.shape[ax] // 2 for ax in axes]

72

73 return roll(x, shift, axes)

74

75

76@array_function_dispatch(_fftshift_dispatcher, module='numpy.fft')

77 def ifftshift(x, axes=None):

78 """

79 The inverse of `fftshift`. Although identical for even-length `x`, the

80 functions differ by one sample for odd-length `x`.

81

82 Parameters

83 ----------

84 x : array_like

85 Input array.

86 axes : int or shape tuple, optional

87 Axes over which to calculate. Defaults to None, which shifts all axes.

88

89 Returns

90 -------

91 y : ndarray

92 The shifted array.

93

94 See Also

95 --------

96 fftshift : Shift zero-frequency component to the center of the spectrum.

97

98 Examples

99 --------

100 >>> freqs = np.fft.fftfreq(9, d=1./9).reshape(3, 3)

101 >>> freqs

102 array([[ 0., 1., 2.],

103 [ 3., 4., -4.],

104 [-3., -2., -1.]])

105 >>> np.fft.ifftshift(np.fft.fftshift(freqs))

106 array([[ 0., 1., 2.],

107 [ 3., 4., -4.],

108 [-3., -2., -1.]])

109

110 """

111 x = asarray(x)

112 if axes is None:

113 axes = tuple(range(x.ndim))

114 shift = [-(dim // 2) for dim in x.shape]

115 elif isinstance(axes, integer_types):

116 shift = -(x.shape[axes] // 2)

117 else:

118 shift = [-(x.shape[ax] // 2) for ax in axes]

119

120 return roll(x, shift, axes)

121

122

123@set_module('numpy.fft')

124 def fftfreq(n, d=1.0):

125 """

126 Return the Discrete Fourier Transform sample frequencies.

127

128 The returned float array `f` contains the frequency bin centers in cycles

129 per unit of the sample spacing (with zero at the start). For instance, if

130 the sample spacing is in seconds, then the frequency unit is cycles/second.

131

132 Given a window length `n` and a sample spacing `d`::

133

134 f = [0, 1, ..., n/2-1, -n/2, ..., -1] / (d*n) if n is even

135 f = [0, 1, ..., (n-1)/2, -(n-1)/2, ..., -1] / (d*n) if n is odd

136

137 Parameters

138 ----------

139 n : int

140 Window length.

141 d : scalar, optional

142 Sample spacing (inverse of the sampling rate). Defaults to 1.

143

144 Returns

145 -------

146 f : ndarray

147 Array of length `n` containing the sample frequencies.

148

149 Examples

150 --------

151 >>> signal = np.array([-2, 8, 6, 4, 1, 0, 3, 5], dtype=float)

152 >>> fourier = np.fft.fft(signal)

153 >>> n = signal.size

154 >>> timestep = 0.1

155 >>> freq = np.fft.fftfreq(n, d=timestep)

156 >>> freq

157 array([ 0. , 1.25, 2.5 , ..., -3.75, -2.5 , -1.25])

158

159 """

160 if not isinstance(n, integer_types):

161 raise ValueError("n should be an integer")

162 val = 1.0 / (n * d)

163 results = empty(n, int)

164 N = (n-1)//2 + 1

165 p1 = arange(0, N, dtype=int)

166 results[:N] = p1

167 p2 = arange(-(n//2), 0, dtype=int)

168 results[N:] = p2

169 return results * val

170

171

172@set_module('numpy.fft')

173 def rfftfreq(n, d=1.0):

174 """

175 Return the Discrete Fourier Transform sample frequencies

176 (for usage with rfft, irfft).

177

178 The returned float array `f` contains the frequency bin centers in cycles

179 per unit of the sample spacing (with zero at the start). For instance, if

180 the sample spacing is in seconds, then the frequency unit is cycles/second.

181

182 Given a window length `n` and a sample spacing `d`::

183

184 f = [0, 1, ..., n/2-1, n/2] / (d*n) if n is even

185 f = [0, 1, ..., (n-1)/2-1, (n-1)/2] / (d*n) if n is odd

186

187 Unlike `fftfreq` (but like `scipy.fftpack.rfftfreq`)

188 the Nyquist frequency component is considered to be positive.

189

190 Parameters

191 ----------

192 n : int

193 Window length.

194 d : scalar, optional

195 Sample spacing (inverse of the sampling rate). Defaults to 1.

196

197 Returns

198 -------

199 f : ndarray

200 Array of length ``n//2 + 1`` containing the sample frequencies.

201

202 Examples

203 --------

204 >>> signal = np.array([-2, 8, 6, 4, 1, 0, 3, 5, -3, 4], dtype=float)

205 >>> fourier = np.fft.rfft(signal)

206 >>> n = signal.size

207 >>> sample_rate = 100

208 >>> freq = np.fft.fftfreq(n, d=1./sample_rate)

209 >>> freq

210 array([ 0., 10., 20., ..., -30., -20., -10.])

211 >>> freq = np.fft.rfftfreq(n, d=1./sample_rate)

212 >>> freq

213 array([ 0., 10., 20., 30., 40., 50.])

214

215 """

216 if not isinstance(n, integer_types):

217 raise ValueError("n should be an integer")

218 val = 1.0/(n*d)

219 N = n//2 + 1

220 results = arange(0, N, dtype=int)

221 return results * val