rfft2.hpp

Go to the documentation of this file.

#pragma once
 
#include <complex>
 
#include "NumCpp/Core/Internal/StaticAsserts.hpp"
#include "NumCpp/Core/Types.hpp"
#include "NumCpp/Functions/complex.hpp"
#include "NumCpp/NdArray.hpp"
 
namespace nc::fft
{
    namespace detail
    {
        //===========================================================================
        // Method Description:
        inline NdArray<std::complex<double>> rfft2_internal(const NdArray<std::complex<double>>& x, const Shape& shape)
        {
            if (shape.rows == 0 || shape.cols == 0)
            {
                return {};
            }
 
            const auto realN  = shape.cols / 2 + 1;
            auto       result = NdArray<std::complex<double>>(shape.rows, realN);
 
            stl_algorithms::for_each(result.begin(),
                                     result.end(),
                                     [&](auto& resultElement)
                                     {
                                         const auto i  = &resultElement - result.data();
                                         const auto k  = static_cast<double>(i / realN);
                                         const auto l  = static_cast<double>(i % realN);
                                         resultElement = std::complex<double>{ 0., 0. };
                                         for (auto m = 0u; m < std::min(shape.rows, x.numRows()); ++m)
                                         {
                                             for (auto n = 0u; n < std::min(shape.cols, x.numCols()); ++n)
                                             {
                                                 const auto angle =
                                                     -constants::twoPi *
                                                     (((static_cast<double>(m) * k) / static_cast<double>(shape.rows)) +
                                                      ((static_cast<double>(n) * l) / static_cast<double>(shape.cols)));
                                                 resultElement += (x(m, n) * std::polar(1., angle));
                                             }
                                         }
                                     });
 
            return result;
        }
    } // namespace detail
 
    //============================================================================
    // Method Description:
    template<typename dtype>
    NdArray<std::complex<double>> rfft2(const NdArray<dtype>& inArray, const Shape& inShape)
    {
        STATIC_ASSERT_ARITHMETIC(dtype);
 
        const auto data = nc::complex<dtype, double>(inArray);
        return detail::rfft2_internal(data, inShape);
    }
 
    //============================================================================
    // Method Description:
    template<typename dtype>
    NdArray<std::complex<double>> rfft2(const NdArray<dtype>& inArray)
    {
        STATIC_ASSERT_ARITHMETIC(dtype);
 
        return rfft2(inArray, inArray.shape());
    }
} // namespace nc::fft