1 '''
  2  Radix-2 DIF FFT in Python 2.7.3
  3  John Bryan, 2016
  4 '''
  5 
  6 import numpy as np
  7 import matplotlib.pyplot as plt
  8 
  9 
 10 def bracewell_buneman(xarray, length, log2length):
 11     ''' 
 12     bracewell-buneman bit reversal function
 13     inputs: xarray is array; length is array length; log2length=log2(length).
 14     output: bit reversed array xarray. 
 15     '''
 16     muplus = int((log2length+1)/2)
 17     mvar = 1
 18     reverse = np.zeros(length, dtype = int)
 19     upper_range = muplus+1
 20     for _ in np.arange(1, upper_range):
 21         for kvar in np.arange(0, mvar):
 22             tvar = 2*reverse[kvar]
 23             reverse[kvar] = tvar
 24             reverse[kvar+mvar] = tvar+1
 25         mvar = mvar+mvar
 26     if (log2length & 0x01):
 27         mvar = mvar/2
 28     for qvar in np.arange(1, mvar):
 29         nprime = qvar-mvar
 30         rprimeprime = reverse[qvar]*mvar
 31         for pvar in np.arange(0, reverse[qvar]):
 32             nprime = nprime+mvar
 33             rprime = rprimeprime+reverse[pvar]
 34             temp = xarray[nprime]
 35             xarray[nprime] = xarray[rprime]
 36             xarray[rprime] = temp
 37     return xarray
 38 
 39 
 40 
 41 def dif_fft0 (xarray, twiddle, log2length):
 42     ''' 
 43     radix-2 dif fft 
 44     '''
 45     xarray = xarray.astype(np.complex_)
 46     b_p = 1
 47     nvar_p = xarray.size
 48     twiddle_step_size = 1
 49     for _ in range(0,  log2length):           # pass loop
 50         nvar_pp =  nvar_p/2
 51         base_e = 0
 52         for _ in range(0,  b_p):       # block loop
 53             base_o = base_e+nvar_pp
 54             for nvar in range(0,  nvar_pp):   # butterfly loop
 55                 evar =  xarray[base_e+nvar]+xarray[base_o+nvar]
 56                 if nvar == 0:
 57                     ovar = xarray[base_e+nvar]-xarray[base_o+nvar]
 58                 else:
 59                     twiddle_factor =  nvar*twiddle_step_size
 60                     ovar = (xarray[base_e+nvar] \
 61                         -xarray[base_o+nvar])*twiddle[twiddle_factor]
 62                 xarray[base_e+nvar] = evar
 63                 xarray[base_o+nvar] = ovar
 64             base_e = base_e+nvar_p
 65         b_p = b_p*2
 66         nvar_p = nvar_p/2
 67         twiddle_step_size = 2*twiddle_step_size
 68     xarray = bracewell_buneman(xarray, xarray.size, log2length)
 69     return xarray
 70 
 71 
 72 def test(time, yarray, samplefreq):
 73     '''
 74     Set up plot, call FFT function, plot result.
 75     Called  from testbench function.
 76     Inputs time:time vector, yarray: array, samplefreq: sampling rate. 
 77     Outputs: none.
 78     '''
 79     plt.subplot(2, 1, 1)
 80     plt.title('Test of DIF FFT with 10 Hz Sine Input')
 81     plt.plot(time, yarray ,'k-')
 82     plt.xlabel('time')
 83     plt.ylabel('amplitude')
 84     plt.subplot(2, 1, 2)
 85     ylength = len(yarray)                       # length of the signal
 86     kvar = np.arange(ylength)
 87     tvar = ylength/samplefreq
 88     frq = kvar/tvar                        # two-sided frequency range
 89     freq = frq[range(ylength/2)]           # one-sided frequency range
 90     twiddle = np.exp(-2j*np.pi*np.arange(0, 0.5, 1./ylength, dtype=np.complex_))
 91     y2array = abs(dif_fft0(yarray, twiddle, \
 92               int(np.log2(ylength)))/ylength)   # fft normalized magnitude
 93     y3array = y2array[range(ylength/2)]
 94     markerline, stemlines, baseline = plt.stem(freq, y3array, '--')
 95     plt.xlabel('freq (Hz)')
 96     plt.ylabel('|Y(freq)|')
 97     plt.ylim((0.0, 0.55))
 98     plt.setp(markerline, 'markerfacecolor', 'b')
 99     plt.setp(baseline, 'color', 'b', 'linewidth', 2)
100     plt.show()
101     return None
102 
103 
104 def testbench():
105     '''
106     no inputs, no outputs, calls test function.  
107     '''
108     samplefreq = 128                                # sampling rate
109     samplinginterval = 1.0/samplefreq               # sampling interval
110     time = np.arange(0, 1, samplinginterval)        # time vector
111     frequency = 10                                  # frequency of the signal
112     yarray = np.sin(2 * np.pi * frequency * time)   # 10 hz sine wave signal
113     test(time, yarray, samplefreq)                 # send sine to test 
114     return None
115 
116 
117 testbench()