thrust/complex.h

File members: thrust/complex.h

/*
 *  Copyright 2008-2019 NVIDIA Corporation
 *  Copyright 2013 Filipe RNC Maia
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

#pragma once

#include <thrust/detail/config.h>

#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
#  pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
#  pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
#  pragma system_header
#endif // no system header

#include <thrust/detail/type_traits.h>
#include <thrust/type_traits/is_trivially_relocatable.h>

#include <cmath>
#include <complex>
#include <sstream>

#define THRUST_STD_COMPLEX_REAL(z) \
  reinterpret_cast<const typename ::cuda::std::remove_reference_t<decltype(z)>::value_type(&)[2]>(z)[0]
#define THRUST_STD_COMPLEX_IMAG(z) \
  reinterpret_cast<const typename ::cuda::std::remove_reference_t<decltype(z)>::value_type(&)[2]>(z)[1]
#define THRUST_STD_COMPLEX_DEVICE _CCCL_DEVICE

THRUST_NAMESPACE_BEGIN

/*
 *  Calls to the standard math library from inside the thrust namespace
 *  with real arguments require explicit scope otherwise they will fail
 *  to resolve as it will find the equivalent complex function but then
 *  fail to match the template, and give up looking for other scopes.
 */

template <typename T>
struct complex
{
public:
  using value_type = T;

  /* --- Constructors --- */

  _CCCL_HOST_DEVICE complex(const T& re);

  _CCCL_HOST_DEVICE complex(const T& re, const T& im);

  complex() = default;

  complex(const complex<T>& z) = default;

  template <typename U>
  _CCCL_HOST_DEVICE complex(const complex<U>& z);

  _CCCL_HOST THRUST_STD_COMPLEX_DEVICE complex(const std::complex<T>& z);

  template <typename U>
  _CCCL_HOST THRUST_STD_COMPLEX_DEVICE complex(const std::complex<U>& z);

  /* --- Assignment Operators --- */

  _CCCL_HOST_DEVICE complex& operator=(const T& re);

  complex& operator=(const complex<T>& z) = default;

  template <typename U>
  _CCCL_HOST_DEVICE complex& operator=(const complex<U>& z);

  _CCCL_HOST THRUST_STD_COMPLEX_DEVICE complex& operator=(const std::complex<T>& z);

  template <typename U>
  _CCCL_HOST THRUST_STD_COMPLEX_DEVICE complex& operator=(const std::complex<U>& z);

  /* --- Compound Assignment Operators --- */

  template <typename U>
  _CCCL_HOST_DEVICE complex<T>& operator+=(const complex<U>& z);

  template <typename U>
  _CCCL_HOST_DEVICE complex<T>& operator-=(const complex<U>& z);

  template <typename U>
  _CCCL_HOST_DEVICE complex<T>& operator*=(const complex<U>& z);

  template <typename U>
  _CCCL_HOST_DEVICE complex<T>& operator/=(const complex<U>& z);

  template <typename U>
  _CCCL_HOST_DEVICE complex<T>& operator+=(const U& z);

  template <typename U>
  _CCCL_HOST_DEVICE complex<T>& operator-=(const U& z);

  template <typename U>
  _CCCL_HOST_DEVICE complex<T>& operator*=(const U& z);

  template <typename U>
  _CCCL_HOST_DEVICE complex<T>& operator/=(const U& z);

  /* --- Getter functions ---
   * The volatile ones are there to help for example
   * with certain reductions optimizations
   */

  _CCCL_HOST_DEVICE T real() const volatile
  {
    return data.x;
  }

  _CCCL_HOST_DEVICE T imag() const volatile
  {
    return data.y;
  }

  _CCCL_HOST_DEVICE T real() const
  {
    return data.x;
  }

  _CCCL_HOST_DEVICE T imag() const
  {
    return data.y;
  }

  /* --- Setter functions ---
   * The volatile ones are there to help for example
   * with certain reductions optimizations
   */

  _CCCL_HOST_DEVICE void real(T re) volatile
  {
    data.x = re;
  }

  _CCCL_HOST_DEVICE void imag(T im) volatile
  {
    data.y = im;
  }

  _CCCL_HOST_DEVICE void real(T re)
  {
    data.x = re;
  }

  _CCCL_HOST_DEVICE void imag(T im)
  {
    data.y = im;
  }

  /* --- Casting functions --- */

  _CCCL_HOST operator std::complex<T>() const
  {
    return std::complex<T>(real(), imag());
  }

private:
#if defined(_CCCL_CUDA_COMPILER_NVCC) && _CCCL_CUDACC_BELOW(11, 7)
  struct __align__(sizeof(T) * 2) storage
#elif _CCCL_COMPILER(ICC)
  struct storage
#else // !(_CCCL_COMPILER(ICC) || (defined(_CCCL_CUDA_COMPILER_NVCC) && _CCCL_CUDACC_BELOW(11, 7)))
  struct alignas(sizeof(T) * 2) storage
#endif // !(_CCCL_COMPILER(ICC) || (defined(_CCCL_CUDA_COMPILER_NVCC) && _CCCL_CUDACC_BELOW(11, 7)))
  {
    T x;
    T y;
  }
#if _CCCL_COMPILER(ICC)
  __attribute__((aligned(sizeof(T) * 2)))
#endif // _CCCL_COMPILER(ICC)
  ;
  storage data;
};

/* --- General Functions --- */

template <typename T>
_CCCL_HOST_DEVICE T abs(const complex<T>& z);

template <typename T>
_CCCL_HOST_DEVICE T arg(const complex<T>& z);

template <typename T>
_CCCL_HOST_DEVICE T norm(const complex<T>& z);

template <typename T>
_CCCL_HOST_DEVICE complex<T> conj(const complex<T>& z);

template <typename T0, typename T1>
_CCCL_HOST_DEVICE complex<::cuda::std::common_type_t<T0, T1>> polar(const T0& m, const T1& theta = T1());

template <typename T>
_CCCL_HOST_DEVICE complex<T> proj(const T& z);

/* --- Binary Arithmetic operators --- */

template <typename T0, typename T1>
_CCCL_HOST_DEVICE complex<::cuda::std::common_type_t<T0, T1>> operator+(const complex<T0>& x, const complex<T1>& y);

template <typename T0, typename T1>
_CCCL_HOST_DEVICE complex<::cuda::std::common_type_t<T0, T1>> operator+(const complex<T0>& x, const T1& y);

template <typename T0, typename T1>
_CCCL_HOST_DEVICE complex<::cuda::std::common_type_t<T0, T1>> operator+(const T0& x, const complex<T1>& y);

template <typename T0, typename T1>
_CCCL_HOST_DEVICE complex<::cuda::std::common_type_t<T0, T1>> operator-(const complex<T0>& x, const complex<T1>& y);

template <typename T0, typename T1>
_CCCL_HOST_DEVICE complex<::cuda::std::common_type_t<T0, T1>> operator-(const complex<T0>& x, const T1& y);

template <typename T0, typename T1>
_CCCL_HOST_DEVICE complex<::cuda::std::common_type_t<T0, T1>> operator-(const T0& x, const complex<T1>& y);

template <typename T0, typename T1>
_CCCL_HOST_DEVICE complex<::cuda::std::common_type_t<T0, T1>> operator*(const complex<T0>& x, const complex<T1>& y);

template <typename T0, typename T1>
_CCCL_HOST_DEVICE complex<::cuda::std::common_type_t<T0, T1>> operator*(const complex<T0>& x, const T1& y);

template <typename T0, typename T1>
_CCCL_HOST_DEVICE complex<::cuda::std::common_type_t<T0, T1>> operator*(const T0& x, const complex<T1>& y);

template <typename T0, typename T1>
_CCCL_HOST_DEVICE complex<::cuda::std::common_type_t<T0, T1>> operator/(const complex<T0>& x, const complex<T1>& y);

template <typename T0, typename T1>
_CCCL_HOST_DEVICE complex<::cuda::std::common_type_t<T0, T1>> operator/(const complex<T0>& x, const T1& y);

template <typename T0, typename T1>
_CCCL_HOST_DEVICE complex<::cuda::std::common_type_t<T0, T1>> operator/(const T0& x, const complex<T1>& y);

/* --- Unary Arithmetic operators --- */

template <typename T>
_CCCL_HOST_DEVICE complex<T> operator+(const complex<T>& y);

template <typename T>
_CCCL_HOST_DEVICE complex<T> operator-(const complex<T>& y);

/* --- Exponential Functions --- */

template <typename T>
_CCCL_HOST_DEVICE complex<T> exp(const complex<T>& z);

template <typename T>
_CCCL_HOST_DEVICE complex<T> log(const complex<T>& z);

template <typename T>
_CCCL_HOST_DEVICE complex<T> log10(const complex<T>& z);

/* --- Power Functions --- */

template <typename T0, typename T1>
_CCCL_HOST_DEVICE complex<::cuda::std::common_type_t<T0, T1>> pow(const complex<T0>& x, const complex<T1>& y);

template <typename T0, typename T1>
_CCCL_HOST_DEVICE complex<::cuda::std::common_type_t<T0, T1>> pow(const complex<T0>& x, const T1& y);

template <typename T0, typename T1>
_CCCL_HOST_DEVICE complex<::cuda::std::common_type_t<T0, T1>> pow(const T0& x, const complex<T1>& y);

template <typename T>
_CCCL_HOST_DEVICE complex<T> sqrt(const complex<T>& z);

/* --- Trigonometric Functions --- */

template <typename T>
_CCCL_HOST_DEVICE complex<T> cos(const complex<T>& z);

template <typename T>
_CCCL_HOST_DEVICE complex<T> sin(const complex<T>& z);

template <typename T>
_CCCL_HOST_DEVICE complex<T> tan(const complex<T>& z);

/* --- Hyperbolic Functions --- */

template <typename T>
_CCCL_HOST_DEVICE complex<T> cosh(const complex<T>& z);

template <typename T>
_CCCL_HOST_DEVICE complex<T> sinh(const complex<T>& z);

template <typename T>
_CCCL_HOST_DEVICE complex<T> tanh(const complex<T>& z);

/* --- Inverse Trigonometric Functions --- */

template <typename T>
_CCCL_HOST_DEVICE complex<T> acos(const complex<T>& z);

template <typename T>
_CCCL_HOST_DEVICE complex<T> asin(const complex<T>& z);

template <typename T>
_CCCL_HOST_DEVICE complex<T> atan(const complex<T>& z);

/* --- Inverse Hyperbolic Functions --- */

template <typename T>
_CCCL_HOST_DEVICE complex<T> acosh(const complex<T>& z);

template <typename T>
_CCCL_HOST_DEVICE complex<T> asinh(const complex<T>& z);

template <typename T>
_CCCL_HOST_DEVICE complex<T> atanh(const complex<T>& z);

/* --- Stream Operators --- */

template <typename T, typename CharT, typename Traits>
std::basic_ostream<CharT, Traits>& operator<<(std::basic_ostream<CharT, Traits>& os, const complex<T>& z);

template <typename T, typename CharT, typename Traits>
_CCCL_HOST std::basic_istream<CharT, Traits>& operator>>(std::basic_istream<CharT, Traits>& is, complex<T>& z);

/* --- Equality Operators --- */

template <typename T0, typename T1>
_CCCL_HOST_DEVICE bool operator==(const complex<T0>& x, const complex<T1>& y);

template <typename T0, typename T1>
_CCCL_HOST THRUST_STD_COMPLEX_DEVICE bool operator==(const complex<T0>& x, const std::complex<T1>& y);

template <typename T0, typename T1>
_CCCL_HOST THRUST_STD_COMPLEX_DEVICE bool operator==(const std::complex<T0>& x, const complex<T1>& y);

template <typename T0, typename T1>
_CCCL_HOST_DEVICE bool operator==(const T0& x, const complex<T1>& y);

template <typename T0, typename T1>
_CCCL_HOST_DEVICE bool operator==(const complex<T0>& x, const T1& y);

template <typename T0, typename T1>
_CCCL_HOST_DEVICE bool operator!=(const complex<T0>& x, const complex<T1>& y);

template <typename T0, typename T1>
_CCCL_HOST THRUST_STD_COMPLEX_DEVICE bool operator!=(const complex<T0>& x, const std::complex<T1>& y);

template <typename T0, typename T1>
_CCCL_HOST THRUST_STD_COMPLEX_DEVICE bool operator!=(const std::complex<T0>& x, const complex<T1>& y);

template <typename T0, typename T1>
_CCCL_HOST_DEVICE bool operator!=(const T0& x, const complex<T1>& y);

template <typename T0, typename T1>
_CCCL_HOST_DEVICE bool operator!=(const complex<T0>& x, const T1& y);

THRUST_NAMESPACE_END

#include <thrust/detail/complex/complex.inl>

#undef THRUST_STD_COMPLEX_REAL
#undef THRUST_STD_COMPLEX_IMAG
#undef THRUST_STD_COMPLEX_DEVICE