Merge pull request #2787 from stan-dev/feature/owens-t-varmat

Allow varmat for owens_t

Author: andrjohns

stan/math/prim/fun/owens_t.hpp

@@ -67,10 +67,11 @@ inline double owens_t(double h, double a) { return boost::math::owens_t(h, a); }
  * @param b Second input
  * @return owens_t function applied to the two inputs.
  */
-template <typename T1, typename T2, require_any_container_t<T1, T2>* = nullptr>
+template <typename T1, typename T2, require_any_container_t<T1, T2>* = nullptr,
+          require_all_not_var_and_matrix_types<T1, T2>* = nullptr>
 inline auto owens_t(const T1& a, const T2& b) {
   return apply_scalar_binary(
-      a, b, [&](const auto& c, const auto& d) { return owens_t(c, d); });
+      a, b, [](const auto& c, const auto& d) { return owens_t(c, d); });
 }
 
 }  // namespace math

stan/math/prim/meta.hpp

@@ -207,6 +207,7 @@
 #include <stan/math/prim/meta/is_string_convertible.hpp>
 #include <stan/math/prim/meta/is_tuple.hpp>
 #include <stan/math/prim/meta/is_var.hpp>
+#include <stan/math/prim/meta/is_var_and_matrix_types.hpp>
 #include <stan/math/prim/meta/is_var_matrix.hpp>
 #include <stan/math/prim/meta/is_var_dense_dynamic.hpp>
 #include <stan/math/prim/meta/is_var_eigen.hpp>
@@ -220,6 +221,7 @@
 #include <stan/math/prim/meta/partials_return_type.hpp>
 #include <stan/math/prim/meta/partials_type.hpp>
 #include <stan/math/prim/meta/plain_type.hpp>
+#include <stan/math/prim/meta/possibly_sum.hpp>
 #include <stan/math/prim/meta/promote_args.hpp>
 #include <stan/math/prim/meta/promote_scalar_type.hpp>
 #include <stan/math/prim/meta/ref_type.hpp>

stan/math/prim/meta/is_var_and_matrix_types.hpp

@@ -0,0 +1,32 @@
+#ifndef STAN_MATH_PRIM_META_IS_VAR_AND_MATRIX_TYPES_HPP
+#define STAN_MATH_PRIM_META_IS_VAR_AND_MATRIX_TYPES_HPP
+
+#include <stan/math/prim/meta/disjunction.hpp>
+#include <stan/math/prim/meta/is_var.hpp>
+#include <stan/math/prim/meta/is_matrix.hpp>
+#include <stan/math/prim/meta/require_helpers.hpp>
+#include <stan/math/prim/meta/return_type.hpp>
+
+namespace stan {
+
+/** \ingroup type_trait
+ * Extends std::true_type when instantiated with at least one type that has a
+ * var `scalar_type` and at least one type is a matrix. Extends std::false_type
+ * otherwise.
+ * @tparam Types Types to test
+ */
+template <typename... Types>
+using is_var_and_matrix_types
+    = bool_constant<is_var<return_type_t<Types...>>::value
+                    && stan::math::disjunction<is_matrix<Types>...>::value>;
+
+template <typename... Types>
+using require_all_var_and_matrix_types
+    = require_t<is_var_and_matrix_types<Types...>>;
+
+template <typename... Types>
+using require_all_not_var_and_matrix_types
+    = require_not_t<is_var_and_matrix_types<Types...>>;
+
+}  // namespace stan
+#endif

stan/math/prim/meta/possibly_sum.hpp

@@ -0,0 +1,37 @@
+#ifndef STAN_MATH_PRIM_META_POSSIBLY_SUM_HPP
+#define STAN_MATH_PRIM_META_POSSIBLY_SUM_HPP
+
+#include <stan/math/prim/fun/sum.hpp>
+#include <stan/math/prim/meta/require_helpers.hpp>
+
+namespace stan {
+namespace math {
+
+/**
+ * Conditionally sum the input at compile time.
+ * @tparam CondSum A struct with a static boolean member `value` which if true
+ *  will allow the input value to be summed
+ * @tparam T A scalar, Eigen type, or standard vector with inner scalar type.
+ * @param x The value to be summed.
+ */
+template <typename CondSum, typename T, require_t<CondSum>* = nullptr>
+inline auto possibly_sum(T&& x) {
+  return sum(std::forward<T>(x));
+}
+
+/**
+ * Conditionally sum the input at compile time. This overload does not sum.
+ * @tparam CondSum A struct with a static boolean member `value` which if false
+ *  will pass the input to the output.
+ * @tparam T A scalar, Eigen type, or standard vector with inner scalar type.
+ * @param x The value to be passed trhough.
+ */
+template <typename CondSum, typename T1, require_not_t<CondSum>* = nullptr>
+inline auto possibly_sum(T1&& x) {
+  return std::forward<T1>(x);
+}
+
+}  // namespace math
+}  // namespace stan
+
+#endif

stan/math/rev/fun/owens_t.hpp

@@ -5,64 +5,49 @@
 #include <stan/math/rev/core.hpp>
 #include <stan/math/prim/fun/constants.hpp>
 #include <stan/math/prim/fun/erf.hpp>
+#include <stan/math/prim/fun/eval.hpp>
 #include <stan/math/prim/fun/owens_t.hpp>
 #include <stan/math/prim/fun/square.hpp>
 #include <cmath>
 
 namespace stan {
 namespace math {
 
-namespace internal {
-class owens_t_vv_vari : public op_vv_vari {
- public:
-  owens_t_vv_vari(vari* avi, vari* bvi)
-      : op_vv_vari(owens_t(avi->val_, bvi->val_), avi, bvi) {}
-  void chain() {
-    const double neg_avi_sq_div_2 = -square(avi_->val_) * 0.5;
-    const double one_p_bvi_sq = 1.0 + square(bvi_->val_);
-
-    avi_->adj_ += adj_ * erf(bvi_->val_ * avi_->val_ * INV_SQRT_TWO)
-                  * std::exp(neg_avi_sq_div_2) * INV_SQRT_TWO_PI * -0.5;
-    bvi_->adj_ += adj_ * std::exp(neg_avi_sq_div_2 * one_p_bvi_sq)
-                  / (one_p_bvi_sq * TWO_PI);
-  }
-};
-
-class owens_t_vd_vari : public op_vd_vari {
- public:
-  owens_t_vd_vari(vari* avi, double b)
-      : op_vd_vari(owens_t(avi->val_, b), avi, b) {}
-  void chain() {
-    avi_->adj_ += adj_ * erf(bd_ * avi_->val_ * INV_SQRT_TWO)
-                  * std::exp(-square(avi_->val_) * 0.5) * INV_SQRT_TWO_PI
-                  * -0.5;
-  }
-};
-
-class owens_t_dv_vari : public op_dv_vari {
- public:
-  owens_t_dv_vari(double a, vari* bvi)
-      : op_dv_vari(owens_t(a, bvi->val_), a, bvi) {}
-  void chain() {
-    const double one_p_bvi_sq = 1.0 + square(bvi_->val_);
-    bvi_->adj_ += adj_ * std::exp(-0.5 * square(ad_) * one_p_bvi_sq)
-                  / (one_p_bvi_sq * TWO_PI);
-  }
-};
-}  // namespace internal
-
 /**
  * The Owen's T function of h and a.
  *
  * Used to compute the cumulative density function for the skew normal
  * distribution.
  *
+ * @tparam Var1 A scalar or Eigen type whose `scalar_type` is an var.
+ * @tparam Var2 A scalar or Eigen type whose `scalar_type` is an var.
  * @param h var parameter.
  * @param a var parameter.
  * @return The Owen's T function.
  */
-inline var owens_t(const var& h, const var& a) {
-  return var(new internal::owens_t_vv_vari(h.vi_, a.vi_));
+template <typename Var1, typename Var2,
+          require_all_st_var<Var1, Var2>* = nullptr,
+          require_all_not_std_vector_t<Var1, Var2>* = nullptr>
+inline auto owens_t(const Var1& h, const Var2& a) {
+  auto h_arena = to_arena(h);
+  auto a_arena = to_arena(a);
+  using return_type
+      = return_var_matrix_t<decltype(owens_t(h_arena.val(), a_arena.val())),
+                            Var1, Var2>;
+  arena_t<return_type> ret = owens_t(h_arena.val(), a_arena.val());
+  reverse_pass_callback([h_arena, a_arena, ret]() mutable {
+    const auto& h_val = as_value_array_or_scalar(h_arena);
+    const auto& a_val = as_value_array_or_scalar(a_arena);
+    const auto neg_h_sq_div_2 = stan::math::eval(-square(h_val) * 0.5);
+    const auto one_p_a_sq = stan::math::eval(1.0 + square(a_val));
+    as_array_or_scalar(h_arena).adj() += possibly_sum<is_stan_scalar<Var1>>(
+        as_array_or_scalar(ret.adj()) * erf(a_val * h_val * INV_SQRT_TWO)
+        * exp(neg_h_sq_div_2) * INV_SQRT_TWO_PI * -0.5);
+    as_array_or_scalar(a_arena).adj() += possibly_sum<is_stan_scalar<Var2>>(
+        as_array_or_scalar(ret.adj()) * exp(neg_h_sq_div_2 * one_p_a_sq)
+        / (one_p_a_sq * TWO_PI));
+  });
+  return return_type(ret);
 }
 
 /**
@@ -71,12 +56,30 @@ inline var owens_t(const var& h, const var& a) {
  * Used to compute the cumulative density function for the skew normal
  * distribution.
  *
+ * @tparam Var A scalar or Eigen type whose `scalar_type` is an var.
+ * @tparam Arith A scalar or Eigen type with an inner arirthmetic scalar value.
  * @param h var parameter.
  * @param a double parameter.
  * @return The Owen's T function.
  */
-inline var owens_t(const var& h, double a) {
-  return var(new internal::owens_t_vd_vari(h.vi_, a));
+template <typename Var, typename Arith, require_st_arithmetic<Arith>* = nullptr,
+          require_all_not_std_vector_t<Var, Arith>* = nullptr,
+          require_st_var<Var>* = nullptr>
+inline auto owens_t(const Var& h, const Arith& a) {
+  auto h_arena = to_arena(h);
+  auto a_arena = to_arena(a);
+  using return_type
+      = return_var_matrix_t<decltype(owens_t(h_arena.val(), a_arena)), Var,
+                            Arith>;
+  arena_t<return_type> ret = owens_t(h_arena.val(), a_arena);
+  reverse_pass_callback([h_arena, a_arena, ret]() mutable {
+    const auto& h_val = as_value_array_or_scalar(h_arena);
+    as_array_or_scalar(h_arena).adj() += possibly_sum<is_stan_scalar<Var>>(
+        as_array_or_scalar(ret.adj())
+        * erf(as_array_or_scalar(a_arena) * h_val * INV_SQRT_TWO)
+        * exp(-square(h_val) * 0.5) * INV_SQRT_TWO_PI * -0.5);
+  });
+  return return_type(ret);
 }
 
 /**
@@ -85,12 +88,31 @@ inline var owens_t(const var& h, double a) {
  * Used to compute the cumulative density function for the skew normal
  * distribution.
  *
+ * @tparam Var A scalar or Eigen type whose `scalar_type` is an var.
+ * @tparam Arith A scalar or Eigen type with an inner arithmetic scalar value.
  * @param h double parameter.
  * @param a var parameter.
  * @return The Owen's T function.
  */
-inline var owens_t(double h, const var& a) {
-  return var(new internal::owens_t_dv_vari(h, a.vi_));
+template <typename Arith, typename Var, require_st_arithmetic<Arith>* = nullptr,
+          require_all_not_std_vector_t<Var, Arith>* = nullptr,
+          require_st_var<Var>* = nullptr>
+inline auto owens_t(const Arith& h, const Var& a) {
+  auto h_arena = to_arena(h);
+  auto a_arena = to_arena(a);
+  using return_type
+      = return_var_matrix_t<decltype(owens_t(h_arena, a_arena.val())), Var,
+                            Arith>;
+  arena_t<return_type> ret = owens_t(h_arena, a_arena.val());
+  reverse_pass_callback([h_arena, a_arena, ret]() mutable {
+    const auto one_p_a_sq
+        = eval(1.0 + square(as_value_array_or_scalar(a_arena)));
+    as_array_or_scalar(a_arena).adj() += possibly_sum<is_stan_scalar<Var>>(
+        as_array_or_scalar(ret.adj())
+        * exp(-0.5 * square(as_array_or_scalar(h_arena)) * one_p_a_sq)
+        / (one_p_a_sq * TWO_PI));
+  });
+  return return_type(ret);
 }
 
 }  // namespace math

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

@@ -15,6 +15,21 @@ TEST(mathMixScalFun, owensT) {
   stan::test::expect_ad(f, nan, nan);
 }
 
+TEST(mathMixScalFun, owensT_varmat) {
+  auto f = [](const auto& x1, const auto& x2) {
+    return stan::math::owens_t(x1, x2);
+  };
+  double scal = 2.0;
+  Eigen::MatrixXd mat = Eigen::MatrixXd::Random(2, 2);
+  Eigen::VectorXd vec = Eigen::VectorXd::Random(2);
+  stan::test::expect_ad_matvar(f, mat, mat);
+  stan::test::expect_ad_matvar(f, mat, scal);
+  stan::test::expect_ad_matvar(f, scal, mat);
+  stan::test::expect_ad_matvar(f, vec, vec);
+  stan::test::expect_ad_matvar(f, vec, scal);
+  stan::test::expect_ad_matvar(f, scal, vec);
+}
+
 TEST(mathMixScalFun, owensT_vec) {
   auto f = [](const auto& x1, const auto& x2) {
     using stan::math::owens_t;
@@ -27,3 +42,16 @@ TEST(mathMixScalFun, owensT_vec) {
   in2 << 3.0, 4.0;
   stan::test::expect_ad_vectorized_binary(f, in1, in2);
 }
+
+TEST(mathMixScalFun, owensT_vec_matvar) {
+  auto f = [](const auto& x1, const auto& x2) {
+    using stan::math::owens_t;
+    return owens_t(x1, x2);
+  };
+
+  Eigen::MatrixXd in1(2, 2);
+  in1 << 0.5, 3.4, 5.2, 0.5;
+  Eigen::MatrixXd in2(2, 2);
+  in2 << 3.3, 0.9, 6.7, 3.3;
+  stan::test::expect_ad_vectorized_matvar(f, in1, in2);
+}