Allow for integer arrays with vectorisation

Author: Andrew Johnson

stan/math/prim/functor/apply_scalar_binary.hpp

@@ -6,6 +6,7 @@
 #include <stan/math/prim/meta/is_container.hpp>
 #include <stan/math/prim/meta/is_eigen.hpp>
 #include <stan/math/prim/meta/require_generics.hpp>
+#include <stan/math/prim/fun/num_elements.hpp>
 #include <vector>
 
 namespace stan {
@@ -55,6 +56,116 @@ inline auto apply_scalar_binary(const T1& x, const T2& y, const F& f) {
   return x.binaryExpr(y, f).eval();
 }
 
+/**
+ * Specialisation for use with one Eigen vector (row or column) and
+ * a one-dimensional std::vector of integer types
+ *
+ * @tparam T1 Type of first argument to which functor is applied.
+ * @tparam T2 Type of second argument to which functor is applied.
+ * @tparam F Type of functor to apply.
+ * @param x Eigen input to which operation is applied.
+ * @param y Integer std::vector input to which operation is applied.
+ * @param f functor to apply to inputs.
+ * @return Eigen object with result of applying functor to inputs.
+ */
+template <typename T1, typename T2, typename F,
+          require_eigen_vector_vt<is_stan_scalar, T1>* = nullptr,
+          require_std_vector_vt<std::is_integral, T2>* = nullptr>
+inline auto apply_scalar_binary(const T1& x, const T2& y, const F& f) {
+  check_matching_sizes("Binary function", "x", x, "y", y);
+  using int_vec_t = promote_scalar_t<value_type_t<T2>, plain_type_t<T1>>;
+  Eigen::Map<const int_vec_t> y_map(y.data(), y.size());
+  return x.binaryExpr(y_map, f).eval();
+}
+
+/**
+ * Specialisation for use with a one-dimensional std::vector of integer types
+ * and one Eigen vector (row or column).
+ *
+ * @tparam T1 Type of first argument to which functor is applied.
+ * @tparam T2 Type of second argument to which functor is applied.
+ * @tparam F Type of functor to apply.
+ * @param x Integer std::vector input to which operation is applied.
+ * @param y Eigen input to which operation is applied.
+ * @param f functor to apply to inputs.
+ * @return Eigen object with result of applying functor to inputs.
+ */
+template <typename T1, typename T2, typename F,
+          require_std_vector_vt<std::is_integral, T1>* = nullptr,
+          require_eigen_vector_vt<is_stan_scalar, T2>* = nullptr>
+inline auto apply_scalar_binary(const T1& x, const T2& y, const F& f) {
+  check_matching_sizes("Binary function", "x", x, "y", y);
+  using int_vec_t = promote_scalar_t<value_type_t<T1>, plain_type_t<T2>>;
+  Eigen::Map<const int_vec_t> x_map(x.data(), x.size());
+  return x_map.binaryExpr(y, f).eval();
+}
+
+/**
+ * Specialisation for use with one Eigen matrix and
+ * a two-dimensional std::vector of integer types
+ *
+ * @tparam T1 Type of first argument to which functor is applied.
+ * @tparam T2 Type of second argument to which functor is applied.
+ * @tparam F Type of functor to apply.
+ * @param x Eigen matrix input to which operation is applied.
+ * @param y Nested integer std::vector input to which operation is applied.
+ * @param f functor to apply to inputs.
+ * @return Eigen object with result of applying functor to inputs.
+ */
+template <typename T1, typename T2, typename F,
+          require_eigen_matrix_vt<is_stan_scalar, T1>* = nullptr,
+          require_std_vector_vt<is_std_vector, T2>* = nullptr,
+          require_std_vector_st<std::is_integral, T2>* = nullptr>
+inline auto apply_scalar_binary(const T1& x, const T2& y, const F& f) {
+  if (num_elements(x) != num_elements(y)) {
+    std::ostringstream msg;
+    msg << "Inputs to vectorized binary function must match in"
+        << " size if one is not a scalar";
+    throw std::invalid_argument(msg.str());
+  }
+  using return_st = decltype(f(x(0), y[0][0]));
+  Eigen::Matrix<return_st, Eigen::Dynamic, Eigen::Dynamic>
+    result(x.rows(), x.cols());
+  for (size_t i = 0; i < y.size(); ++i) {
+    result.row(i) = apply_scalar_binary(x.row(i).transpose(),
+                                        as_column_vector_or_scalar(y[i]), f);
+  }
+  return result;
+}
+
+/**
+ * Specialisation for use with a two-dimensional std::vector of integer types
+ * and one Eigen matrix.
+ *
+ * @tparam T1 Type of first argument to which functor is applied.
+ * @tparam T2 Type of second argument to which functor is applied.
+ * @tparam F Type of functor to apply.
+ * @param x Nested integer std::vector input to which operation is applied.
+ * @param y Eigen matrix input to which operation is applied.
+ * @param f functor to apply to inputs.
+ * @return Eigen object with result of applying functor to inputs.
+ */
+template <typename T1, typename T2, typename F,
+          require_std_vector_vt<is_std_vector, T1>* = nullptr,
+          require_std_vector_st<std::is_integral, T1>* = nullptr,
+          require_eigen_matrix_vt<is_stan_scalar, T2>* = nullptr>
+inline auto apply_scalar_binary(const T1& x, const T2& y, const F& f) {
+  if (num_elements(x) != num_elements(y)) {
+    std::ostringstream msg;
+    msg << "Inputs to vectorized binary function must match in"
+        << " size if one is not a scalar";
+    throw std::invalid_argument(msg.str());
+  }
+  using return_st = decltype(f(x[0][0], y(0)));
+  Eigen::Matrix<return_st, Eigen::Dynamic, Eigen::Dynamic>
+    result(y.rows(), y.cols());
+  for (size_t i = 0; i < x.size(); ++i) {
+    result.row(i) = apply_scalar_binary(as_column_vector_or_scalar(x[i]),
+                                        y.row(i).transpose(), f);
+  }
+  return result;
+}
+
 /**
  * Specialisation for use when the first input is an Eigen type and the second
  * is a scalar. Eigen's unaryExpr framework is used for more efficient indexing

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

@@ -20,9 +20,12 @@ TEST(mathMixScalFun, besselFirstKind_vec) {
     return bessel_first_kind(x1, x2);
   };
 
-  Eigen::VectorXi in1(2);
-  in1 << 3, 1;
+  std::vector<int> std_in1{3, 1};
   Eigen::VectorXd in2(2);
   in2 << 0.5, 3.4;
-  stan::test::expect_ad_vectorized_binary(f, in1, in2);
+  stan::test::expect_ad_vectorized_binary(f, std_in1, in2);
+
+  std::vector<std::vector<int>> std_std_in1{std_in1, std_in1};
+  Eigen::MatrixXd mat_in2 = in2.replicate(1, 2);
+  stan::test::expect_ad_vectorized_binary(f, std_std_in1, mat_in2);
 }

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

@@ -29,7 +29,11 @@ TEST(mathMixScalFun, fallingFactorial_vec) {
 
   Eigen::VectorXd in1(2);
   in1 << 0.5, 3.4;
-  Eigen::VectorXi in2(2);
-  in2 << 3, 1;
-  stan::test::expect_ad_vectorized_binary(f, in1, in2);
+  std::vector<int> std_in2{3, 1};
+  stan::test::expect_ad_vectorized_binary(f, in1, std_in2);
+
+
+  Eigen::MatrixXd mat_in1 = in1.replicate(1, 2);
+  std::vector<std::vector<int>> std_std_in2{std_in2, std_in2};
+  stan::test::expect_ad_vectorized_binary(f, mat_in1, std_std_in2);
 }

test/unit/math/prim/fun/bessel_first_kind_test.cpp

@@ -26,9 +26,13 @@ TEST(MathFunctions, bessel_first_kind_vec) {
     return bessel_first_kind(x1, x2);
   };
 
-  Eigen::VectorXi in1(3);
-  in1 << 1, 3, 1;
+  std::vector<int> std_in1{1, 3, 1};
   Eigen::VectorXd in2(3);
   in2 << -1.3, 0.7, 2.8;
-  stan::test::binary_scalar_tester(f, in1, in2);
+  stan::test::binary_scalar_tester(f, std_in1, in2);
+
+  Eigen::MatrixXd mat_in2 = in2.replicate(1, 3);
+  std::vector<std::vector<int>> std_std_in1{std_in1, std_in1, std_in1};
+  stan::test::binary_scalar_tester(f, std_std_in1, mat_in2);
+
 }

test/unit/math/prim/fun/binary_scalar_tester.hpp

@@ -17,7 +17,7 @@ namespace test {
  * @param f functor to apply to inputs.
  */
 template <typename F, typename T1, typename T2,
-          require_all_not_std_vector_t<T1, T2>* = nullptr>
+          require_all_not_vector_t<T1, T2>* = nullptr>
 void binary_scalar_tester_impl(const F& f, const T1& x, const T2& y) {
   auto vec_vec = f(x, y);
   auto vec_scal = f(x, y(0));
@@ -87,7 +87,7 @@ void binary_scalar_tester_impl(const F& f, const T1& x, const T2& y) {
  * @param f functor to apply to inputs.
  */
 template <typename F, typename T1, typename T2,
-          require_all_std_vector_t<T1, T2>* = nullptr>
+          require_all_vector_t<T1, T2>* = nullptr>
 void binary_scalar_tester_impl(const F& f, const T1& x, const T2& y) {
   auto vec_vec = f(x, y);
   auto vec_scal = f(x, y[0]);
@@ -97,13 +97,13 @@ void binary_scalar_tester_impl(const F& f, const T1& x, const T2& y) {
     EXPECT_FLOAT_EQ(f(x[i], y[0]), vec_scal[i]);
     EXPECT_FLOAT_EQ(f(x[0], y[i]), scal_vec[i]);
   }
-  T1 x_zero;
-  T2 y_zero;
+  plain_type_t<T1> x_zero;
+  plain_type_t<T2> y_zero;
   EXPECT_THROW(f(x_zero, y), std::invalid_argument);
   EXPECT_THROW(f(x, y_zero), std::invalid_argument);
 
-  std::vector<T1> nest_x{x, x, x};
-  std::vector<T2> nest_y{y, y, y};
+  std::vector<plain_type_t<T1>> nest_x{x, x, x};
+  std::vector<plain_type_t<T2>> nest_y{y, y, y};
   auto nestvec_nestvec = f(nest_x, nest_y);
   auto nestvec_scal = f(nest_x, y[0]);
   auto scal_nestvec = f(x[0], nest_y);
@@ -114,13 +114,13 @@ void binary_scalar_tester_impl(const F& f, const T1& x, const T2& y) {
       EXPECT_FLOAT_EQ(f(x[0], nest_y[i][j]), scal_nestvec[i][j]);
     }
   }
-  std::vector<T1> nest_x_small{x, x};
-  std::vector<T2> nest_y_small{y, y};
+  std::vector<plain_type_t<T1>> nest_x_small{x, x};
+  std::vector<plain_type_t<T2>> nest_y_small{y, y};
   EXPECT_THROW(f(nest_x, nest_y_small), std::invalid_argument);
   EXPECT_THROW(f(nest_x_small, nest_y), std::invalid_argument);
 
-  std::vector<std::vector<T1>> nest_nest_x{nest_x, nest_x, nest_x};
-  std::vector<std::vector<T2>> nest_nest_y{nest_y, nest_y, nest_y};
+  std::vector<std::vector<plain_type_t<T1>>> nest_nest_x{nest_x, nest_x, nest_x};
+  std::vector<std::vector<plain_type_t<T2>>> nest_nest_y{nest_y, nest_y, nest_y};
   auto nestnestvec_nestnestvec = f(nest_nest_x, nest_nest_y);
   auto nestnestvec_scal = f(nest_nest_x, y[0]);
   auto scal_nestnestvec = f(x[0], nest_nest_y);
@@ -136,12 +136,129 @@ void binary_scalar_tester_impl(const F& f, const T1& x, const T2& y) {
       }
     }
   }
-  std::vector<std::vector<T1>> nest_nest_x_small{nest_x, nest_x};
+  std::vector<std::vector<plain_type_t<T1>>> nest_nest_x_small{nest_x, nest_x};
+  std::vector<std::vector<plain_type_t<T2>>> nest_nest_y_small{nest_y, nest_y};
+  EXPECT_THROW(f(nest_nest_x, nest_nest_y_small), std::invalid_argument);
+  EXPECT_THROW(f(nest_nest_x_small, nest_nest_y), std::invalid_argument);
+}
+
+
+/**
+ * Implementation function which checks that the binary vectorisation
+ * framework returns the same value as the function with scalar inputs,
+ * for all valid combinations of scalar/vector/nested vector.
+ *
+ * @tparam F Type of functor to apply.
+ * @tparam T1 Type of first vector.
+ * @tparam T2 Type of second vector.
+ * @param x First vector input to which operation is applied.
+ * @param y Second vector input to which operation is applied.
+ * @param f functor to apply to inputs.
+ */
+template <typename F, typename T1, typename T2,
+          typename T1_plain = plain_type_t<T1>,
+          require_eigen_matrix_t<T1>* = nullptr,
+          require_std_vector_t<T2>* = nullptr>
+void binary_scalar_tester_impl(const F& f, const T1& x, const T2& y) {
+  auto vec_vec = f(x, y);
+  for (int r = 0; r < x.rows(); ++r) {
+    for(int c = 0; c < x.cols(); ++c) {
+      EXPECT_FLOAT_EQ(f(x(r, c), y[r][c]), vec_vec(r, c));
+    }
+  }
+
+  T1_plain x_zero;
+  T2 y_zero;
+  EXPECT_THROW(f(x_zero, y), std::invalid_argument);
+  EXPECT_THROW(f(x, y_zero), std::invalid_argument);
+
+  std::vector<T1_plain> nest_x{x, x, x};
+  std::vector<T2> nest_y{y, y, y};
+  auto nestvec_nestvec = f(nest_x, nest_y);
+  for (int i = 0; i < 3; ++i) {
+    for (int r = 0; r < x.rows(); ++r) {
+      for(int c = 0; c < x.cols(); ++c) {
+        EXPECT_FLOAT_EQ(f(nest_x[i](r, c), nest_y[i][r][c]), nestvec_nestvec[i](r, c));
+      }
+    }
+  }
+  std::vector<T1_plain> nest_x_small{x, x};
+  std::vector<T2> nest_y_small{y, y};
+  EXPECT_THROW(f(nest_x, nest_y_small), std::invalid_argument);
+  EXPECT_THROW(f(nest_x_small, nest_y), std::invalid_argument);
+
+  std::vector<std::vector<T1_plain>> nest_nest_x{nest_x, nest_x, nest_x};
+  std::vector<std::vector<T2>> nest_nest_y{nest_y, nest_y, nest_y};
+
+  auto nestnestvec_nestnestvec = f(nest_nest_x, nest_nest_y);
+  for (int i = 0; i < 3; ++i) {
+    for (int j = 0; j < 3; ++j) {
+      for (int r = 0; r < x.rows(); ++r) {
+        for(int c = 0; c < x.cols(); ++c) {
+          EXPECT_FLOAT_EQ(f(nest_nest_x[i][j](r, c), nest_nest_y[i][j][r][c]),
+                          nestnestvec_nestnestvec[i][j](r, c));
+        }
+      }
+    }
+  }
+  std::vector<std::vector<T1_plain>> nest_nest_x_small{nest_x, nest_x};
   std::vector<std::vector<T2>> nest_nest_y_small{nest_y, nest_y};
   EXPECT_THROW(f(nest_nest_x, nest_nest_y_small), std::invalid_argument);
   EXPECT_THROW(f(nest_nest_x_small, nest_nest_y), std::invalid_argument);
 }
 
+template <typename F, typename T1, typename T2,
+          typename T2_plain = plain_type_t<T2>,
+          require_std_vector_t<T1>* = nullptr,
+          require_eigen_matrix_t<T2>* = nullptr>
+void binary_scalar_tester_impl(const F& f, const T1& x, const T2& y) {
+  auto vec_vec = f(x, y);
+  for (int r = 0; r < y.rows(); ++r) {
+    for(int c = 0; c < y.cols(); ++c) {
+      EXPECT_FLOAT_EQ(f(x[r][c], y(r, c)), vec_vec(r, c));
+    }
+  }
+
+  T1 x_zero;
+  T2_plain y_zero;
+  EXPECT_THROW(f(x_zero, y), std::invalid_argument);
+  EXPECT_THROW(f(x, y_zero), std::invalid_argument);
+
+  std::vector<T1> nest_x{x, x, x};
+  std::vector<T2_plain> nest_y{y, y, y};
+  auto nestvec_nestvec = f(nest_x, nest_y);
+  for (int i = 0; i < 3; ++i) {
+    for (int r = 0; r < y.rows(); ++r) {
+      for(int c = 0; c < y.cols(); ++c) {
+        EXPECT_FLOAT_EQ(f(nest_x[i][r][c], nest_y[i](r, c)), nestvec_nestvec[i](r, c));
+      }
+    }
+  }
+  std::vector<T1> nest_x_small{x, x};
+  std::vector<T2_plain> nest_y_small{y, y};
+  EXPECT_THROW(f(nest_x, nest_y_small), std::invalid_argument);
+  EXPECT_THROW(f(nest_x_small, nest_y), std::invalid_argument);
+
+  std::vector<std::vector<T1>> nest_nest_x{nest_x, nest_x, nest_x};
+  std::vector<std::vector<T2_plain>> nest_nest_y{nest_y, nest_y, nest_y};
+
+  auto nestnestvec_nestnestvec = f(nest_nest_x, nest_nest_y);
+  for (int i = 0; i < 3; ++i) {
+    for (int j = 0; j < 3; ++j) {
+      for (int r = 0; r < y.rows(); ++r) {
+        for(int c = 0; c < y.cols(); ++c) {
+          EXPECT_FLOAT_EQ(f(nest_nest_x[i][j][r][c], nest_nest_y[i][j](r, c)),
+                          nestnestvec_nestnestvec[i][j](r, c));
+        }
+      }
+    }
+  }
+  std::vector<std::vector<T1>> nest_nest_x_small{nest_x, nest_x};
+  std::vector<std::vector<T2_plain>> nest_nest_y_small{nest_y, nest_y};
+  EXPECT_THROW(f(nest_nest_x, nest_nest_y_small), std::invalid_argument);
+  EXPECT_THROW(f(nest_nest_x_small, nest_nest_y), std::invalid_argument);
+}
+
 /**
  * Testing framework for checking that the vectorisation of binary
  * functions returns the same results as the binary function with
@@ -172,5 +289,10 @@ void binary_scalar_tester(const F& f, const T1& x, const T2& y) {
       std::vector<typename T2::Scalar>(y.data(), y.data() + y.size()));
 }
 
+template <typename F, typename T1, typename T2,
+          require_any_std_vector_t<T1, T2>* = nullptr>
+void binary_scalar_tester(const F& f, const T1& x, const T2& y) {
+  binary_scalar_tester_impl(f, x, y);
+}
 }  // namespace test
 }  // namespace stan

test/unit/math/prim/fun/falling_factorial_test.cpp

@@ -25,7 +25,10 @@ TEST(MathFunctions, falling_factorial_vec) {
 
   Eigen::VectorXd in1(3);
   in1 << -1.3, 0.7, 2.8;
-  Eigen::VectorXi in2(3);
-  in2 << 1, 3, 1;
-  stan::test::binary_scalar_tester(f, in1, in2);
+  std::vector<int> std_in2{1, 3, 1};
+  stan::test::binary_scalar_tester(f, in1, std_in2);
+
+  Eigen::MatrixXd mat_in1 = in1.replicate(1, 3);
+  std::vector<std::vector<int>> std_std_in2{std_in2, std_in2, std_in2};
+  stan::test::binary_scalar_tester(f, mat_in1, std_std_in2);
 }