Complex functions now accept batch of real values

JohanMabille · JohanMabille · commit 31cfa5212563 · 2021-10-06T05:48:22.000+02:00
diff --git a/include/xsimd/arch/generic/xsimd_generic_complex.hpp b/include/xsimd/arch/generic/xsimd_generic_complex.hpp
@@ -3,42 +3,94 @@
 
 #include "./xsimd_generic_details.hpp"
 
-
 namespace xsimd {
 
   namespace kernel {
 
     using namespace types;
 
+    namespace detail
+    {
+      template <class B, bool is_complex>
+      struct real_imag_kernel
+      {
+        using return_type = typename B::real_batch;
+
+        static return_type real(B const& z)
+        {
+          return z.real();
+        }
+
+        static return_type imag(B const& z)
+        {
+          return z.imag();
+        }
+      };
+
+      template <class B>
+      struct real_imag_kernel<B, false>
+      {
+        using return_type = B;
+
+        static return_type real(B const& z)
+        {
+          return z;
+        }
+
+        static return_type imag(B const&)
+        {
+          return B(typename B::value_type(0));
+        }
+      };
+    }
+
+    // real
+    template <class A, class T>
+    real_batch_type_t<batch<T, A>> real(batch<T, A> const& self, requires_arch<generic>) {
+      using batch_type = batch<T, A>;
+      constexpr bool is_cplx = xsimd::detail::is_complex<typename batch_type::value_type>::value;
+      return detail::real_imag_kernel<batch_type, is_cplx>::real(self);
+    }
+
+    // imag
+    template <class A, class T>
+    real_batch_type_t<batch<T, A>> imag(batch<T, A> const& self, requires_arch<generic>) {
+      using batch_type = batch<T, A>;
+      constexpr bool is_cplx = xsimd::detail::is_complex<typename batch_type::value_type>::value;
+      return detail::real_imag_kernel<batch_type, is_cplx>::imag(self);
+    }
+
     // arg
     template<class A, class T>
-    batch<T, A> arg(batch<std::complex<T>, A> const& self, requires_arch<generic>) {
-      return atan2(self.imag(), self.real());
+    real_batch_type_t<batch<T, A>> arg(batch<T, A> const& self, requires_arch<generic>) {
+      return atan2(imag(self), real(self));
     }
 
-
     // conj
-    template<class A, class T> batch<std::complex<T>, A> conj(batch<std::complex<T>, A> const& self, requires_arch<generic>) {
-      return {self.real(), - self.imag()};
+    template<class A, class T>
+    complex_batch_type_t<batch<T, A>> conj(batch<T, A> const& self, requires_arch<generic>) {
+      return {real(self), - imag(self)};
     }
 
     // norm
-    template<class A, class T> batch<T, A> norm(batch<std::complex<T>, A> const& self, requires_arch<generic>) {
-      return {fma(self.real(), self.real(), self.imag() * self.imag())};
+    template<class A, class T>
+    real_batch_type_t<batch<T, A>> norm(batch<T, A> const& self, requires_arch<generic>) {
+      return {fma(real(self), real(self), imag(self) * imag(self))};
     }
 
-
     // proj
-    template<class A, class T> batch<std::complex<T>, A> proj(batch<std::complex<T>, A> const& self, requires_arch<generic>) {
-      using batch_type = batch<std::complex<T>, A>;
+    template<class A, class T>
+    complex_batch_type_t<batch<T, A>> proj(batch<T, A> const& self, requires_arch<generic>) {
+      using batch_type = complex_batch_type_t<batch<T, A>>;
       using real_batch = typename batch_type::real_batch;
-      auto cond = xsimd::isinf(self.real()) || xsimd::isinf(self.imag());
-      return select(cond, batch_type(constants::infinity<real_batch>(), copysign(real_batch(0.), self.imag())), self);
+      using real_value_type = typename real_batch::value_type;
+      auto cond = xsimd::isinf(real(self)) || xsimd::isinf(imag(self));
+      return select(cond,
+                    batch_type(constants::infinity<real_batch>(),
+                               copysign(real_batch(real_value_type(0)), imag(self))),
+                    batch_type(self));
     }
-
-
   }
-
 }
 
 #endif
diff --git a/include/xsimd/types/xsimd_api.hpp b/include/xsimd/types/xsimd_api.hpp
@@ -96,11 +96,11 @@ batch<T, A> acosh(batch<T, A> const& x) {
  * @ingroup batch_complex
  *
  * Computes the argument of the batch \c z.
- * @param z batch of complex values.
+ * @param z batch of complex or real values.
  * @return the argument of \c z.
  */
 template<class T, class A>
-batch<T, A> arg(batch<std::complex<T>, A> const& z) {
+real_batch_type_t<batch<T, A>> arg(batch<T, A> const& z) {
   return kernel::arg<A>(z, A{});
 }
 
@@ -358,7 +358,7 @@ batch<T, A> clip(batch<T, A> const& x, batch<T, A> const& lo, batch<T, A> const&
  * @return the argument of \c z.
  */
 template<class A, class T>
-batch<std::complex<T>, A> conj(batch<std::complex<T>, A> const& z) {
+complex_batch_type_t<batch<T, A>> conj(batch<T, A> const& z) {
   return kernel::conj(z, A{});
 }
 
@@ -758,6 +758,17 @@ batch<T, A> hypot(batch<T, A> const& x, batch<T, A> const& y) {
   return kernel::hypot<A>(x, y, A{});
 }
 
+/**
+ * @ingroup batch_complex
+ *
+ * Computes the imaginary part of the batch \c z.
+ * @param z batch of complex or real values.
+ * @return the argument of \c z.
+ */
+template <class T, class A>
+real_batch_type_t<batch<T, A>> imag(batch<T, A> const& x) {
+  return kernel::imag<A>(x, A{});
+}
 
 /**
  * @ingroup batch_constant
@@ -1115,11 +1126,11 @@ batch<T, A> nextafter(batch<T, A> const& x, batch<T, A> const& y) {
  * @ingroup batch_complex
  *
  * Computes the norm of the batch \c x.
- * @param x batch of complex values.
+ * @param x batch of complex or real values.
  * @return the norm of \c x.
  */
 template<class A, class T>
-batch<T, A> norm(batch<std::complex<T>, A> const& x) {
+real_batch_type_t<batch<T, A>> norm(batch<T, A> const& x) {
   return kernel::norm(x, A{});
 }
 
@@ -1167,14 +1178,26 @@ batch<T, A> pow(batch<T, A> const& x, ITy y) {
  * @ingroup batch_complex
  *
  * Computes the projection of the batch \c x.
- * @param x batch of complex values.
+ * @param x batch of complex or real values.
  * @return the projection of \c x.
  */
 template<class A, class T>
-batch<std::complex<T>, A> proj(batch<std::complex<T>, A> const& x) {
+complex_batch_type_t<batch<T, A>> proj(batch<T, A> const& x) {
   return kernel::proj(x, A{});
 }
 
+/**
+ * @ingroup batch_complex
+ *
+ * Computes the real part of the batch \c z.
+ * @param z batch of complex or real values.
+ * @return the argument of \c z.
+ */
+template <class T, class A>
+real_batch_type_t<batch<T, A>> real(batch<T, A> const& x) {
+  return kernel::real<A>(x, A{});
+}
+
 /**
  * @ingroup batch_math
  *
diff --git a/include/xsimd/types/xsimd_batch.hpp b/include/xsimd/types/xsimd_batch.hpp
@@ -350,6 +350,45 @@ struct batch<std::complex<T>, A> {
 template <class T, class A>
 constexpr std::size_t batch<std::complex<T>, A>::size;
 
+        /*******************
+         * real_batch_type *
+         *******************/
+
+        template <class B>
+        struct real_batch_type
+        {
+            using type = B;
+        };
+
+        template <class T, class A>
+        struct real_batch_type<batch<std::complex<T>, A>>
+        {
+            using type = batch<T, A>;
+        };
+
+        template <class B>
+        using real_batch_type_t = typename real_batch_type<B>::type;
+
+        /**********************
+         * complex_batch_type *
+         **********************/
+
+        template <class B>
+        struct complex_batch_type
+        {
+            using real_value_type = typename B::value_type;
+            using arch_type = typename B::arch_type;
+            using type = batch<std::complex<real_value_type>, arch_type>;
+        };
+
+        template <class T, class A>
+        struct complex_batch_type<batch<std::complex<T>, A>>
+        {
+            using type = batch<std::complex<T>, A>;
+        };
+
+        template <class B>
+        using complex_batch_type_t = typename complex_batch_type<B>::type;
 }
 
 #include "../types/xsimd_batch_constant.hpp"
diff --git a/test/test_batch_complex.cpp b/test/test_batch_complex.cpp
@@ -486,6 +486,35 @@ class batch_complex_test : public testing::Test
         }
     }
 
+    void test_conj_norm_proj_real() const
+    {
+        // conj real batch
+        {
+            array_type expected;
+            std::transform(lhs.cbegin(), lhs.cend(), expected.begin(),
+                           [](const value_type& v) { return std::conj(std::real(v)); });
+            batch_type res = conj(real(batch_lhs()));
+            EXPECT_BATCH_EQ(res, expected) << print_function_name("conj(real batch)");
+            
+        }
+        // norm real batch
+        {
+            real_array_type expected;
+            std::transform(lhs.cbegin(), lhs.cend(), expected.begin(),
+                            [](const value_type& v) {  return std::norm(std::real(v)); });
+            real_batch_type res = norm(real(batch_lhs()));
+            EXPECT_BATCH_EQ(res, expected) << print_function_name("norm(real_batch)");
+        }
+        // proj real batch
+        {
+            array_type expected;
+            std::transform(lhs.cbegin(), lhs.cend(), expected.begin(),
+                            [](const value_type& v) { return std::proj(std::real(v)); });
+            batch_type res = proj(real(batch_lhs()));
+            EXPECT_BATCH_EQ(res, expected) << print_function_name("proj(real_batch)");
+        }
+    }
+
     void test_horizontal_operations() const
     {
         // hadd
@@ -599,6 +628,11 @@ TYPED_TEST(batch_complex_test, conj_norm_proj)
     this->test_conj_norm_proj();
 }
 
+TYPED_TEST(batch_complex_test, conj_norm_proj_real)
+{
+    this->test_conj_norm_proj_real();
+}
+
 TYPED_TEST(batch_complex_test, horizontal_operations)
 {
     this->test_horizontal_operations();
