File indexing completed on 2024-05-19 04:54:32

0001 #pragma once
0002 
0003 #include <complex>
0004 #include <vector>
0005 
0006 namespace kissfft_utils {
0007 
0008 template <typename T_scalar>
0009 struct traits
0010 {
0011     typedef T_scalar scalar_type;
0012     typedef std::complex<scalar_type> cpx_type;
0013     void fill_twiddles( std::complex<T_scalar> * dst ,int nfft,bool inverse)
0014     {
0015         T_scalar phinc =  (inverse?2:-2)* acos( (T_scalar) -1)  / nfft;
0016         for (int i=0;i<nfft;++i)
0017             dst[i] = exp( std::complex<T_scalar>(0,i*phinc) );
0018     }
0019 
0020     void prepare(
0021             std::vector< std::complex<T_scalar> > & dst,
0022             int nfft,bool inverse, 
0023             std::vector<int> & stageRadix, 
0024             std::vector<int> & stageRemainder )
0025     {
0026         _twiddles.resize(nfft);
0027         fill_twiddles( &_twiddles[0],nfft,inverse);
0028         dst = _twiddles;
0029 
0030         //factorize
0031         //start factoring out 4's, then 2's, then 3,5,7,9,...
0032         int n= nfft;
0033         int p=4;
0034         do {
0035             while (n % p) {
0036                 switch (p) {
0037                     case 4: p = 2; break;
0038                     case 2: p = 3; break;
0039                     default: p += 2; break;
0040                 }
0041                 if (p*p>n)
0042                     p=n;// no more factors
0043             }
0044             n /= p;
0045             stageRadix.push_back(p);
0046             stageRemainder.push_back(n);
0047         }while(n>1);
0048     }
0049     std::vector<cpx_type> _twiddles;
0050 
0051 
0052     const cpx_type twiddle(int i) { return _twiddles[i]; }
0053 };
0054 
0055 }
0056 
0057 template <typename T_Scalar,
0058          typename T_traits=kissfft_utils::traits<T_Scalar> 
0059          >
0060 class kissfft
0061 {
0062     public:
0063         typedef T_traits traits_type;
0064         typedef typename traits_type::scalar_type scalar_type;
0065         typedef typename traits_type::cpx_type cpx_type;
0066 
0067         kissfft(int nfft,bool inverse,const traits_type & traits=traits_type() ) 
0068             :_nfft(nfft),_inverse(inverse),_traits(traits)
0069         {
0070             _traits.prepare(_twiddles, _nfft,_inverse ,_stageRadix, _stageRemainder);
0071         }
0072 
0073         void transform(const cpx_type * src , cpx_type * dst)
0074         {
0075             kf_work(0, dst, src, 1,1);
0076         }
0077 
0078     private:
0079         void kf_work( int stage,cpx_type * Fout, const cpx_type * f, size_t fstride,size_t in_stride)
0080         {
0081             int p = _stageRadix[stage];
0082             int m = _stageRemainder[stage];
0083             cpx_type * Fout_beg = Fout;
0084             cpx_type * Fout_end = Fout + p*m;
0085 
0086             if (m==1) {
0087                 do{
0088                     *Fout = *f;
0089                     f += fstride*in_stride;
0090                 }while(++Fout != Fout_end );
0091             }else{
0092                 do{
0093                     // recursive call:
0094                     // DFT of size m*p performed by doing
0095                     // p instances of smaller DFTs of size m, 
0096                     // each one takes a decimated version of the input
0097                     kf_work(stage+1, Fout , f, fstride*p,in_stride);
0098                     f += fstride*in_stride;
0099                 }while( (Fout += m) != Fout_end );
0100             }
0101 
0102             Fout=Fout_beg;
0103 
0104             // recombine the p smaller DFTs 
0105             switch (p) {
0106                 case 2: kf_bfly2(Fout,fstride,m); break;
0107                 case 3: kf_bfly3(Fout,fstride,m); break;
0108                 case 4: kf_bfly4(Fout,fstride,m); break;
0109                 case 5: kf_bfly5(Fout,fstride,m); break;
0110                 default: kf_bfly_generic(Fout,fstride,m,p); break;
0111             }
0112         }
0113 
0114         // these were #define macros in the original kiss_fft
0115         void C_ADD( cpx_type & c,const cpx_type & a,const cpx_type & b) { c=a+b;}
0116         void C_MUL( cpx_type & c,const cpx_type & a,const cpx_type & b) { c=a*b;}
0117         void C_SUB( cpx_type & c,const cpx_type & a,const cpx_type & b) { c=a-b;}
0118         void C_ADDTO( cpx_type & c,const cpx_type & a) { c+=a;}
0119         void C_FIXDIV( cpx_type & ,int ) {} // NO-OP for float types
0120         scalar_type S_MUL( const scalar_type & a,const scalar_type & b) { return a*b;}
0121         scalar_type HALF_OF( const scalar_type & a) { return a*.5;}
0122         void C_MULBYSCALAR(cpx_type & c,const scalar_type & a) {c*=a;}
0123 
0124         void kf_bfly2( cpx_type * Fout, const size_t fstride, int m)
0125         {
0126             for (int k=0;k<m;++k) {
0127                 cpx_type t = Fout[m+k] * _traits.twiddle(k*fstride);
0128                 Fout[m+k] = Fout[k] - t;
0129                 Fout[k] += t;
0130             }
0131         }
0132 
0133         void kf_bfly4( cpx_type * Fout, const size_t fstride, const size_t m)
0134         {
0135             cpx_type scratch[7];
0136             int negative_if_inverse = _inverse * -2 +1;
0137             for (size_t k=0;k<m;++k) {
0138                 scratch[0] = Fout[k+m] * _traits.twiddle(k*fstride);
0139                 scratch[1] = Fout[k+2*m] * _traits.twiddle(k*fstride*2);
0140                 scratch[2] = Fout[k+3*m] * _traits.twiddle(k*fstride*3);
0141                 scratch[5] = Fout[k] - scratch[1];
0142 
0143                 Fout[k] += scratch[1];
0144                 scratch[3] = scratch[0] + scratch[2];
0145                 scratch[4] = scratch[0] - scratch[2];
0146                 scratch[4] = cpx_type( scratch[4].imag()*negative_if_inverse , -scratch[4].real()* negative_if_inverse );
0147 
0148                 Fout[k+2*m]  = Fout[k] - scratch[3];
0149                 Fout[k] += scratch[3];
0150                 Fout[k+m] = scratch[5] + scratch[4];
0151                 Fout[k+3*m] = scratch[5] - scratch[4];
0152             }
0153         }
0154 
0155         void kf_bfly3( cpx_type * Fout, const size_t fstride, const size_t m)
0156         {
0157             size_t k=m;
0158             const size_t m2 = 2*m;
0159             cpx_type *tw1,*tw2;
0160             cpx_type scratch[5];
0161             cpx_type epi3;
0162             epi3 = _twiddles[fstride*m];
0163 
0164             tw1=tw2=&_twiddles[0];
0165 
0166             do{
0167                 C_FIXDIV(*Fout,3); C_FIXDIV(Fout[m],3); C_FIXDIV(Fout[m2],3);
0168 
0169                 C_MUL(scratch[1],Fout[m] , *tw1);
0170                 C_MUL(scratch[2],Fout[m2] , *tw2);
0171 
0172                 C_ADD(scratch[3],scratch[1],scratch[2]);
0173                 C_SUB(scratch[0],scratch[1],scratch[2]);
0174                 tw1 += fstride;
0175                 tw2 += fstride*2;
0176 
0177                 Fout[m] = cpx_type( Fout->real() - HALF_OF(scratch[3].real() ) , Fout->imag() - HALF_OF(scratch[3].imag() ) );
0178 
0179                 C_MULBYSCALAR( scratch[0] , epi3.imag() );
0180 
0181                 C_ADDTO(*Fout,scratch[3]);
0182 
0183                 Fout[m2] = cpx_type(  Fout[m].real() + scratch[0].imag() , Fout[m].imag() - scratch[0].real() );
0184 
0185                 C_ADDTO( Fout[m] , cpx_type( -scratch[0].imag(),scratch[0].real() ) );
0186                 ++Fout;
0187             }while(--k);
0188         }
0189 
0190         void kf_bfly5( cpx_type * Fout, const size_t fstride, const size_t m)
0191         {
0192             cpx_type *Fout0,*Fout1,*Fout2,*Fout3,*Fout4;
0193             size_t u;
0194             cpx_type scratch[13];
0195             cpx_type * twiddles = &_twiddles[0];
0196             cpx_type *tw;
0197             cpx_type ya,yb;
0198             ya = twiddles[fstride*m];
0199             yb = twiddles[fstride*2*m];
0200 
0201             Fout0=Fout;
0202             Fout1=Fout0+m;
0203             Fout2=Fout0+2*m;
0204             Fout3=Fout0+3*m;
0205             Fout4=Fout0+4*m;
0206 
0207             tw=twiddles;
0208             for ( u=0; u<m; ++u ) {
0209                 C_FIXDIV( *Fout0,5); C_FIXDIV( *Fout1,5); C_FIXDIV( *Fout2,5); C_FIXDIV( *Fout3,5); C_FIXDIV( *Fout4,5);
0210                 scratch[0] = *Fout0;
0211 
0212                 C_MUL(scratch[1] ,*Fout1, tw[u*fstride]);
0213                 C_MUL(scratch[2] ,*Fout2, tw[2*u*fstride]);
0214                 C_MUL(scratch[3] ,*Fout3, tw[3*u*fstride]);
0215                 C_MUL(scratch[4] ,*Fout4, tw[4*u*fstride]);
0216 
0217                 C_ADD( scratch[7],scratch[1],scratch[4]);
0218                 C_SUB( scratch[10],scratch[1],scratch[4]);
0219                 C_ADD( scratch[8],scratch[2],scratch[3]);
0220                 C_SUB( scratch[9],scratch[2],scratch[3]);
0221 
0222                 C_ADDTO( *Fout0, scratch[7]);
0223                 C_ADDTO( *Fout0, scratch[8]);
0224 
0225                 scratch[5] = scratch[0] + cpx_type(
0226                         S_MUL(scratch[7].real(),ya.real() ) + S_MUL(scratch[8].real() ,yb.real() ),
0227                         S_MUL(scratch[7].imag(),ya.real()) + S_MUL(scratch[8].imag(),yb.real())
0228                         );
0229 
0230                 scratch[6] =  cpx_type( 
0231                         S_MUL(scratch[10].imag(),ya.imag()) + S_MUL(scratch[9].imag(),yb.imag()),
0232                         -S_MUL(scratch[10].real(),ya.imag()) - S_MUL(scratch[9].real(),yb.imag()) 
0233                         );
0234 
0235                 C_SUB(*Fout1,scratch[5],scratch[6]);
0236                 C_ADD(*Fout4,scratch[5],scratch[6]);
0237 
0238                 scratch[11] = scratch[0] + 
0239                     cpx_type(
0240                             S_MUL(scratch[7].real(),yb.real()) + S_MUL(scratch[8].real(),ya.real()),
0241                             S_MUL(scratch[7].imag(),yb.real()) + S_MUL(scratch[8].imag(),ya.real())
0242                             );
0243 
0244                 scratch[12] = cpx_type(
0245                         -S_MUL(scratch[10].imag(),yb.imag()) + S_MUL(scratch[9].imag(),ya.imag()),
0246                         S_MUL(scratch[10].real(),yb.imag()) - S_MUL(scratch[9].real(),ya.imag())
0247                         );
0248 
0249                 C_ADD(*Fout2,scratch[11],scratch[12]);
0250                 C_SUB(*Fout3,scratch[11],scratch[12]);
0251 
0252                 ++Fout0;++Fout1;++Fout2;++Fout3;++Fout4;
0253             }
0254         }
0255 
0256         /* perform the butterfly for one stage of a mixed radix FFT */
0257         void kf_bfly_generic(
0258                 cpx_type * Fout,
0259                 const size_t fstride,
0260                 int m,
0261                 int p
0262                 )
0263         {
0264             int u,k,q1,q;
0265             cpx_type * twiddles = &_twiddles[0];
0266             cpx_type t;
0267             int Norig = _nfft;
0268             cpx_type scratchbuf[p];
0269 
0270             for ( u=0; u<m; ++u ) {
0271                 k=u;
0272                 for ( q1=0 ; q1<p ; ++q1 ) {
0273                     scratchbuf[q1] = Fout[ k  ];
0274                     C_FIXDIV(scratchbuf[q1],p);
0275                     k += m;
0276                 }
0277 
0278                 k=u;
0279                 for ( q1=0 ; q1<p ; ++q1 ) {
0280                     int twidx=0;
0281                     Fout[ k ] = scratchbuf[0];
0282                     for (q=1;q<p;++q ) {
0283                         twidx += fstride * k;
0284                         if (twidx>=Norig) twidx-=Norig;
0285                         C_MUL(t,scratchbuf[q] , twiddles[twidx] );
0286                         C_ADDTO( Fout[ k ] ,t);
0287                     }
0288                     k += m;
0289                 }
0290             }
0291         }
0292 
0293         int _nfft;
0294         bool _inverse;
0295         std::vector<cpx_type> _twiddles;
0296         std::vector<int> _stageRadix;
0297         std::vector<int> _stageRemainder;
0298         traits_type _traits;
0299 };