Tuesday, December 13, 2011

A simple FFT

So, following my PDE's class, I decided to try my hand at writing an FFT algorithm. I ended up doing it in both c and Fortran. Surprisingly, I found it easier to write in Fortran (accounting for the fact that I am more familiar with c, and had to look up the syntax for just about everything in fortran). So, without further ado, here is my code:

FORTRAN:


 module FFT  
   implicit none  
   save  
   public  
   REAL, parameter :: PI = 4*atan(1.0)  
   contains  
   recursive function fft1(f) result(r)  
     !assumes that f is of length n where n is a power of 2  
     !output is zero indexed  
     COMPLEX, intent(in) :: f(:)   
     integer :: N, first, last, k   
     COMPLEX :: r(0:size(f)-1)  
     COMPLEX, dimension(0:size(f)/2-1) :: even, odd  
     N = size(f)  
     first = lbound(f,1)  
     last=ubound(f,1)  
     if( N == 1) then  
       r = f  
       return  
     end if  
     even = fft1(f(first:last-1:2))  
     odd = fft1(f(first+1:last:2))  
     do k=0,N/2-1  
      odd(k) = exp((0,-2)*PI*k/N)*odd(k)  
     end do  
     r(0:N/2-1) = even+odd  
     r(N/2:N-1) = even-odd  
   end function  
   recursive function ifft1(f) result(r)  
     !assumes that f is of length n where n is a power of 2  
     !output is zero indexed  
     COMPLEX, intent(in) :: f(:)   
     integer :: N, first, last, k   
     COMPLEX :: r(0:size(f)-1)  
     COMPLEX, dimension(0:size(f)/2-1) :: even, odd  
     N = size(f)  
     first = lbound(f,1)  
     last=ubound(f,1)  
     if( N == 1) then  
       r = f  
       return  
     end if  
     even = ifft1(f(first:last-1:2))  
     odd = ifft1(f(first+1:last:2))  
     do k=0,N/2-1  
       odd(k) = exp((0,2)*PI*k/N)*odd(k)  
     end do  
     r(0:N/2-1) = even+odd  
     r(N/2:N-1) = even-odd  
     r=r/2  
   end function  
 end module  

C:

 #include <complex.h>  
 #include <math.h>  
 #include <stdlib.h>  
 #include "fft.h"  
 double complex* _fft1( double complex*, size_t, unsigned int);  
 double complex* _ifft1(double complex*, size_t, unsigned int);  
 /** use the Cooley-Tukey method to take the fast fourier transform of an  
  * Array of lenth 2^n   
  *  
  * input and output are both complex arrays of length N, which must be a  
  * power of 2. Returns 0 if succusfull, nonzero if failed.  
  **/  
 double complex* fft1( double complex *input, size_t N) {  
   return _fft1(input,N,1);  
 }  
 double complex* _fft1( double complex *data, size_t N, unsigned int step) {  
   double complex *result;  
   double complex *even;  
   double complex *odd;  
   size_t k;  
   if ( (N > 2) && ( N % 2 != 0 )) return NULL;  
   result = (double complex*) malloc( sizeof(double complex) * N);  
   if ( N ==1 ) {  
     result[0]=data[0];  
     return result;  
   }  
   even = _fft1(data, N/2, 2*step);  
   odd = _fft1(data+step, N/2, 2*step);  
   //multiply the odds by the twiddle factor e^(-i*2*pi*k/N)  
   for( k=0; k < N/2; k++) {  
     odd[k] = cexp(-2.0*I*M_PI*k/N)*odd[k];  
   }  
   //generate result  
   for( k =0; k < N/2; k++) {  
     result[k] = even[k] + odd[k]; //first N/2 results  
     result[k+N/2] = even[k]-odd[k]; // last N/2 results  
   }  
   free(even);  
   free(odd);  
   return result;  
 }  
 double complex* ifft1( double complex* data, size_t N) {  
   return _ifft1(data, N, 1);  
 }  
 double complex* _ifft1( double complex *data, size_t N, unsigned int step) {  
   double complex *result;  
   double complex *even;  
   double complex *odd;  
   size_t k;  
   if ( N % 2 != 0 ) return NULL;  
   result = (double complex*) malloc( sizeof(double complex) * N);  
   if ( N ==1 ) {  
     result[0]==data[0];  
     return result;  
   }  
   even = _ifft1(data, N/2, 2*step);  
   odd = _ifft1(data+1, N/2, 2*step);  
   //multiply the odds by the twiddle factor e^(i*2*pi*k/N)  
   for( k=0; k < N/2; k++) {  
     odd[k] = cexp(I*2*M_PI*k/N)*odd[k];  
   }  
   //generate result  
   for( k =0; k < N/2; k++) {  
     result[k] = (even[k] + odd[k])/2; //first N/2 results  
     result[k+N/2] = (even[k]-odd[k])/2; // last N/2 results  
   }  
   free(even);  
   free(odd);  
   return result;  
 }  

No comments:

Post a Comment