Add matrix/vector array reshaping helpers

stan/math/prim/fun.hpp

@@ -326,7 +326,9 @@
 #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>
+#include <stan/math/prim/fun/to_row_vector_array.hpp>
 #include <stan/math/prim/fun/to_vector.hpp>
+#include <stan/math/prim/fun/to_vector_array.hpp>
 #include <stan/math/prim/fun/trace.hpp>
 #include <stan/math/prim/fun/trace_gen_inv_quad_form_ldlt.hpp>
 #include <stan/math/prim/fun/trace_gen_quad_form.hpp>

stan/math/prim/fun/to_matrix.hpp

@@ -39,6 +39,29 @@ inline auto to_matrix(EigVec&& matrix) {
       std::forward<EigVec>(matrix));
 }
 
+/**
+ * Returns a matrix representation of a standard vector of Eigen
+ * column vectors. Each input vector becomes a column in the result.
+ *
+ * @tparam T type of the elements in the vector
+ * @param x standard vector of Eigen column vectors
+ * @return the matrix representation of the input
+ */
+template <typename T>
+inline Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic> to_matrix(
+    const std::vector<Eigen::Matrix<T, Eigen::Dynamic, 1>>& x) {
+  int cols = x.size();
+  if (cols == 0) {
+    return {};
+  }
+  int rows = x[0].size();
+  Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic> result(rows, cols);
+  for (int j = 0; j < cols; ++j) {
+    result.col(j) = x[j];
+  }
+  return result;
+}
+
 /**
  * Returns a matrix representation of a standard vector of Eigen
  * row vectors with the same dimensions and indexing order.

stan/math/prim/fun/to_row_vector_array.hpp

@@ -0,0 +1,33 @@
+#ifndef STAN_MATH_PRIM_FUN_TO_ROW_VECTOR_ARRAY_HPP
+#define STAN_MATH_PRIM_FUN_TO_ROW_VECTOR_ARRAY_HPP
+
+#include <stan/math/prim/meta.hpp>
+#include <stan/math/prim/fun/Eigen.hpp>
+#include <vector>
+
+namespace stan {
+namespace math {
+
+/**
+ * Returns a standard vector of Eigen row vectors from the rows of the input
+ * matrix.
+ *
+ * @tparam EigMat type of the input matrix
+ * @param matrix input matrix
+ * @return the array of row vectors representation of the input
+ */
+template <typename EigMat, require_eigen_t<EigMat>* = nullptr>
+inline std::vector<Eigen::Matrix<value_type_t<EigMat>, 1, Eigen::Dynamic>>
+to_row_vector_array(const EigMat& matrix) {
+  using T = value_type_t<EigMat>;
+  std::vector<Eigen::Matrix<T, 1, Eigen::Dynamic>> result;
+  result.reserve(matrix.rows());
+  for (int i = 0; i < matrix.rows(); ++i) {
+    result.push_back(matrix.row(i));
+  }
+  return result;
+}
+
+}  // namespace math
+}  // namespace stan
+#endif

stan/math/prim/fun/to_vector_array.hpp

@@ -0,0 +1,33 @@
+#ifndef STAN_MATH_PRIM_FUN_TO_VECTOR_ARRAY_HPP
+#define STAN_MATH_PRIM_FUN_TO_VECTOR_ARRAY_HPP
+
+#include <stan/math/prim/meta.hpp>
+#include <stan/math/prim/fun/Eigen.hpp>
+#include <vector>
+
+namespace stan {
+namespace math {
+
+/**
+ * Returns a standard vector of Eigen column vectors from the columns of
+ * the input matrix.
+ *
+ * @tparam EigMat type of the input matrix
+ * @param matrix input matrix
+ * @return the array of column vectors representation of the input
+ */
+template <typename EigMat, require_eigen_t<EigMat>* = nullptr>
+inline std::vector<Eigen::Matrix<value_type_t<EigMat>, Eigen::Dynamic, 1>>
+to_vector_array(const EigMat& matrix) {
+  using T = value_type_t<EigMat>;
+  std::vector<Eigen::Matrix<T, Eigen::Dynamic, 1>> result;
+  result.reserve(matrix.cols());
+  for (int j = 0; j < matrix.cols(); ++j) {
+    result.push_back(matrix.col(j));
+  }
+  return result;
+}
+
+}  // namespace math
+}  // namespace stan
+#endif

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

@@ -1,6 +1,7 @@
 #include <stan/math/prim.hpp>
 #include <test/unit/util.hpp>
 #include <gtest/gtest.h>
+#include <type_traits>
 #include <vector>
 #include <stdexcept>
 
@@ -205,8 +206,8 @@ TEST(ToMatrixRowVector, answers) {
 // [[T]] -> Matrix
 inline void test_to_matrix_2darray_answers(int m, int n) {
   using stan::math::to_matrix;
-  std::vector<std::vector<double> > vec(m, std::vector<double>(n));
-  std::vector<std::vector<int> > vec_int(m, std::vector<int>(n));
+  std::vector<std::vector<double>> vec(m, std::vector<double>(n));
+  std::vector<std::vector<int>> vec_int(m, std::vector<int>(n));
   // Any vec (0, C) will become (0, 0)
   if (m == 0)
     n = 0;
@@ -229,3 +230,111 @@ TEST(ToMatrix2dArray, answers) {
   test_to_matrix_2darray_answers(3, 0);
   test_to_matrix_2darray_answers(0, 3);
 }
+
+TEST(ToMatrixVectorArray, answers) {
+  using stan::math::to_matrix;
+  using stan::math::to_vector_array;
+  using vector_d = Eigen::Matrix<double, Eigen::Dynamic, 1>;
+
+  std::vector<vector_d> vecs(2, vector_d(3));
+  vecs[0] << 1.1, 2.2, 3.3;
+  vecs[1] << 4.4, 5.5, 6.6;
+
+  Eigen::MatrixXd expected(3, 2);
+  expected << 1.1, 4.4, 2.2, 5.5, 3.3, 6.6;
+
+  Eigen::MatrixXd result = to_matrix(vecs);
+  EXPECT_MATRIX_FLOAT_EQ(expected, result);
+
+  std::vector<vector_d> round_trip = to_vector_array(result);
+  ASSERT_EQ(vecs.size(), round_trip.size());
+  for (size_t i = 0; i < vecs.size(); ++i) {
+    EXPECT_MATRIX_FLOAT_EQ(vecs[i], round_trip[i]);
+  }
+}
+
+TEST(ToMatrixVectorArray, empty) {
+  using stan::math::to_matrix;
+  using stan::math::to_vector_array;
+  using vector_d = Eigen::Matrix<double, Eigen::Dynamic, 1>;
+
+  std::vector<vector_d> empty;
+  Eigen::MatrixXd empty_result = to_matrix(empty);
+  EXPECT_EQ(0, empty_result.rows());
+  EXPECT_EQ(0, empty_result.cols());
+
+  std::vector<vector_d> zero_rows(3, vector_d(0));
+  Eigen::MatrixXd zero_rows_result = to_matrix(zero_rows);
+  EXPECT_EQ(0, zero_rows_result.rows());
+  EXPECT_EQ(3, zero_rows_result.cols());
+
+  std::vector<vector_d> round_trip = to_vector_array(zero_rows_result);
+  ASSERT_EQ(zero_rows.size(), round_trip.size());
+  for (size_t i = 0; i < zero_rows.size(); ++i) {
+    EXPECT_MATRIX_FLOAT_EQ(zero_rows[i], round_trip[i]);
+  }
+}
+
+TEST(ToMatrixVectorArray, preservesScalarType) {
+  using stan::math::to_matrix;
+  using vector_i = Eigen::Matrix<int, Eigen::Dynamic, 1>;
+
+  std::vector<vector_i> vecs(2, vector_i(2));
+  vecs[0] << 1, 2;
+  vecs[1] << 3, 4;
+
+  auto result = to_matrix(vecs);
+  static_assert(
+      std::is_same<decltype(result),
+                   Eigen::Matrix<int, Eigen::Dynamic, Eigen::Dynamic>>::value,
+      "to_matrix should preserve the vector scalar type");
+  EXPECT_EQ(1, result(0, 0));
+  EXPECT_EQ(2, result(1, 0));
+  EXPECT_EQ(3, result(0, 1));
+  EXPECT_EQ(4, result(1, 1));
+}
+
+TEST(ToMatrixRowVectorArray, answers) {
+  using stan::math::to_matrix;
+  using stan::math::to_row_vector_array;
+  using row_vector_d = Eigen::Matrix<double, 1, Eigen::Dynamic>;
+
+  std::vector<row_vector_d> row_vecs(3, row_vector_d(2));
+  row_vecs[0] << 1.1, 4.4;
+  row_vecs[1] << 2.2, 5.5;
+  row_vecs[2] << 3.3, 6.6;
+
+  Eigen::MatrixXd expected(3, 2);
+  expected << 1.1, 4.4, 2.2, 5.5, 3.3, 6.6;
+
+  Eigen::MatrixXd result = to_matrix(row_vecs);
+  EXPECT_MATRIX_FLOAT_EQ(expected, result);
+
+  std::vector<row_vector_d> round_trip = to_row_vector_array(result);
+  ASSERT_EQ(row_vecs.size(), round_trip.size());
+  for (size_t i = 0; i < row_vecs.size(); ++i) {
+    EXPECT_MATRIX_FLOAT_EQ(row_vecs[i], round_trip[i]);
+  }
+}
+
+TEST(ToMatrixRowVectorArray, empty) {
+  using stan::math::to_matrix;
+  using stan::math::to_row_vector_array;
+  using row_vector_d = Eigen::Matrix<double, 1, Eigen::Dynamic>;
+
+  std::vector<row_vector_d> empty;
+  Eigen::MatrixXd empty_result = to_matrix(empty);
+  EXPECT_EQ(0, empty_result.rows());
+  EXPECT_EQ(0, empty_result.cols());
+
+  std::vector<row_vector_d> zero_cols(3, row_vector_d(0));
+  Eigen::MatrixXd zero_cols_result = to_matrix(zero_cols);
+  EXPECT_EQ(3, zero_cols_result.rows());
+  EXPECT_EQ(0, zero_cols_result.cols());
+
+  std::vector<row_vector_d> round_trip = to_row_vector_array(zero_cols_result);
+  ASSERT_EQ(zero_cols.size(), round_trip.size());
+  for (size_t i = 0; i < zero_cols.size(); ++i) {
+    EXPECT_MATRIX_FLOAT_EQ(zero_cols[i], round_trip[i]);
+  }
+}

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

@@ -0,0 +1,66 @@
+#include <stan/math/prim.hpp>
+#include <test/unit/util.hpp>
+#include <gtest/gtest.h>
+#include <type_traits>
+#include <vector>
+
+TEST(ToRowVectorArray, values) {
+  using stan::math::to_matrix;
+  using stan::math::to_row_vector_array;
+
+  Eigen::MatrixXd matrix(3, 2);
+  matrix << 1.1, 4.4, 2.2, 5.5, 3.3, 6.6;
+
+  auto result = to_row_vector_array(matrix);
+  ASSERT_EQ(3, result.size());
+
+  Eigen::RowVectorXd expected_row_0(2);
+  expected_row_0 << 1.1, 4.4;
+  Eigen::RowVectorXd expected_row_1(2);
+  expected_row_1 << 2.2, 5.5;
+  Eigen::RowVectorXd expected_row_2(2);
+  expected_row_2 << 3.3, 6.6;
+
+  EXPECT_MATRIX_FLOAT_EQ(expected_row_0, result[0]);
+  EXPECT_MATRIX_FLOAT_EQ(expected_row_1, result[1]);
+  EXPECT_MATRIX_FLOAT_EQ(expected_row_2, result[2]);
+  EXPECT_MATRIX_FLOAT_EQ(matrix, to_matrix(result));
+}
+
+TEST(ToRowVectorArray, emptyShapes) {
+  using stan::math::to_matrix;
+  using stan::math::to_row_vector_array;
+
+  Eigen::MatrixXd zero_by_zero(0, 0);
+  auto zero_by_zero_result = to_row_vector_array(zero_by_zero);
+  EXPECT_EQ(0, zero_by_zero_result.size());
+
+  Eigen::MatrixXd zero_by_three(0, 3);
+  auto zero_by_three_result = to_row_vector_array(zero_by_three);
+  EXPECT_EQ(0, zero_by_three_result.size());
+
+  Eigen::MatrixXd three_by_zero(3, 0);
+  auto three_by_zero_result = to_row_vector_array(three_by_zero);
+  ASSERT_EQ(3, three_by_zero_result.size());
+  for (const auto& result : three_by_zero_result) {
+    EXPECT_EQ(0, result.size());
+  }
+  EXPECT_MATRIX_FLOAT_EQ(three_by_zero, to_matrix(three_by_zero_result));
+}
+
+TEST(ToRowVectorArray, preservesScalarType) {
+  using stan::math::to_row_vector_array;
+
+  Eigen::Matrix<int, Eigen::Dynamic, Eigen::Dynamic> matrix(2, 2);
+  matrix << 1, 3, 2, 4;
+
+  auto result = to_row_vector_array(matrix);
+  static_assert(std::is_same<decltype(result)::value_type,
+                             Eigen::Matrix<int, 1, Eigen::Dynamic>>::value,
+                "to_row_vector_array should preserve the matrix scalar type");
+  ASSERT_EQ(2, result.size());
+  EXPECT_EQ(1, result[0](0));
+  EXPECT_EQ(3, result[0](1));
+  EXPECT_EQ(2, result[1](0));
+  EXPECT_EQ(4, result[1](1));
+}

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

@@ -0,0 +1,63 @@
+#include <stan/math/prim.hpp>
+#include <test/unit/util.hpp>
+#include <gtest/gtest.h>
+#include <type_traits>
+#include <vector>
+
+TEST(ToVectorArray, values) {
+  using stan::math::to_matrix;
+  using stan::math::to_vector_array;
+
+  Eigen::MatrixXd matrix(3, 2);
+  matrix << 1.1, 4.4, 2.2, 5.5, 3.3, 6.6;
+
+  auto result = to_vector_array(matrix);
+  ASSERT_EQ(2, result.size());
+
+  Eigen::VectorXd expected_col_0(3);
+  expected_col_0 << 1.1, 2.2, 3.3;
+  Eigen::VectorXd expected_col_1(3);
+  expected_col_1 << 4.4, 5.5, 6.6;
+
+  EXPECT_MATRIX_FLOAT_EQ(expected_col_0, result[0]);
+  EXPECT_MATRIX_FLOAT_EQ(expected_col_1, result[1]);
+  EXPECT_MATRIX_FLOAT_EQ(matrix, to_matrix(result));
+}
+
+TEST(ToVectorArray, emptyShapes) {
+  using stan::math::to_matrix;
+  using stan::math::to_vector_array;
+
+  Eigen::MatrixXd zero_by_zero(0, 0);
+  auto zero_by_zero_result = to_vector_array(zero_by_zero);
+  EXPECT_EQ(0, zero_by_zero_result.size());
+
+  Eigen::MatrixXd zero_by_three(0, 3);
+  auto zero_by_three_result = to_vector_array(zero_by_three);
+  ASSERT_EQ(3, zero_by_three_result.size());
+  for (const auto& result : zero_by_three_result) {
+    EXPECT_EQ(0, result.size());
+  }
+  EXPECT_MATRIX_FLOAT_EQ(zero_by_three, to_matrix(zero_by_three_result));
+
+  Eigen::MatrixXd three_by_zero(3, 0);
+  auto three_by_zero_result = to_vector_array(three_by_zero);
+  EXPECT_EQ(0, three_by_zero_result.size());
+}
+
+TEST(ToVectorArray, preservesScalarType) {
+  using stan::math::to_vector_array;
+
+  Eigen::Matrix<int, Eigen::Dynamic, Eigen::Dynamic> matrix(2, 2);
+  matrix << 1, 3, 2, 4;
+
+  auto result = to_vector_array(matrix);
+  static_assert(std::is_same<decltype(result)::value_type,
+                             Eigen::Matrix<int, Eigen::Dynamic, 1>>::value,
+                "to_vector_array should preserve the matrix scalar type");
+  ASSERT_EQ(2, result.size());
+  EXPECT_EQ(1, result[0](0));
+  EXPECT_EQ(2, result[0](1));
+  EXPECT_EQ(3, result[1](0));
+  EXPECT_EQ(4, result[1](1));
+}