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)26 const 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)48 const 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