fft_expr.hpp

//=======================================================================
// Copyright (c) 2014-2023 Baptiste Wicht
// Distributed under the terms of the MIT License.
// (See accompanying file LICENSE or copy at
//  http://opensource.org/licenses/MIT)
//=======================================================================

#pragma once

#include "etl/expr/base_temporary_expr.hpp"

namespace etl {

template <etl_expr A, typename T, typename Impl>
struct fft_expr : base_temporary_expr_un<fft_expr<A, T, Impl>, A> {
    using value_type = T;                                    
    using this_type  = fft_expr<A, T, Impl>;                 
    using base_type  = base_temporary_expr_un<this_type, A>; 
    using sub_traits = decay_traits<A>;                      

    static constexpr auto storage_order = sub_traits::storage_order; 

    static constexpr bool gpu_computable = Impl::template gpu_computable<A>;

    explicit fft_expr(A a) : base_type(a) {
        //Nothing else to init
    }

    // Assignment functions

    template <same_dimensions<A> C>
    void assign_to(C&& c) const {
        inc_counter("temp:assign");

        Impl::apply(this->a(), c);
    }

    template <same_dimensions<A> L>
    void assign_add_to(L&& lhs) const {
        std_add_evaluate(*this, lhs);
    }

    template <same_dimensions<A> L>
    void assign_sub_to(L&& lhs) const {
        std_sub_evaluate(*this, lhs);
    }

    template <same_dimensions<A> L>
    void assign_mul_to(L&& lhs) const {
        std_mul_evaluate(*this, lhs);
    }

    template <same_dimensions<A> L>
    void assign_div_to(L&& lhs) const {
        std_div_evaluate(*this, lhs);
    }

    template <same_dimensions<A> L>
    void assign_mod_to(L&& lhs) const {
        std_mod_evaluate(*this, lhs);
    }

    friend std::ostream& operator<<(std::ostream& os, const fft_expr& expr) {
        return os << "fft(" << expr._a << ")";
    }
};

template <typename A, typename T, typename Impl>
struct etl_traits<etl::fft_expr<A, T, Impl>> {
    using expr_t     = etl::fft_expr<A, T, Impl>; 
    using sub_expr_t = std::decay_t<A>;           
    using sub_traits = etl_traits<sub_expr_t>;    
    using value_type = T;                         

    static constexpr bool is_etl         = true;                                 
    static constexpr bool is_transformer = false;                                
    static constexpr bool is_view        = false;                                
    static constexpr bool is_magic_view  = false;                                
    static constexpr bool is_fast        = sub_traits::is_fast;                  
    static constexpr bool is_linear      = false;                                
    static constexpr bool is_thread_safe = true;                                 
    static constexpr bool is_value       = false;                                
    static constexpr bool is_direct      = true;                                 
    static constexpr bool is_generator   = false;                                
    static constexpr bool is_padded      = false;                                
    static constexpr bool is_aligned     = true;                                 
    static constexpr bool is_temporary   = true;                                 
    static constexpr bool gpu_computable = is_gpu_t<value_type> && cuda_enabled; 
    static constexpr order storage_order = sub_traits::storage_order;            

    template <vector_mode_t V>
    static constexpr bool vectorizable = true;

    template <size_t DD>
    static constexpr size_t dim() {
        return decay_traits<A>::template dim<DD>();
    }

    static size_t dim(const expr_t& e, size_t d) {
        return etl::dim(e._a, d);
    }

    static size_t size(const expr_t& e) {
        return sub_traits::size(e._a);
    }

    static constexpr size_t size() {
        return sub_traits::size();
    }

    static constexpr size_t dimensions() {
        return sub_traits::dimensions();
    }

    static constexpr int complexity() noexcept {
        return -1;
    }
};

//Helpers to compute the type of the result

namespace detail {

template <typename A>
using fft_value_type = std::conditional_t<is_complex<A>, value_t<A>, std::complex<value_t<A>>>;

template <typename A>
using ifft_value_type = std::conditional_t<is_complex<A>, value_t<A>, std::complex<value_t<A>>>;

template <typename A>
using ifft_real_value_type = std::conditional_t<is_complex<A>, typename value_t<A>::value_type, value_t<A>>;

} //end of namespace detail

template <etl_expr A>
fft_expr<detail::build_type<A>, detail::fft_value_type<A>, detail::fft1_impl> fft_1d(A&& a) {
    return fft_expr<detail::build_type<A>, detail::fft_value_type<A>, detail::fft1_impl>{a};
}

template <etl_expr A, etl_expr C>
auto fft_1d(A&& a, C&& c) {
    validate_assign(c, a);

    c = fft_1d(a);
    return c;
}

template <etl_expr A>
fft_expr<detail::build_type<A>, detail::ifft_value_type<A>, detail::ifft1_impl> ifft_1d(A&& a) {
    return fft_expr<detail::build_type<A>, detail::ifft_value_type<A>, detail::ifft1_impl>{a};
}

template <etl_expr A, etl_expr C>
auto ifft_1d(A&& a, C&& c) {
    validate_assign(c, a);

    c = ifft_1d(a);
    return c;
}

template <etl_expr A>
fft_expr<detail::build_type<A>, detail::ifft_real_value_type<A>, detail::ifft1_real_impl> ifft_1d_real(A&& a) {
    return fft_expr<detail::build_type<A>, detail::ifft_real_value_type<A>, detail::ifft1_real_impl>{a};
}

template <etl_expr A, etl_expr C>
auto ifft_1d_real(A&& a, C&& c) {
    validate_assign(c, a);

    c = ifft_1d_real(a);
    return c;
}

template <etl_expr A>
fft_expr<detail::build_type<A>, detail::fft_value_type<A>, detail::fft2_impl> fft_2d(A&& a) {
    return fft_expr<detail::build_type<A>, detail::fft_value_type<A>, detail::fft2_impl>{a};
}

template <etl_expr A, etl_expr C>
auto fft_2d(A&& a, C&& c) {
    validate_assign(c, a);

    c = fft_2d(a);
    return c;
}

template <etl_expr A>
fft_expr<detail::build_type<A>, detail::ifft_value_type<A>, detail::ifft2_impl> ifft_2d(A&& a) {
    return fft_expr<detail::build_type<A>, detail::ifft_value_type<A>, detail::ifft2_impl>{a};
}

template <etl_expr A, etl_expr C>
auto ifft_2d(A&& a, C&& c) {
    validate_assign(c, a);

    c = ifft_2d(a);
    return c;
}

template <etl_expr A>
fft_expr<detail::build_type<A>, detail::ifft_real_value_type<A>, detail::ifft2_real_impl> ifft_2d_real(A&& a) {
    return fft_expr<detail::build_type<A>, detail::ifft_real_value_type<A>, detail::ifft2_real_impl>{a};
}

template <etl_expr A, etl_expr C>
auto ifft_2d_real(A&& a, C&& c) {
    validate_assign(c, a);

    c = ifft_2d_real(a);
    return c;
}

template <matrix A>
fft_expr<detail::build_type<A>, detail::fft_value_type<A>, detail::fft1_many_impl> fft_1d_many(A&& a) {
    return fft_expr<detail::build_type<A>, detail::fft_value_type<A>, detail::fft1_many_impl>{a};
}

template <matrix A, matrix C>
auto fft_1d_many(A&& a, C&& c) {
    validate_assign(c, a);

    c = fft_1d_many(a);
    return c;
}

template <matrix A>
fft_expr<detail::build_type<A>, detail::ifft_value_type<A>, detail::ifft1_many_impl> ifft_1d_many(A&& a) {
    return fft_expr<detail::build_type<A>, detail::ifft_value_type<A>, detail::ifft1_many_impl>{a};
}

template <matrix A, matrix C>
auto ifft_1d_many(A&& a, C&& c) {
    validate_assign(c, a);

    c = ifft_1d_many(a);
    return c;
}

template <deep_mat A>
fft_expr<detail::build_type<A>, detail::fft_value_type<A>, detail::fft2_many_impl> fft_2d_many(A&& a) {
    return fft_expr<detail::build_type<A>, detail::fft_value_type<A>, detail::fft2_many_impl>{a};
}

template <deep_mat A, deep_mat C>
auto fft_2d_many(A&& a, C&& c) {
    validate_assign(c, a);

    c = fft_2d_many(a);
    return c;
}

template <deep_mat A>
fft_expr<detail::build_type<A>, detail::ifft_value_type<A>, detail::ifft2_many_impl> ifft_2d_many(A&& a) {
    return fft_expr<detail::build_type<A>, detail::ifft_value_type<A>, detail::ifft2_many_impl>{a};
}

template <deep_mat A, deep_mat C>
auto ifft_2d_many(A&& a, C&& c) {
    validate_assign(c, a);

    c = ifft_2d_many(a);
    return c;
}

} //end of namespace etl