diff --git a/stan/math/mix/functor/laplace_marginal_density.hpp b/stan/math/mix/functor/laplace_marginal_density.hpp index cf505417cc4..e462e2476fd 100644 --- a/stan/math/mix/functor/laplace_marginal_density.hpp +++ b/stan/math/mix/functor/laplace_marginal_density.hpp @@ -133,6 +133,35 @@ inline constexpr void copy_compute_s2(Output&& output, Input&& input) { std::forward(output), std::forward(input)); } +/** + * Returns a `theta_size` x `theta_size` matrix that holds `entry(r, c)` in + * its diagonal blocks of size `hessian_block_size` and zero elsewhere. + * `laplace_likelihood::compute_s2()` reads only these entries of its matrix + * argument, so they are computed in O(n^2 b) instead of forming the full + * matrix with O(n^3) products. + * @tparam EntryFun Type of a functor `(Eigen::Index r, Eigen::Index c)` + * returning `double` + * @param theta_size Number of rows and columns + * @param hessian_block_size Size of the diagonal blocks + * @param entry Functor that returns the entry at row `r` and column `c` + */ +template +inline Eigen::MatrixXd block_diagonal_entries( + const Eigen::Index theta_size, const Eigen::Index hessian_block_size, + EntryFun&& entry) { + Eigen::MatrixXd A = Eigen::MatrixXd::Zero(theta_size, theta_size); + const Eigen::Index n_blocks = theta_size / hessian_block_size; + for (Eigen::Index j = 0; j < n_blocks; ++j) { + const Eigen::Index start = j * hessian_block_size; + for (Eigen::Index c = start; c < start + hessian_block_size; ++c) { + for (Eigen::Index r = start; r < start + hessian_block_size; ++r) { + A(r, c) = entry(r, c); + } + } + } + return A; +} + } // namespace internal /** @@ -218,18 +247,23 @@ inline auto laplace_marginal_density(LLFun&& ll_fun, LLTupleArgs&& ll_args, = md_est.L.template triangularView().solve( md_est.W_r * covariance_val); if constexpr (ll_args_contain_var) { - arena_t A = covariance_val - C.transpose() * C; + // only the diagonal blocks of A = K - C' C, see compute_s2() + arena_t A = internal::block_diagonal_entries( + covariance_val.rows(), options.hessian_block_size, + [&](Eigen::Index r, Eigen::Index c) { + return covariance_val(r, c) - C.col(r).dot(C.col(c)); + }); auto s2_tmp = laplace_likelihood::compute_s2( ll_fun, md_est.theta, A, options.hessian_block_size, ll_args_copy, msgs); s2.deep_copy(s2_tmp); internal::copy_compute_s2(partial_parm, ll_args_filter); } else { - s2.deep_copy( - (0.5 - * (covariance_val.diagonal() - (C.transpose() * C).diagonal()) - .cwiseProduct(laplace_likelihood::third_diff( - ll_fun, md_est.theta, ll_args_vals, msgs)))); + s2.deep_copy((0.5 + * (covariance_val.diagonal() + - C.colwise().squaredNorm().transpose()) + .cwiseProduct(laplace_likelihood::third_diff( + ll_fun, md_est.theta, ll_args_vals, msgs)))); } } else { @@ -239,7 +273,12 @@ inline auto laplace_marginal_density(LLFun&& ll_fun, LLTupleArgs&& ll_args, arena_t C = md_est.L.template triangularView().solve( md_est.W_r * covariance_val); - arena_t A = covariance_val - C.transpose() * C; + // only the diagonal blocks of A = K - C' C, see compute_s2() + arena_t A = internal::block_diagonal_entries( + covariance_val.rows(), options.hessian_block_size, + [&](Eigen::Index r, Eigen::Index c) { + return covariance_val(r, c) - C.col(r).dot(C.col(c)); + }); auto s2_tmp = laplace_likelihood::compute_s2(ll_fun, md_est.theta, A, options.hessian_block_size, ll_args_copy, msgs); @@ -258,8 +297,14 @@ inline auto laplace_marginal_density(LLFun&& ll_fun, LLTupleArgs&& ll_args, arena_t C = md_est.L.template triangularView().solve( md_est.K_root.transpose()); + // only the diagonal blocks of C' C, see compute_s2() auto s2_tmp = laplace_likelihood::compute_s2( - ll_fun, md_est.theta, (C.transpose() * C).eval(), + ll_fun, md_est.theta, + internal::block_diagonal_entries(md_est.theta.size(), + options.hessian_block_size, + [&](Eigen::Index r, Eigen::Index c) { + return C.col(r).dot(C.col(c)); + }), options.hessian_block_size, ll_args_copy, msgs); s2.deep_copy(s2_tmp); internal::copy_compute_s2(partial_parm, ll_args_filter); @@ -269,8 +314,14 @@ inline auto laplace_marginal_density(LLFun&& ll_fun, LLTupleArgs&& ll_args, = Eigen::MatrixXd::Identity(md_est.W_r.rows(), md_est.W_r.cols()) - LU_solve_covariance * md_est.W_r; R = md_est.W_r * I_minus_BinvKW; // == W - W B^{-1} K W - arena_t A - = covariance_val - covariance_val * md_est.W_r * LU_solve_covariance; + // only the diagonal blocks of A = K - K W B^-1 K, see compute_s2() + const Eigen::MatrixXd KW = covariance_val * md_est.W_r; + arena_t A = internal::block_diagonal_entries( + covariance_val.rows(), options.hessian_block_size, + [&](Eigen::Index r, Eigen::Index c) { + return covariance_val(r, c) + - KW.row(r).dot(LU_solve_covariance.col(c)); + }); auto s2_tmp = laplace_likelihood::compute_s2(ll_fun, md_est.theta, A, options.hessian_block_size, ll_args_copy, msgs); @@ -296,7 +347,7 @@ inline auto laplace_marginal_density(LLFun&& ll_fun, LLTupleArgs&& ll_args, internal::collect_adjoints(partial_parm, ll_args_filter); arena_t v; if (md_est.solver_used == 1 || md_est.solver_used == 2) { - v = covariance_val * s2 - covariance_val * R * covariance_val * s2; + v = covariance_val * s2 - covariance_val * (R * (covariance_val * s2)); } else { v = LU_solve_covariance * s2; } diff --git a/test/unit/math/laplace/laplace_block_diagonal_entries_test.cpp b/test/unit/math/laplace/laplace_block_diagonal_entries_test.cpp new file mode 100644 index 00000000000..3c2cf40dd2b --- /dev/null +++ b/test/unit/math/laplace/laplace_block_diagonal_entries_test.cpp @@ -0,0 +1,72 @@ +#include +#include +#include +#include +#include + +namespace { + +// symmetric test matrix without zero entries +Eigen::MatrixXd dense_test_matrix(int n) { + Eigen::MatrixXd m(n, n); + for (int c = 0; c < n; ++c) { + for (int r = 0; r < n; ++r) { + m(r, c) = 1.0 / (1.0 + r + c) + (r == c ? 2.0 : 0.0); + } + } + return m; +} + +} // namespace + +TEST(laplace_block_diagonal_entries, keeps_only_the_diagonal_blocks) { + const int n = 6; + const Eigen::MatrixXd full = dense_test_matrix(n); + for (int block_size : {1, 2, 3}) { + const Eigen::MatrixXd blocks = stan::math::internal::block_diagonal_entries( + n, block_size, + [&](Eigen::Index r, Eigen::Index c) { return full(r, c); }); + ASSERT_EQ(n, blocks.rows()); + ASSERT_EQ(n, blocks.cols()); + for (int c = 0; c < n; ++c) { + for (int r = 0; r < n; ++r) { + if (r / block_size == c / block_size) { + EXPECT_EQ(full(r, c), blocks(r, c)); + } else { + EXPECT_EQ(0.0, blocks(r, c)); + } + } + } + } +} + +// laplace_marginal_density() passes only the diagonal blocks of A to +// compute_s2(), because compute_s2() reads no other entry of A: the result +// is the same for the full matrix and for its diagonal blocks. +TEST(laplace_block_diagonal_entries, + compute_s2_reads_only_the_diagonal_blocks) { + const int n = 6; + const Eigen::MatrixXd full = dense_test_matrix(n); + Eigen::VectorXd theta(n); + theta << 0.3, -0.2, 0.5, 1.0, -0.7, 0.1; + const std::vector y{1, 0, 2, 3, 0, 1, 2}; + const std::vector y_index{1, 2, 3, 4, 5, 6, 3}; + const Eigen::VectorXd mean = Eigen::VectorXd::Constant(n, 0.4); + std::ostream* msgs = nullptr; + for (int block_size : {1, 2, 3}) { + const Eigen::MatrixXd blocks = stan::math::internal::block_diagonal_entries( + n, block_size, + [&](Eigen::Index r, Eigen::Index c) { return full(r, c); }); + const Eigen::VectorXd s2_full = stan::math::laplace_likelihood::compute_s2( + stan::math::poisson_log_likelihood{}, theta, full, block_size, + std::forward_as_tuple(y, y_index, mean), msgs); + const Eigen::VectorXd s2_blocks + = stan::math::laplace_likelihood::compute_s2( + stan::math::poisson_log_likelihood{}, theta, blocks, block_size, + std::forward_as_tuple(y, y_index, mean), msgs); + ASSERT_EQ(n, s2_full.size()); + for (int i = 0; i < n; ++i) { + EXPECT_EQ(s2_full(i), s2_blocks(i)) << "block size " << block_size; + } + } +}