FFT.py 3.99 KB
Newer Older
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's avatar
s.gretchko committed
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