xmipp_fftw.h

Go to the documentation of this file.

/***************************************************************************
 *
 * Authors:    Roberto Marabini                 (roberto@cnb.csic.es)
 *             Carlos Oscar S. Sorzano          (coss@cnb.csic.es)
 *
 * Unidad de  Bioinformatica of Centro Nacional de Biotecnologia , CSIC
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA
 * 02111-1307  USA
 *
 *  All comments concerning this program package may be sent to the
 *  e-mail address 'xmipp@cnb.csic.es'
 ***************************************************************************/

#ifndef _CORE__XmippFFTW_H
#define _CORE__XmippFFTW_H

#include <complex>
#include "fftw3.h"
#include "multidim_array.h"
#include "multidim_array_generic.h"
#include "xmipp_fft.h"


class FourierTransformer
{
public:
    MultidimArray<double> *fReal;

    MultidimArray<std::complex<double> > *fComplex;

    MultidimArray< std::complex<double> > fFourier;

    /* fftw Forawrd plan */
    fftw_plan fPlanForward;

    /* fftw Backward plan */
    fftw_plan fPlanBackward;

    /* number of threads*/
    int nthreads;

    /* Threads has been used in this program*/
    bool threadsSetOn;

    /* Sign where the normalization is applied */
    int normSign;

    // Public methods
public:
    FourierTransformer();

    FourierTransformer(const FourierTransformer& fTransform);

    FourierTransformer(int _normSign);

    FourierTransformer & operator= (const FourierTransformer & other);

    ~FourierTransformer();

    void setThreadsNumber(int tNumber)
    {
        if (tNumber!=1)
        {
            threadsSetOn=true;
            nthreads = tNumber;
            if(fftw_init_threads()==0)
                REPORT_ERROR(ERR_THREADS_NOTINIT, (std::string)"FFTW cannot init threads (setThreadsNumber)");
            fftw_plan_with_nthreads(nthreads);
        }
    }
    void changeThreadsNumber(int tNumber)
    {
        nthreads = tNumber;
        fftw_plan_with_nthreads(nthreads);
    }

    void destroyThreads(void )
    {
        nthreads = 1;
        if(threadsSetOn)
            fftw_cleanup_threads();

        threadsSetOn=false;
    }

    template <typename T, typename T1>
    void FourierTransform(T& v, T1& V, bool getCopy=true)
    {
        setReal(v);
        Transform(FFTW_FORWARD);
        if (getCopy)
            getFourierCopy(V);
        else
            getFourierAlias(V);
    }

    void FourierTransform();

    void enforceHermitianSymmetry();

    void inverseFourierTransform();

    template <typename T, typename T1>
    void inverseFourierTransform(T& V, T1& v)
    {
        setReal(v);
        setFourier(V);
        Transform(FFTW_BACKWARD);
    }

    template <typename T>
    void getFourierAlias(T& V)
    {
        V.alias(fFourier);
        return;
    }

    template <typename T>
    void setFourierAlias(T& V)
    {
        fFourier.alias(V);
        return;
    }

    template <typename T>
    void getFourierCopy(T& V)
    {
        V.resizeNoCopy(fFourier);
        memcpy(MULTIDIM_ARRAY(V),MULTIDIM_ARRAY(fFourier),
               MULTIDIM_SIZE(fFourier)*2*sizeof(double));
    }


    template <typename T>
    void getCompleteFourier(T& V)
    {
        V.resizeNoCopy(*fReal);
        int ndim=3;
        if (ZSIZE(*fReal)==1)
        {
            ndim=2;
            if (YSIZE(*fReal)==1)
                ndim=1;
        }
        double *ptrSource=NULL;
        double *ptrDest=NULL;
        switch (ndim)
        {
        case 1:
            FOR_ALL_DIRECT_ELEMENTS_IN_ARRAY1D(V)
            {
                ptrDest=(double*)&DIRECT_A1D_ELEM(V,i);
                if (i<XSIZE(fFourier))
                {
                    ptrSource=(double*)&DIRECT_A1D_ELEM(fFourier,i);
                    *ptrDest=*ptrSource;
                    *(ptrDest+1)=*(ptrSource+1);
                }
                else
                {
                    ptrSource=(double*)&DIRECT_A1D_ELEM(fFourier,XSIZE(*fReal)-i);
                    *ptrDest=*ptrSource;
                    *(ptrDest+1)=-(*(ptrSource+1));
                }
            }
            break;
        case 2:
            FOR_ALL_DIRECT_ELEMENTS_IN_ARRAY2D(V)
            {
                ptrDest=(double*)&DIRECT_A2D_ELEM(V,i,j);
                if (j<XSIZE(fFourier))
                {
                    ptrSource=(double*)&DIRECT_A2D_ELEM(fFourier,i,j);
                    *ptrDest=*ptrSource;
                    *(ptrDest+1)=*(ptrSource+1);
                }
                else
                {
                    ptrSource=(double*)&DIRECT_A2D_ELEM(fFourier,
                                                        (YSIZE(*fReal)-i)%YSIZE(*fReal),
                                                        XSIZE(*fReal)-j);
                    *ptrDest=*ptrSource;
                    *(ptrDest+1)=-(*(ptrSource+1));
                }
            }
            break;
        case 3:
            FOR_ALL_DIRECT_ELEMENTS_IN_ARRAY3D(V)
            {
                ptrDest=(double*)&DIRECT_A3D_ELEM(V,k,i,j);
                if (j<XSIZE(fFourier))
                {
                    ptrSource=(double*)&DIRECT_A3D_ELEM(fFourier,k,i,j);
                    *ptrDest=*ptrSource;
                    *(ptrDest+1)=*(ptrSource+1);
                }
                else
                {
                    ptrSource=(double*)&DIRECT_A3D_ELEM(fFourier,
                                                        (ZSIZE(*fReal)-k)%ZSIZE(*fReal),
                                                        (YSIZE(*fReal)-i)%YSIZE(*fReal),
                                                        XSIZE(*fReal)-j);
                    *ptrDest=*ptrSource;
                    *(ptrDest+1)=-(*(ptrSource+1));
                }
            }
            break;
        }
    }


    template <typename T, typename T1>
    void completeFourierTransform(T& v, T1& V)
    {
        setReal(v);
        Transform(FFTW_FORWARD);
        getCompleteFourier(V);
    }

    template <typename T>
    void setFromCompleteFourier(T& V)
    {
        int ndim=3;
        if (ZSIZE(*fReal)==1)
        {
            ndim=2;
            if (YSIZE(*fReal)==1)
                ndim=1;
        }
        switch (ndim)
        {
        case 1:
            FOR_ALL_DIRECT_ELEMENTS_IN_ARRAY1D(fFourier)
            DIRECT_A1D_ELEM(fFourier,i)=DIRECT_A1D_ELEM(V,i);
            break;
        case 2:
            FOR_ALL_DIRECT_ELEMENTS_IN_ARRAY2D(fFourier)
            DIRECT_A2D_ELEM(fFourier,i,j) = DIRECT_A2D_ELEM(V,i,j);
            break;
        case 3:
            FOR_ALL_DIRECT_ELEMENTS_IN_ARRAY3D(fFourier)
            DIRECT_A3D_ELEM(fFourier,k,i,j) = DIRECT_A3D_ELEM(V,k,i,j);
            break;
        }
    }

    // Internal methods
public:
    /* Pointer to the array of doubles with which the plan was computed */
    double * dataPtr;

    /* Pointer to the array of complex<double> with which the plan was computed */
    std::complex<double> * complexDataPtr;

    /* Init object*/
    void init();
    void clear();
    void cleanup(void)
    {
        fftw_cleanup();
    }

    void recomputePlanR2C();

    void Transform(int sign);

    const MultidimArray<double> &getReal() const;
    const MultidimArray<std::complex<double> > &getComplex() const;

    void setReal(MultidimArray<double> &img);

    void setReal(MultidimArray<std::complex<double> > &img);

    void setFourier(const MultidimArray<std::complex<double> > &imgFourier);

    /* Set normalization sign.
     * It defines when the normalization must be applied, when doing
     * FFTW_FORWARD OR FFTW_BACKWARD. By default, FFTW_FORWARD.*/
    void setNormalizationSign(int _normSign)
    {
        normSign = _normSign;
    }

};

void FFT_magnitude(const MultidimArray< std::complex< double > > & v,
                   MultidimArray< double >& mag);

void FFT_phase(const MultidimArray< std::complex< double > > & v,
               MultidimArray< double >& phase);

void radial_magnitude(const MultidimArray<double> & v, MultidimArray< std::complex< double > > &V,
                      MultidimArray< double >& radialMagnitude);

template <typename T>
void auto_correlation_vector(const MultidimArray< T > & Img, MultidimArray< double >& R)
{
    Img.checkDimension(1);

    // Compute the Fourier Transform
    MultidimArray< std::complex< double > > FFT1;
    FourierTransformer transformer1;
    R=Img;
    transformer1.FourierTransform(R, FFT1, false);

    // Multiply FFT1 * FFT1'
    double dSize=XSIZE(Img);
    FOR_ALL_DIRECT_ELEMENTS_IN_MULTIDIMARRAY(FFT1)
    {
        double *ptr=(double*)&DIRECT_MULTIDIM_ELEM(FFT1,n);
        double &realPart=*ptr;
        double &imagPart=*(ptr+1);
        realPart=dSize*(realPart*realPart+imagPart*imagPart);
        imagPart=0;
    }

    // Invert the product, in order to obtain the correlation image
    transformer1.inverseFourierTransform();

    // Center the resulting image to obtain a centered autocorrelation
    CenterFFT(R, true);
}

template <typename T>
void correlation_vector(const MultidimArray< T > & m1,
                        const MultidimArray< T > & m2,
                        MultidimArray< double >& R)
{
    m1.checkDimension(2);
    m2.checkDimension(2);

    // Compute the Fourier Transforms
    MultidimArray< std::complex< double > > FFT1, FFT2;
    FourierTransformer transformer1, transformer2;
    R=m1;
    transformer1.FourierTransform(R, FFT1, false);
    transformer2.FourierTransform((MultidimArray<T> &)m2, FFT2, false);

    // Multiply FFT1 * FFT2'
    double dSize=XSIZE(m1);
    FOR_ALL_ELEMENTS_IN_ARRAY1D(FFT1)
    FFT1(i) *= dSize * conj(FFT2(i));

    // Invert the product, in order to obtain the correlation image
    transformer1.inverseFourierTransform();

    // Center the resulting image to obtain a centered autocorrelation
    CenterFFT(R, true);
}

void fast_correlation_vector(const MultidimArray< std::complex<double> > & FFT1,
                        const MultidimArray< std::complex<double> > & FFT2,
                        MultidimArray< double >& R,
                        FourierTransformer &transformer);

template <class T>
void correlation_vector_no_Fourier(const MultidimArray<T> &v1, const MultidimArray<T> &v2,
                                   MultidimArray<T> &result)
{
    v1.checkDimension(1);
    v2.checkDimension(1);

    result.initZeros(v1);
    result.setXmippOrigin();
    int N=XSIZE(v1)-1;
    FOR_ALL_ELEMENTS_IN_ARRAY1D(result)
    for (int k=0; k<XSIZE(v1); ++k)
        A1D_ELEM(result,i)+=DIRECT_A1D_ELEM(v1,intWRAP(k+i,0,N))*
                            DIRECT_A1D_ELEM(v2,k);
    STARTINGX(result)=0;
}

class CorrelationAux
{
public:
    MultidimArray< std::complex< double > > FFT1, FFT2;
    FourierTransformer transformer1, transformer2;
};

void correlation_matrix(const MultidimArray<double> & m1,
                        const MultidimArray<double> & m2,
                        MultidimArray< double >& R,
                        CorrelationAux &aux,
                        bool center=true);

void correlation_matrix(const MultidimArray< std::complex< double > > & FFT1,
                        const MultidimArray<double> & m2,
                        MultidimArray<double>& R,
                        CorrelationAux &aux,
                        bool center=true);

void correlation_matrix(const MultidimArray< std::complex< double > > & FFT1,
                        const MultidimArray< std::complex< double > > & FFT2,
                        MultidimArray<double>& R,
                        CorrelationAux &aux,
                        bool center=true);

template <typename T>
void auto_correlation_matrix(const MultidimArray< T > & Img, MultidimArray< double >& R)
{
    // Compute the Fourier Transform
    MultidimArray< std::complex< double > > FFT1;
    FourierTransformer transformer1;
    R=Img;
    transformer1.FourierTransform(R, FFT1, false);

    // Multiply FFT1 * FFT1'
    double dSize=MULTIDIM_SIZE(Img);
    FOR_ALL_DIRECT_ELEMENTS_IN_MULTIDIMARRAY(FFT1)
    {
        double *ptr=(double*)&DIRECT_MULTIDIM_ELEM(FFT1,n);
        double &realPart=*ptr;
        double &imagPart=*(ptr+1);
        realPart=dSize*(realPart*realPart+imagPart*imagPart);
        imagPart=0;
    }

    // Invert the product, in order to obtain the correlation image
    transformer1.inverseFourierTransform();

    // Center the resulting image to obtain a centered autocorrelation
    CenterFFT(R, true);
}

void auto_correlation_matrix(const MultidimArray<double> & Img, MultidimArray< double >& R, CorrelationAux &aux);

void convolutionFFTStack(const MultidimArray<double> &img,
                    const MultidimArray<double> &kernel,
                    MultidimArray<double> &result);

void convolutionFFT(MultidimArray<double> &img,
                    MultidimArray<double> &kernel,
                    MultidimArray<double> &result);

void frc_dpr(MultidimArray< double > & m1,
             MultidimArray< double > & m2,
             double sampling_rate,
             MultidimArray< double >& freq,
             MultidimArray< double >& frc,
             MultidimArray< double >& frc_noise,
             MultidimArray< double >& dpr,
             MultidimArray< double >& error_l2,
             bool skipdpr=false,
             bool doRfactor = false,
             double minFreq = -1,
             double maxFreq = 0.5,
                         double * rFactor= NULL);

void scaleToSizeFourier(int Zdim, int Ydim, int Xdim, MultidimArray<double> &mdaIn, MultidimArray<double> &mdaOut, int nThreads=1);

template<typename T>
void scaleToSizeFourier(MultidimArray<T> &mdaIn, MultidimArray<T> &mdaOut,
    MultidimArray<std::complex<T> > &inFourier, MultidimArray<std::complex<T> > &outFourier);

void selfScaleToSizeFourier(int Zdim, int Ydim, int Xdim, MultidimArray<double> &mda, int nthreads=1);

void selfScaleToSizeFourier(int Ydim, int Xdim, MultidimArray<double> &mda, int nthreads=1);
void selfScaleToSizeFourier(int Zdim, int Ydim, int Xdim, MultidimArrayGeneric &mda, int nthreads=1);
void selfScaleToSizeFourier(int Ydim, int Xdim, MultidimArrayGeneric &mda, int nthreads=1);

#define POWER_SPECTRUM 0
#define AMPLITUDE_SPECTRUM 1

void getSpectrum(MultidimArray<double> &Min,
                 MultidimArray<double> &spectrum,
                 int spectrum_type=AMPLITUDE_SPECTRUM);

void divideBySpectrum(MultidimArray<double> &Min,
                      MultidimArray<double> &spectrum,
                      bool leave_origin_intact=false);

void multiplyBySpectrum(MultidimArray<double> &Min,
                        MultidimArray<double> &spectrum,
                        bool leave_origin_intact=false);

void whitenSpectrum(MultidimArray<double> &Min,
                    MultidimArray<double> &Mout,
                    int spectrum_type=AMPLITUDE_SPECTRUM,
                    bool leave_origin_intact=false);

void adaptSpectrum(MultidimArray<double> &Min,
                   MultidimArray<double> &Mout,
                   const MultidimArray<double> &spectrum_ref,
                   int spectrum_type=AMPLITUDE_SPECTRUM,
                   bool leave_origin_intact=false);

void randomizePhases(MultidimArray<double> &Min, double wRandom);
#endif