Kokkos_Complex.hpp

/*
//@HEADER
// ************************************************************************
//
//                        Kokkos v. 3.0
//       Copyright (2020) National Technology & Engineering
//               Solutions of Sandia, LLC (NTESS).
//
// Under the terms of Contract DE-NA0003525 with NTESS,
// the U.S. Government retains certain rights in this software.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
// 1. Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright
// notice, this list of conditions and the following disclaimer in the
// documentation and/or other materials provided with the distribution.
//
// 3. Neither the name of the Corporation nor the names of the
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY NTESS "AS IS" AND ANY
// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NTESS OR THE
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
//
// Questions? Contact Christian R. Trott (crtrott@sandia.gov)
//
// ************************************************************************
//@HEADER
*/
#ifndef KOKKOS_COMPLEX_HPP
#define KOKKOS_COMPLEX_HPP
 
#include <Kokkos_Atomic.hpp>
#include <Kokkos_MathematicalFunctions.hpp>
#include <Kokkos_NumericTraits.hpp>
#include <impl/Kokkos_Error.hpp>
#include <complex>
#include <type_traits>
#include <iosfwd>
 
namespace Kokkos {
 
template <class RealType>
class
#ifdef KOKKOS_ENABLE_COMPLEX_ALIGN
    alignas(2 * sizeof(RealType))
#endif
        complex {
 private:
  RealType re_{};
  RealType im_{};
 
 public:
  using value_type = RealType;
 
  KOKKOS_DEFAULTED_FUNCTION
  complex() noexcept = default;
 
  KOKKOS_DEFAULTED_FUNCTION
  complex(const complex&) noexcept = default;
 
  KOKKOS_DEFAULTED_FUNCTION
  complex& operator=(const complex&) noexcept = default;
 
  template <class RType,
            typename std::enable_if<std::is_convertible<RType, RealType>::value,
                                    int>::type = 0>
  KOKKOS_INLINE_FUNCTION complex(const complex<RType>& other) noexcept
      // Intentionally do the conversions implicitly here so that users don't
      // get any warnings about narrowing, etc., that they would expect to get
      // otherwise.
      : re_(other.real()), im_(other.imag()) {}
 
  KOKKOS_INLINE_FUNCTION
  complex(const std::complex<RealType>& src) noexcept
      // We can use this aspect of the standard to avoid calling
      // non-device-marked functions `std::real` and `std::imag`: "For any
      // object z of type complex<T>, reinterpret_cast<T(&)[2]>(z)[0] is the
      // real part of z and reinterpret_cast<T(&)[2]>(z)[1] is the imaginary
      // part of z." Now we don't have to provide a whole bunch of the overloads
      // of things taking either Kokkos::complex or std::complex
      : re_(reinterpret_cast<const RealType (&)[2]>(src)[0]),
        im_(reinterpret_cast<const RealType (&)[2]>(src)[1]) {}
 
  // TODO: make explicit.  DJS 2019-08-28
  operator std::complex<RealType>() const noexcept {
    return std::complex<RealType>(re_, im_);
  }
 
  KOKKOS_INLINE_FUNCTION complex(const RealType& val) noexcept
      : re_(val), im_(static_cast<RealType>(0)) {}
 
  KOKKOS_INLINE_FUNCTION
  complex(const RealType& re, const RealType& im) noexcept : re_(re), im_(im) {}
 
  KOKKOS_INLINE_FUNCTION complex& operator=(const RealType& val) noexcept {
    re_ = val;
    im_ = RealType(0);
    return *this;
  }
 
  complex& operator=(const std::complex<RealType>& src) noexcept {
    *this = complex(src);
    return *this;
  }
 
  KOKKOS_INLINE_FUNCTION
  KOKKOS_CONSTEXPR_14 RealType& imag() noexcept { return im_; }
 
  KOKKOS_INLINE_FUNCTION
  KOKKOS_CONSTEXPR_14 RealType& real() noexcept { return re_; }
 
  KOKKOS_INLINE_FUNCTION
  constexpr RealType imag() const noexcept { return im_; }
 
  KOKKOS_INLINE_FUNCTION
  constexpr RealType real() const noexcept { return re_; }
 
  KOKKOS_INLINE_FUNCTION
  KOKKOS_CONSTEXPR_14
  void imag(RealType v) noexcept { im_ = v; }
 
  KOKKOS_INLINE_FUNCTION
  KOKKOS_CONSTEXPR_14
  void real(RealType v) noexcept { re_ = v; }
 
  KOKKOS_CONSTEXPR_14 KOKKOS_INLINE_FUNCTION complex& operator+=(
      const complex<RealType>& src) noexcept {
    re_ += src.re_;
    im_ += src.im_;
    return *this;
  }
 
  KOKKOS_CONSTEXPR_14 KOKKOS_INLINE_FUNCTION complex& operator+=(
      const RealType& src) noexcept {
    re_ += src;
    return *this;
  }
 
  KOKKOS_CONSTEXPR_14 KOKKOS_INLINE_FUNCTION complex& operator-=(
      const complex<RealType>& src) noexcept {
    re_ -= src.re_;
    im_ -= src.im_;
    return *this;
  }
 
  KOKKOS_CONSTEXPR_14 KOKKOS_INLINE_FUNCTION complex& operator-=(
      const RealType& src) noexcept {
    re_ -= src;
    return *this;
  }
 
  KOKKOS_CONSTEXPR_14 KOKKOS_INLINE_FUNCTION complex& operator*=(
      const complex<RealType>& src) noexcept {
    const RealType realPart = re_ * src.re_ - im_ * src.im_;
    const RealType imagPart = re_ * src.im_ + im_ * src.re_;
    re_                     = realPart;
    im_                     = imagPart;
    return *this;
  }
 
  KOKKOS_CONSTEXPR_14 KOKKOS_INLINE_FUNCTION complex& operator*=(
      const RealType& src) noexcept {
    re_ *= src;
    im_ *= src;
    return *this;
  }
 
  // Conditional noexcept, just in case RType throws on divide-by-zero
  KOKKOS_CONSTEXPR_14 KOKKOS_INLINE_FUNCTION complex& operator/=(
      const complex<RealType>& y) noexcept(noexcept(RealType{} / RealType{})) {
    using Kokkos::Experimental::fabs;
    // Scale (by the "1-norm" of y) to avoid unwarranted overflow.
    // If the real part is +/-Inf and the imaginary part is -/+Inf,
    // this won't change the result.
    const RealType s = fabs(y.real()) + fabs(y.imag());
 
    // If s is 0, then y is zero, so x/y == real(x)/0 + i*imag(x)/0.
    // In that case, the relation x/y == (x/s) / (y/s) doesn't hold,
    // because y/s is NaN.
    // TODO mark this branch unlikely
    if (s == RealType(0)) {
      this->re_ /= s;
      this->im_ /= s;
    } else {
      const complex x_scaled(this->re_ / s, this->im_ / s);
      const complex y_conj_scaled(y.re_ / s, -(y.im_) / s);
      const RealType y_scaled_abs =
          y_conj_scaled.re_ * y_conj_scaled.re_ +
          y_conj_scaled.im_ * y_conj_scaled.im_;  // abs(y) == abs(conj(y))
      *this = x_scaled * y_conj_scaled;
      *this /= y_scaled_abs;
    }
    return *this;
  }
 
  KOKKOS_CONSTEXPR_14
  KOKKOS_INLINE_FUNCTION complex& operator/=(
      const std::complex<RealType>& y) noexcept(noexcept(RealType{} /
                                                         RealType{})) {
    using Kokkos::Experimental::fabs;
    // Scale (by the "1-norm" of y) to avoid unwarranted overflow.
    // If the real part is +/-Inf and the imaginary part is -/+Inf,
    // this won't change the result.
    const RealType s = fabs(y.real()) + fabs(y.imag());
 
    // If s is 0, then y is zero, so x/y == real(x)/0 + i*imag(x)/0.
    // In that case, the relation x/y == (x/s) / (y/s) doesn't hold,
    // because y/s is NaN.
    if (s == RealType(0)) {
      this->re_ /= s;
      this->im_ /= s;
    } else {
      const complex x_scaled(this->re_ / s, this->im_ / s);
      const complex y_conj_scaled(y.re_ / s, -(y.im_) / s);
      const RealType y_scaled_abs =
          y_conj_scaled.re_ * y_conj_scaled.re_ +
          y_conj_scaled.im_ * y_conj_scaled.im_;  // abs(y) == abs(conj(y))
      *this = x_scaled * y_conj_scaled;
      *this /= y_scaled_abs;
    }
    return *this;
  }
 
  KOKKOS_CONSTEXPR_14 KOKKOS_INLINE_FUNCTION complex& operator/=(
      const RealType& src) noexcept(noexcept(RealType{} / RealType{})) {
    re_ /= src;
    im_ /= src;
    return *this;
  }
 
  //---------------------------------------------------------------------------
  // TODO: refactor Kokkos reductions to remove dependency on
  // volatile member overloads since they are being deprecated in c++20
  //---------------------------------------------------------------------------
 
  template <class RType,
            typename std::enable_if<std::is_convertible<RType, RealType>::value,
                                    int>::type = 0>
  KOKKOS_INLINE_FUNCTION complex(const volatile complex<RType>& src) noexcept
      // Intentionally do the conversions implicitly here so that users don't
      // get any warnings about narrowing, etc., that they would expect to get
      // otherwise.
      : re_(src.re_), im_(src.im_) {}
 
  //
  // Templated, so as not to be a copy assignment operator (Kokkos issue #2577)
  // Intended to behave as
  //    void operator=(const complex&) volatile noexcept
  //
  // Use cases:
  //    complex r;
  //    const complex cr;
  //    volatile complex vl;
  //    vl = r;
  //    vl = cr;
  template <class Complex,
            typename std::enable_if<std::is_same<Complex, complex>::value,
                                    int>::type = 0>
  KOKKOS_INLINE_FUNCTION void operator=(const Complex& src) volatile noexcept {
    re_ = src.re_;
    im_ = src.im_;
    // We deliberately do not return anything here.  See explanation
    // in public documentation above.
  }
 
  // TODO Should this return void like the other volatile assignment operators?
  //
  // Templated, so as not to be a copy assignment operator (Kokkos issue #2577)
  // Intended to behave as
  //    volatile complex& operator=(const volatile complex&) volatile noexcept
  //
  // Use cases:
  //    volatile complex vr;
  //    const volatile complex cvr;
  //    volatile complex vl;
  //    vl = vr;
  //    vl = cvr;
  template <class Complex,
            typename std::enable_if<std::is_same<Complex, complex>::value,
                                    int>::type = 0>
  KOKKOS_INLINE_FUNCTION volatile complex& operator=(
      const volatile Complex& src) volatile noexcept {
    re_ = src.re_;
    im_ = src.im_;
    return *this;
  }
 
  //
  // Templated, so as not to be a copy assignment operator (Kokkos issue #2577)
  // Intended to behave as
  //    complex& operator=(const volatile complex&) noexcept
  //
  // Use cases:
  //    volatile complex vr;
  //    const volatile complex cvr;
  //    complex l;
  //    l = vr;
  //    l = cvr;
  //
  template <class Complex,
            typename std::enable_if<std::is_same<Complex, complex>::value,
                                    int>::type = 0>
  KOKKOS_INLINE_FUNCTION complex& operator=(
      const volatile Complex& src) noexcept {
    re_ = src.re_;
    im_ = src.im_;
    return *this;
  }
 
  // Mirroring the behavior of the assignment operators from complex RHS in the
  // RealType RHS versions.
 
  KOKKOS_INLINE_FUNCTION void operator=(const volatile RealType& val) noexcept {
    re_ = val;
    im_ = RealType(0);
    // We deliberately do not return anything here.  See explanation
    // in public documentation above.
  }
 
  KOKKOS_INLINE_FUNCTION complex& operator=(
      const RealType& val) volatile noexcept {
    re_ = val;
    im_ = RealType(0);
    return *this;
  }
 
  // TODO Should this return void like the other volatile assignment operators?
  KOKKOS_INLINE_FUNCTION complex& operator=(
      const volatile RealType& val) volatile noexcept {
    re_ = val;
    im_ = RealType(0);
    return *this;
  }
 
  KOKKOS_INLINE_FUNCTION
  volatile RealType& imag() volatile noexcept { return im_; }
 
  KOKKOS_INLINE_FUNCTION
  volatile RealType& real() volatile noexcept { return re_; }
 
  KOKKOS_INLINE_FUNCTION
  RealType imag() const volatile noexcept { return im_; }
 
  KOKKOS_INLINE_FUNCTION
  RealType real() const volatile noexcept { return re_; }
 
  KOKKOS_INLINE_FUNCTION void operator+=(
      const volatile complex<RealType>& src) volatile noexcept {
    re_ += src.re_;
    im_ += src.im_;
  }
 
  KOKKOS_INLINE_FUNCTION void operator+=(
      const volatile RealType& src) volatile noexcept {
    re_ += src;
  }
 
  KOKKOS_INLINE_FUNCTION void operator*=(
      const volatile complex<RealType>& src) volatile noexcept {
    const RealType realPart = re_ * src.re_ - im_ * src.im_;
    const RealType imagPart = re_ * src.im_ + im_ * src.re_;
 
    re_ = realPart;
    im_ = imagPart;
  }
 
  KOKKOS_INLINE_FUNCTION void operator*=(
      const volatile RealType& src) volatile noexcept {
    re_ *= src;
    im_ *= src;
  }
 
  // TODO DSH 2019-10-7 why are there no volatile /= and friends?
};
 
//==============================================================================
// <editor-fold desc="Equality and inequality"> {{{1
 
// Note that this is not the same behavior as std::complex, which doesn't allow
// implicit conversions, but since this is the way we had it before, we have
// to do it this way now.
 
template <class RealType1, class RealType2>
KOKKOS_INLINE_FUNCTION bool operator==(complex<RealType1> const& x,
                                       complex<RealType2> const& y) noexcept {
  using common_type = typename std::common_type<RealType1, RealType2>::type;
  return common_type(x.real()) == common_type(y.real()) &&
         common_type(x.imag()) == common_type(y.imag());
}
 
// TODO (here and elsewhere) decide if we should convert to a Kokkos::complex
//      and do the comparison in a device-marked function
template <class RealType1, class RealType2>
inline bool operator==(std::complex<RealType1> const& x,
                       complex<RealType2> const& y) noexcept {
  using common_type = typename std::common_type<RealType1, RealType2>::type;
  return common_type(x.real()) == common_type(y.real()) &&
         common_type(x.imag()) == common_type(y.imag());
}
 
template <class RealType1, class RealType2>
inline bool operator==(complex<RealType1> const& x,
                       std::complex<RealType2> const& y) noexcept {
  using common_type = typename std::common_type<RealType1, RealType2>::type;
  return common_type(x.real()) == common_type(y.real()) &&
         common_type(x.imag()) == common_type(y.imag());
}
 
template <
    class RealType1, class RealType2,
    // Constraints to avoid participation in oparator==() for every possible RHS
    typename std::enable_if<std::is_convertible<RealType2, RealType1>::value,
                            int>::type = 0>
KOKKOS_INLINE_FUNCTION bool operator==(complex<RealType1> const& x,
                                       RealType2 const& y) noexcept {
  using common_type = typename std::common_type<RealType1, RealType2>::type;
  return common_type(x.real()) == common_type(y) &&
         common_type(x.imag()) == common_type(0);
}
 
template <
    class RealType1, class RealType2,
    // Constraints to avoid participation in oparator==() for every possible RHS
    typename std::enable_if<std::is_convertible<RealType1, RealType2>::value,
                            int>::type = 0>
KOKKOS_INLINE_FUNCTION bool operator==(RealType1 const& x,
                                       complex<RealType2> const& y) noexcept {
  using common_type = typename std::common_type<RealType1, RealType2>::type;
  return common_type(x) == common_type(y.real()) &&
         common_type(0) == common_type(y.imag());
}
 
template <class RealType1, class RealType2>
KOKKOS_INLINE_FUNCTION bool operator!=(complex<RealType1> const& x,
                                       complex<RealType2> const& y) noexcept {
  using common_type = typename std::common_type<RealType1, RealType2>::type;
  return common_type(x.real()) != common_type(y.real()) ||
         common_type(x.imag()) != common_type(y.imag());
}
 
template <class RealType1, class RealType2>
inline bool operator!=(std::complex<RealType1> const& x,
                       complex<RealType2> const& y) noexcept {
  using common_type = typename std::common_type<RealType1, RealType2>::type;
  return common_type(x.real()) != common_type(y.real()) ||
         common_type(x.imag()) != common_type(y.imag());
}
 
template <class RealType1, class RealType2>
inline bool operator!=(complex<RealType1> const& x,
                       std::complex<RealType2> const& y) noexcept {
  using common_type = typename std::common_type<RealType1, RealType2>::type;
  return common_type(x.real()) != common_type(y.real()) ||
         common_type(x.imag()) != common_type(y.imag());
}
 
template <
    class RealType1, class RealType2,
    // Constraints to avoid participation in oparator==() for every possible RHS
    typename std::enable_if<std::is_convertible<RealType2, RealType1>::value,
                            int>::type = 0>
KOKKOS_INLINE_FUNCTION bool operator!=(complex<RealType1> const& x,
                                       RealType2 const& y) noexcept {
  using common_type = typename std::common_type<RealType1, RealType2>::type;
  return common_type(x.real()) != common_type(y) ||
         common_type(x.imag()) != common_type(0);
}
 
template <
    class RealType1, class RealType2,
    // Constraints to avoid participation in oparator==() for every possible RHS
    typename std::enable_if<std::is_convertible<RealType1, RealType2>::value,
                            int>::type = 0>
KOKKOS_INLINE_FUNCTION bool operator!=(RealType1 const& x,
                                       complex<RealType2> const& y) noexcept {
  using common_type = typename std::common_type<RealType1, RealType2>::type;
  return common_type(x) != common_type(y.real()) ||
         common_type(0) != common_type(y.imag());
}
 
// </editor-fold> end Equality and inequality }}}1
//==============================================================================
 
template <class RealType1, class RealType2>
KOKKOS_INLINE_FUNCTION
    complex<typename std::common_type<RealType1, RealType2>::type>
    operator+(const complex<RealType1>& x,
              const complex<RealType2>& y) noexcept {
  return complex<typename std::common_type<RealType1, RealType2>::type>(
      x.real() + y.real(), x.imag() + y.imag());
}
 
template <class RealType1, class RealType2>
KOKKOS_INLINE_FUNCTION
    complex<typename std::common_type<RealType1, RealType2>::type>
    operator+(const complex<RealType1>& x, const RealType2& y) noexcept {
  return complex<typename std::common_type<RealType1, RealType2>::type>(
      x.real() + y, x.imag());
}
 
template <class RealType1, class RealType2>
KOKKOS_INLINE_FUNCTION
    complex<typename std::common_type<RealType1, RealType2>::type>
    operator+(const RealType1& x, const complex<RealType2>& y) noexcept {
  return complex<typename std::common_type<RealType1, RealType2>::type>(
      x + y.real(), y.imag());
}
 
template <class RealType>
KOKKOS_INLINE_FUNCTION complex<RealType> operator+(
    const complex<RealType>& x) noexcept {
  return complex<RealType>{+x.real(), +x.imag()};
}
 
template <class RealType1, class RealType2>
KOKKOS_INLINE_FUNCTION
    complex<typename std::common_type<RealType1, RealType2>::type>
    operator-(const complex<RealType1>& x,
              const complex<RealType2>& y) noexcept {
  return complex<typename std::common_type<RealType1, RealType2>::type>(
      x.real() - y.real(), x.imag() - y.imag());
}
 
template <class RealType1, class RealType2>
KOKKOS_INLINE_FUNCTION
    complex<typename std::common_type<RealType1, RealType2>::type>
    operator-(const complex<RealType1>& x, const RealType2& y) noexcept {
  return complex<typename std::common_type<RealType1, RealType2>::type>(
      x.real() - y, x.imag());
}
 
template <class RealType1, class RealType2>
KOKKOS_INLINE_FUNCTION
    complex<typename std::common_type<RealType1, RealType2>::type>
    operator-(const RealType1& x, const complex<RealType2>& y) noexcept {
  return complex<typename std::common_type<RealType1, RealType2>::type>(
      x - y.real(), -y.imag());
}
 
template <class RealType>
KOKKOS_INLINE_FUNCTION complex<RealType> operator-(
    const complex<RealType>& x) noexcept {
  return complex<RealType>(-x.real(), -x.imag());
}
 
template <class RealType1, class RealType2>
KOKKOS_INLINE_FUNCTION
    complex<typename std::common_type<RealType1, RealType2>::type>
    operator*(const complex<RealType1>& x,
              const complex<RealType2>& y) noexcept {
  return complex<typename std::common_type<RealType1, RealType2>::type>(
      x.real() * y.real() - x.imag() * y.imag(),
      x.real() * y.imag() + x.imag() * y.real());
}
 
template <class RealType1, class RealType2>
inline complex<typename std::common_type<RealType1, RealType2>::type> operator*(
    const std::complex<RealType1>& x, const complex<RealType2>& y) {
  return complex<typename std::common_type<RealType1, RealType2>::type>(
      x.real() * y.real() - x.imag() * y.imag(),
      x.real() * y.imag() + x.imag() * y.real());
}
 
template <class RealType1, class RealType2>
KOKKOS_INLINE_FUNCTION
    complex<typename std::common_type<RealType1, RealType2>::type>
    operator*(const RealType1& x, const complex<RealType2>& y) noexcept {
  return complex<typename std::common_type<RealType1, RealType2>::type>(
      x * y.real(), x * y.imag());
}
 
template <class RealType1, class RealType2>
KOKKOS_INLINE_FUNCTION
    complex<typename std::common_type<RealType1, RealType2>::type>
    operator*(const complex<RealType1>& y, const RealType2& x) noexcept {
  return complex<typename std::common_type<RealType1, RealType2>::type>(
      x * y.real(), x * y.imag());
}
 
template <class RealType>
KOKKOS_INLINE_FUNCTION RealType imag(const complex<RealType>& x) noexcept {
  return x.imag();
}
 
template <class RealType>
KOKKOS_INLINE_FUNCTION RealType real(const complex<RealType>& x) noexcept {
  return x.real();
}
 
template <class T>
KOKKOS_INLINE_FUNCTION complex<T> polar(const T& r, const T& theta = T()) {
  using Kokkos::Experimental::cos;
  using Kokkos::Experimental::sin;
  KOKKOS_EXPECTS(r >= 0);
  return complex<T>(r * cos(theta), r * sin(theta));
}
 
template <class RealType>
KOKKOS_INLINE_FUNCTION RealType abs(const complex<RealType>& x) {
  using Kokkos::Experimental::hypot;
  return hypot(x.real(), x.imag());
}
 
template <class T>
KOKKOS_INLINE_FUNCTION complex<T> pow(const complex<T>& x, const T& y) {
  using Kokkos::Experimental::atan2;
  using Kokkos::Experimental::pow;
  T r     = abs(x);
  T theta = atan2(x.imag(), x.real());
  return polar(pow(r, y), y * theta);
}
 
template <class T>
KOKKOS_INLINE_FUNCTION complex<T> pow(const T& x, const complex<T>& y) {
  return pow(complex<T>(x), y);
}
 
template <class T>
KOKKOS_INLINE_FUNCTION complex<T> pow(const complex<T>& x,
                                      const complex<T>& y) {
  using Kokkos::Experimental::log;
 
  return x == T() ? T() : exp(y * log(x));
}
 
namespace Impl {
// NOTE promote would also be useful for math functions
template <class T, bool = std::is_integral<T>::value>
struct promote {
  using type = double;
};
template <class T>
struct promote<T, false> {};
template <>
struct promote<long double> {
  using type = long double;
};
template <>
struct promote<double> {
  using type = double;
};
template <>
struct promote<float> {
  using type = float;
};
template <class T>
using promote_t = typename promote<T>::type;
template <class T, class U>
struct promote_2 {
  using type = decltype(promote_t<T>() + promote_t<U>());
};
template <class T, class U>
using promote_2_t = typename promote_2<T, U>::type;
}  // namespace Impl
 
template <class T, class U,
          class = std::enable_if_t<std::is_arithmetic<T>::value>>
KOKKOS_INLINE_FUNCTION complex<Impl::promote_2_t<T, U>> pow(
    const T& x, const complex<U>& y) {
  using type = Impl::promote_2_t<T, U>;
  return pow(type(x), complex<type>(y));
}
 
template <class T, class U,
          class = std::enable_if_t<std::is_arithmetic<U>::value>>
KOKKOS_INLINE_FUNCTION complex<Impl::promote_2_t<T, U>> pow(const complex<T>& x,
                                                            const U& y) {
  using type = Impl::promote_2_t<T, U>;
  return pow(complex<type>(x), type(y));
}
 
template <class T, class U>
KOKKOS_INLINE_FUNCTION complex<Impl::promote_2_t<T, U>> pow(
    const complex<T>& x, const complex<U>& y) {
  using type = Impl::promote_2_t<T, U>;
  return pow(complex<type>(x), complex<type>(y));
}
 
template <class RealType>
KOKKOS_INLINE_FUNCTION Kokkos::complex<RealType> sqrt(
    const complex<RealType>& x) {
  using Kokkos::Experimental::fabs;
  using Kokkos::Experimental::sqrt;
 
  RealType r = x.real();
  RealType i = x.imag();
 
  if (r == RealType()) {
    RealType t = sqrt(fabs(i) / 2);
    return Kokkos::complex<RealType>(t, i < RealType() ? -t : t);
  } else {
    RealType t = sqrt(2 * (abs(x) + fabs(r)));
    RealType u = t / 2;
    return r > RealType() ? Kokkos::complex<RealType>(u, i / t)
                          : Kokkos::complex<RealType>(fabs(i) / t,
                                                      i < RealType() ? -u : u);
  }
}
 
template <class RealType>
KOKKOS_INLINE_FUNCTION complex<RealType> conj(
    const complex<RealType>& x) noexcept {
  return complex<RealType>(real(x), -imag(x));
}
 
template <class RealType>
KOKKOS_INLINE_FUNCTION complex<RealType> exp(const complex<RealType>& x) {
  using Kokkos::Experimental::cos;
  using Kokkos::Experimental::exp;
  using Kokkos::Experimental::sin;
  return exp(x.real()) * complex<RealType>(cos(x.imag()), sin(x.imag()));
}
 
template <class RealType>
KOKKOS_INLINE_FUNCTION Kokkos::complex<RealType> log(
    const complex<RealType>& x) {
  using Kokkos::Experimental::atan2;
  using Kokkos::Experimental::log;
  RealType phi = atan2(x.imag(), x.real());
  return Kokkos::complex<RealType>(log(abs(x)), phi);
}
 
template <class RealType>
KOKKOS_INLINE_FUNCTION Kokkos::complex<RealType> sin(
    const complex<RealType>& x) {
  using Kokkos::Experimental::cos;
  using Kokkos::Experimental::cosh;
  using Kokkos::Experimental::sin;
  using Kokkos::Experimental::sinh;
  return Kokkos::complex<RealType>(sin(x.real()) * cosh(x.imag()),
                                   cos(x.real()) * sinh(x.imag()));
}
 
template <class RealType>
KOKKOS_INLINE_FUNCTION Kokkos::complex<RealType> cos(
    const complex<RealType>& x) {
  using Kokkos::Experimental::cos;
  using Kokkos::Experimental::cosh;
  using Kokkos::Experimental::sin;
  using Kokkos::Experimental::sinh;
  return Kokkos::complex<RealType>(cos(x.real()) * cosh(x.imag()),
                                   -sin(x.real()) * sinh(x.imag()));
}
 
template <class RealType>
KOKKOS_INLINE_FUNCTION Kokkos::complex<RealType> tan(
    const complex<RealType>& x) {
  return sin(x) / cos(x);
}
 
template <class RealType>
KOKKOS_INLINE_FUNCTION Kokkos::complex<RealType> sinh(
    const complex<RealType>& x) {
  using Kokkos::Experimental::cos;
  using Kokkos::Experimental::cosh;
  using Kokkos::Experimental::sin;
  using Kokkos::Experimental::sinh;
  return Kokkos::complex<RealType>(sinh(x.real()) * cos(x.imag()),
                                   cosh(x.real()) * sin(x.imag()));
}
 
template <class RealType>
KOKKOS_INLINE_FUNCTION Kokkos::complex<RealType> cosh(
    const complex<RealType>& x) {
  using Kokkos::Experimental::cos;
  using Kokkos::Experimental::cosh;
  using Kokkos::Experimental::sin;
  using Kokkos::Experimental::sinh;
  return Kokkos::complex<RealType>(cosh(x.real()) * cos(x.imag()),
                                   sinh(x.real()) * sin(x.imag()));
}
 
template <class RealType>
KOKKOS_INLINE_FUNCTION Kokkos::complex<RealType> tanh(
    const complex<RealType>& x) {
  return sinh(x) / cosh(x);
}
 
template <class RealType>
KOKKOS_INLINE_FUNCTION Kokkos::complex<RealType> asinh(
    const complex<RealType>& x) {
  return log(x + sqrt(x * x + RealType(1.0)));
}
 
template <class RealType>
KOKKOS_INLINE_FUNCTION Kokkos::complex<RealType> acosh(
    const complex<RealType>& x) {
  return RealType(2.0) * log(sqrt(RealType(0.5) * (x + RealType(1.0))) +
                             sqrt(RealType(0.5) * (x - RealType(1.0))));
}
 
template <class RealType>
KOKKOS_INLINE_FUNCTION Kokkos::complex<RealType> atanh(
    const complex<RealType>& x) {
  using Kokkos::Experimental::atan2;
  using Kokkos::Experimental::log;
 
  const RealType i2 = x.imag() * x.imag();
  const RealType r  = RealType(1.0) - i2 - x.real() * x.real();
 
  RealType p = RealType(1.0) + x.real();
  RealType m = RealType(1.0) - x.real();
 
  p = i2 + p * p;
  m = i2 + m * m;
 
  RealType phi = atan2(RealType(2.0) * x.imag(), r);
  return Kokkos::complex<RealType>(RealType(0.25) * (log(p) - log(m)),
                                   RealType(0.5) * phi);
}
 
template <class RealType>
KOKKOS_INLINE_FUNCTION Kokkos::complex<RealType> asin(
    const complex<RealType>& x) {
  Kokkos::complex<RealType> t =
      asinh(Kokkos::complex<RealType>(-x.imag(), x.real()));
  return Kokkos::complex<RealType>(t.imag(), -t.real());
}
 
template <class RealType>
KOKKOS_INLINE_FUNCTION Kokkos::complex<RealType> acos(
    const complex<RealType>& x) {
  using Kokkos::Experimental::acos;
  Kokkos::complex<RealType> t = asin(x);
  RealType pi_2               = acos(RealType(0.0));
  return Kokkos::complex<RealType>(pi_2 - t.real(), -t.imag());
}
 
template <class RealType>
KOKKOS_INLINE_FUNCTION Kokkos::complex<RealType> atan(
    const complex<RealType>& x) {
  using Kokkos::Experimental::atan2;
  using Kokkos::Experimental::log;
  const RealType r2 = x.real() * x.real();
  const RealType i  = RealType(1.0) - r2 - x.imag() * x.imag();
 
  RealType p = x.imag() + RealType(1.0);
  RealType m = x.imag() - RealType(1.0);
 
  p = r2 + p * p;
  m = r2 + m * m;
 
  return Kokkos::complex<RealType>(
      RealType(0.5) * atan2(RealType(2.0) * x.real(), i),
      RealType(0.25) * log(p / m));
}
 
template <class RealType>
inline complex<RealType> exp(const std::complex<RealType>& c) {
  return complex<RealType>(std::exp(c.real()) * std::cos(c.imag()),
                           std::exp(c.real()) * std::sin(c.imag()));
}
 
template <class RealType1, class RealType2>
KOKKOS_INLINE_FUNCTION
    complex<typename std::common_type<RealType1, RealType2>::type>
    operator/(const complex<RealType1>& x,
              const RealType2& y) noexcept(noexcept(RealType1{} /
                                                    RealType2{})) {
  return complex<typename std::common_type<RealType1, RealType2>::type>(
      real(x) / y, imag(x) / y);
}
 
template <class RealType1, class RealType2>
KOKKOS_INLINE_FUNCTION
    complex<typename std::common_type<RealType1, RealType2>::type>
    operator/(const complex<RealType1>& x,
              const complex<RealType2>& y) noexcept(noexcept(RealType1{} /
                                                             RealType2{})) {
  using Kokkos::Experimental::fabs;
  // Scale (by the "1-norm" of y) to avoid unwarranted overflow.
  // If the real part is +/-Inf and the imaginary part is -/+Inf,
  // this won't change the result.
  using common_real_type =
      typename std::common_type<RealType1, RealType2>::type;
  const common_real_type s = fabs(real(y)) + fabs(imag(y));
 
  // If s is 0, then y is zero, so x/y == real(x)/0 + i*imag(x)/0.
  // In that case, the relation x/y == (x/s) / (y/s) doesn't hold,
  // because y/s is NaN.
  if (s == 0.0) {
    return complex<common_real_type>(real(x) / s, imag(x) / s);
  } else {
    const complex<common_real_type> x_scaled(real(x) / s, imag(x) / s);
    const complex<common_real_type> y_conj_scaled(real(y) / s, -imag(y) / s);
    const RealType1 y_scaled_abs =
        real(y_conj_scaled) * real(y_conj_scaled) +
        imag(y_conj_scaled) * imag(y_conj_scaled);  // abs(y) == abs(conj(y))
    complex<common_real_type> result = x_scaled * y_conj_scaled;
    result /= y_scaled_abs;
    return result;
  }
}
 
template <class RealType1, class RealType2>
KOKKOS_INLINE_FUNCTION
    complex<typename std::common_type<RealType1, RealType2>::type>
    operator/(const RealType1& x,
              const complex<RealType2>& y) noexcept(noexcept(RealType1{} /
                                                             RealType2{})) {
  return complex<typename std::common_type<RealType1, RealType2>::type>(x) / y;
}
 
template <class RealType>
std::ostream& operator<<(std::ostream& os, const complex<RealType>& x) {
  const std::complex<RealType> x_std(Kokkos::real(x), Kokkos::imag(x));
  os << x_std;
  return os;
}
 
template <class RealType>
std::istream& operator>>(std::istream& is, complex<RealType>& x) {
  std::complex<RealType> x_std;
  is >> x_std;
  x = x_std;  // only assigns on success of above
  return is;
}
 
template <class T>
struct reduction_identity<Kokkos::complex<T>> {
  using t_red_ident = reduction_identity<T>;
  KOKKOS_FORCEINLINE_FUNCTION constexpr static Kokkos::complex<T>
  sum() noexcept {
    return Kokkos::complex<T>(t_red_ident::sum(), t_red_ident::sum());
  }
  KOKKOS_FORCEINLINE_FUNCTION constexpr static Kokkos::complex<T>
  prod() noexcept {
    return Kokkos::complex<T>(t_red_ident::prod(), t_red_ident::sum());
  }
};
 
}  // namespace Kokkos
 
#endif  // KOKKOS_COMPLEX_HPP