Merge pull request #2256 from stan-dev/feature/varmat-a-to-c-unary

var matrix specializations for a-c unary functions in rev

Author: bbbales2

stan/math/prim/fun/abs.hpp

@@ -3,21 +3,57 @@
 
 #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>
 
 namespace stan {
 namespace math {
 
 /**
- * Return floating-point absolute value.
+ * Structure to wrap `abs()` so it can be vectorized.
  *
- * Delegates to <code>fabs(double)</code> rather than
- * <code>std::abs(int)</code>.
+ * @tparam T type of variable
+ * @param x argument
+ * @return Absolute value of variable.
+ */
+struct abs_fun {
+  template <typename T>
+  static inline T fun(const T& x) {
+    using std::fabs;
+    return fabs(x);
+  }
+};
+
+/**
+ * Returns the elementwise `abs()` of the input,
+ * which may be a scalar or any Stan container of numeric scalars.
  *
- * @param x scalar
- * @return absolute value of scalar
+ * @tparam Container type of container
+ * @param x argument
+ * @return Absolute value of each variable in the container.
  */
-inline double abs(double x) { return std::fabs(x); }
+template <typename Container,
+          require_not_container_st<std::is_arithmetic, Container>* = nullptr,
+          require_not_var_matrix_t<Container>* = nullptr>
+inline auto abs(const Container& x) {
+  return apply_scalar_unary<abs_fun, Container>::apply(x);
+}
+
+/**
+ * Version of `abs()` that accepts std::vectors, Eigen Matrix/Array objects
+ *  or expressions, and containers of these.
+ *
+ * @tparam Container Type of x
+ * @param x argument
+ * @return Absolute value of each variable in the container.
+ */
+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(); });
+}
 
 namespace internal {
 /**

stan/math/prim/fun/acos.hpp

@@ -43,6 +43,7 @@ struct acos_fun {
  */
 template <typename Container,
           require_not_container_st<std::is_arithmetic, Container>* = nullptr,
+          require_not_var_matrix_t<Container>* = nullptr,
           require_all_not_nonscalar_prim_or_rev_kernel_expression_t<
               Container>* = nullptr>
 inline auto acos(const Container& x) {

stan/math/prim/fun/acosh.hpp

@@ -83,8 +83,9 @@ struct acosh_fun {
  * @param x container
  * @return Elementwise acosh of members of container.
  */
-template <typename T, require_all_not_nonscalar_prim_or_rev_kernel_expression_t<
-                          T>* = nullptr>
+template <
+    typename T, require_not_var_matrix_t<T>* = nullptr,
+    require_all_not_nonscalar_prim_or_rev_kernel_expression_t<T>* = nullptr>
 inline auto acosh(const T& x) {
   return apply_scalar_unary<acosh_fun, T>::apply(x);
 }

stan/math/prim/fun/asin.hpp

@@ -41,6 +41,7 @@ struct asin_fun {
  */
 template <typename Container,
           require_not_container_st<std::is_arithmetic, Container>* = nullptr,
+          require_not_var_matrix_t<Container>* = nullptr,
           require_all_not_nonscalar_prim_or_rev_kernel_expression_t<
               Container>* = nullptr>
 inline auto asin(const Container& x) {

stan/math/prim/fun/asinh.hpp

@@ -43,8 +43,9 @@ struct asinh_fun {
  * @param x container
  * @return Inverse hyperbolic sine of each value in the container.
  */
-template <typename T, require_all_not_nonscalar_prim_or_rev_kernel_expression_t<
-                          T>* = nullptr>
+template <
+    typename T, require_not_var_matrix_t<T>* = nullptr,
+    require_all_not_nonscalar_prim_or_rev_kernel_expression_t<T>* = nullptr>
 inline auto asinh(const T& x) {
   return apply_scalar_unary<asinh_fun, T>::apply(x);
 }

stan/math/prim/fun/atan.hpp

@@ -39,6 +39,7 @@ struct atan_fun {
  */
 template <typename Container,
           require_not_container_st<std::is_arithmetic, Container>* = nullptr,
+          require_not_var_matrix_t<Container>* = nullptr,
           require_all_not_nonscalar_prim_or_rev_kernel_expression_t<
               Container>* = nullptr>
 inline auto atan(const Container& x) {

stan/math/prim/fun/atanh.hpp

@@ -72,8 +72,9 @@ struct atanh_fun {
  * @param x container
  * @return Elementwise atanh of members of container.
  */
-template <typename T, require_all_not_nonscalar_prim_or_rev_kernel_expression_t<
-                          T>* = nullptr>
+template <
+    typename T, require_not_var_matrix_t<T>* = nullptr,
+    require_all_not_nonscalar_prim_or_rev_kernel_expression_t<T>* = nullptr>
 inline auto atanh(const T& x) {
   return apply_scalar_unary<atanh_fun, T>::apply(x);
 }

stan/math/prim/fun/block.hpp

@@ -18,7 +18,7 @@ namespace math {
  * @param ncols Number of columns in block.
  * @throw std::out_of_range if either index is out of range.
  */
-template <typename T, require_eigen_t<T>* = nullptr>
+template <typename T, require_matrix_t<T>* = nullptr>
 inline auto block(const T& m, size_t i, size_t j, size_t nrows, size_t ncols) {
   check_row_index("block", "i", m, i);
   check_row_index("block", "i+nrows-1", m, i + nrows - 1);

stan/math/prim/fun/cbrt.hpp

@@ -31,7 +31,7 @@ struct cbrt_fun {
  * @param x container
  * @return Cube root of each value in x.
  */
-template <typename T>
+template <typename T, require_not_var_matrix_t<T>* = nullptr>
 inline auto cbrt(const T& x) {
   return apply_scalar_unary<cbrt_fun, T>::apply(x);
 }

stan/math/prim/fun/cos.hpp

@@ -38,6 +38,7 @@ struct cos_fun {
  */
 template <typename Container,
           require_not_container_st<std::is_arithmetic, Container>* = nullptr,
+          require_not_var_matrix_t<Container>* = nullptr,
           require_all_not_nonscalar_prim_or_rev_kernel_expression_t<
               Container>* = nullptr>
 inline auto cos(const Container& x) {

stan/math/prim/fun/cosh.hpp

@@ -37,6 +37,7 @@ struct cosh_fun {
  */
 template <typename Container,
           require_not_container_st<std::is_arithmetic, Container>* = nullptr,
+          require_not_var_matrix_t<Container>* = nullptr,
           require_all_not_nonscalar_prim_or_rev_kernel_expression_t<
               Container>* = nullptr>
 inline auto cosh(const Container& x) {

stan/math/prim/fun/fabs.hpp

@@ -34,7 +34,8 @@ struct fabs_fun {
  * @return Absolute value of each value in x.
  */
 template <typename Container,
-          require_not_container_st<std::is_arithmetic, Container>* = nullptr>
+          require_not_container_st<std::is_arithmetic, Container>* = nullptr,
+          require_not_var_matrix_t<Container>* = nullptr>
 inline auto fabs(const Container& x) {
   return apply_scalar_unary<fabs_fun, Container>::apply(x);
 }

stan/math/prim/fun/sin.hpp

@@ -37,6 +37,7 @@ struct sin_fun {
  */
 template <
     typename T, require_not_container_st<std::is_arithmetic, T>* = nullptr,
+    require_not_var_matrix_t<T>* = nullptr,
     require_all_not_nonscalar_prim_or_rev_kernel_expression_t<T>* = nullptr>
 inline auto sin(const T& x) {
   return apply_scalar_unary<sin_fun, T>::apply(x);

stan/math/prim/fun/sinh.hpp

@@ -35,6 +35,7 @@ struct sinh_fun {
  */
 template <typename Container,
           require_not_container_st<std::is_arithmetic, Container>* = nullptr,
+          require_not_var_matrix_t<Container>* = nullptr,
           require_all_not_nonscalar_prim_or_rev_kernel_expression_t<
               Container>* = nullptr>
 inline auto sinh(const Container& x) {

stan/math/prim/fun/tan.hpp

@@ -37,6 +37,7 @@ struct tan_fun {
  */
 template <typename Container,
           require_not_container_st<std::is_arithmetic, Container>* = nullptr,
+          require_not_var_matrix_t<Container>* = nullptr,
           require_all_not_nonscalar_prim_or_rev_kernel_expression_t<
               Container>* = nullptr>
 inline auto tan(const Container& x) {

stan/math/prim/fun/tanh.hpp

@@ -37,6 +37,7 @@ struct tanh_fun {
  */
 template <typename Container,
           require_not_container_st<std::is_arithmetic, Container>* = nullptr,
+          require_not_var_matrix_t<Container>* = nullptr,
           require_all_not_nonscalar_prim_or_rev_kernel_expression_t<
               Container>* = nullptr>
 inline auto tanh(const Container& x) {

stan/math/prim/functor/apply_vector_unary.hpp

@@ -159,6 +159,12 @@ struct apply_vector_unary<T, require_std_vector_vt<is_stan_scalar, T>> {
   }
 };
 
+namespace internal {
+template <typename T>
+using is_container_or_var_matrix
+    = disjunction<is_container<T>, is_var_matrix<T>>;
+}
+
 /**
  * Specialisation for use with nested containers (std::vectors).
  * For each of the member functions, an std::vector with the appropriate
@@ -170,7 +176,8 @@ struct apply_vector_unary<T, require_std_vector_vt<is_stan_scalar, T>> {
  *
  */
 template <typename T>
-struct apply_vector_unary<T, require_std_vector_vt<is_container, T>> {
+struct apply_vector_unary<
+    T, require_std_vector_vt<internal::is_container_or_var_matrix, T>> {
   using T_vt = value_type_t<T>;
 
   /**

stan/math/prim/meta/is_container.hpp

@@ -20,8 +20,7 @@ namespace stan {
  */
 template <typename Container>
 using is_container = bool_constant<
-    math::disjunction<is_eigen<Container>, is_std_vector<Container>,
-                      is_var_matrix<Container>>::value>;
+    math::disjunction<is_eigen<Container>, is_std_vector<Container>>::value>;
 
 STAN_ADD_REQUIRE_UNARY(container, is_container, general_types);
 STAN_ADD_REQUIRE_CONTAINER(container, is_container, general_types);

stan/math/rev/fun/abs.hpp

@@ -33,10 +33,14 @@ namespace math {
    \end{cases}
    \f]
  *
+ * @tparam T A floating point type or an Eigen type with floating point scalar.
  * @param a Variable input.
  * @return Absolute value of variable.
  */
-inline var abs(const var& a) { return fabs(a); }
+template <typename T>
+inline auto abs(const var_value<T>& a) {
+  return fabs(a);
+}
 
 /**
  * Return the absolute value of the complex argument.

stan/math/rev/fun/acos.hpp

@@ -16,16 +16,6 @@
 namespace stan {
 namespace math {
 
-namespace internal {
-class acos_vari : public op_v_vari {
- public:
-  explicit acos_vari(vari* avi) : op_v_vari(std::acos(avi->val_), avi) {}
-  void chain() {
-    avi_->adj_ -= adj_ / std::sqrt(1.0 - (avi_->val_ * avi_->val_));
-  }
-};
-}  // namespace internal
-
 /**
  * Return the principal value of the arc cosine of a variable,
  * in radians (cmath).
@@ -62,7 +52,27 @@ class acos_vari : public op_v_vari {
  * @param x argument
  * @return Arc cosine of variable, in radians.
  */
-inline var acos(const var& x) { return var(new internal::acos_vari(x.vi_)); }
+inline var acos(const var& x) {
+  return make_callback_var(std::acos(x.val()), [x](const auto& vi) mutable {
+    x.adj() -= vi.adj() / std::sqrt(1.0 - (x.val() * x.val()));
+  });
+}
+
+/**
+ * Return the principal value of the arc cosine of a variable,
+ * in radians (cmath).
+ *
+ * @param x a `var_value` with inner Eigen type
+ * @return Arc cosine of variable, in radians.
+ */
+template <typename VarMat, require_var_matrix_t<VarMat>* = nullptr>
+inline auto acos(const VarMat& x) {
+  return make_callback_var(
+      x.val().array().acos().matrix(), [x](const auto& vi) mutable {
+        x.adj().array()
+            -= vi.adj().array() / (1.0 - (x.val().array().square())).sqrt();
+      });
+}
 
 /**
  * Return the arc cosine of the complex argument.

stan/math/rev/fun/acosh.hpp

@@ -21,16 +21,6 @@
 namespace stan {
 namespace math {
 
-namespace internal {
-class acosh_vari : public op_v_vari {
- public:
-  acosh_vari(double val, vari* avi) : op_v_vari(val, avi) {}
-  void chain() {
-    avi_->adj_ += adj_ / std::sqrt(avi_->val_ * avi_->val_ - 1.0);
-  }
-};
-}  // namespace internal
-
 /**
  * The inverse hyperbolic cosine function for variables (C99).
  *
@@ -67,11 +57,31 @@ class acosh_vari : public op_v_vari {
    \frac{\partial \, \cosh^{-1}(x)}{\partial x} = \frac{1}{\sqrt{x^2-1}}
    \f]
  *
- * @param a The variable.
+ * @param x The variable.
+ * @return Inverse hyperbolic cosine of the variable.
+ */
+inline var acosh(const var& x) {
+  return make_callback_var(acosh(x.val()), [x](const auto& vi) mutable {
+    x.adj() += vi.adj() / std::sqrt(x.val() * x.val() - 1.0);
+  });
+}
+/**
+ * The inverse hyperbolic cosine function for variables (C99).
+ *
+ * For non-variable function, see ::acosh().
+ *
+ * @tparam Varmat a `var_value` with inner Eigen type
+ * @param x The variable
  * @return Inverse hyperbolic cosine of the variable.
  */
-inline var acosh(const var& a) {
-  return var(new internal::acosh_vari(acosh(a.val()), a.vi_));
+template <typename VarMat, require_var_matrix_t<VarMat>* = nullptr>
+inline auto acosh(const VarMat& x) {
+  return make_callback_var(
+      x.val().unaryExpr([](const auto x) { return acosh(x); }),
+      [x](const auto& vi) mutable {
+        x.adj().array()
+            += vi.adj().array() / (x.val().array().square() - 1.0).sqrt();
+      });
 }
 
 /**

stan/math/rev/fun/asin.hpp

@@ -14,16 +14,6 @@
 namespace stan {
 namespace math {
 
-namespace internal {
-class asin_vari : public op_v_vari {
- public:
-  explicit asin_vari(vari* avi) : op_v_vari(std::asin(avi->val_), avi) {}
-  void chain() {
-    avi_->adj_ += adj_ / std::sqrt(1.0 - (avi_->val_ * avi_->val_));
-  }
-};
-}  // namespace internal
-
 /**
  * Return the principal value of the arc sine, in radians, of the
  * specified variable (cmath).
@@ -57,10 +47,31 @@ class asin_vari : public op_v_vari {
    \frac{\partial \, \arcsin(x)}{\partial x} = \frac{1}{\sqrt{1-x^2}}
    \f]
  *
- * @param a Variable in range [-1, 1].
+ * @param x Variable in range [-1, 1].
+ * @return Arc sine of variable, in radians.
+ */
+inline var asin(const var& x) {
+  return make_callback_var(std::asin(x.val()), [x](const auto& vi) mutable {
+    x.adj() += vi.adj() / std::sqrt(1.0 - (x.val() * x.val()));
+  });
+}
+
+/**
+ * Return the principal value of the arc sine, in radians, of the
+ * specified variable (cmath).
+ *
+ * @tparam Varmat a `var_value` with inner Eigen type
+ * @param x Variable with cells in range [-1, 1].
  * @return Arc sine of variable, in radians.
  */
-inline var asin(const var& a) { return var(new internal::asin_vari(a.vi_)); }
+template <typename VarMat, require_var_matrix_t<VarMat>* = nullptr>
+inline auto asin(const VarMat& x) {
+  return make_callback_var(
+      x.val().array().asin().matrix(), [x](const auto& vi) mutable {
+        x.adj().array()
+            += vi.adj().array() / (1.0 - (x.val().array().square())).sqrt();
+      });
+}
 
 /**
  * Return the arc sine of the complex argument.