diff --git a/stan/math/prim/functor.hpp b/stan/math/prim/functor.hpp index 2f9de59c706..22a3934f99d 100644 --- a/stan/math/prim/functor.hpp +++ b/stan/math/prim/functor.hpp @@ -23,6 +23,7 @@ #include #include #include +#include #include #include #include diff --git a/stan/math/prim/functor/map.hpp b/stan/math/prim/functor/map.hpp new file mode 100644 index 00000000000..ba87d604402 --- /dev/null +++ b/stan/math/prim/functor/map.hpp @@ -0,0 +1,369 @@ +#ifndef STAN_MATH_PRIM_FUNCTOR_MAP_HPP +#define STAN_MATH_PRIM_FUNCTOR_MAP_HPP + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace stan { +namespace math { + +namespace internal { + +template +inline void check_all_matching_sizes(const char* function, const T& x, + const Types&... xs) { + std::size_t arg_idx = 2; + ( + [&](const auto& y) { + if (x.size() != y.size()) { + [&]() STAN_COLD_PATH { + const std::string name = "x" + std::to_string(arg_idx); + check_matching_sizes(function, "x1", x, name.c_str(), y); + }(); + } + ++arg_idx; + }(xs), + ...); +} + +template +inline void check_all_matching_dims(const char* function, const T& x, + const Types&... xs) { + std::size_t arg_idx = 2; + ( + [&](const auto& y) { + if (x.rows() != y.rows() || x.cols() != y.cols()) { + [&]() STAN_COLD_PATH { + const std::string name = "x" + std::to_string(arg_idx); + check_matching_dims(function, "x1", x, name.c_str(), y); + }(); + } + ++arg_idx; + }(xs), + ...); +} + +template +inline void assign_matrix_row( + const char* function, + Eigen::Matrix& result, + Eigen::Index i, T&& x) { + if (i == 0) { + result = Eigen::Matrix( + result.rows(), x.cols()); + } else { + check_size_match(function, "columns of result", result.cols(), + "columns of returned row", x.cols()); + } + result.row(i) = std::forward(x); +} + +template +inline void assign_matrix_col( + const char* function, + Eigen::Matrix& result, + Eigen::Index j, T&& x) { + if (j == 0) { + result = Eigen::Matrix( + x.rows(), result.cols()); + } else { + check_size_match(function, "rows of result", result.rows(), + "rows of returned column", x.rows()); + } + result.col(j) = std::forward(x); +} + +} // namespace internal + +/** + * Apply a functor to each element of a std::vector and collect the results. + * + * For a vector `x` of length `n`, returns a std::vector whose `i`-th element + * is `f(x[i], args...)`. Additional arguments are shared across calls, which + * allows passing parameters without closures. If `x` is empty, returns an + * empty std::vector. + * + * The element type of the returned std::vector is deduced from the return type + * of `f` to allow for cases where the input and return scalar types differ + * (e.g., functions implicitly promoting integers). + * + * @tparam F Type of functor to apply. + * @tparam T Type of input std::vector. + * @tparam Args Types of shared arguments passed to each call. + * @param f functor to apply to each element. + * @param x std::vector input to which operation is applied. + * @param args shared arguments passed to each call. + * @return std::vector with result of applying functor to each element of `x`. + */ +template * = nullptr> +inline auto map(F&& f, T&& x, Args&&... args) { + using T_return = std::decay_t; + std::vector result; + result.reserve(x.size()); + + for (auto&& xi : x) { + result.push_back(f(xi, args...)); + } + + return result; +} + +/** + * Apply a functor elementwise to N std::vector inputs of the same size. + * + * For vectors `args...` of length `n`, returns a std::vector whose `i`-th + * element is `f(args[0][i], args[1][i], ...)`. If the inputs are empty, + * returns an empty std::vector. + * + * Unlike `map`, this overload only zips containers and does not accept + * additional shared trailing arguments. Pass shared state through the + * functor (e.g. a capturing lambda) when needed. + * + * The element type of the returned std::vector is deduced from the return type + * of `f` to allow for cases where the input and return scalar types differ + * (e.g., functions implicitly promoting integers). + * + * @tparam F Type of functor to apply. + * @tparam Types Types of input std::vector containers. + * @param f functor to apply to each tuple of elements. + * @param args std::vector inputs to which operation is applied. + * @return std::vector with result of applying functor to each tuple of + * elements. + * @throw std::invalid_argument if the inputs have different sizes. + */ +template * = nullptr> +inline auto mapN(F&& f, Types&&... args) { + static_assert(sizeof...(Types) >= 2, + "mapN requires at least two std::vector inputs."); + static constexpr const char* function = "mapN"; + internal::check_all_matching_sizes(function, args...); + + const std::size_t n = std::get<0>(std::forward_as_tuple(args...)).size(); + using T_return = std::decay_t; + std::vector result; + result.reserve(n); + + for (std::size_t i = 0; i < n; ++i) { + result.push_back(f((args[i])...)); + } + + return result; +} + +/** + * Apply a functor to each row of an Eigen matrix and collect the results. + * + * For a matrix `m` with `n` rows, returns a matrix whose `i`-th row is + * `f(m.row(i), args...)`. Additional arguments are shared across calls. If + * `m` has zero rows, returns a 0x0 matrix (output column count is unknown + * when the functor is never called). + * + * Expensive Eigen expressions passed as `m` are evaluated once via `to_ref` + * before the row loop, so compound expressions (e.g. `A * B`) are not + * re-evaluated for every row. + * + * The functor must return a type assignable to a matrix row. All returned + * rows must have the same number of columns. + * + * @tparam F Type of functor to apply. + * @tparam T Eigen matrix type. + * @tparam Args Types of shared arguments passed to each call. + * @param f functor to apply to each row. + * @param m Eigen matrix input to which operation is applied. + * @param args shared arguments passed to each call. + * @return Eigen matrix with result of applying functor to each row of `m`. + * @throw std::invalid_argument if returned rows have inconsistent sizes. + */ +template * = nullptr> +inline auto row_map(F&& f, T&& m, Args&&... args) { + static constexpr const char* function = "row_map"; + decltype(auto) m_ref = to_ref(std::forward(m)); + using T_return + = scalar_type_t>; + using matrix_t = Eigen::Matrix; + + const Eigen::Index n_rows = m_ref.rows(); + if (n_rows == 0) { + return matrix_t(0, 0); + } + + matrix_t result(n_rows, 0); + for (Eigen::Index i = 0; i < n_rows; ++i) { + internal::assign_matrix_row(function, result, i, f(m_ref.row(i), args...)); + } + return result; +} + +/** + * Apply a functor to each column of an Eigen matrix and collect the results. + * + * For a matrix `m` with `k` columns, returns a matrix whose `j`-th column is + * `f(m.col(j), args...)`. Additional arguments are shared across calls. If + * `m` has zero columns, returns a 0x0 matrix (output row count is unknown + * when the functor is never called). + * + * Expensive Eigen expressions passed as `m` are evaluated once via `to_ref` + * before the column loop, so compound expressions (e.g. `A * B`) are not + * re-evaluated for every column. + * + * The functor must return a type assignable to a matrix column. All returned + * columns must have the same number of rows. + * + * @tparam F Type of functor to apply. + * @tparam T Eigen matrix type. + * @tparam Args Types of shared arguments passed to each call. + * @param f functor to apply to each column. + * @param m Eigen matrix input to which operation is applied. + * @param args shared arguments passed to each call. + * @return Eigen matrix with result of applying functor to each column of `m`. + * @throw std::invalid_argument if returned columns have inconsistent sizes. + */ +template * = nullptr> +inline auto col_map(F&& f, T&& m, Args&&... args) { + static constexpr const char* function = "col_map"; + decltype(auto) m_ref = to_ref(std::forward(m)); + using T_return + = scalar_type_t>; + using matrix_t = Eigen::Matrix; + + const Eigen::Index n_cols = m_ref.cols(); + if (n_cols == 0) { + return matrix_t(0, 0); + } + + matrix_t result(0, n_cols); + for (Eigen::Index j = 0; j < n_cols; ++j) { + internal::assign_matrix_col(function, result, j, f(m_ref.col(j), args...)); + } + return result; +} + +/** + * Apply a functor rowwise to N Eigen matrices with identical dimensions. + * + * For matrices `args...` with `n` rows, returns a matrix whose `i`-th row is + * `f(args[0].row(i), args[1].row(i), ...)`. If the inputs have zero rows, + * returns a 0x0 matrix. + * + * Unlike `row_map`, this overload only zips matrices and does not accept + * additional shared trailing arguments. + * + * Expensive Eigen expressions in `args...` are each evaluated once via + * `to_ref` before the row loop. + * + * The functor must return a type assignable to a matrix row. All returned + * rows must have the same number of columns. + * + * @tparam F Type of functor to apply. + * @tparam Types Eigen matrix types. + * @param f functor to apply to each tuple of rows. + * @param args Eigen matrix inputs to which operation is applied. + * @return Eigen matrix with result of applying functor to each tuple of rows. + * @throw std::invalid_argument if the inputs have different dimensions or + * returned rows have inconsistent sizes. + */ +template * = nullptr> +inline auto row_mapN(F&& f, Types&&... args) { + static_assert(sizeof...(Types) >= 2, + "row_mapN requires at least two Eigen matrix inputs."); + static constexpr const char* function = "row_mapN"; + internal::check_all_matching_dims(function, args...); + + std::tuple...> m_refs{ + to_ref(std::forward(args))...}; + const Eigen::Index n_rows = std::get<0>(m_refs).rows(); + using T_return = scalar_type_t>().row(0))...))>>; + using matrix_t = Eigen::Matrix; + + if (n_rows == 0) { + return matrix_t(0, 0); + } + + matrix_t result(n_rows, 0); + apply( + [&](auto&&... ms) { + for (Eigen::Index i = 0; i < n_rows; ++i) { + internal::assign_matrix_row(function, result, i, f((ms.row(i))...)); + } + }, + m_refs); + return result; +} + +/** + * Apply a functor columnwise to N Eigen matrices with identical dimensions. + * + * For matrices `args...` with `k` columns, returns a matrix whose `j`-th + * column is `f(args[0].col(j), args[1].col(j), ...)`. If the inputs have + * zero columns, returns a 0x0 matrix. + * + * Unlike `col_map`, this overload only zips matrices and does not accept + * additional shared trailing arguments. + * + * Expensive Eigen expressions in `args...` are each evaluated once via + * `to_ref` before the column loop. + * + * The functor must return a type assignable to a matrix column. All returned + * columns must have the same number of rows. + * + * @tparam F Type of functor to apply. + * @tparam Types Eigen matrix types. + * @param f functor to apply to each tuple of columns. + * @param args Eigen matrix inputs to which operation is applied. + * @return Eigen matrix with result of applying functor to each tuple of + * columns. + * @throw std::invalid_argument if the inputs have different dimensions or + * returned columns have inconsistent sizes. + */ +template * = nullptr> +inline auto col_mapN(F&& f, Types&&... args) { + static_assert(sizeof...(Types) >= 2, + "col_mapN requires at least two Eigen matrix inputs."); + static constexpr const char* function = "col_mapN"; + internal::check_all_matching_dims(function, args...); + + std::tuple...> m_refs{ + to_ref(std::forward(args))...}; + const Eigen::Index n_cols = std::get<0>(m_refs).cols(); + using T_return = scalar_type_t>().col(0))...))>>; + using matrix_t = Eigen::Matrix; + + if (n_cols == 0) { + return matrix_t(0, 0); + } + + matrix_t result(0, n_cols); + apply( + [&](auto&&... ms) { + for (Eigen::Index j = 0; j < n_cols; ++j) { + internal::assign_matrix_col(function, result, j, f((ms.col(j))...)); + } + }, + m_refs); + return result; +} + +} // namespace math +} // namespace stan + +#endif diff --git a/test/unit/math/prim/functor/map_test.cpp b/test/unit/math/prim/functor/map_test.cpp new file mode 100644 index 00000000000..e11d86148e0 --- /dev/null +++ b/test/unit/math/prim/functor/map_test.cpp @@ -0,0 +1,389 @@ +#include +#include +#include +#include +#include +#include + +TEST(MathFunctor, map_square) { + std::vector x{1.0, 2.0, 3.0}; + auto y = stan::math::map([](double v) { return v * v; }, x); + ASSERT_EQ(y.size(), 3u); + EXPECT_FLOAT_EQ(y[0], 1.0); + EXPECT_FLOAT_EQ(y[1], 4.0); + EXPECT_FLOAT_EQ(y[2], 9.0); +} + +TEST(MathFunctor, map_type_change) { + std::vector x{1, 2, 3}; + auto y = stan::math::map([](int v) { return v + 0.5; }, x); + EXPECT_FLOAT_EQ(y[0], 1.5); + EXPECT_FLOAT_EQ(y[1], 2.5); + EXPECT_FLOAT_EQ(y[2], 3.5); +} + +TEST(MathFunctor, map_empty) { + std::vector x; + auto y = stan::math::map([](double v) { return v; }, x); + EXPECT_EQ(y.size(), 0u); +} + +TEST(MathFunctor, map_const_vector) { + const std::vector x{1.0, 2.0, 3.0}; + auto y = stan::math::map([](double v) { return 2.0 * v; }, x); + ASSERT_EQ(y.size(), 3u); + EXPECT_FLOAT_EQ(y[0], 2.0); + EXPECT_FLOAT_EQ(y[1], 4.0); + EXPECT_FLOAT_EQ(y[2], 6.0); +} + +TEST(MathFunctor, map_variadic_args) { + std::vector x{1.0, 2.0, 3.0}; + double offset = 10.0; + double scale = 2.0; + auto y = stan::math::map( + [](double v, double mu, double sigma) { return mu + sigma * v; }, x, + offset, scale); + ASSERT_EQ(y.size(), 3u); + EXPECT_FLOAT_EQ(y[0], 12.0); + EXPECT_FLOAT_EQ(y[1], 14.0); + EXPECT_FLOAT_EQ(y[2], 16.0); +} + +TEST(MathFunctor, map_shared_arg_reused) { + std::vector x{1.0, 2.0, 3.0}; + auto offset = std::make_unique(10.0); + auto y = stan::math::map( + [](double v, const std::unique_ptr& mu) { return v + *mu; }, x, + std::move(offset)); + ASSERT_EQ(y.size(), 3u); + EXPECT_FLOAT_EQ(y[0], 11.0); + EXPECT_FLOAT_EQ(y[1], 12.0); + EXPECT_FLOAT_EQ(y[2], 13.0); +} + +TEST(MathFunctor, mapN_add) { + std::vector a{1.0, 2.0}; + std::vector b{10.0, 20.0}; + auto y = stan::math::mapN([](double u, double v) { return u + v; }, a, b); + ASSERT_EQ(y.size(), 2u); + EXPECT_FLOAT_EQ(y[0], 11.0); + EXPECT_FLOAT_EQ(y[1], 22.0); +} + +TEST(MathFunctor, mapN_mixed_scalar_types) { + std::vector a{1, 2, 3}; + std::vector b{0.5, 1.5, 2.5}; + auto y = stan::math::mapN([](int u, double v) { return u + v; }, a, b); + ASSERT_EQ(y.size(), 3u); + EXPECT_FLOAT_EQ(y[0], 1.5); + EXPECT_FLOAT_EQ(y[1], 3.5); + EXPECT_FLOAT_EQ(y[2], 5.5); +} + +TEST(MathFunctor, mapN_three_vectors) { + std::vector a{1.0, 2.0}; + std::vector b{10.0, 20.0}; + std::vector c{100.0, 200.0}; + auto y = stan::math::mapN( + [](double u, double v, double w) { return u + v + w; }, a, b, c); + ASSERT_EQ(y.size(), 2u); + EXPECT_FLOAT_EQ(y[0], 111.0); + EXPECT_FLOAT_EQ(y[1], 222.0); +} + +TEST(MathFunctor, mapN_empty) { + std::vector a; + std::vector b; + auto y = stan::math::mapN([](double u, double v) { return u + v; }, a, b); + EXPECT_EQ(y.size(), 0u); +} + +TEST(MathFunctor, mapN_size_mismatch) { + std::vector a{1.0, 2.0}; + std::vector b{10.0}; + EXPECT_THROW(stan::math::mapN([](double u, double v) { return u + v; }, a, b), + std::invalid_argument); +} + +TEST(MathFunctor, row_map_standardize) { + Eigen::MatrixXd x(2, 3); + x << 1.0, 2.0, 3.0, 4.0, 5.0, 6.0; + Eigen::RowVectorXd mu(3); + mu << 1.0, 2.0, 3.0; + Eigen::RowVectorXd sigma(3); + sigma << 1.0, 2.0, 3.0; + + auto y = stan::math::row_map( + [](const Eigen::RowVectorXd& row, const Eigen::RowVectorXd& m, + const Eigen::RowVectorXd& s) { + return (row.array() - m.array()) / s.array(); + }, + x, mu, sigma); + + Eigen::MatrixXd expected(2, 3); + expected << 0.0, 0.0, 0.0, 3.0, 1.5, 1.0; + EXPECT_MATRIX_FLOAT_EQ(y, expected); +} + +TEST(MathFunctor, row_map_empty) { + Eigen::MatrixXd x(0, 3); + auto y = stan::math::row_map( + [](const Eigen::RowVectorXd& row) { return row; }, x); + EXPECT_EQ(y.rows(), 0); + EXPECT_EQ(y.cols(), 0); +} + +TEST(MathFunctor, row_map_inconsistent_row_size) { + Eigen::MatrixXd x(2, 2); + x << 1.0, 2.0, 3.0, 4.0; + int call = 0; + EXPECT_THROW(stan::math::row_map( + [&call](const Eigen::RowVectorXd&) { + ++call; + if (call == 1) { + return Eigen::RowVectorXd::Ones(2); + } + return Eigen::RowVectorXd::Ones(3); + }, + x), + std::invalid_argument); +} + +TEST(MathFunctor, row_map_shared_arg_reused) { + Eigen::MatrixXd x(2, 2); + x << 1.0, 2.0, 3.0, 4.0; + auto offset = std::make_unique(10.0); + auto y = stan::math::row_map( + [](const Eigen::RowVectorXd& row, const std::unique_ptr& mu) { + return row.array() + *mu; + }, + x, std::move(offset)); + + Eigen::MatrixXd expected(2, 2); + expected << 11.0, 12.0, 13.0, 14.0; + EXPECT_MATRIX_FLOAT_EQ(y, expected); +} + +TEST(MathFunctor, col_map_sum_rows) { + Eigen::MatrixXd x(2, 2); + x << 1.0, 2.0, 3.0, 4.0; + + auto y = stan::math::col_map( + [](const Eigen::VectorXd& col) { + return Eigen::VectorXd::Constant(1, col.sum()); + }, + x); + + Eigen::MatrixXd expected(1, 2); + expected << 4.0, 6.0; + EXPECT_MATRIX_FLOAT_EQ(y, expected); +} + +TEST(MathFunctor, col_map_empty) { + Eigen::MatrixXd x(3, 0); + auto y + = stan::math::col_map([](const Eigen::VectorXd& col) { return col; }, x); + EXPECT_EQ(y.rows(), 0); + EXPECT_EQ(y.cols(), 0); +} + +TEST(MathFunctor, col_map_inconsistent_col_size) { + Eigen::MatrixXd x(2, 2); + x << 1.0, 2.0, 3.0, 4.0; + int call = 0; + EXPECT_THROW(stan::math::col_map( + [&call](const Eigen::VectorXd&) { + ++call; + if (call == 1) { + return Eigen::VectorXd::Ones(2); + } + return Eigen::VectorXd::Ones(3); + }, + x), + std::invalid_argument); +} + +TEST(MathFunctor, col_map_shared_arg_reused) { + Eigen::MatrixXd x(2, 2); + x << 1.0, 2.0, 3.0, 4.0; + auto scale = std::make_unique(2.0); + auto y = stan::math::col_map( + [](const Eigen::VectorXd& col, const std::unique_ptr& s) { + return (*s) * col; + }, + x, std::move(scale)); + + Eigen::MatrixXd expected(2, 2); + expected << 2.0, 4.0, 6.0, 8.0; + EXPECT_MATRIX_FLOAT_EQ(y, expected); +} + +TEST(MathFunctor, row_mapN_add) { + Eigen::MatrixXd a(2, 2); + a << 1.0, 2.0, 3.0, 4.0; + Eigen::MatrixXd b(2, 2); + b << 10.0, 20.0, 30.0, 40.0; + + auto y + = stan::math::row_mapN([](const Eigen::RowVectorXd& u, + const Eigen::RowVectorXd& v) { return u + v; }, + a, b); + + Eigen::MatrixXd expected(2, 2); + expected << 11.0, 22.0, 33.0, 44.0; + EXPECT_MATRIX_FLOAT_EQ(y, expected); +} + +TEST(MathFunctor, row_mapN_dim_mismatch) { + Eigen::MatrixXd a(2, 2); + a << 1.0, 2.0, 3.0, 4.0; + Eigen::MatrixXd b(2, 3); + b << 1.0, 2.0, 3.0, 4.0, 5.0, 6.0; + + EXPECT_THROW( + stan::math::row_mapN([](const Eigen::RowVectorXd& u, + const Eigen::RowVectorXd& v) { return u + v; }, + a, b), + std::invalid_argument); +} + +TEST(MathFunctor, row_mapN_empty) { + Eigen::MatrixXd a(0, 2); + Eigen::MatrixXd b(0, 2); + auto y + = stan::math::row_mapN([](const Eigen::RowVectorXd& u, + const Eigen::RowVectorXd& v) { return u + v; }, + a, b); + EXPECT_EQ(y.rows(), 0); + EXPECT_EQ(y.cols(), 0); +} + +TEST(MathFunctor, row_mapN_mixed_scalar_types) { + Eigen::MatrixXi a(2, 2); + a << 1, 2, 3, 4; + Eigen::MatrixXd b(2, 2); + b << 0.5, 1.5, 2.5, 3.5; + auto y = stan::math::row_mapN( + [](const auto& u, const auto& v) { + return u.template cast() + v; + }, + a, b); + + Eigen::MatrixXd expected(2, 2); + expected << 1.5, 3.5, 5.5, 7.5; + EXPECT_MATRIX_FLOAT_EQ(y, expected); +} + +TEST(MathFunctor, row_mapN_inconsistent_row_size) { + Eigen::MatrixXd a(2, 2); + a << 1.0, 2.0, 3.0, 4.0; + Eigen::MatrixXd b(2, 2); + b << 10.0, 20.0, 30.0, 40.0; + int call = 0; + EXPECT_THROW( + stan::math::row_mapN( + [&call](const Eigen::RowVectorXd&, const Eigen::RowVectorXd&) { + ++call; + if (call == 1) { + return Eigen::RowVectorXd::Ones(2); + } + return Eigen::RowVectorXd::Ones(3); + }, + a, b), + std::invalid_argument); +} + +TEST(MathFunctor, col_mapN_add) { + Eigen::MatrixXd a(2, 2); + a << 1.0, 2.0, 3.0, 4.0; + Eigen::MatrixXd b(2, 2); + b << 10.0, 20.0, 30.0, 40.0; + + auto y = stan::math::col_mapN( + [](const Eigen::VectorXd& u, const Eigen::VectorXd& v) { return u + v; }, + a, b); + + Eigen::MatrixXd expected(2, 2); + expected << 11.0, 22.0, 33.0, 44.0; + EXPECT_MATRIX_FLOAT_EQ(y, expected); +} + +TEST(MathFunctor, col_mapN_dim_mismatch) { + Eigen::MatrixXd a(2, 2); + a << 1.0, 2.0, 3.0, 4.0; + Eigen::MatrixXd b(3, 2); + b << 1.0, 2.0, 3.0, 4.0, 5.0, 6.0; + + EXPECT_THROW( + stan::math::col_mapN([](const Eigen::VectorXd& u, + const Eigen::VectorXd& v) { return u + v; }, + a, b), + std::invalid_argument); +} + +TEST(MathFunctor, col_mapN_empty) { + Eigen::MatrixXd a(2, 0); + Eigen::MatrixXd b(2, 0); + auto y = stan::math::col_mapN( + [](const Eigen::VectorXd& u, const Eigen::VectorXd& v) { return u + v; }, + a, b); + EXPECT_EQ(y.rows(), 0); + EXPECT_EQ(y.cols(), 0); +} + +TEST(MathFunctor, col_mapN_mixed_scalar_types) { + Eigen::MatrixXi a(2, 2); + a << 1, 2, 3, 4; + Eigen::MatrixXd b(2, 2); + b << 0.5, 1.5, 2.5, 3.5; + auto y = stan::math::col_mapN( + [](const auto& u, const auto& v) { + return u.template cast() + v; + }, + a, b); + + Eigen::MatrixXd expected(2, 2); + expected << 1.5, 3.5, 5.5, 7.5; + EXPECT_MATRIX_FLOAT_EQ(y, expected); +} + +TEST(MathFunctor, col_mapN_inconsistent_col_size) { + Eigen::MatrixXd a(2, 2); + a << 1.0, 2.0, 3.0, 4.0; + Eigen::MatrixXd b(2, 2); + b << 10.0, 20.0, 30.0, 40.0; + int call = 0; + EXPECT_THROW(stan::math::col_mapN( + [&call](const Eigen::VectorXd&, const Eigen::VectorXd&) { + ++call; + if (call == 1) { + return Eigen::VectorXd::Ones(2); + } + return Eigen::VectorXd::Ones(3); + }, + a, b), + std::invalid_argument); +} + +TEST(MathFunctor, row_map_block_expression) { + Eigen::MatrixXd x(3, 3); + x << 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0; + auto y = stan::math::row_map([](const auto& row) { return 2.0 * row; }, + x.block(0, 0, 2, 3)); + + Eigen::MatrixXd expected(2, 3); + expected << 2.0, 4.0, 6.0, 8.0, 10.0, 12.0; + EXPECT_MATRIX_FLOAT_EQ(y, expected); +} + +TEST(MathFunctor, col_map_block_expression) { + Eigen::MatrixXd x(3, 3); + x << 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0; + auto y = stan::math::col_map([](const auto& col) { return 2.0 * col; }, + x.block(0, 0, 3, 2)); + + Eigen::MatrixXd expected(3, 2); + expected << 2.0, 4.0, 8.0, 10.0, 14.0, 16.0; + EXPECT_MATRIX_FLOAT_EQ(y, expected); +}