fftw.h

Go to the documentation of this file.

// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
#ifndef SP_FFTW_H
#define SP_FFTW_H
#include <fftw3.h>
 
namespace sp
{
 
class FFTW
{
  private:
    fftw_plan    pl_fft;     
    fftw_plan    pl_ifft;    
    fftw_plan    pl_fft_cx;  
    fftw_plan    pl_ifft_cx; 
    unsigned int N;          
    unsigned int R, C;       
    int          alg;        
    int export_alg; 
  public:
    FFTW( unsigned int _N, int _alg = FFTW_ESTIMATE )
    {
        N          = _N;
        R          = 0;
        C          = 0;
        alg        = _alg;
        export_alg = FFTW_PATIENT;
        pl_fft     = NULL;
        pl_ifft    = NULL;
        pl_fft_cx  = NULL;
        pl_ifft_cx = NULL;
    }
 
    FFTW( unsigned int _R, unsigned int _C, int _alg )
    {
        R          = _R;
        C          = _C;
        N          = 0;
        alg        = _alg;
        export_alg = FFTW_PATIENT;
        pl_fft     = NULL;
        pl_ifft    = NULL;
        pl_fft_cx  = NULL;
        pl_ifft_cx = NULL;
    }
 
    ~FFTW()
    {
        if( pl_fft != NULL )
            fftw_destroy_plan( pl_fft );
        if( pl_ifft != NULL )
            fftw_destroy_plan( pl_ifft );
        if( pl_fft_cx != NULL )
            fftw_destroy_plan( pl_fft_cx );
        if( pl_ifft_cx != NULL )
            fftw_destroy_plan( pl_ifft_cx );
        fftw_cleanup();
    }
 
    void fft_cx( arma::cx_vec& x, arma::cx_vec& Pxx )
    {
        fftw_complex* in  = reinterpret_cast<fftw_complex*>( x.memptr() );
        fftw_complex* out = reinterpret_cast<fftw_complex*>( Pxx.memptr() );
        if( pl_fft_cx == NULL )
        {
            pl_fft_cx = fftw_plan_dft_1d( N, in, out, FFTW_FORWARD, alg );
            if( pl_fft_cx == NULL )
            {
                err_handler( "Unable to create complex data FFTW plan" );
            }
        }
        fftw_execute_dft( pl_fft_cx, in, out );
    }
 
    arma::cx_vec fft_cx( arma::cx_vec& x )
    {
        arma::cx_vec Pxx( N );
        fft_cx( x, Pxx );
        return Pxx;
    }
 
    void ifft_cx( arma::cx_vec& Pxx, arma::cx_vec& x )
    {
        fftw_complex* in  = reinterpret_cast<fftw_complex*>( Pxx.memptr() );
        fftw_complex* out = reinterpret_cast<fftw_complex*>( x.memptr() );
        if( pl_ifft_cx == NULL )
        {
            pl_ifft_cx = fftw_plan_dft_1d( N, in, out, FFTW_BACKWARD, alg );
            if( pl_ifft_cx == NULL )
            {
                err_handler( "Unable to create complex data IFFTW plan" );
            }
        }
        fftw_execute_dft( pl_ifft_cx, in, out );
        x /= N;
    }
 
    arma::cx_vec ifft_cx( arma::cx_vec& Pxx )
    {
        arma::cx_vec x( N );
        ifft_cx( Pxx, x );
        return x;
    }
 
    void fft( arma::vec& x, arma::cx_vec& Pxx )
    {
        double*       in  = x.memptr();
        fftw_complex* out = reinterpret_cast<fftw_complex*>( Pxx.memptr() );
        if( pl_fft == NULL )
        {
            pl_fft = fftw_plan_dft_r2c_1d( N, in, out, alg );
            if( pl_fft == NULL )
            {
                err_handler( "Unable to create real data FFTW plan" );
            }
        }
 
        fftw_execute_dft_r2c( pl_fft, in, out );
        int offset = static_cast<int>( ceil( N / 2.0 ) );
        int n_elem = N - offset;
        for( int i = 0; i < n_elem; ++i )
        {
            Pxx( offset + i ) = std::conj( Pxx( n_elem - i ) );
        }
    }
 
    arma::cx_vec fft( arma::vec& x )
    {
        arma::cx_vec Pxx( N );
        fft( x, Pxx );
        return Pxx;
    }
 
    void ifft( arma::cx_vec& Pxx, arma::vec& x )
    {
        fftw_complex* in  = reinterpret_cast<fftw_complex*>( Pxx.memptr() );
        double*       out = x.memptr();
        if( pl_ifft == NULL )
        {
            pl_ifft = fftw_plan_dft_c2r_1d( N, in, out, alg );
            if( pl_ifft == NULL )
            {
                err_handler( "Unable to create real data IFFTW plan" );
            }
        }
        fftw_execute_dft_c2r( pl_ifft, in, out );
        x /= N;
    }
 
    arma::vec ifft( arma::cx_vec& Pxx )
    {
        arma::vec x( N );
        ifft( Pxx, x );
        return x;
    }
 
    void fft2( arma::mat& x, arma::cx_mat& Pxx )
    {
        arma::cx_mat  Ptmp( R / 2 + 1, C, arma::fill::ones );
        double*       in  = x.memptr();
        fftw_complex* out = reinterpret_cast<fftw_complex*>( Ptmp.memptr() );
        if( pl_fft == NULL )
        {
            pl_fft = fftw_plan_dft_r2c_2d( C, R, in, out,
                                           alg ); // Column to row-major order trick: switch C and R
            if( pl_fft == NULL )
            {
                err_handler( "Unable to create real data FFTW plan" );
            }
        }
 
        fftw_execute_dft_r2c( pl_fft, in, out );
 
        // Reshape Pxx - upper half
        std::complex<double>* ptr  = reinterpret_cast<std::complex<double>*>( out );
        const unsigned int    Roff = R / 2 + 1;
        for( unsigned int r = 0; r < Roff; r++ )
        {
            for( unsigned int c = 0; c < C; c++ )
            {
                Pxx( r, c ) = ptr[r + c * Roff];
            }
        }
        // Reshape Pxx - conj symmetry
        for( unsigned int r = Roff; r < R; r++ )
        {
            Pxx( r, 0 ) = conj( ptr[R - r] );
            for( unsigned int c = 1; c < C; c++ )
            {
                Pxx( r, c ) = conj( ptr[R - r + ( C - c ) * Roff] );
            }
        }
    }
 
    arma::cx_mat fft2( arma::mat& x )
    {
        arma::cx_mat Pxx( R, C, arma::fill::ones );
        fft2( x, Pxx );
        return Pxx;
    }
 
    void ifft2( arma::cx_mat& Pxx, arma::mat& x )
    {
        // Reshape to row-major format
        unsigned int          Roff = R / 2 + 1;
        arma::cx_mat          Ptmp( Roff, C );
        std::complex<double>* ptr = reinterpret_cast<std::complex<double>*>( Ptmp.memptr() );
        for( unsigned int r = 0; r < Roff; r++ )
        {
            for( unsigned int c = 0; c < C; c++ )
            {
                ptr[r + c * Roff] = Pxx( r, c );
            }
        }
 
        fftw_complex* in  = reinterpret_cast<fftw_complex*>( Ptmp.memptr() );
        double*       out = x.memptr();
        if( pl_ifft == NULL )
        {
            pl_ifft = fftw_plan_dft_c2r_2d( C, R, in, out, alg );
            if( pl_ifft == NULL )
            {
                err_handler( "Unable to create real data IFFTW plan" );
            }
        }
        fftw_execute_dft_c2r( pl_ifft, in, out );
        x /= ( R * C );
    }
 
    arma::mat ifft2( arma::cx_mat& Pxx )
    {
        arma::mat x( R, C );
        ifft2( Pxx, x );
        return x;
    }
 
    void import_wisdom_string( const std::string wisd )
    {
        int res = fftw_import_wisdom_from_string( wisd.c_str() );
        if( res == 0 )
            err_handler( "Unable to import wisdom from string!" );
    }
 
    void import_wisdom_file( const std::string fname )
    {
        int res = fftw_import_wisdom_from_filename( fname.c_str() );
        if( res == 0 )
            err_handler( "Unable to import wisdom from file!" );
    }
 
    void export_wisdom_fft( const std::string fname )
    {
        fftw_plan     pl_w = NULL;
        double*       x_r;
        fftw_complex* x_cx1;
 
        if( R == 0 || C == 0 ) // 1D
        {
            x_r   = fftw_alloc_real( N );
            x_cx1 = fftw_alloc_complex( N );
 
            // Replan using wisdom
            pl_w = fftw_plan_dft_r2c_1d( N, x_r, x_cx1, export_alg );
        }
        else // 2D
        {
            x_r   = fftw_alloc_real( R * C );
            x_cx1 = fftw_alloc_complex( R * C );
 
            // Replan using wisdom
            pl_w = fftw_plan_dft_r2c_2d( C, R, x_r, x_cx1, export_alg );
        }
        if( pl_w == NULL )
        {
            err_handler( "Unable to create real data FFTW plan" );
        }
 
        // Export
        if( fftw_export_wisdom_to_filename( fname.c_str() ) == 0 )
        {
            err_handler( "Could not export wisdom to file!" );
        }
 
        fftw_destroy_plan( pl_w );
        fftw_free( x_r );
        fftw_free( x_cx1 );
    }
 
    void export_wisdom_ifft( const std::string fname )
    {
        fftw_plan     pl_w = NULL;
        double*       x_r;
        fftw_complex* x_cx1;
 
        if( R == 0 || C == 0 ) // 1D
        {
            x_r   = fftw_alloc_real( N );
            x_cx1 = fftw_alloc_complex( N );
 
            // Replan using wisdom
            pl_w = fftw_plan_dft_c2r_1d( N, x_cx1, x_r, export_alg );
        }
        else // 2D
        {
            x_r   = fftw_alloc_real( R * C );
            x_cx1 = fftw_alloc_complex( R * C );
 
            // Replan using wisdom
            pl_w = fftw_plan_dft_c2r_2d( C, R, x_cx1, x_r, export_alg );
        }
 
        if( pl_w == NULL )
        {
            err_handler( "Unable to create real data FFTW plan" );
        }
 
        // Export
        if( fftw_export_wisdom_to_filename( fname.c_str() ) == 0 )
        {
            err_handler( "Could not export wisdom to file!" );
        }
 
        fftw_destroy_plan( pl_w );
        fftw_free( x_r );
        fftw_free( x_cx1 );
    }
 
    void export_wisdom_fft_cx( const std::string fname )
    {
        fftw_plan     pl_w = NULL;
        fftw_complex *x_cx1, *x_cx2;
 
        if( R == 0 || C == 0 ) // 1D
        {
            x_cx1 = fftw_alloc_complex( N );
            x_cx2 = fftw_alloc_complex( N );
 
            // Replan using wisdom
            pl_w = fftw_plan_dft_1d( N, x_cx1, x_cx2, FFTW_FORWARD, export_alg );
        }
        else
        {
            x_cx1 = fftw_alloc_complex( R * C );
            x_cx2 = fftw_alloc_complex( R * C );
 
            // Replan using wisdom
            pl_w = fftw_plan_dft_2d( C, R, x_cx1, x_cx2, FFTW_FORWARD, export_alg );
        }
 
        if( pl_w == NULL )
        {
            err_handler( "Unable to create complex data FFTW plan" );
        }
 
        // Export
        if( fftw_export_wisdom_to_filename( fname.c_str() ) == 0 )
        {
            err_handler( "Could not export wisdom to file!" );
        }
 
        fftw_destroy_plan( pl_w );
        fftw_free( x_cx1 );
        fftw_free( x_cx2 );
    }
 
    void export_wisdom_ifft_cx( const std::string fname )
    {
        fftw_plan     pl_w = NULL;
        fftw_complex *x_cx1, *x_cx2;
 
        if( R == 0 || C == 0 ) // 1D
        {
            x_cx1 = fftw_alloc_complex( N );
            x_cx2 = fftw_alloc_complex( N );
 
            // Replan using wisdom
            pl_w = fftw_plan_dft_1d( N, x_cx2, x_cx1, FFTW_BACKWARD, export_alg );
        }
        else
        {
            x_cx1 = fftw_alloc_complex( R * C );
            x_cx2 = fftw_alloc_complex( R * C );
 
            // Replan using wisdom
            pl_w = fftw_plan_dft_2d( C, R, x_cx2, x_cx1, FFTW_BACKWARD, export_alg );
        }
        if( pl_w == NULL )
        {
            err_handler( "Unable to create complex data IFFTW plan" );
        }
 
        // Export
        if( fftw_export_wisdom_to_filename( fname.c_str() ) == 0 )
        {
            err_handler( "Could not export wisdom to file!" );
        }
 
        fftw_destroy_plan( pl_w );
        fftw_free( x_cx1 );
        fftw_free( x_cx2 );
    }
};
 
 
} // namespace sp
#endif