Homec4science
Diffusion µSpectre prototype implementation (master) Browse

muSpectre_proto/src/system/fftw_engine_r2c.cc
8353160ae1bdmaster

Last Change
Raw File

fftw_engine_r2c.cc
View Options

/**
* file fftw_engine_r2c.cc
*
* @author Till Junge <till.junge@epfl.ch>
*
* @date 13 May 2017
*
* @brief fft_engine using real-data fftw transform
*
* @section LICENCE
*
* Copyright (C) 2017 Till Junge
*
* µSpectre 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 3, or (at
* your option) any later version.
*
* µSpectre 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 GNU Emacs; see the file COPYING. If not, write to the
* Free Software Foundation, Inc., 59 Temple Place - Suite 330,
* Boston, MA 02111-1307, USA.
*/
#include <complex>
#include <array>
#include <numeric>
#include <unsupported/Eigen/CXX11/Tensor>
#include "system/fftw_engine_r2c.hh"
#include <stdexcept>
namespace muSpectre {
/** here, I'm doing everything out-of plas, so that the real input arrays
* can be contiguous */
template<Dim_t DimS, Dim_t DimM>
FFTW_EngineR2C<DimS, DimM>::FFTW_EngineR2C(Ccoord nb_pixels, FFT_PlanFlags flags)
:FFT_Engine<DimS, DimM>(nb_pixels),
nb_hermitian(this->tot_nb_pixels/nb_pixels.back() * (nb_pixels.back()/2+1))
{
this->work_space =
fftw_alloc_complex(this->nb_hermitian* DimM*DimM);
this->init(flags);
}
//----------------------------------------------------------------------------//
template<Dim_t DimS, Dim_t DimM>
FFTW_EngineR2C<DimS, DimM>::~FFTW_EngineR2C()
{
fftw_destroy_plan(this->plan_fft);
fftw_destroy_plan(this->plan_ifft);
fftw_free(this->work_space);
fftw_cleanup();
}
//----------------------------------------------------------------------------//
template<Dim_t DimS, Dim_t DimM>
void FFTW_EngineR2C<DimS, DimM>::init(FFT_PlanFlags inflags) {
const int & rank = DimS;
std::array<int, DimS> narr;
const int * const n = &narr[0];
std::copy(this->nb_pixels.begin(), this->nb_pixels.end(), narr.begin());
int howmany = DimM*DimM;
//temporary buffer for plan
Real * r_work_space = fftw_alloc_real(this->tot_nb_pixels*howmany);
Real * in = r_work_space;
const int * const inembed = nullptr;//nembed are tricky: they refer to physical layout
int istride = howmany;
int idist = 1;
fftw_complex * out = this->work_space;
const int * const onembed = nullptr;
int ostride = howmany;
int odist = idist;
unsigned int flags;
if (inflags == FFT_PlanFlags::estimate) {
flags = FFTW_ESTIMATE;
}
if (inflags == FFT_PlanFlags::measure) {
flags = FFTW_MEASURE;
}
if (inflags == FFT_PlanFlags::patient) {
flags = FFTW_PATIENT;
}
this->plan_fft = fftw_plan_many_dft_r2c(rank, n, howmany, in, inembed, istride,
idist, out, onembed, ostride, odist,
FFTW_PRESERVE_INPUT | flags);
if (this->plan_fft == nullptr) {
throw std::runtime_error("Plan failed");
}
fftw_complex * i_in = this->work_space;
Real * i_out = r_work_space;
this->plan_ifft = fftw_plan_many_dft_c2r(rank, n, howmany, i_in, inembed,
istride, idist, i_out, onembed,
ostride, odist, flags);
if (this->plan_ifft == nullptr) {
throw std::runtime_error("Plan failed");
}
fftw_free(r_work_space);
}
//----------------------------------------------------------------------------//
template<Dim_t DimS, Dim_t DimM>
FFTW_EngineR2C<DimS, DimM>::index_iterator::index_iterator(const FFTW_EngineR2C<DimS, DimM> & parent, bool begin)
:parent(parent), index(begin ? 0: parent.nb_hermitian)
{
this->half_strides.back() = this->full_strides.back() = 1;
this->full_strides[DimS-2] = this->parent.nb_pixels[DimS-1];
this->half_strides[DimS-2]= this->full_strides[DimS-2]/2+1;
for (Dim_t i = DimS-3; i >= 0; --i) {
//the full strides are cumulative
this->full_strides[i] = this->parent.nb_pixels[i+1] * this->full_strides[i+1];
// the half strides are not (operator* divided tmp_index successively)
this->half_strides[i] = this->parent.nb_pixels[i+1];
}
}
//----------------------------------------------------------------------------//
template <Dim_t DimS, Dim_t DimM>
void FFTW_EngineR2C<DimS, DimM>::convolve(const Real *const Ghat,
const Real * const in,
Real * out) {
fftw_execute_dft_r2c(this->plan_fft, const_cast<Real*>(in),
this->work_space);
//----------------------------------------------------------------------------//
const std::array<Eigen::IndexPair<int>,2>
prod_dims {Eigen::IndexPair<int>{2, 1},
Eigen::IndexPair<int>{3, 0}};
for (auto pix_id:*this) {
const size_t & real_pix_id = std::get<0>(pix_id);
const size_t & complex_pix_id = std::get<1>(pix_id);
using T4shape = Eigen::Sizes<DimM, DimM,
DimM,DimM>;
using T2shape = Eigen::Sizes<DimM, DimM>;
using T4type = Eigen::TensorFixedSize<Real, T4shape>;
using T2type = Eigen::TensorFixedSize<Complex, T2shape>;
using T4map = Eigen::TensorMap<T4type>;
using T2map = Eigen::TensorMap<T2type>;
const Real* const G_ptr = &Ghat[real_pix_id*DimM*DimM*DimM*DimM];
const T4map G(const_cast<Real*>(G_ptr),DimM, DimM, DimM,DimM);
Complex * P_ptr =
reinterpret_cast<Complex*>(&this->work_space[complex_pix_id*DimM*DimM]);
T2map P(P_ptr, DimM, DimM);
P = 1./static_cast<Complex>(this->tot_nb_pixels)*G.template cast<Complex>().contract(P, prod_dims).eval();
}
fftw_execute_dft_c2r(this->plan_ifft, this->work_space, out);
}
//----------------------------------------------------------------------------//
template <Dim_t DimS, Dim_t DimM>
typename FFTW_EngineR2C<DimS, DimM>::index_iterator
FFTW_EngineR2C<DimS, DimM>::begin() {
return index_iterator(*this);
}
//----------------------------------------------------------------------------//
template <Dim_t DimS, Dim_t DimM>
typename FFTW_EngineR2C<DimS, DimM>::index_iterator
FFTW_EngineR2C<DimS, DimM>::end() {
return index_iterator(*this, false);
}
template class FFTW_EngineR2C<1, 1>;
template class FFTW_EngineR2C<2, 1>;
template class FFTW_EngineR2C<2, 2>;
template class FFTW_EngineR2C<3, 3>;
} // muSpectre
c4science · Help