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)25const 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)44const 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