Subversion Repositories – MLAB_E

(root)//projekty/Laserovy_projektor/SW/mando/fourier.tcc @ 515

Rev 178

#include <fftw3.h>
#include <boost/array.hpp>

namespace mimas {

// Helping function-object for fft and invfft.
template< typename Real >
struct dft_
{
  void operator()( int rank, const int *n, const std::complex< Real > *in,
                   std::complex< Real > *out, int sign );
};

template< typename Real, std::size_t Dim, typename TPtr >
boost::multi_array< std::complex< Real >, Dim > fft_t< Real, Dim, TPtr >::operator()
  ( const boost::const_multi_array_ref< std::complex< Real >, Dim, TPtr > &field )
{
  boost::array< int, Dim > shape;
  std::copy( field.shape(), field.shape() + Dim, shape.begin() );
  
  boost::multi_array< std::complex< Real >, Dim > retVal( shape );

  dft_< Real > dft;
  dft( Dim, shape.begin(), field.data(), retVal.data(), -1 );

  return retVal;
}

template< typename Real, std::size_t Dim, typename TPtr >
boost::multi_array< std::complex< Real >, Dim > invfft_t< Real, Dim, TPtr >::operator()
  ( const boost::const_multi_array_ref< std::complex< Real >, Dim, TPtr > &field )
{
  boost::array< int, Dim > shape;
  std::copy( field.shape(), field.shape() + Dim, shape.begin() );

  boost::multi_array< std::complex< Real >, Dim > retVal( shape );
    
  dft_< Real > dft;
  dft( Dim, shape.begin(), field.data(), retVal.data(), +1 );

  return retVal * std::complex< Real >( 1.0 / retVal.num_elements() );
}

// Helping function-object for rfft.
template< typename Real >
struct dft_r2c_
{
  void operator()( int rank, const int *n, const Real *in,
                   std::complex< Real > *out );
};

template< typename Real, std::size_t Dim, typename TPtr >
boost::multi_array< std::complex< Real >, Dim > rfft_t< Real, Dim, TPtr >::operator()
  ( const boost::const_multi_array_ref< Real, Dim, TPtr > &field )
{
  assert( field.shape()[ Dim - 1 ] % 2 == 0 );

  boost::array< int, Dim > shape1, shape2;
  std::copy( field.shape(), field.shape() + Dim, shape1.begin() );
  std::copy( field.shape(), field.shape() + Dim, shape2.begin() );
  shape1[ Dim - 1 ] = shape1[ Dim - 1 ] / 2 + 1;
  
  boost::multi_array< std::complex< Real >, Dim > retVal( shape1 );

  dft_r2c_< Real > dft_r2c;
  dft_r2c( Dim, shape2.begin(), field.data(), retVal.data() );
  
  return retVal;
}

// Helping function-object for invrfft.
template< typename Real >
struct dft_c2r_
{
  void operator()( int rank, const int *n, const std::complex< Real > *in,
                   Real *out );
};

  template< typename Real, std::size_t Dim, typename TPtr >
boost::multi_array< Real, Dim > invrfft_t< Real, Dim, TPtr >::operator()
  ( const boost::const_multi_array_ref< std::complex< Real >, Dim, TPtr > &field )
{
  // Doesn't preserve input.
  boost::multi_array< std::complex< Real >, Dim > copy( field );

  boost::array< int, Dim > shape;
  std::copy( field.shape(), field.shape() + Dim, shape.begin() );
  shape[ Dim - 1 ] = ( shape[ Dim - 1 ] - 1 ) * 2;

  boost::multi_array< Real, Dim > retVal( shape );

  dft_c2r_< Real > dft_c2r;
  dft_c2r( Dim, shape.begin(), copy.data(), retVal.data() );

  return retVal * (Real)( 1.0 / retVal.num_elements() );
}

}