1 #ifndef STAN_MATH_PRIM_FUN_CSR_EXTRACT_V_HPP 2 #define STAN_MATH_PRIM_FUN_CSR_EXTRACT_V_HPP 3 4 #include <stan/math/prim/meta.hpp> 5 #include <stan/math/prim/fun/Eigen.hpp> 6 #include <stan/math/prim/fun/to_ref.hpp> 7 #include <vector> 8 #include <numeric> 9 10 namespace stan { 11 namespace math { 12 13 /** \addtogroup csr_format 14 * @{ 15 */ 16 17 /** 18 * Extract the column indexes for non-zero value from a sparse 19 * matrix. 20 * 21 * @tparam T type of elements in the matrix 22 * @param A Sparse matrix. 23 * @return Array of column indexes for non-zero entries of A. 24 */ 25 template <typename T> csr_extract_v(const Eigen::SparseMatrix<T,Eigen::RowMajor> & A)26const std::vector<int> csr_extract_v( 27 const Eigen::SparseMatrix<T, Eigen::RowMajor>& A) { 28 std::vector<int> v(A.nonZeros()); 29 for (int nze = 0; nze < A.nonZeros(); ++nze) { 30 v[nze] = *(A.innerIndexPtr() + nze) + stan::error_index::value; 31 } 32 return v; 33 } 34 35 /** 36 * Extract the column indexes for non-zero values from a dense 37 * matrix by converting to sparse and calling the sparse matrix 38 * extractor. 39 * 40 * @tparam T type of elements in the matrix 41 * @tparam R number of rows, can be Eigen::Dynamic 42 * @tparam C number of columns, can be Eigen::Dynamic 43 * 44 * @param[in] A dense matrix. 45 * @return Array of column indexes to non-zero entries of A. 46 */ 47 template <typename T, require_eigen_dense_base_t<T>* = nullptr> csr_extract_v(const T & A)48const std::vector<int> csr_extract_v(const T& A) { 49 // conversion to sparse seems to touch data twice, so we need to call to_ref 50 Eigen::SparseMatrix<scalar_type_t<T>, Eigen::RowMajor> B 51 = to_ref(A).sparseView(); 52 std::vector<int> v = csr_extract_v(B); 53 return v; 54 } 55 56 /** @} */ // end of csr_format group 57 58 } // namespace math 59 } // namespace stan 60 61 #endif 62