FFT/fft.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <complex>
 
#include "NumCpp/Core/Constants.hpp"
#include "NumCpp/Core/Internal/StaticAsserts.hpp"
#include "NumCpp/Core/Internal/StlAlgorithms.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>> fft_internal(const NdArray<std::complex<double>>& x, uint32 n)
        {
            if (n == 0)
            {
                return {};
            }
 
            auto result = NdArray<std::complex<double>>(1, n);
 
            stl_algorithms::for_each(result.begin(),
                                     result.end(),
                                     [&](auto& resultElement)
                                     {
                                         const auto k = static_cast<double>(&resultElement - result.data());
                                         const auto minusTwoPiKOverN = -constants::twoPi * k / static_cast<double>(n);
                                         resultElement               = std::complex<double>{ 0., 0. };
                                         for (auto m = 0u; m < std::min(n, x.size()); ++m)
                                         {
                                             const auto angle = minusTwoPiKOverN * static_cast<double>(m);
                                             resultElement += (x[m] * std::polar(1., angle));
                                         }
                                     });
 
            return result;
        }
    } // namespace detail
 
    //===========================================================================
    // Method Description:
    template<typename dtype>
    NdArray<std::complex<double>> fft(const NdArray<dtype>& inArray, uint32 inN, Axis inAxis = Axis::NONE)
    {
        STATIC_ASSERT_ARITHMETIC(dtype);
 
        switch (inAxis)
        {
            case Axis::NONE:
            {
                const auto data = nc::complex<dtype, double>(inArray);
                return detail::fft_internal(data, inN);
            }
            case Axis::COL:
            {
                auto        data           = nc::complex<dtype, double>(inArray);
                const auto& shape          = inArray.shape();
                auto        result         = NdArray<std::complex<double>>(shape.rows, inN);
                const auto  dataColSlice   = data.cSlice();
                const auto  resultColSlice = result.cSlice();
 
                for (uint32 row = 0; row < data.numRows(); ++row)
                {
                    const auto rowData   = data(row, dataColSlice);
                    const auto rowResult = detail::fft_internal(rowData, inN);
                    result.put(row, resultColSlice, rowResult);
                }
 
                return result;
            }
            case Axis::ROW:
            {
                return fft(inArray.transpose(), inN, Axis::COL).transpose();
            }
            default:
            {
                THROW_INVALID_ARGUMENT_ERROR("Unimplemented axis type.");
                return {};
            }
        }
    }
 
    //===========================================================================
    // Method Description:
    template<typename dtype>
    NdArray<std::complex<double>> fft(const NdArray<dtype>& inArray, Axis inAxis = Axis::NONE)
    {
        STATIC_ASSERT_ARITHMETIC(dtype);
 
        switch (inAxis)
        {
            case Axis::NONE:
            {
                return fft(inArray, inArray.size(), inAxis);
            }
            case Axis::COL:
            {
                return fft(inArray, inArray.numCols(), inAxis);
            }
            case Axis::ROW:
            {
                return fft(inArray, inArray.numRows(), inAxis);
            }
            default:
            {
                THROW_INVALID_ARGUMENT_ERROR("Unimplemented axis type.");
                return {};
            }
        }
    }
 
    //============================================================================
    // Method Description:
    template<typename dtype>
    NdArray<std::complex<double>> fft(const NdArray<std::complex<dtype>>& inArray, uint32 inN, Axis inAxis = Axis::NONE)
    {
        STATIC_ASSERT_ARITHMETIC(dtype);
 
        switch (inAxis)
        {
            case Axis::NONE:
            {
                const auto data = nc::complex<dtype, double>(inArray);
                return detail::fft_internal(data, inN);
            }
            case Axis::COL:
            {
                const auto  data           = nc::complex<dtype, double>(inArray);
                const auto& shape          = inArray.shape();
                auto        result         = NdArray<std::complex<double>>(shape.rows, inN);
                const auto  dataColSlice   = data.cSlice();
                const auto  resultColSlice = result.cSlice();
 
                for (uint32 row = 0; row < data.numRows(); ++row)
                {
                    const auto rowData   = data(row, dataColSlice);
                    const auto rowResult = detail::fft_internal(rowData, inN);
                    result.put(row, resultColSlice, rowResult);
                }
 
                return result;
            }
            case Axis::ROW:
            {
                return fft(inArray.transpose(), inN, Axis::COL).transpose();
            }
            default:
            {
                THROW_INVALID_ARGUMENT_ERROR("Unimplemented axis type.");
                return {};
            }
        }
    }
 
    //============================================================================
    // Method Description:
    template<typename dtype>
    NdArray<std::complex<double>> fft(const NdArray<std::complex<dtype>>& inArray, Axis inAxis = Axis::NONE)
    {
        STATIC_ASSERT_ARITHMETIC(dtype);
 
        switch (inAxis)
        {
            case Axis::NONE:
            {
                return fft(inArray, inArray.size(), inAxis);
            }
            case Axis::COL:
            {
                return fft(inArray, inArray.numCols(), inAxis);
            }
            case Axis::ROW:
            {
                return fft(inArray, inArray.numRows(), inAxis);
            }
            default:
            {
                THROW_INVALID_ARGUMENT_ERROR("Unimplemented axis type.");
                return {};
            }
        }
    }
} // namespace nc::fft