Specialized reduction of sparse arrays

Hi All,

I have created a specialized version of reduction for sparse arrays, as the default falls back to generic arrays, which is too slow.

The code is pretty simple as

function Base.reduce(::typeof(hcat), A::AbstractVector{T}) where {T<: SparseMatrixCSC{Tv, Ti} where {Tv, Ti}}
  length(A) <= 1 && return(A)
  Tv, Ti = eltype(A[1].nzval), eltype(A[1].colptr)
  coloffset = 0
  Is, Js, Vs = Vector{Vector{Ti}}(undef, length(A)), Vector{Vector{Ti}}(undef, length(A)), Vector{Vector{Tv}}(undef, length(A))
  for (i, a) in enumerate(A)
    (I, J, V) = findnz(a, coloffset)
    Is[i] =  I
    Js[i] =  J
    Vs[i] =  V
    coloffset += a.n
  sparse(reduce(vcat, Is), reduce(vcat, Js), reduce(vcat, Vs))

where I have slightly modified the default findnz function

function SparseArrays.findnz(S::SparseMatrixCSC{Tv,Ti}, coloffset) where {Tv,Ti}
    numnz = nnz(S)
    I = Vector{Ti}(undef, numnz)
    J = Vector{Ti}(undef, numnz)
    V = Vector{Tv}(undef, numnz)

    count = 1
    @inbounds for col = 1 : S.n, k = S.colptr[col] : (S.colptr[col+1]-1)
        I[count] = S.rowval[k]
        J[count] = col + coloffset
        V[count] = S.nzval[k]
        count += 1

    return (I, J, V)

My question is, if it is of sufficient interest to be pushed to julia’s Sparse Arrays package. Also, if it possible to controbute only to SparseArrays and not to the whole Julia base.

Thanks for answer.

1 Like