1 #ifndef STAN_MATH_PRIM_FUN_CSR_EXTRACT_U_HPP
2 #define STAN_MATH_PRIM_FUN_CSR_EXTRACT_U_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 NZE index for each entry from a sparse matrix.
19  *
20  * @tparam T type of elements in the matrix
21  * @param A Sparse matrix.
22  * @return Array of indexes into non-zero entries of A.
23  */
24 template <typename T>
csr_extract_u(const Eigen::SparseMatrix<T,Eigen::RowMajor> & A)25 const std::vector<int> csr_extract_u(
26     const Eigen::SparseMatrix<T, Eigen::RowMajor>& A) {
27   std::vector<int> u(A.outerSize() + 1);  // last entry is garbage.
28   for (int nze = 0; nze <= A.outerSize(); ++nze) {
29     u[nze] = *(A.outerIndexPtr() + nze) + stan::error_index::value;
30   }
31   return u;
32 }
33 
34 /**
35  * Extract the NZE index for each entry from a sparse matrix.
36  *
37  * @tparam T type of elements in the matrix
38  * @tparam R number of rows, can be Eigen::Dynamic
39  * @tparam C number of columns, can be Eigen::Dynamic
40  * @param A Dense matrix.
41  * @return Array of indexes into non-zero entries of A.
42  */
43 template <typename T, require_eigen_dense_base_t<T>* = nullptr>
csr_extract_u(const T & A)44 const std::vector<int> csr_extract_u(const T& A) {
45   // conversion to sparse seems to touch data twice, so we need to call to_ref
46   Eigen::SparseMatrix<scalar_type_t<T>, Eigen::RowMajor> B
47       = to_ref(A).sparseView();
48   std::vector<int> u = csr_extract_u(B);
49   return u;
50 }
51 
52 /** @} */  // end of csr_format group
53 
54 }  // namespace math
55 }  // namespace stan
56 
57 #endif
58