xcomplex.hpp

/***************************************************************************
 * Copyright (c) Johan Mabille, Sylvain Corlay and Wolf Vollprecht          *
 * Copyright (c) QuantStack                                                 *
 *                                                                          *
 * Distributed under the terms of the BSD 3-Clause License.                 *
 *                                                                          *
 * The full license is in the file LICENSE, distributed with this software. *
 ****************************************************************************/
 
#ifndef XTENSOR_COMPLEX_HPP
#define XTENSOR_COMPLEX_HPP
 
#include <type_traits>
#include <utility>
 
#include <xtl/xcomplex.hpp>
 
#include "xtensor/xbuilder.hpp"
#include "xtensor/xexpression.hpp"
#include "xtensor/xoffset_view.hpp"
 
namespace xt
{
 
    /******************************
     * real and imag declarations *
     ******************************/
 
    template <class E>
    decltype(auto) real(E&& e) noexcept;
 
    template <class E>
    decltype(auto) imag(E&& e) noexcept;
 
    /********************************
     * real and imag implementation *
     ********************************/
 
    namespace detail
    {
        template <bool iscomplex = true>
        struct complex_helper
        {
            template <class E>
            inline static auto real(E&& e) noexcept
            {
                using real_type = typename std::decay_t<E>::value_type::value_type;
                return xoffset_view<xclosure_t<E>, real_type, 0>(std::forward<E>(e));
            }
 
            template <class E>
            inline static auto imag(E&& e) noexcept
            {
                using real_type = typename std::decay_t<E>::value_type::value_type;
                return xoffset_view<xclosure_t<E>, real_type, sizeof(real_type)>(std::forward<E>(e));
            }
        };
 
        template <>
        struct complex_helper<false>
        {
            template <class E>
            inline static decltype(auto) real(E&& e) noexcept
            {
                return std::forward<E>(e);
            }
 
            template <class E>
            inline static auto imag(E&& e) noexcept
            {
                return zeros<typename std::decay_t<E>::value_type>(e.shape());
            }
        };
 
        template <bool isexpression = true>
        struct complex_expression_helper
        {
            template <class E>
            inline static decltype(auto) real(E&& e) noexcept
            {
                return detail::complex_helper<xtl::is_complex<typename std::decay_t<E>::value_type>::value>::real(
                    std::forward<E>(e)
                );
            }
 
            template <class E>
            inline static decltype(auto) imag(E&& e) noexcept
            {
                return detail::complex_helper<xtl::is_complex<typename std::decay_t<E>::value_type>::value>::imag(
                    std::forward<E>(e)
                );
            }
        };
 
        template <>
        struct complex_expression_helper<false>
        {
            template <class E>
            inline static decltype(auto) real(E&& e) noexcept
            {
                return xtl::forward_real(std::forward<E>(e));
            }
 
            template <class E>
            inline static decltype(auto) imag(E&& e) noexcept
            {
                return xtl::forward_imag(std::forward<E>(e));
            }
        };
    }
 
    template <class E>
    inline decltype(auto) real(E&& e) noexcept
    {
        return detail::complex_expression_helper<is_xexpression<std::decay_t<E>>::value>::real(std::forward<E>(e
        ));
    }
 
    template <class E>
    inline decltype(auto) imag(E&& e) noexcept
    {
        return detail::complex_expression_helper<is_xexpression<std::decay_t<E>>::value>::imag(std::forward<E>(e
        ));
    }
 
#define UNARY_COMPLEX_FUNCTOR(NS, NAME)             \
    struct NAME##_fun                               \
    {                                               \
        template <class T>                          \
        constexpr auto operator()(const T& t) const \
        {                                           \
            using NS::NAME;                         \
            return NAME(t);                         \
        }                                           \
                                                    \
        template <class B>                          \
        constexpr auto simd_apply(const B& t) const \
        {                                           \
            using NS::NAME;                         \
            return NAME(t);                         \
        }                                           \
    }
 
    namespace math
    {
        namespace detail
        {
            template <class T>
            constexpr std::complex<T> conj_impl(const std::complex<T>& c)
            {
                return std::complex<T>(c.real(), -c.imag());
            }
 
            template <class T>
            constexpr std::complex<T> conj_impl(const T& real)
            {
                return std::complex<T>(real, 0);
            }
 
#ifdef XTENSOR_USE_XSIMD
            template <class T, class A>
            xsimd::complex_batch_type_t<xsimd::batch<T, A>> conj_impl(const xsimd::batch<T, A>& z)
            {
                return xsimd::conj(z);
            }
#endif
        }
 
        UNARY_COMPLEX_FUNCTOR(std, norm);
        UNARY_COMPLEX_FUNCTOR(std, arg);
        UNARY_COMPLEX_FUNCTOR(detail, conj_impl);
    }
 
#undef UNARY_COMPLEX_FUNCTOR
 
    template <class E>
    inline auto conj(E&& e) noexcept
    {
        using functor = math::conj_impl_fun;
        using type = xfunction<functor, const_xclosure_t<E>>;
        return type(functor(), std::forward<E>(e));
    }
 
    template <class E>
    inline auto arg(E&& e) noexcept
    {
        using functor = math::arg_fun;
        using type = xfunction<functor, const_xclosure_t<E>>;
        return type(functor(), std::forward<E>(e));
    }
 
    template <class E>
    inline auto angle(E&& e, bool deg = false) noexcept
    {
        using value_type = xtl::complex_value_type_t<typename std::decay_t<E>::value_type>;
        value_type multiplier = 1.0;
        if (deg)
        {
            multiplier = value_type(180) / numeric_constants<value_type>::PI;
        }
        return arg(std::forward<E>(e)) * std::move(multiplier);
    }
 
    template <class E>
    inline auto norm(E&& e) noexcept
    {
        using functor = math::norm_fun;
        using type = xfunction<functor, const_xclosure_t<E>>;
        return type(functor(), std::forward<E>(e));
    }
}
#endif