From aa929c989303e702b6bee7fbbb536d43d729cdb4 Mon Sep 17 00:00:00 2001 From: Adam Haber Date: Sun, 18 Jul 2021 11:02:46 +0300 Subject: [PATCH 1/5] initial implementation --- stan/math/prim/prob/multi_normal_lpdf.hpp | 21 ++++++++++++++++++- stan/math/prim/prob/multi_normal_rng.hpp | 19 ++++++++++++++++- .../unit/math/prim/prob/multi_normal_test.cpp | 16 ++++++++++++++ 3 files changed, 54 insertions(+), 2 deletions(-) diff --git a/stan/math/prim/prob/multi_normal_lpdf.hpp b/stan/math/prim/prob/multi_normal_lpdf.hpp index 38f3f41da82..b055c7ac8c0 100644 --- a/stan/math/prim/prob/multi_normal_lpdf.hpp +++ b/stan/math/prim/prob/multi_normal_lpdf.hpp @@ -15,7 +15,8 @@ namespace stan { namespace math { -template +template * = nullptr> return_type_t multi_normal_lpdf(const T_y& y, const T_loc& mu, const T_covar& Sigma) { @@ -102,6 +103,24 @@ return_type_t multi_normal_lpdf(const T_y& y, return lp; } +template * = nullptr> +return_type_t multi_normal_lpdf(const T_y& y, + const T_loc& mu, + const T_covar& Sigma) { + using lp_type = return_type_t; + + const size_t N = y.rows(); + lp_type lp(0.0); + + for (size_t n = 0; n < N; ++n) { + auto current_y = y.row(n); + auto current_mu = mu.row(n); + lp += multi_normal_lpdf(current_y, current_mu, Sigma); + } + return lp; +} + template inline return_type_t multi_normal_lpdf( const T_y& y, const T_loc& mu, const T_covar& Sigma) { diff --git a/stan/math/prim/prob/multi_normal_rng.hpp b/stan/math/prim/prob/multi_normal_rng.hpp index 8f342306164..e1394192306 100644 --- a/stan/math/prim/prob/multi_normal_rng.hpp +++ b/stan/math/prim/prob/multi_normal_rng.hpp @@ -29,7 +29,7 @@ namespace math { * std::invalid_argument if the length of (each) mu is not equal to the number * of rows and columns in S */ -template +template * = nullptr> inline typename StdVectorBuilder::type multi_normal_rng(const T_loc& mu, const Eigen::Matrix& S, @@ -80,6 +80,23 @@ multi_normal_rng(const T_loc& mu, return output.data(); } +template * = nullptr> +inline typename Eigen::MatrixXd multi_normal_rng( + const T_loc& mu, + const Eigen::Matrix& S, RNG& rng) { + using Eigen::MatrixXd; + + const size_t N = mu.rows(); + const size_t M = mu.cols(); + + MatrixXd output(N, M); + for (int n = 0; n < N; ++n) { + auto current_mu = mu.row(n); + output.row(n) = multi_normal_rng(current_mu, S, rng); + } + return output; +} + } // namespace math } // namespace stan #endif diff --git a/test/unit/math/prim/prob/multi_normal_test.cpp b/test/unit/math/prim/prob/multi_normal_test.cpp index 4cf0f4bcf3d..d254be8fbcd 100644 --- a/test/unit/math/prim/prob/multi_normal_test.cpp +++ b/test/unit/math/prim/prob/multi_normal_test.cpp @@ -21,6 +21,22 @@ TEST(ProbDistributionsMultiNormal, NotVectorized) { EXPECT_NO_THROW(stan::math::multi_normal_rng(mu, Sigma, rng)); } +TEST(ProbDistributionsMultiNormal, MatrixNotVectorized) { + using Eigen::Dynamic; + using Eigen::Matrix; + using std::vector; + + boost::random::mt19937 rng; + Matrix y(2, 3); + y << 2.0, -2.0, 11.0, 2.0, -2.0, 11.0; + Matrix mu(2, 3); + mu << 1.0, -1.0, 3.0, 1.0, -1.0, 3.0; + Matrix Sigma(3, 3); + Sigma << 9.0, -3.0, 0.0, -3.0, 4.0, 0.0, 0.0, 0.0, 5.0; + EXPECT_FLOAT_EQ(-23.47816, stan::math::multi_normal_lpdf(y, mu, Sigma)); + EXPECT_NO_THROW(stan::math::multi_normal_rng(mu, Sigma, rng)); +} + TEST(ProbDistributionsMultiNormal, Vectorized) { using Eigen::Dynamic; using Eigen::Matrix; From 44b939b5c56c2ccab6ea3d9b61a96090b1909aca Mon Sep 17 00:00:00 2001 From: Steve Bronder Date: Tue, 20 Jul 2021 13:01:17 -0400 Subject: [PATCH 2/5] setup mvn to take a matrix but fails rn --- stan/math/prim/prob/multi_normal_lpdf.hpp | 50 ++++++++++++++++--- .../unit/math/prim/prob/multi_normal_test.cpp | 11 +++- 2 files changed, 53 insertions(+), 8 deletions(-) diff --git a/stan/math/prim/prob/multi_normal_lpdf.hpp b/stan/math/prim/prob/multi_normal_lpdf.hpp index b055c7ac8c0..5625769bf22 100644 --- a/stan/math/prim/prob/multi_normal_lpdf.hpp +++ b/stan/math/prim/prob/multi_normal_lpdf.hpp @@ -108,15 +108,53 @@ template multi_normal_lpdf(const T_y& y, const T_loc& mu, const T_covar& Sigma) { + using T_covar_elem = typename scalar_type::type; using lp_type = return_type_t; - - const size_t N = y.rows(); + using Eigen::Dynamic; + static const char* function = "multi_normal_lpdf"; + check_positive(function, "Covariance matrix rows", Sigma.rows()); + auto&& Sigma_ref = to_ref(Sigma); + const Eigen::Index number_of_y = y.rows(); + const Eigen::Index number_of_mu = mu.rows(); + if (number_of_y == 0 || number_of_mu == 0) { + return 0.0; + } + check_size_match(function, "Rows of random variable", y.rows(), + "rows of location parameter", mu.rows()); + check_size_match(function, "Columns of random variable", y.cols(), + "columns of location parameter", mu.cols()); + check_size_match(function, "Columns of random variable", y.cols(), + "rows of covariance parameter", Sigma.rows()); + if (unlikely(number_of_mu == 0 || number_of_y == 0)) { + return 0.0; + } lp_type lp(0.0); + auto&& y_vec = to_ref(y); + auto&& mu_vec = to_ref(mu); + const Eigen::Index size_y = y_vec.cols(); + const Eigen::Index size_mu = mu_vec.cols(); + + const size_t size_vec = size_y; + + check_finite(function, "Location parameter", mu_vec); + check_not_nan(function, "Random variable", y_vec); + check_symmetric(function, "Covariance matrix", Sigma_ref); + + auto ldlt_Sigma = make_ldlt_factor(Sigma_ref); + check_ldlt_factor(function, "LDLT_Factor of covariance parameter", + ldlt_Sigma); - for (size_t n = 0; n < N; ++n) { - auto current_y = y.row(n); - auto current_mu = mu.row(n); - lp += multi_normal_lpdf(current_y, current_mu, Sigma); + if (include_summand::value) { + lp += NEG_LOG_SQRT_TWO_PI * y.rows() * y.cols(); + } + + if (include_summand::value) { + lp -= 0.5 * log_determinant_ldlt(ldlt_Sigma) * y.cols(); + } + + if (include_summand::value) { + lp_type sum_lp_vec(0.0); + lp -= 0.5 * trace_inv_quad_form_ldlt(ldlt_Sigma, (y_vec - mu_vec).transpose()); } return lp; } diff --git a/test/unit/math/prim/prob/multi_normal_test.cpp b/test/unit/math/prim/prob/multi_normal_test.cpp index d254be8fbcd..7e450d28fde 100644 --- a/test/unit/math/prim/prob/multi_normal_test.cpp +++ b/test/unit/math/prim/prob/multi_normal_test.cpp @@ -24,7 +24,6 @@ TEST(ProbDistributionsMultiNormal, NotVectorized) { TEST(ProbDistributionsMultiNormal, MatrixNotVectorized) { using Eigen::Dynamic; using Eigen::Matrix; - using std::vector; boost::random::mt19937 rng; Matrix y(2, 3); @@ -33,7 +32,15 @@ TEST(ProbDistributionsMultiNormal, MatrixNotVectorized) { mu << 1.0, -1.0, 3.0, 1.0, -1.0, 3.0; Matrix Sigma(3, 3); Sigma << 9.0, -3.0, 0.0, -3.0, 4.0, 0.0, 0.0, 0.0, 5.0; - EXPECT_FLOAT_EQ(-23.47816, stan::math::multi_normal_lpdf(y, mu, Sigma)); + std::vector test_y; + std::vector test_mu; + test_y.push_back(y.row(0).transpose()); + test_y.push_back(y.row(1).transpose()); + test_mu.push_back(mu.row(0).transpose()); + test_mu.push_back(mu.row(1).transpose()); + auto test_eval = stan::math::multi_normal_lpdf(y, mu, Sigma); + auto test_orig = stan::math::multi_normal_lpdf(test_y, test_mu, Sigma); + EXPECT_FLOAT_EQ(test_orig, test_eval); EXPECT_NO_THROW(stan::math::multi_normal_rng(mu, Sigma, rng)); } From dc6034be3c1f834cfa889aa5fded42d5309caa67 Mon Sep 17 00:00:00 2001 From: Steve Bronder Date: Tue, 20 Jul 2021 13:20:53 -0400 Subject: [PATCH 3/5] fix so tests pass --- stan/math/prim/prob/multi_normal_lpdf.hpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/stan/math/prim/prob/multi_normal_lpdf.hpp b/stan/math/prim/prob/multi_normal_lpdf.hpp index 5625769bf22..ab862e2295b 100644 --- a/stan/math/prim/prob/multi_normal_lpdf.hpp +++ b/stan/math/prim/prob/multi_normal_lpdf.hpp @@ -149,7 +149,7 @@ return_type_t multi_normal_lpdf(const T_y& y, } if (include_summand::value) { - lp -= 0.5 * log_determinant_ldlt(ldlt_Sigma) * y.cols(); + lp -= 0.5 * log_determinant_ldlt(ldlt_Sigma) * y.rows(); } if (include_summand::value) { From cc3f2f5aff0a3fe429d9c05fdd4dc68b6c4620a1 Mon Sep 17 00:00:00 2001 From: Steve Bronder Date: Tue, 20 Jul 2021 13:25:12 -0400 Subject: [PATCH 4/5] cleanup a bit --- stan/math/prim/prob/multi_normal_lpdf.hpp | 25 ++++++++--------------- 1 file changed, 9 insertions(+), 16 deletions(-) diff --git a/stan/math/prim/prob/multi_normal_lpdf.hpp b/stan/math/prim/prob/multi_normal_lpdf.hpp index ab862e2295b..51f11f1a287 100644 --- a/stan/math/prim/prob/multi_normal_lpdf.hpp +++ b/stan/math/prim/prob/multi_normal_lpdf.hpp @@ -104,7 +104,7 @@ return_type_t multi_normal_lpdf(const T_y& y, } template * = nullptr> + require_all_dense_dynamic_t* = nullptr> return_type_t multi_normal_lpdf(const T_y& y, const T_loc& mu, const T_covar& Sigma) { @@ -116,7 +116,7 @@ return_type_t multi_normal_lpdf(const T_y& y, auto&& Sigma_ref = to_ref(Sigma); const Eigen::Index number_of_y = y.rows(); const Eigen::Index number_of_mu = mu.rows(); - if (number_of_y == 0 || number_of_mu == 0) { + if (unlikely(number_of_y == 0 || number_of_mu == 0)) { return 0.0; } check_size_match(function, "Rows of random variable", y.rows(), @@ -125,21 +125,14 @@ return_type_t multi_normal_lpdf(const T_y& y, "columns of location parameter", mu.cols()); check_size_match(function, "Columns of random variable", y.cols(), "rows of covariance parameter", Sigma.rows()); - if (unlikely(number_of_mu == 0 || number_of_y == 0)) { - return 0.0; - } lp_type lp(0.0); - auto&& y_vec = to_ref(y); - auto&& mu_vec = to_ref(mu); - const Eigen::Index size_y = y_vec.cols(); - const Eigen::Index size_mu = mu_vec.cols(); - - const size_t size_vec = size_y; - - check_finite(function, "Location parameter", mu_vec); - check_not_nan(function, "Random variable", y_vec); + auto&& y_ref = to_ref(y); + auto&& mu_ref = to_ref(mu); + const Eigen::Index size_y = y_ref.cols(); + const Eigen::Index size_mu = mu_ref.cols(); + check_finite(function, "Location parameter", mu_ref); + check_not_nan(function, "Random variable", y_ref); check_symmetric(function, "Covariance matrix", Sigma_ref); - auto ldlt_Sigma = make_ldlt_factor(Sigma_ref); check_ldlt_factor(function, "LDLT_Factor of covariance parameter", ldlt_Sigma); @@ -154,7 +147,7 @@ return_type_t multi_normal_lpdf(const T_y& y, if (include_summand::value) { lp_type sum_lp_vec(0.0); - lp -= 0.5 * trace_inv_quad_form_ldlt(ldlt_Sigma, (y_vec - mu_vec).transpose()); + lp -= 0.5 * trace_inv_quad_form_ldlt(ldlt_Sigma, (y_ref - mu_ref).transpose()); } return lp; } From 0a293e4e74c00d16fb8d8bff8a4b8bf9760a6bde Mon Sep 17 00:00:00 2001 From: Steve Bronder Date: Tue, 20 Jul 2021 13:32:54 -0400 Subject: [PATCH 5/5] make test in mix for multi normal with matrix inputs --- stan/math/prim/prob/multi_normal_lpdf.hpp | 5 +---- test/unit/math/mix/prob/multi_normal_test.cpp | 16 ++++++++++++++++ 2 files changed, 17 insertions(+), 4 deletions(-) diff --git a/stan/math/prim/prob/multi_normal_lpdf.hpp b/stan/math/prim/prob/multi_normal_lpdf.hpp index 51f11f1a287..ec04aa25a40 100644 --- a/stan/math/prim/prob/multi_normal_lpdf.hpp +++ b/stan/math/prim/prob/multi_normal_lpdf.hpp @@ -109,8 +109,6 @@ return_type_t multi_normal_lpdf(const T_y& y, const T_loc& mu, const T_covar& Sigma) { using T_covar_elem = typename scalar_type::type; - using lp_type = return_type_t; - using Eigen::Dynamic; static const char* function = "multi_normal_lpdf"; check_positive(function, "Covariance matrix rows", Sigma.rows()); auto&& Sigma_ref = to_ref(Sigma); @@ -125,7 +123,7 @@ return_type_t multi_normal_lpdf(const T_y& y, "columns of location parameter", mu.cols()); check_size_match(function, "Columns of random variable", y.cols(), "rows of covariance parameter", Sigma.rows()); - lp_type lp(0.0); + return_type_t lp(0.0); auto&& y_ref = to_ref(y); auto&& mu_ref = to_ref(mu); const Eigen::Index size_y = y_ref.cols(); @@ -146,7 +144,6 @@ return_type_t multi_normal_lpdf(const T_y& y, } if (include_summand::value) { - lp_type sum_lp_vec(0.0); lp -= 0.5 * trace_inv_quad_form_ldlt(ldlt_Sigma, (y_ref - mu_ref).transpose()); } return lp; diff --git a/test/unit/math/mix/prob/multi_normal_test.cpp b/test/unit/math/mix/prob/multi_normal_test.cpp index bb1fc347d70..ddbd6c37f90 100644 --- a/test/unit/math/mix/prob/multi_normal_test.cpp +++ b/test/unit/math/mix/prob/multi_normal_test.cpp @@ -47,6 +47,22 @@ TEST(ProbDistributionsMultiNormal, matvar) { stan::test::expect_ad_matvar(f, y1, mu1, Sigma00); } +TEST(ProbDistributionsMultiNormal, matrix) { + auto f = [](const auto& y, const auto& mu, const auto& sigma) { + auto sigma_sym = stan::math::multiply(0.5, sigma + sigma.transpose()); + return stan::math::multi_normal_lpdf(y, mu, sigma_sym); + }; + + Eigen::MatrixXd y(2, 3); + y << 2.0, -2.0, 11.0, 2.0, -2.0, 11.0; + Eigen::MatrixXd mu(2, 3); + mu << 1.0, -1.0, 3.0, 1.0, -1.0, 3.0; + Eigen::MatrixXd Sigma(3, 3); + Sigma << 9.0, -3.0, 0.0, -3.0, 4.0, 0.0, 0.0, 0.0, 5.0; + stan::test::expect_ad(f, y, mu, Sigma); + stan::test::expect_ad_matvar(f, y, mu, Sigma); +} + TEST(ProbDistributionsMultiNormal, fvar_var) { using Eigen::Dynamic; using Eigen::Matrix;