FFT.py 3.99 KB
 alexander.dornheim committed Apr 05, 2017 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 ``````# Fourier transform processor. FFT's data in the physical domain and # corrects for the peculiarities of numpy's fft algorithm. # First we fft the data, then we take the negative sign of every other # mode in each direction (x and y), then we # divide the bad boy by (sqrt(res))^2 to get the correct intensity scaling, # finally we fftshift everything. from numpy.random import rand from numpy import exp, sqrt, log, tan, pi, floor, zeros, ones from scipy.special import gammaln import numpy as np def FFT(f): shape = f.shape; res=max(shape); if shape[0] == shape[1]: F= np.fft.fft2(f); F=F/(res); F[1:res:2,:]=-F[1:res:2,:]; F[:,1:res:2]=-F[:,1:res:2]; else: F = np.fft.fft(f,axis=0); F = F.ravel('F'); #create 1-d array, colum-major order (matlab style), not really nice F[1:res:2] =-F[1:res:2]; F = F.reshape(shape[1],shape[0]).T; #back to original shape F=F/np.sqrt(res); return np.fft.fftshift(F); `````` s.gretchko committed Jan 25, 2020 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 `````` def fft(a): """ fft(a) Compute the one-dimensional discrete Fourier transform of a the way Matlab does. When a is a vector, the Fourier transform of the vector is returned. When a is a matrix, each column vector of a is transformed individually, and a new matrix containing the transformed column vectors of a is returned. Parameters ---------- a : array_like 1-D or 2-D input array (can be complex) Returns ------- result : ndarray 1-D or 2-D array of similar shape and type containing the discrete fourier transform of a See Also -------- ifft Notes ----- Using the Numpy function fft on a matrix does not produce results similar to what Matlab does. This helper function uses the Numpy functions to produce a resut that agrees with what Matlab does. """ return transformVectors(a, np.fft.fft) def ifft(a): """ ifft(a) Compute the one-dimensional inverse discrete Fourier transform the way Matlab does. When a is a vector, the inverse Fourier transform of the vector is returned. When a is a matrix, each column vector of a is transformed individually, and a new matrix containing the transformed column vectors of a is returned. Parameters ---------- a : array_like 1-D or 2-D input array (can be complex) Returns ------- out : ndarray 1-D or 2-D array of similar shape and type containing the inverse discrete fourier transform of a See Also -------- fft Notes ----- Using the Numpy function ifft on a matrix does not produce results similar to what Matlab does. This helper function uses the Numpy functions to produce a resut that agrees with what Matlab does. """ return transformVectors(a, np.fft.ifft) def transformVectors(a, transform): """ transformVectors(a, transform) Transform a according to the given transform function. When a is a vector, it applies the transform function to a and returns the result. When a is a matrix, each column vector of a is transformed individually using the given transform function, and a new matrix containing the transformed column vectors of a is returned. Parameters ---------- a : array_like 1-D or 2-D input array transform : callable function This function takes a 1-D array as argument and returns a 1-D array of same size. This function is applied to a if it is a vector or to the column vectors of a if a is a matrix. Returns ------- out : ndarray 1-D or 2-D array of similar shape and type containing the transformed data See Also -------- fft, ifft Notes ----- This function is used by fft(a) and ifft(a). """ if a.ndim == 1: # a is a vector out = transform(a) else: # assume a is a matrix (2d array) shape = a.shape colCount = shape[1] #result = np.empty_like(a) out = np.zeros_like(a) # for each column vector in a for i in range(0, colCount): col = a[:,i] fft_col = transform(col) out[:,i] = fft_col return out ``````