FourierTransform.cpp

Go to the documentation of this file.

//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Base/Math/FourierTransform.cpp
//! @brief     Implements class FourierTransform.
//!
//! @homepage  http://www.bornagainproject.org
//! @license   GNU General Public License v3 or higher (see COPYING)
//! @copyright Forschungszentrum Jülich GmbH 2015
//! @authors   Scientific Computing Group at MLZ Garching
//
//  ************************************************************************************************
 
#include "Base/Math/FourierTransform.h"
#include "Base/Util/Assert.h"
#include <algorithm>
#include <cmath>
 
FourierTransform::FourierTransform() = default;
 
FourierTransform::Workspace::~Workspace()
{
    clear();
}
 
void FourierTransform::Workspace::clear()
{
    // rows (h) and columns (w) of the input 'source'
    h_src = 0;
    w_src = 0;
 
    delete[] in_src;
    in_src = nullptr;
 
    if (out_fftw)
        fftw_free((fftw_complex*)out_fftw);
    out_fftw = nullptr;
 
    if (p_forw_src != nullptr)
        fftw_destroy_plan(p_forw_src);
 
    fftw_cleanup();
}
 
/* ************************************************************************* */
// Fourier Transform in 2D
/* ************************************************************************* */
void FourierTransform::fft(const double2d_t& source, double2d_t& result)
{
    // rows (h) and columns (w) of the input 'source'
    int h_src = static_cast<int>(source.size());
    int w_src = static_cast<int>((!source.empty() ? source[0].size() : 0));
 
    // initialization
    init(h_src, w_src);
 
    // Compute the forward FT
    fftw_forward_FT(source);
 
    double* ptr = ws.out_fftw;
    result.clear();
 
    const size_t nh = ws.h_fftw;
 
    // Resize the array for holding the FT output to correct dimensions
    result.resize(nh, std::vector<double>(static_cast<size_t>(ws.w_fftw)));
 
    // Storing FT output
    for (size_t i = 0; i < nh; i++) {
        size_t k = nh - i;
        if (i == 0)
            k -= nh;
        for (size_t j = 0; j < static_cast<size_t>(ws.w_fftw / 2 + 1); j++) {
            result[i][j] = *ptr;
            size_t l = static_cast<size_t>(ws.w_fftw) - j;
            if (j != 0)
                result[k][l] = result[i][j];
            ptr += 2; // Only interested in the magnitudes of the complex Fourier coefficients
        }
    }
}
 
/* ************************************************************************* */
// Fourier Transform 2D shift - center array around zero frequency
/* ************************************************************************* */
void FourierTransform::fftshift(double2d_t& result) const
{
    // Centering FT around zero frequency
    size_t row_shift;
    if (result.size() % 2 == 0)
        row_shift = result.size() / 2;
    else
        row_shift = result.size() / 2 + 1;
 
    size_t col_shift;
    if (result[0].size() % 2 == 0)
        col_shift = result[0].size() / 2;
    else
        col_shift = result[0].size() / 2 + 1;
 
    // First, shifting the rows
    std::rotate(result.begin(), result.begin() + static_cast<int>(row_shift), result.end());
 
    // Second, shifting the cols
    for (size_t i = 0; i < static_cast<size_t>(ws.h_fftw); i++)
        std::rotate(result[i].begin(), result[i].begin() + static_cast<int>(col_shift),
                    result[i].end());
}
 
/* ************************************************************************* */
// Fourier Transform in 1D
/* ************************************************************************* */
void FourierTransform::fft(const double1d_t& source, double1d_t& result)
{
    // we simply create 2d arrays with length of first dimension equal to 1, and call 2d FT
    double2d_t source2d;
    source2d.push_back(source);
 
    double2d_t result2d;
    fft(source2d, result2d);
 
    ASSERT(result2d.size() == 1);
    result = result2d[0];
}
 
/* ************************************************************************* */
// Fourier Transform 1D shift - center around zero frequency
/* ************************************************************************* */
void FourierTransform::fftshift(double1d_t& result) const
{
    // Centering FT around zero frequency
    size_t col_shift;
    if (result.size() % 2 == 0)
        col_shift = result.size() / 2;
    else
        col_shift = result.size() / 2 + 1;
 
    std::rotate(result.begin(), result.begin() + static_cast<int>(col_shift), result.end());
}
 
/* ************************************************************************************ */
// initialise input and output arrays in workspace for fast Fourier transformation (fftw)
/* ************************************************************************************ */
void FourierTransform::init(int h_src, int w_src)
{
    ASSERT(h_src);
    ASSERT(w_src);
 
    ws.clear();
    ws.h_src = h_src;
    ws.w_src = w_src;
 
    ws.h_fftw = h_src;
    ws.w_fftw = w_src;
 
    ws.in_src = new double[static_cast<size_t>(ws.h_fftw * ws.w_fftw)];
    ws.out_fftw = static_cast<double*>(
        fftw_malloc(sizeof(fftw_complex) * static_cast<size_t>(ws.h_fftw * (ws.w_fftw / 2 + 1))));
 
    // Initialization of the plans
    ws.p_forw_src = fftw_plan_dft_r2c_2d(ws.h_fftw, ws.w_fftw, ws.in_src,
                                         (fftw_complex*)ws.out_fftw, FFTW_ESTIMATE);
    // ws.p_forw_src = fftw_plan_dft_r2c_2d(ws.h_fftw, ws.w_fftw, ws.in_src,
    //                                     static_cast<double*>(ws.out_src), FFTW_ESTIMATE);
 
    ASSERT(ws.p_forw_src);
}
 
/* ************************************************************************* */
// initialise input and output arrays for fast Fourier transformation
/* ************************************************************************* */
 
void FourierTransform::fftw_forward_FT(const double2d_t& src)
{
    ASSERT(ws.h_fftw > 0);
    ASSERT(ws.w_fftw > 0);
 
    double *ptr, *ptr_end;
 
    // Initializing the content of ws.in_src to zero for all elements
    for (ptr = ws.in_src, ptr_end = ws.in_src + ws.h_fftw * ws.w_fftw; ptr != ptr_end; ++ptr)
        *ptr = 0.0;
 
    // Building the input signal for fftw algorithm
    for (size_t row = 0; row < static_cast<size_t>(ws.h_src); ++row)
        for (size_t col = 0; col < static_cast<size_t>(ws.w_src); ++col)
            ws.in_src[(static_cast<int>(row) % ws.h_fftw) * ws.w_fftw
                      + (static_cast<int>(col) % ws.w_fftw)] += src[row][col];
 
    // Computing the FFT with fftw plan
    fftw_execute(ws.p_forw_src);
 
    for (ptr = ws.out_fftw, ptr_end = ws.out_fftw + 2 * ws.h_fftw * (ws.w_fftw / 2 + 1);
         ptr != ptr_end; ++ptr) {
        double re_out = *ptr;
        double im_out = *(++ptr);
 
        // magnitude
        *(ptr - 1) = sqrt(re_out * re_out + im_out * im_out);
        // phase
        *ptr = atan2(im_out, re_out);
    }
}