Merge pull request #2554 from nickdidio/complex-numbers

Complex numbers get_real(), get_imag(), to_complex

Author: bob-carpenter

stan/math/prim/fun.hpp

@@ -111,7 +111,9 @@
 #include <stan/math/prim/fun/get.hpp>
 #include <stan/math/prim/fun/get_base1.hpp>
 #include <stan/math/prim/fun/get_base1_lhs.hpp>
+#include <stan/math/prim/fun/get_imag.hpp>
 #include <stan/math/prim/fun/get_lp.hpp>
+#include <stan/math/prim/fun/get_real.hpp>
 #include <stan/math/prim/fun/gp_dot_prod_cov.hpp>
 #include <stan/math/prim/fun/gp_exponential_cov.hpp>
 #include <stan/math/prim/fun/gp_matern32_cov.hpp>
@@ -325,6 +327,7 @@
 #include <stan/math/prim/fun/tgamma.hpp>
 #include <stan/math/prim/fun/to_array_1d.hpp>
 #include <stan/math/prim/fun/to_array_2d.hpp>
+#include <stan/math/prim/fun/to_complex.hpp>
 #include <stan/math/prim/fun/to_matrix.hpp>
 #include <stan/math/prim/fun/to_ref.hpp>
 #include <stan/math/prim/fun/to_row_vector.hpp>

stan/math/prim/fun/abs.hpp

@@ -1,11 +1,13 @@
 #ifndef STAN_MATH_PRIM_FUN_ABS_HPP
 #define STAN_MATH_PRIM_FUN_ABS_HPP
 
+#include <stan/math/prim/core.hpp>
 #include <stan/math/prim/meta.hpp>
 #include <stan/math/prim/fun/hypot.hpp>
 #include <stan/math/prim/functor/apply_scalar_unary.hpp>
 #include <stan/math/prim/functor/apply_vector_unary.hpp>
 #include <cmath>
+#include <complex>
 
 namespace stan {
 namespace math {
@@ -62,7 +64,7 @@ template <typename Container,
           require_container_st<std::is_arithmetic, Container>* = nullptr>
 inline auto abs(const Container& x) {
   return apply_vector_unary<Container>::apply(
-      x, [](const auto& v) { return v.array().abs(); });
+      x, [&](const auto& v) { return v.array().abs(); });
 }
 
 namespace internal {

stan/math/prim/fun/cos.hpp

@@ -13,16 +13,6 @@
 namespace stan {
 namespace math {
 
-/**
- * Arithmetic version of `cos()`
- * @tparam T An `Arithmetic` type.
- * @param x Arithmetic scalar.
- */
-template <typename T, require_arithmetic_t<T>* = nullptr>
-inline auto cos(T x) {
-  return std::cos(x);
-}
-
 /**
  * Structure to wrap `cos()` so it can be vectorized.
  *
@@ -33,6 +23,7 @@ inline auto cos(T x) {
 struct cos_fun {
   template <typename T>
   static inline T fun(const T& x) {
+    using std::cos;
     return cos(x);
   }
 };
@@ -47,7 +38,6 @@ struct cos_fun {
  */
 template <typename Container,
           require_not_container_st<std::is_arithmetic, Container>* = nullptr,
-          require_not_stan_scalar_t<Container>* = nullptr,
           require_not_var_matrix_t<Container>* = nullptr,
           require_all_not_nonscalar_prim_or_rev_kernel_expression_t<
               Container>* = nullptr>
@@ -67,7 +57,7 @@ template <typename Container,
           require_container_st<std::is_arithmetic, Container>* = nullptr>
 inline auto cos(const Container& x) {
   return apply_vector_unary<Container>::apply(
-      x, [](const auto& v) { return v.array().cos(); });
+      x, [&](const auto& v) { return v.array().cos(); });
 }
 
 namespace internal {

stan/math/prim/fun/get_imag.hpp

@@ -0,0 +1,25 @@
+#ifndef STAN_MATH_PRIM_FUN_GET_IMAG_HPP
+#define STAN_MATH_PRIM_FUN_GET_IMAG_HPP
+
+#include <stan/math/prim/meta.hpp>
+#include <complex>
+
+namespace stan {
+namespace math {
+
+/**
+ * Return the imaginary component of the complex argument.
+ *
+ * @tparam T value type of complex argument
+ * @param[in] z complex value whose imaginary component is extracted
+ * @return imaginary component of argument
+ */
+template <typename T>
+T get_imag(const std::complex<T>& z) {
+  return z.imag();
+}
+
+}  // namespace math
+}  // namespace stan
+
+#endif

stan/math/prim/fun/get_real.hpp

@@ -0,0 +1,25 @@
+#ifndef STAN_MATH_PRIM_FUN_GET_REAL_HPP
+#define STAN_MATH_PRIM_FUN_GET_REAL_HPP
+
+#include <stan/math/prim/meta.hpp>
+#include <complex>
+
+namespace stan {
+namespace math {
+
+/**
+ * Return the real component of the complex argument.
+ *
+ * @tparam T value type of complex argument
+ * @param[in] z complex value whose real component is extracted
+ * @return real component of argument
+ */
+template <typename T>
+T get_real(const std::complex<T>& z) {
+  return z.real();
+}
+
+}  // namespace math
+}  // namespace stan
+
+#endif

stan/math/prim/fun/imag.hpp

@@ -8,11 +8,11 @@ namespace stan {
 namespace math {
 
 /**
- * Return the imaginary part of the complex argument.
+ * Return the imaginary component of the complex argument.
  *
- * @tparam T value type of argument
- * @param[in] z argument
- * @return imaginary part of argument
+ * @tparam T value type of complex argument
+ * @param[in] z complex value whose imaginary component is extracted
+ * @return imaginary component of argument
  */
 template <typename T, require_autodiff_t<T>>
 T imag(const std::complex<T>& z) {

stan/math/prim/fun/real.hpp

@@ -8,11 +8,11 @@ namespace stan {
 namespace math {
 
 /**
- * Return the real part of the complex argument.
+ * Return the real component of the complex argument.
  *
- * @tparam T value type of argument
- * @param[in] z argument
- * @return real part of argument
+ * @tparam T value type of complex argument
+ * @param[in] z complex value whose real component is extracted
+ * @return real component of argument
  */
 template <typename T, require_autodiff_t<T>>
 T real(const std::complex<T>& z) {

stan/math/prim/fun/to_complex.hpp

@@ -0,0 +1,30 @@
+#ifndef STAN_MATH_PRIM_FUN_TO_COMPLEX_HPP
+#define STAN_MATH_PRIM_FUN_TO_COMPLEX_HPP
+
+#include <stan/math/prim/meta.hpp>
+#include <stan/math/prim/meta/return_type.hpp>
+#include <complex>
+
+namespace stan {
+namespace math {
+
+/**
+ * Return a complex value from a real component and an imaginary component.
+ * Default values for both components is 0.
+ *
+ * @tparam T type of real component
+ * @tparam S type of imaginary component
+ * @param[in] re real component (default = 0)
+ * @param[in] im imaginary component (default = 0)
+ * @return complex value with specified real and imaginary components
+ */
+template <typename T = double, typename S = double>
+inline std::complex<stan::real_return_t<T, S>> to_complex(const T& re = 0,
+                                                          const S& im = 0) {
+  return std::complex<stan::real_return_t<T, S>>(re, im);
+}
+
+}  // namespace math
+}  // namespace stan
+
+#endif

test/unit/math/mix/fun/get_imag_test.cpp

@@ -0,0 +1,7 @@
+#include <test/unit/math/test_ad.hpp>
+#include <complex>
+
+TEST(mathMixMatFun, get_imag) {
+  auto f = [](const auto& z) { return stan::math::get_imag(z); };
+  stan::test::expect_complex_common(f);
+}

test/unit/math/mix/fun/get_real_test.cpp

@@ -0,0 +1,7 @@
+#include <test/unit/math/test_ad.hpp>
+#include <complex>
+
+TEST(mathMixMatFun, get_real) {
+  auto f = [](const auto& z) { return stan::math::get_real(z); };
+  stan::test::expect_complex_common(f);
+}

test/unit/math/mix/fun/to_complex_test.cpp

@@ -0,0 +1,14 @@
+#include <test/unit/math/test_ad.hpp>
+#include <complex>
+
+TEST(mathMixMatFun, to_complex) {
+  auto f0 = []() { return stan::math::to_complex(); };
+  auto f1 = [](const auto& x) { return stan::math::to_complex(x); };
+  auto f2 = [](const auto& x, const auto& y) {
+    return stan::math::to_complex(x, y);
+  };
+
+  EXPECT_EQ(std::complex<double>(), stan::math::to_complex());
+  stan::test::expect_common_unary(f1);
+  stan::test::expect_common_binary(f2);
+}

test/unit/math/test_ad_test.cpp

@@ -233,6 +233,7 @@ int baz_fvar = 0;
 int baz_complex = 0;
 int baz_complex_var = 0;
 int baz_complex_fvar = 0;
+
 double baz(int x) {
   ++baz_int;
   return x / 2.0;