include/DGVFourierTransform.h

Go to the documentation of this file.

#ifndef DGVFOURIERTRANSFORM_H
#define DGVFOURIERTRANSFORM_H

#include <complex>
#include <fftw3.h>

#include "DGVAliases.h"

template<typename type, int rank>
class DGV_TEMPLATEDLL DGVFourierTransform
{
public:
    // -------------------------------------------------------
    DGVFourierTransform();
    virtual ~DGVFourierTransform();

        void FFT(Array<complex<type>,rank> &field, Array<complex<type>,rank> &fftOfField);
        void iFFT(Array<complex<type>,rank> &fftOfField, Array<complex<type>,rank> &field);

        void FFT_Real2Complex(Array<type,rank> &field, Array<complex<type>,rank> &fftOfField);
        void FFT_Complex2Real(Array<complex<type>,rank> &fftOfField, Array<type,rank> &field);

    void normalise(Array<complex<type>,rank> &field);
    void normalise(Array<type,rank> &field);
    Array<type,rank> shiftDC(Array<type,rank> &field);
    Array<complex<type>,rank> shiftDC(Array<complex<type>,rank> &field);
    inline void reset()
    {   mode = FFTW_ESTIMATE;   fftw_cleanup(); }

protected:
    unsigned mode; 
    fftw_plan plan; 

    void execute();
};

template<typename type, int rank>
DGVFourierTransform<type,rank>::DGVFourierTransform()
{
    mode = FFTW_ESTIMATE;
}

template<typename type, int rank>
DGVFourierTransform<type,rank>::~DGVFourierTransform()
{
    fftw_cleanup();
}

template<typename type, int rank>
void DGVFourierTransform<type,rank>::FFT(Array<complex<type>,rank> &field, Array<complex<type>,rank> &fftOfField)
{
    if(rank < 4 && rank > 0) //Protect FFTW
    {
        const int extent[3] = {field.rows(),field.cols(),field.depth()};
        //cerr << "FFT Extent: " << extent[0] << ", " << extent[1] << ", " << extent[2] << endl;

        plan = fftw_plan_dft(rank,extent,reinterpret_cast<fftw_complex*>(field.data()),reinterpret_cast<fftw_complex*>(fftOfField.data()),FFTW_FORWARD,mode);

        execute();
    }
}

template<typename type, int rank>
void DGVFourierTransform<type,rank>::iFFT(Array<complex<type>,rank> &fftOfField, Array<complex<type>,rank> &field)
{
    if(rank < 4 && rank > 0) //Protect FFTW
    {
        const int extent[3] = {fftOfField.rows(),fftOfField.cols(),fftOfField.depth()};
        //cerr << "iFFT Extent: " << extent[0] << ", " << extent[1] << ", " << extent[2] << endl;

        plan = fftw_plan_dft(rank,extent,reinterpret_cast<fftw_complex*>(fftOfField.data()),reinterpret_cast<fftw_complex*>(field.data()),FFTW_BACKWARD,mode);

        execute();
    }
}

template<typename type, int rank>
void DGVFourierTransform<type,rank>::FFT_Real2Complex(Array<type,rank> &field, Array<complex<type>,rank> &fftOfField)
{
    if(rank < 4 && rank > 0) //Protect FFTW
    {
        const int extent[3] = {field.rows(),field.cols(),field.depth()};
        //cerr << "FFT Extent: " << extent[0] << ", " << extent[1] << ", " << extent[2] << endl;

        plan = fftw_plan_dft_r2c(rank,extent,field.data(),reinterpret_cast<fftw_complex*>(fftOfField.data()),mode);

        execute();
    }
}

template<typename type, int rank>
void DGVFourierTransform<type,rank>::FFT_Complex2Real(Array<complex<type>,rank> &fftOfField, Array<type,rank> &field)
{
    if(rank < 4 && rank > 0) //Protect FFTW
    {
        const int extent[3] = {fftOfField.rows(),fftOfField.cols(),fftOfField.depth()};
        //cerr << "FFT Extent: " << extent[0] << ", " << extent[1] << ", " << extent[2] << endl;

        plan = fftw_plan_dft_c2r(rank,extent,reinterpret_cast<fftw_complex*>(fftOfField.data()),field.data(),mode);

        execute();
    }
}

template<typename type, int rank>
void DGVFourierTransform<type,rank>::normalise(Array<complex<type>,rank> &field)
{
    type factor;

    if(rank == 1)
        factor = static_cast<type>(field.size());
    else if(rank == 2)
        factor = static_cast<type>(field.rows()*field.cols());
    else
        factor = static_cast<type>(field.rows()*field.cols()*field.depth());

    //cerr << "Normalising by " << factor << endl;
    field /= factor;
}

template<typename type, int rank>
void DGVFourierTransform<type,rank>::normalise(Array<type,rank> &field)
{
    type factor;

    if(rank == 1)
        factor = static_cast<type>(field.size());
    else if(rank == 2)
        factor = static_cast<type>(field.rows()*field.cols());
    else
        factor = static_cast<type>(field.rows()*field.cols()*field.depth());

    //cerr << "Normalising by " << factor << endl;
    field /= factor;
}

template<typename type, int rank>
Array<type,rank> DGVFourierTransform<type,rank>::shiftDC(Array<type,rank> &field)
{
    int rowsHalf = static_cast<int>(field.rows()/2.0), colsHalf = static_cast<int>(field.cols()/2.0);
    Array<type,rank> tmpField(field.shape());

    for(int j = 0; j < field.rows(); j ++)
        for(int k = 0; k < field.cols(); k ++)
            tmpField( (j+rowsHalf)%field.rows() , (k+colsHalf)%field.cols() ) = field(j,k);

    return tmpField;
}

template<typename type, int rank>
Array<complex<type>,rank> DGVFourierTransform<type,rank>::shiftDC(Array<complex<type>,rank> &field)
{
    int rowsHalf = static_cast<int>(field.rows()/2.0), colsHalf = static_cast<int>(field.cols()/2.0);
    Array<complex<type>,rank> tmpField(field.shape());

    for(int j = 0; j < field.rows(); j ++)
        for(int k = 0; k < field.cols(); k ++)
            tmpField( (j+rowsHalf)%field.rows() , (k+colsHalf)%field.cols() ) = field(j,k);

    return tmpField;
}

//Private
template<typename type, int rank>
void DGVFourierTransform<type,rank>::execute()
{
    if(plan != NULL)
    {
        fftw_execute(plan);
        fftw_destroy_plan(plan);
    }
}

#endif // DGVFOURIERTRANSFORM_H