diff --git a/include/armadillo_bits/SpSubview_meat.hpp b/include/armadillo_bits/SpSubview_meat.hpp index 750f98ac..a31e7014 100644 --- a/include/armadillo_bits/SpSubview_meat.hpp +++ b/include/armadillo_bits/SpSubview_meat.hpp @@ -44,18 +44,28 @@ SpSubview::SpSubview(const SpMat& in_m, const uword in_row1, const uword m.sync_csc(); - // There must be a O(1) way to do this - uword lend = m.col_ptrs[in_col1 + in_n_cols]; - uword lend_row = in_row1 + in_n_rows; - uword count = 0; + // count the number of non-zeros in the subview + uword count = 0; - for(uword i = m.col_ptrs[in_col1]; i < lend; ++i) + if( (n_cols == 1) && (n_rows == m.n_rows) ) { - const uword m_row_indices_i = m.row_indices[i]; + count = m.col_ptrs[aux_col1 + 1] - m.col_ptrs[aux_col1]; + } + else + { + arma_debug_print("counting non-zeros in sparse subview"); - const bool condition = (m_row_indices_i >= in_row1) && (m_row_indices_i < lend_row); + uword lend = m.col_ptrs[in_col1 + in_n_cols]; + uword lend_row = in_row1 + in_n_rows; - count += condition ? uword(1) : uword(0); + for(uword i = m.col_ptrs[in_col1]; i < lend; ++i) + { + const uword m_row_indices_i = m.row_indices[i]; + + const bool condition = (m_row_indices_i >= in_row1) && (m_row_indices_i < lend_row); + + count += condition ? uword(1) : uword(0); + } } access::rw(n_nonzero) = count; diff --git a/include/armadillo_bits/fn_dot.hpp b/include/armadillo_bits/fn_dot.hpp index c2294787..f22cd45a 100644 --- a/include/armadillo_bits/fn_dot.hpp +++ b/include/armadillo_bits/fn_dot.hpp @@ -289,13 +289,31 @@ dot { arma_debug_sigprint(); + typedef typename T1::elem_type eT; + + if(is_SpSubview_col::value) + { + // TODO: refactor to use C++17 "if constexpr" to avoid reinterpret_cast shenanigans + + const SpSubview_col& yy = reinterpret_cast< const SpSubview_col& >(y); + + if(yy.n_rows == yy.m.n_rows) + { + arma_debug_print("using sparse column vector specialisation"); + + const quasi_unwrap U(x); + + arma_conform_assert_same_size(U.M.n_elem, uword(1), yy.n_elem, uword(1), "dot()"); + + return dense_sparse_helper::dot(U.M.memptr(), yy.m, yy.aux_col1); + } + } + const Proxy pa(x); const SpProxy pb(y); arma_conform_assert_same_size(pa.get_n_rows(), pa.get_n_cols(), pb.get_n_rows(), pb.get_n_cols(), "dot()"); - typedef typename T1::elem_type eT; - eT result = eT(0); typename SpProxy::const_iterator_type it = pb.begin();