1 #ifndef STAN_MATH_PRIM_FUN_ADD_DIAG_HPP
2 #define STAN_MATH_PRIM_FUN_ADD_DIAG_HPP
3
4 #include <stan/math/prim/meta.hpp>
5 #include <stan/math/prim/err.hpp>
6 #include <stan/math/prim/fun/Eigen.hpp>
7 #include <stan/math/prim/fun/as_array_or_scalar.hpp>
8 #include <stan/math/prim/fun/as_column_vector_or_scalar.hpp>
9 #include <algorithm>
10
11 namespace stan {
12 namespace math {
13
14 /**
15 * Returns a Matrix with values added along the main diagonal
16 *
17 * @tparam T_m type of element in Eigen::Matrix
18 * @tparam T_a Type of element(s) to add along the diagonal
19 *
20 * @param mat a matrix
21 * @param to_add scalar value or column vector or row vector to add along the
22 * diagonal
23 * @return a matrix with to_add added along main diagonal
24 */
25 template <typename T_m, typename T_a, typename = require_eigen_t<T_m>,
26 typename = require_any_t<is_eigen_vector<T_a>, is_stan_scalar<T_a>>>
27 inline typename Eigen::Matrix<return_type_t<T_m, T_a>, Eigen::Dynamic,
28 Eigen::Dynamic>
add_diag(const T_m & mat,const T_a & to_add)29 add_diag(const T_m &mat, const T_a &to_add) {
30 if (is_vector<T_a>::value) {
31 const size_t length_diag = std::min(mat.rows(), mat.cols());
32 check_consistent_size("add_diag", "number of elements of to_add", to_add,
33 length_diag);
34 }
35 Eigen::Matrix<return_type_t<T_m, T_a>, Eigen::Dynamic, Eigen::Dynamic> out
36 = mat;
37 out.diagonal().array()
38 += as_array_or_scalar(as_column_vector_or_scalar(to_add));
39 return out;
40 }
41
42 } // namespace math
43 } // namespace stan
44
45 #endif
46