moved scaling to Eigen::FFT

1 files changed, 61 insertions, 12 deletions

diff --git a/unsupported/Eigen/FFT b/unsupported/Eigen/FFT
index 36afdde8d..e2705abf6 100644
--- a/unsupported/Eigen/FFT
+++ b/unsupported/Eigen/FFT
@@ -28,6 +28,7 @@
 #include <complex>
 #include <vector>
 #include <map>
+#include <Eigen/Core>
 
 #ifdef EIGEN_FFTW_DEFAULT
 // FFTW: faster, GPL -- incompatible with Eigen in LGPL form, bigger code size
@@ -65,10 +66,31 @@ class FFT
     typedef typename impl_type::Scalar Scalar;
     typedef typename impl_type::Complex Complex;
 
-    FFT(const impl_type & impl=impl_type() ) :m_impl(impl) { }
+    enum Flag {
+      Default=0, // goof proof
+      Unscaled=1,
+      HalfSpectrum=2,
+      // SomeOtherSpeedOptimization=4
+      Speedy=32767
+    };
 
-    template <typename _Input>
-    void fwd( Complex * dst, const _Input * src, int nfft)
+    FFT( const impl_type & impl=impl_type() , Flag flags=Default ) :m_impl(impl),m_flag(flags) { }
+
+    inline
+    bool HasFlag(Flag f) const { return (m_flag & (int)f) == f;}
+
+    inline
+    void SetFlag(Flag f) { m_flag |= (int)f;}
+
+    inline
+    void ClearFlag(Flag f) { m_flag &= (~(int)f);}
+
+    void fwd( Complex * dst, const Scalar * src, int nfft)
+    {
+        m_impl.fwd(dst,src,nfft);
+    }
+
+    void fwd( Complex * dst, const Complex * src, int nfft)
     {
         m_impl.fwd(dst,src,nfft);
     }
@@ -76,8 +98,11 @@ class FFT
     template <typename _Input>
     void fwd( std::vector<Complex> & dst, const std::vector<_Input> & src) 
     {
-        dst.resize( src.size() );
-        fwd( &dst[0],&src[0],src.size() );
+      if ( NumTraits<_Input>::IsComplex == 0 && HasFlag(HalfSpectrum) )
+        dst.resize( (src.size()>>1)+1);
+      else
+        dst.resize(src.size());
+      fwd(&dst[0],&src[0],src.size());
     }
 
     template<typename InputDerived, typename ComplexDerived>
@@ -94,17 +119,18 @@ class FFT
         fwd( &dst[0],&src[0],src.size() );
     }
 
-    template <typename _Output>
-    void inv( _Output * dst, const Complex * src, int nfft)
+    void inv( Complex * dst, const Complex * src, int nfft)
     {
         m_impl.inv( dst,src,nfft );
+        if ( HasFlag( Unscaled ) == false)
+          scale(dst,1./nfft,nfft);
     }
 
-    template <typename _Output>
-    void inv( std::vector<_Output> & dst, const std::vector<Complex> & src) 
+    void inv( Scalar * dst, const Complex * src, int nfft)
     {
-        dst.resize( src.size() );
-        inv( &dst[0],&src[0],src.size() );
+        m_impl.inv( dst,src,nfft );
+        if ( HasFlag( Unscaled ) == false)
+          scale(dst,1./nfft,nfft);
     }
 
     template<typename OutputDerived, typename ComplexDerived>
@@ -117,10 +143,24 @@ class FFT
                             YOU_MIXED_DIFFERENT_NUMERIC_TYPES__YOU_NEED_TO_USE_THE_CAST_METHOD_OF_MATRIXBASE_TO_CAST_NUMERIC_TYPES_EXPLICITLY)
         EIGEN_STATIC_ASSERT(int(OutputDerived::Flags)&int(ComplexDerived::Flags)&DirectAccessBit,
                             THIS_METHOD_IS_ONLY_FOR_EXPRESSIONS_WITH_DIRECT_MEMORY_ACCESS_SUCH_AS_MAP_OR_PLAIN_MATRICES)
-        dst.derived().resize( src.size() );
+
+        int nfft = src.size();
+        int nout = HasFlag(HalfSpectrum) ? ((nfft>>1)+1) : nfft;
+        dst.derived().resize( nout );
         inv( &dst[0],&src[0],src.size() );
     }
 
+    template <typename _Output>
+    void inv( std::vector<_Output> & dst, const std::vector<Complex> & src)
+    {
+      if ( NumTraits<_Output>::IsComplex == 0 && HasFlag(HalfSpectrum) )
+        dst.resize( 2*(src.size()-1) );
+      else
+        dst.resize( src.size() );
+      inv( &dst[0],&src[0],dst.size() );
+    }
+
+
     // TODO: multi-dimensional FFTs
     
     // TODO: handle Eigen MatrixBase
@@ -128,7 +168,16 @@ class FFT
 
     impl_type & impl() {return m_impl;}
   private:
+
+    template <typename _It,typename _Val>
+    void scale(_It x,_Val s,int nx)
+    {
+      for (int k=0;k<nx;++k)
+        *x++ *= s;
+    }
+
     impl_type m_impl;
+    int m_flag;
 };
 }
 #endif