diff --git a/src/SparseArrays.jl b/src/SparseArrays.jl index 2688a449..aec895ae 100644 --- a/src/SparseArrays.jl +++ b/src/SparseArrays.jl @@ -32,7 +32,7 @@ export AbstractSparseArray, AbstractSparseMatrix, AbstractSparseVector, SparseMatrixCSC, SparseVector, blockdiag, droptol!, dropzeros!, dropzeros, issparse, nonzeros, nzrange, rowvals, sparse, sparsevec, spdiagm, sprand, sprandn, spzeros, nnz, permute, findnz, fkeep!, ftranspose!, - sparse_hcat, sparse_vcat, sparse_hvcat + sparse_hcat, sparse_vcat, sparse_hvcat, colvals const LinAlgLeftQs = Union{HessenbergQ,QRCompactWYQ,QRPackedQ} diff --git a/src/abstractsparse.jl b/src/abstractsparse.jl index 37f79e25..79210bbb 100644 --- a/src/abstractsparse.jl +++ b/src/abstractsparse.jl @@ -41,6 +41,14 @@ Supertype for matrix with compressed sparse column (CSC). """ abstract type AbstractSparseMatrixCSC{Tv,Ti<:Integer} <: AbstractSparseMatrix{Tv,Ti} end +""" + AbstractSparseMatrixCSR{Tv,Ti<:Integer} <: AbstractSparseMatrix{Tv,Ti} + +Supertype for matrix with compressed sparse row (CSR). +""" +abstract type AbstractSparseMatrixCSR{Tv,Ti<:Integer} <: AbstractSparseMatrix{Tv,Ti} end + +const AbstractSparseMatrixCSCOrCSR{Tv,Ti} = Union{AbstractSparseMatrixCSR{Tv,Ti}, AbstractSparseMatrixCSC{Tv,Ti}} """ issparse(S) diff --git a/src/linalg.jl b/src/linalg.jl index 131a21bc..7d7c2773 100644 --- a/src/linalg.jl +++ b/src/linalg.jl @@ -12,6 +12,7 @@ const DenseTriangular = UpperOrLowerTriangular{<:Any,<:DenseMatrixUnion} const DenseInputVector = Union{StridedVector, BitVector} const DenseVecOrMat = Union{DenseMatrixUnion, DenseInputVector} +# CSC matprod_dest(A::SparseMatrixCSCUnion2, B::DenseTriangular, TS) = similar(B, TS, (size(A, 1), size(B, 2))) matprod_dest(A::AdjOrTrans{<:Any,<:SparseMatrixCSCUnion2}, B::DenseTriangular, TS) = @@ -29,6 +30,24 @@ matprod_dest(A::Union{BitMatrix,AdjOrTrans{<:Any,BitMatrix}}, B::AdjOrTrans{<:An matprod_dest(A::DenseTriangular, B::AdjOrTrans{<:Any,<:SparseMatrixCSCUnion2}, TS) = similar(A, TS, (size(A, 1), size(B, 2))) +# CSR +matprod_dest(A::SparseMatrixCSRUnion2, B::DenseTriangular, TS) = + similar(B, TS, (size(A, 1), size(B, 2))) +matprod_dest(A::AdjOrTrans{<:Any,<:SparseMatrixCSRUnion2}, B::DenseTriangular, TS) = + similar(B, TS, (size(A, 1), size(B, 2))) +matprod_dest(A::StridedMaybeAdjOrTransMat, B::SparseMatrixCSRUnion2, TS) = + similar(A, TS, (size(A, 1), size(B, 2))) +matprod_dest(A::Union{BitMatrix,AdjOrTrans{<:Any,BitMatrix}}, B::SparseMatrixCSRUnion2, TS) = + similar(A, TS, (size(A, 1), size(B, 2))) +matprod_dest(A::DenseTriangular, B::SparseMatrixCSRUnion2, TS) = + similar(A, TS, (size(A, 1), size(B, 2))) +matprod_dest(A::StridedMaybeAdjOrTransMat, B::AdjOrTrans{<:Any,<:SparseMatrixCSRUnion2}, TS) = + similar(A, TS, (size(A, 1), size(B, 2))) +matprod_dest(A::Union{BitMatrix,AdjOrTrans{<:Any,BitMatrix}}, B::AdjOrTrans{<:Any,<:SparseMatrixCSRUnion2}, TS) = + similar(A, TS, (size(A, 1), size(B, 2))) +matprod_dest(A::DenseTriangular, B::AdjOrTrans{<:Any,<:SparseMatrixCSRUnion2}, TS) = + similar(A, TS, (size(A, 1), size(B, 2))) + for op ∈ (:+, :-), Wrapper ∈ (:Hermitian, :Symmetric) @eval begin $op(A::AbstractSparseMatrix, B::$Wrapper{<:Any,<:AbstractSparseMatrix}) = $op(A, sparse(B)) @@ -54,6 +73,13 @@ generic_matmatmul!(C::StridedMatrix, tA, tB, A::SparseMatrixCSCUnion2, B::Abstra generic_matvecmul!(C::StridedVecOrMat, tA, A::SparseMatrixCSCUnion2, B::DenseInputVector, alpha::Number, beta::Number) = spdensemul!(C, tA, 'N', A, B, alpha, beta) +generic_matmatmul!(C::StridedMatrix, tA, tB, A::SparseMatrixCSRUnion2, B::DenseMatrixUnion, alpha::Number, beta::Number) = + spdensemul!(C, tA, tB, A, B, alpha, beta) +generic_matmatmul!(C::StridedMatrix, tA, tB, A::SparseMatrixCSRUnion2, B::AbstractTriangular, alpha::Number, beta::Number) = + spdensemul!(C, tA, tB, A, B, alpha, beta) +generic_matvecmul!(C::StridedVecOrMat, tA, A::SparseMatrixCSRUnion2, B::DenseInputVector, alpha::Number, beta::Number) = + spdensemul!(C, tA, 'N', A, B, alpha, beta) + Base.@constprop :aggressive function spdensemul!(C, tA, tB, A, B, alpha, beta) tA_uc, tB_uc = uppercase(tA), uppercase(tB) if tA_uc == 'N' @@ -74,7 +100,7 @@ Base.@constprop :aggressive function spdensemul!(C, tA, tB, A, B, alpha, beta) return C end -function _spmatmul!(C, A, B, α, β) +function _spmatmul!(C, A::AbstractSparseMatrixCSC, B, α, β) size(A, 2) == size(B, 1) || throw(DimensionMismatch("second dimension of A, $(size(A,2)), does not match the first dimension of B, $(size(B,1))")) size(A, 1) == size(C, 1) || @@ -95,7 +121,28 @@ function _spmatmul!(C, A, B, α, β) C end -function _At_or_Ac_mul_B!(tfun::Function, C, A, B, α, β) +function _spmatmul!(C, A::AbstractSparseMatrixCSR, B, α, β) + size(A, 2) == size(B, 1) || + throw(DimensionMismatch("second dimension of A, $(size(A,2)), does not match the first dimension of B, $(size(B,1))")) + size(A, 1) == size(C, 1) || + throw(DimensionMismatch("first dimension of A, $(size(A,1)), does not match the first dimension of C, $(size(C,1))")) + size(B, 2) == size(C, 2) || + throw(DimensionMismatch("second dimension of B, $(size(B,2)), does not match the second dimension of C, $(size(C,2))")) + nzv = nonzeros(A) + cv = colvals(A) + β != one(β) && LinearAlgebra._rmul_or_fill!(C, β) + for k in 1:size(C, 1) + @inbounds for row in 1:size(A, 1) + αxj = B[k, row] * α + for j in nzrange(A, row) + C[k, cv[j]] += nzv[j]*αxj + end + end + end + C +end + +function _At_or_Ac_mul_B!(tfun::Function, C, A::AbstractSparseMatrixCSC, B, α, β) size(A, 2) == size(C, 1) || throw(DimensionMismatch("second dimension of A, $(size(A,2)), does not match the first dimension of C, $(size(C,1))")) size(A, 1) == size(B, 1) || @@ -117,6 +164,28 @@ function _At_or_Ac_mul_B!(tfun::Function, C, A, B, α, β) C end +function _At_or_Ac_mul_B!(tfun::Function, C, A::AbstractSparseMatrixCSR, B, α, β) + size(A, 2) == size(C, 1) || + throw(DimensionMismatch("second dimension of A, $(size(A,2)), does not match the first dimension of C, $(size(C,1))")) + size(A, 1) == size(B, 1) || + throw(DimensionMismatch("first dimension of A, $(size(A,1)), does not match the first dimension of B, $(size(B,1))")) + size(B, 2) == size(C, 2) || + throw(DimensionMismatch("second dimension of B, $(size(B,2)), does not match the second dimension of C, $(size(C,2))")) + nzv = nonzeros(A) + cv = colvals(A) + β != one(β) && LinearAlgebra._rmul_or_fill!(C, β) + for k in 1:size(C, 1) + @inbounds for row in 1:size(A, 1) + tmp = zero(eltype(C)) + for j in nzrange(A, row) + tmp += tfun(nzv[j])*B[k,cv[j]] + end + C[k,row] += tmp * α + end + end + C +end + Base.@constprop :aggressive function generic_matmatmul!(C::StridedMatrix, tA, tB, A::DenseMatrixUnion, B::SparseMatrixCSCUnion2, alpha::Number, beta::Number) transA = tA == 'N' ? identity : tA == 'T' ? transpose : adjoint if tB == 'N' @@ -167,6 +236,45 @@ function _spmul!(C::StridedMatrix, X::AdjOrTrans{<:Any,<:DenseMatrixUnion}, A::S C end +function _spmul!(C::StridedMatrix, X::DenseMatrixUnion, A::SparseMatrixCSRUnion2, α::Number, β::Number) + mX, nX = size(X) + nX == size(A, 1) || + throw(DimensionMismatch("second dimension of X, $nX, does not match the first dimension of A, $(size(A,1))")) + mX == size(C, 1) || + throw(DimensionMismatch("first dimension of X, $mX, does not match the first dimension of C, $(size(C,1))")) + size(A, 2) == size(C, 2) || + throw(DimensionMismatch("second dimension of A, $(size(A,2)), does not match the second dimension of C, $(size(C,2))")) + cv = colvals(A) + nzv = nonzeros(A) + β != one(β) && LinearAlgebra._rmul_or_fill!(C, β) + @inbounds for row in 1:size(A, 1), k in nzrange(A, row) + Aiα = nzv[k] * α + cvk = cv[k] + @simd for multivec_col in 1:nX + C[row, multivec_col] += X[cvk, multivec_col] * Aiα + end + end + C +end +function _spmul!(C::StridedMatrix, X::AdjOrTrans{<:Any,<:DenseMatrixUnion}, A::SparseMatrixCSRUnion2, α::Number, β::Number) + mX, nX = size(X) + nX == size(A, 1) || + throw(DimensionMismatch("second dimension of X, $nX, does not match the first dimension of A, $(size(A,1))")) + mX == size(C, 1) || + throw(DimensionMismatch("first dimension of X, $mX, does not match the first dimension of C, $(size(C,1))")) + size(A, 2) == size(C, 2) || + throw(DimensionMismatch("second dimension of A, $(size(A,2)), does not match the second dimension of C, $(size(C,2))")) + cv = colvals(A) + nzv = nonzeros(A) + β != one(β) && LinearAlgebra._rmul_or_fill!(C, β) + for multivec_col in 1:nX, row in 1:size(A, 1) + @inbounds for k in nzrange(A, row) + C[row, multivec_col] += X[cv[k], multivec_col] * nzv[k] * α + end + end + C +end + function _A_mul_Bt_or_Bc!(tfun::Function, C::StridedMatrix, A::AbstractMatrix, B::SparseMatrixCSCUnion2, α::Number, β::Number) mA, nA = size(A) nA == size(B, 2) || @@ -188,22 +296,53 @@ function _A_mul_Bt_or_Bc!(tfun::Function, C::StridedMatrix, A::AbstractMatrix, B C end +function _A_mul_Bt_or_Bc!(tfun::Function, C::StridedMatrix, A::AbstractMatrix, B::SparseMatrixCSRUnion2, α::Number, β::Number) + mA, nA = size(A) + nA == size(B, 2) || + throw(DimensionMismatch("second dimension of A, $nA, does not match the second dimension of B, $(size(B,2))")) + mA == size(C, 1) || + throw(DimensionMismatch("first dimension of A, $mA, does not match the first dimension of C, $(size(C,1))")) + size(B, 1) == size(C, 2) || + throw(DimensionMismatch("first dimension of B, $(size(B,2)), does not match the second dimension of C, $(size(C,2))")) + cv = colvals(B) + nzv = nonzeros(B) + β != one(β) && LinearAlgebra._rmul_or_fill!(C, β) + @inbounds for row in 1:size(B, 1), k in nzrange(B, row) + Biα = tfun(nzv[k]) * α + cvk = cv[k] + @simd for multivec_row in 1:nA + C[cvk, multivec_row] += A[row, multivec_row] * Biα + end + end + C +end + # Sparse matrix multiplication as described in [Gustavson, 1978]: # http://dl.acm.org/citation.cfm?id=355796 -const SparseTriangular{Tv,Ti} = Union{UpperTriangular{Tv,<:SparseMatrixCSCUnion{Tv,Ti}},LowerTriangular{Tv,<:SparseMatrixCSCUnion{Tv,Ti}}} -const SparseOrTri{Tv,Ti} = Union{SparseMatrixCSCUnion{Tv,Ti},SparseTriangular{Tv,Ti}} +const SparseTriangularCSC{Tv,Ti} = Union{UpperTriangular{Tv,<:SparseMatrixCSCUnion{Tv,Ti}},LowerTriangular{Tv,<:SparseMatrixCSCUnion{Tv,Ti}}} +const SparseTriangularCSR{Tv,Ti} = Union{UpperTriangular{Tv,<:SparseMatrixCSRUnion{Tv,Ti}},LowerTriangular{Tv,<:SparseMatrixCSRUnion{Tv,Ti}}} +const SparseTriangular{Tv,Ti} = Union{SparseTriangularCSC{Tv,Ti}, SparseTriangularCSR{Tv,Ti}} +const SparseOrTriCSC{Tv,Ti} = Union{SparseMatrixCSCUnion{Tv,Ti},SparseTriangularCSC{Tv,Ti}} +const SparseOrTriCSR{Tv,Ti} = Union{SparseMatrixCSRUnion{Tv,Ti},SparseTriangularCSR{Tv,Ti}} +const SparseOrTri{Tv,Ti} = Union{SparseOrTriCSC{Tv,Ti}, SparseOrTriCSR{Tv,Ti}} *(A::SparseOrTri, B::AbstractSparseVector) = spmatmulv(A, B) *(A::SparseOrTri, B::SparseColumnView) = spmatmulv(A, B) *(A::SparseOrTri, B::SparseVectorView) = spmatmulv(A, B) *(A::SparseMatrixCSCUnion, B::SparseMatrixCSCUnion) = spmatmul(A,B) +*(A::SparseMatrixCSRUnion, B::SparseMatrixCSRUnion) = spmatmul(A,B) *(A::SparseTriangular, B::SparseMatrixCSCUnion) = spmatmul(A,B) +*(A::SparseTriangular, B::SparseMatrixCSRUnion) = spmatmul(A,B) *(A::SparseMatrixCSCUnion, B::SparseTriangular) = spmatmul(A,B) +*(A::SparseMatrixCSRUnion, B::SparseTriangular) = spmatmul(A,B) *(A::SparseTriangular, B::SparseTriangular) = spmatmul1(A,B) *(A::SparseOrTri, B::AdjOrTrans{<:Any,<:AbstractSparseMatrixCSC}) = spmatmul(A, copy(B)) +*(A::SparseOrTri, B::AdjOrTrans{<:Any,<:AbstractSparseMatrixCSR}) = spmatmul(A, copy(B)) *(A::AdjOrTrans{<:Any,<:AbstractSparseMatrixCSC}, B::SparseOrTri) = spmatmul(copy(A), B) +*(A::AdjOrTrans{<:Any,<:AbstractSparseMatrixCSR}, B::SparseOrTri) = spmatmul(copy(A), B) *(A::AdjOrTrans{<:Any,<:AbstractSparseMatrixCSC}, B::AdjOrTrans{<:Any,<:AbstractSparseMatrixCSC}) = spmatmul(copy(A), copy(B)) +*(A::AdjOrTrans{<:Any,<:AbstractSparseMatrixCSR}, B::AdjOrTrans{<:Any,<:AbstractSparseMatrixCSR}) = spmatmul(copy(A), copy(B)) # Gustavson's matrix multiplication algorithm revisited. # The result rowval vector is already sorted by construction. @@ -213,7 +352,7 @@ const SparseOrTri{Tv,Ti} = Union{SparseMatrixCSCUnion{Tv,Ti},SparseTriangular{Tv # done by a quicksort of the row indices or by a full scan of the dense result vector. # The last is faster, if more than ≈ 1/32 of the result column is nonzero. # TODO: extend to SparseMatrixCSCUnion to allow for SubArrays (view(X, :, r)). -function spmatmul(A::SparseOrTri, B::Union{SparseOrTri,AbstractCompressedVector,SubArray{<:Any,<:Any,<:AbstractSparseArray}}) +function spmatmul(A::SparseOrTriCSC, B::Union{SparseOrTriCSC,AbstractCompressedVector,SubArray{<:Any,<:Any,<:AbstractSparseArray}}) Tv = promote_op(matprod, eltype(A), eltype(B)) Ti = promote_type(indtype(A), indtype(B)) mA, nA = size(A) @@ -248,6 +387,41 @@ function spmatmul(A::SparseOrTri, B::Union{SparseOrTri,AbstractCompressedVector, C = SparseMatrixCSC(mA, nB, colptrC, rowvalC, nzvalC) return C end +function spmatmul(A::MatrixType, B::Union{MatrixType,AbstractCompressedVector,SubArray{<:Any,<:Any,<:AbstractSparseArray}}) where MatrixType <: AbstractSparseMatrixCSR + Tv = promote_op(matprod, eltype(A), eltype(B)) + Ti = promote_type(indtype(A), indtype(B)) + mA, nA = size(A) + nB = size(B, 2) + mB = size(B, 1) + nA == mB || throw(DimensionMismatch("second dimension of A, $nA, does not match the first dimension of B, $mB")) + + nnzC = min(estimate_mulsize(mA, nnz(A), nA, nnz(B), nB) * 11 ÷ 10 + mA, mA*nB) + rowptrC = Vector{Ti}(undef, mA+1) + colvalC = Vector{Ti}(undef, nnzC) + nzvalC = Vector{Tv}(undef, nnzC) + + @inbounds begin + jp = 1 + xb = fill(false, nB) + for j in 1:mA + if jp + nB - 1 > nnzC + nnzC += max(nB, nnzC>>2) + resize!(colvalC, nnzC) + resize!(nzvalC, nnzC) + end + rowptrC[j] = jp + jp = sprowmul!(colvalC, nzvalC, xb, j, jp, A, B) + end + rowptrC[mA+1] = jp + end + + resize!(colvalC, jp - 1) + resize!(nzvalC, jp - 1) + + # This modification of Gustavson algorithm has sorted row indices + C = MatrixType(mA, nB, rowptrC, colvalC, nzvalC) + return C +end # process single rhs column function spcolmul!(rowvalC, nzvalC, xb, i, ip, A, B) @@ -297,6 +471,54 @@ function spcolmul!(rowvalC, nzvalC, xb, i, ip, A, B) end return ip end +# process single rhs row +function sprowmul!(colvalC, nzvalC, xb, j, jp, A, B) + colvalA = colvals(A); nzvalA = nonzeros(A) + colvalB = colvals(B); nzvalB = nonzeros(B) + nB = size(B, 2) + jp0 = jp + k0 = jp - 1 + @inbounds begin + for ip in nzrange(A, j) + nzA = nzvalA[ip] + i = colvalA[ip] + for kp in nzrange(B, i) + nzC = nzvalB[kp] * nzA + k = colvalB[kp] + if xb[k] + nzvalC[k+k0] += nzC + else + nzvalC[k+k0] = nzC + xb[k] = true + colvalC[jp] = k + jp += 1 + end + end + end + if jp > jp0 + if prefer_sort(jp-k0, nB) + # in-place sort of indices. Effort: O(nnz*ln(nnz)). + sort!(colvalC, jp0, jp-1, QuickSort, Base.Order.Forward) + for vp = jp0:jp-1 + k = colvalC[vp] + xb[k] = false + nzvalC[vp] = nzvalC[k+k0] + end + else + # scan result vector (effort O(mA)) + for k = 1:nB + if xb[k] + xb[k] = false + colvalC[jp0] = k + nzvalC[jp0] = nzvalC[k+k0] + jp0 += 1 + end + end + end + end + end + return jp +end # special cases of same twin Upper/LowerTriangular spmatmul1(A, B) = spmatmul(A, B) @@ -1104,17 +1326,27 @@ function _mul!(nzrang::Function, diagop::Function, odiagop::Function, C::Strided end # row range up to (and including if excl=false) diagonal -function nzrangeup(A, i, excl=false) +function nzrangeup(A::SparseMatrixCSCUnion3, i, excl=false) r = nzrange(A, i); r1 = r.start; r2 = r.stop rv = rowvals(A) @inbounds r2 < r1 || rv[r2] <= i - excl ? r : r1:(searchsortedlast(view(rv, r1:r2), i - excl) + r1-1) end # row range from diagonal (included if excl=false) to end -function nzrangelo(A, i, excl=false) +function nzrangelo(A::SparseMatrixCSCUnion3, i, excl=false) r = nzrange(A, i); r1 = r.start; r2 = r.stop rv = rowvals(A) @inbounds r2 < r1 || rv[r1] >= i + excl ? r : (searchsortedfirst(view(rv, r1:r2), i + excl) + r1-1):r2 end +function nzrangeup(A::SparseMatrixCSRUnion3, i, excl=false) + c = nzrange(A, i); c1 = c.start; c2 = c.stop + cv = colvals(A) + @inbounds c2 < c1 || cv[c1] >= i + excl ? c : (searchsortedfirst(view(cv, c1:c2), i + excl) + c1-1):c2 +end +function nzrangelo(A::SparseMatrixCSRUnion3, i, excl=false) + c = nzrange(A, i); c1 = c.start; c2 = c.stop + cv = colvals(A) + @inbounds c2 < c1 || cv[c2] <= i - excl ? c : c1:(searchsortedlast(view(cv, c1:c2), i - excl) + c1-1) +end dot(x::AbstractVector, A::RealHermSymComplexHerm{<:Any,<:AbstractSparseMatrixCSC}, y::AbstractVector) = _dot(x, parent(A), y, A.uplo == 'U' ? nzrangeup : nzrangelo, A isa Symmetric ? identity : real, A isa Symmetric ? transpose : adjoint) diff --git a/src/sparseconvert.jl b/src/sparseconvert.jl index 479654b7..96270ddc 100644 --- a/src/sparseconvert.jl +++ b/src/sparseconvert.jl @@ -7,8 +7,13 @@ """ const SparseMatrixCSCSymmHerm{Tv,Ti} = Union{Symmetric{Tv,<:SparseMatrixCSCUnion{Tv,Ti}}, Hermitian{Tv,<:SparseMatrixCSCUnion{Tv,Ti}}} +const SparseMatrixCSRSymmHerm{Tv,Ti,MatrixType} = Union{ + Symmetric{Tv,MatrixType}, + Hermitian{Tv,MatrixType} +} where MatrixType <: SparseMatrixCSRUnion{Tv,Ti} -const AbstractTriangularSparse{Tv,Ti} = UpperOrLowerTriangular{Tv,<:SparseMatrixCSCUnion{Tv,Ti}} +const AbstractTriangularSparseCSC{Tv,Ti} = UpperOrLowerTriangular{Tv,<:SparseMatrixCSCUnion{Tv,Ti}} +const AbstractTriangularSparseCSR{Tv,Ti,MatrixType} = UpperOrLowerTriangular{Tv,MatrixType} where MatrixType <: MatrixType<:SparseMatrixCSRUnion{Tv,Ti} # converting Symmetric/Hermitian/Triangular/SubArray of SparseMatrixCSC # and Transpose/Adjoint of Triangular of SparseMatrixCSC to SparseMatrixCSC @@ -115,6 +120,7 @@ end # Symmetric/Hermitian of sparse matrix _sparsem(A::SparseMatrixCSCSymmHerm) = _sparsem(A.uplo == 'U' ? nzrangeup : nzrangelo, A) +_sparsem(A::SparseMatrixCSRSymmHerm) = _sparsem(A.uplo == 'U' ? nzrangeup : nzrangelo, A) # Triangular of sparse matrix _sparsem(A::UpperTriangular{T,<:AbstractSparseMatrix}) where T = triu(A.data) _sparsem(A::LowerTriangular{T,<:AbstractSparseMatrix}) where T = tril(A.data) @@ -169,11 +175,60 @@ function _sparsem(fnzrange::Function, sA::SparseMatrixCSCSymmHerm{Tv}) where {Tv end newcolptr[j] = newk end - _sparse_gen(m, n, newcolptr, newrowval, newnzval) + _sparse_gen_csc(SparseMatrixCSC, m, n, newcolptr, newrowval, newnzval) +end +function _sparsem(fnzrange::Function, sA::SparseMatrixCSRSymmHerm{Tv,MatrixType}) where {Tv,MatrixType} + A = sA.data + colval = colvals(A) + nzval = nonzeros(A) + m, n = size(A) + Ti = eltype(colval) + fadj = sA isa Symmetric ? transpose : adjoint + newrowptr = Vector{Ti}(undef, n+1) + diagmap = fadj == transpose ? identity : real + + newrowptr[1] = 1 + rowrange = fnzrange === nzrangeup ? (1:n) : (n:-1:1) + @inbounds for j = colrange + r = fnzrange(A, j); r1 = r.start; r2 = r.stop + newrowptr[j+1] = r2 - r1 + 1 + for k = r1:r2 + row = colval[k] + if row != j + newrowptr[row+1] += 1 + end + end + end + cumsum!(newrowptr, newrowptr) + nz = newrowptr[n+1] - 1 + newcolval = Vector{Ti}(undef, nz) + newnzval = Vector{Tv}(undef, nz) + @inbounds for j = 1:n + newk = newrowptr[j] + for k = fnzrange(A, j) + i = colval[k] + nzv = nzval[k] + if i != j + newcolval[newk] = i + newnzval[newk] = nzv + newk += 1 + ni = newrowptr[i] + newcolval[ni] = j + newnzval[ni] = fadj(nzv) + newrowptr[i] = ni + 1 + else + newcolval[newk] = i + newnzval[newk] = diagmap(nzv) + newk += 1 + end + end + newrowptr[j] = newk + end + _sparse_gen_csr(MatrixType, m, n, newrowptr, newcolval, newnzval) end # 2 cases: Unit(Upper|Lower)Triangular{Tv,AbstractSparseMatrixCSC} -function _sparsem(A::AbstractTriangularSparse{Tv}) where Tv +function _sparsem(A::AbstractTriangularSparseCSC{Tv}) where Tv S = A.data rowval = rowvals(S) nzval = nonzeros(S) @@ -215,10 +270,51 @@ function _sparsem(A::AbstractTriangularSparse{Tv}) where Tv resize!(newnzval, nz) SparseMatrixCSC(m, n, newcolptr, newrowval, newnzval) end +function _sparsem(A::AbstractTriangularSparseCSR{Tv,<:Any,MatrixType}) where {Tv, MatrixType} + S = A.data + colval = colvals(S) + nzval = nonzeros(S) + m, n = size(S) + Ti = eltype(colval) + fnzrange = A isa Union{UpperTriangular,UnitUpperTriangular} ? nzrangeup : nzrangelo + unit = A isa Union{UnitUpperTriangular,UnitLowerTriangular} + nz = nnz(S) + n * unit + newrowptr = Vector{Ti}(undef, n+1) + newcolval = Vector{Ti}(undef, nz) + newnzval = Vector{Tv}(undef, nz) + newrowptr[1] = 1 + uplo = fnzrange == nzrangeup + newk = 1 + @inbounds for j = 1:n + newkk = newk + if unit + newk += !uplo + end + r = fnzrange(S, j); r1 = r.start; r2 = r.stop + for k = r1:r2 + i = colval[k] + if i != j || i == j && !unit + newcolval[newk] = i + newnzval[newk] = nzval[k] + newk += 1 + end + end + if unit + uplo && (newkk = newk) + newcolval[newkk] = j + newnzval[newkk] = one(Tv) + newk += uplo + end + newrowptr[j+1] = newk + end + nz = newrowptr[n+1] - 1 + resize!(newcolval, nz) + resize!(newnzval, nz) + MatrixType(m, n, newrowptr, newcolval, newnzval) +end # 8 cases: (Transpose|Adjoint){Tv,[Unit](Upper|Lower)Triangular} -function _sparsem(taA::AdjOrTrans{Tv,<:AbstractTriangularSparse}) where {Tv} - +function _sparsem(taA::AdjOrTrans{Tv,<:AbstractTriangularSparseCSC}) where {Tv} sA = taA.parent A = sA.data rowval = rowvals(A) @@ -270,14 +366,76 @@ function _sparsem(taA::AdjOrTrans{Tv,<:AbstractTriangularSparse}) where {Tv} newcolptr[j] = ni + 1 end end - _sparse_gen(n, m, newcolptr, newrowval, newnzval) + _sparse_gen_csc(SparseMatrixCSC, n, m, newcolptr, newrowval, newnzval) +end + +function _sparsem(taA::AdjOrTrans{Tv,<:AbstractTriangularSparseCSR{Tv,<:Any,MatrixType}}) where {Tv, MatrixType} + sA = taA.parent + A = sA.data + colval = colvals(A) + nzval = nonzeros(A) + m, n = size(A) + Ti = eltype(colval) + fnzrange = sA isa Union{UpperTriangular,UnitUpperTriangular} ? nzrangeup : nzrangelo + fadj = taA isa Transpose ? transpose : adjoint + unit = sA isa Union{UnitUpperTriangular,UnitLowerTriangular} + uplo = A isa Union{UpperTriangular,UnitUpperTriangular} + + newrowptr = Vector{Ti}(undef, n+1) + fill!(newrowptr, unit) + newrowptr[1] = 1 + @inbounds for j = 1:n + for k = fnzrange(A, j) + i = colval[k] + if i != j || i == j && !unit + newrowptr[i+1] += 1 + end + end + end + cumsum!(newrowptr, newrowptr) + nz = newrowptr[n+1] - 1 + newcolval = Vector{Ti}(undef, nz) + newnzval = Vector{Tv}(undef, nz) + + @inbounds for j = 1:n + if !uplo && unit + ni = newrowptr[j] + newcolval[ni] = j + newnzval[ni] = fadj(one(Tv)) + newrowptr[j] = ni + 1 + end + for k = fnzrange(A, j) + i = colval[k] + nzv = nzval[k] + if i != j || i == j && !unit + ni = newrowptr[i] + newcolval[ni] = j + newnzval[ni] = fadj(nzv) + newrowptr[i] = ni + 1 + end + end + if uplo && unit + ni = newrowptr[j] + newcolval[ni] = j + newnzval[ni] = fadj(one(Tv)) + newrowptr[j] = ni + 1 + end + end + _sparse_gen_csr(MatrixType, n, m, newrowptr, newcolval, newnzval) end -function _sparse_gen(m, n, newcolptr, newrowval, newnzval) +function _sparse_gen_csc(T, m, n, newcolptr, newrowval, newnzval) @inbounds for j = n:-1:1 newcolptr[j+1] = newcolptr[j] end newcolptr[1] = 1 - SparseMatrixCSC(m, n, newcolptr, newrowval, newnzval) + T(m, n, newcolptr, newrowval, newnzval) end +function _sparse_gen_csr(T, m, n, newcolptr, newrowval, newnzval) + @inbounds for j = n:-1:1 + newcolptr[j+1] = newcolptr[j] + end + newcolptr[1] = 1 + T(m, n, newcolptr, newrowval, newnzval) +end diff --git a/src/sparsematrix.jl b/src/sparsematrix.jl index 4fa2adf9..fe287469 100644 --- a/src/sparsematrix.jl +++ b/src/sparsematrix.jl @@ -26,7 +26,7 @@ struct SparseMatrixCSC{Tv,Ti<:Integer} <: AbstractSparseMatrixCSC{Tv,Ti} function SparseMatrixCSC{Tv,Ti}(m::Integer, n::Integer, colptr::Vector{Ti}, rowval::Vector{Ti}, nzval::Vector{Tv}) where {Tv,Ti<:Integer} sparse_check_Ti(m, n, Ti) - _goodbuffers(Int(m), Int(n), colptr, rowval, nzval) || + _goodbufferscsc(Int(m), Int(n), colptr, rowval, nzval) || throw(ArgumentError("Invalid buffers for SparseMatrixCSC construction n=$n, colptr=$(summary(colptr)), rowval=$(summary(rowval)), nzval=$(summary(nzval))")) new(Int(m), Int(n), colptr, rowval, nzval) end @@ -73,7 +73,7 @@ struct FixedSparseCSC{Tv,Ti<:Integer} <: AbstractSparseMatrixCSC{Tv,Ti} rowval::ReadOnly{Ti,1,Vector{Ti}}, nzval::Vector{Tv}) where {Tv,Ti<:Integer} sparse_check_Ti(m, n, Ti) - _goodbuffers(Int(m), Int(n), parent(colptr), parent(rowval), nzval) || + _goodbufferscsc(Int(m), Int(n), parent(colptr), parent(rowval), nzval) || throw(ArgumentError("Invalid buffers for FixedSparseCSC construction n=$n, colptr=$(summary(colptr)), rowval=$(summary(rowval)), nzval=$(summary(nzval))")) new(Int(m), Int(n), colptr, rowval, nzval) end @@ -164,15 +164,21 @@ end size(S::SorF) = (getfield(S, :m), getfield(S, :n)) -_goodbuffers(S::AbstractSparseMatrixCSC) = _goodbuffers(size(S)..., getcolptr(S), getrowval(S), nonzeros(S)) -_checkbuffers(S::AbstractSparseMatrixCSC) = (@assert _goodbuffers(S); S) +_goodbuffers(S::AbstractSparseMatrixCSC) = _goodbufferscsc(size(S)..., getcolptr(S), getrowval(S), nonzeros(S)) +_goodbuffers(S::AbstractSparseMatrixCSR) = _goodbufferscsr(size(S)..., getrowptr(S), getcolval(S), nonzeros(S)) +_checkbuffers(S::AbstractSparseMatrixCSCOrCSR) = (@assert _goodbuffers(S); S) _checkbuffers(S::Union{Adjoint, Transpose}) = (_checkbuffers(parent(S)); S) -function _goodbuffers(m, n, colptr, rowval, nzval) +function _goodbufferscsc(m, n, colptr, rowval, nzval) (length(colptr) == n + 1 && colptr[end] - 1 == length(rowval) == length(nzval)) # stronger check for debugging purposes # && all(issorted(@view rowval[colptr[i]:colptr[i+1]-1]) for i=1:n) end +function _goodbufferscsr(m, n, rowptr, colval, nzval) + (length(rowptr) == m + 1 && rowptr[end] - 1 == length(colval) == length(nzval)) + # stronger check for debugging purposes + # && all(issorted(@view rowval[colptr[i]:colptr[i+1]-1]) for i=1:n) +end # Define an alias for views of a SparseMatrixCSC which include all rows and a unit range of the columns. # Also define a union of SparseMatrixCSC and this view since many methods can be defined efficiently for @@ -190,6 +196,7 @@ const SparseMatrixCSCColumnSubset{Tv,Ti} = SubArray{Tv,2,<:AbstractSparseMatrixCSC{Tv,Ti}, Tuple{Base.Slice{Base.OneTo{Int}},I}} where {I<:AbstractVector{<:Integer}} const SparseMatrixCSCUnion2{Tv,Ti} = Union{AbstractSparseMatrixCSC{Tv,Ti}, SparseMatrixCSCColumnSubset{Tv,Ti}} +const SparseMatrixCSCUnion3{Tv,Ti} = Union{AbstractSparseMatrixCSC{Tv,Ti}, SparseMatrixCSCColumnSubset{Tv,Ti}, SparseMatrixCSCView{Tv, Ti}} getcolptr(S::SorF) = getfield(S, :colptr) getcolptr(S::SparseMatrixCSCView) = view(getcolptr(parent(S)), first(S.indices[2]):(last(S.indices[2]) + 1)) @@ -200,6 +207,25 @@ getnzval( S::AbstractSparseMatrixCSC) = nonzeros(S) getnzval( S::SparseMatrixCSCColumnSubset) = nonzeros(parent(S)) nzvalview(S::AbstractSparseMatrixCSC) = view(nonzeros(S), 1:nnz(S)) +const SparseMatrixCSRView{Tv,Ti} = + SubArray{Tv,2,<:AbstractSparseMatrixCSR{Tv,Ti}, + Tuple{I,Base.Slice{Base.OneTo{Int}}}} where {I<:AbstractUnitRange{<:Integer}} +const SparseMatrixCSRRowSubset{Tv,Ti} = + SubArray{Tv,2,<:AbstractSparseMatrixCSR{Tv,Ti}, + Tuple{I,Base.Slice{Base.OneTo{Int}}}} where {I<:AbstractVector{<:Integer}} +const SparseMatrixCSRUnion{Tv,Ti} = Union{AbstractSparseMatrixCSR{Tv,Ti}, SparseMatrixCSRView{Tv,Ti}} +const SparseMatrixCSRUnion2{Tv,Ti} = Union{AbstractSparseMatrixCSR{Tv,Ti}, SparseMatrixCSRRowSubset{Tv,Ti}} +const SparseMatrixCSRUnion3{Tv,Ti} = Union{AbstractSparseMatrixCSR{Tv,Ti}, SparseMatrixCSRRowSubset{Tv,Ti}, SparseMatrixCSRView{Tv, Ti}} + +getrowptr(S::AbstractSparseMatrixCSR) = getfield(S, :rowptr) +getrowptr(S::SparseMatrixCSRView) = view(getrowptr(parent(S)), first(S.indices[2]):(last(S.indices[2]) + 1)) +getrowptr(S::SparseMatrixCSRRowSubset) = error("getcolptr not well-defined for $(typeof(S))") +getcolval(S::AbstractSparseMatrixCSR) = colvals(S) +getcolval(S::SparseMatrixCSRRowSubset) = colvals(parent(S)) +getnzval( S::AbstractSparseMatrixCSR) = nonzeros(S) +getnzval( S::SparseMatrixCSRRowSubset) = nonzeros(parent(S)) +nzvalview(S::AbstractSparseMatrixCSR) = view(nonzeros(S), 1:nnz(S)) + """ nnz(A) @@ -220,15 +246,27 @@ julia> nnz(A) nnz(S::AbstractSparseMatrixCSC) = Int(getcolptr(S)[size(S, 2) + 1]) - 1 nnz(S::ReshapedArray{<:Any,1,<:AbstractSparseMatrixCSC}) = nnz(parent(S)) nnz(S::AdjOrTrans{<:Any,<:AbstractSparseMatrixCSC}) = nnz(parent(S)) -nnz(S::UpperTriangular{<:Any,<:AbstractSparseMatrixCSC}) = nnz1(S) -nnz(S::LowerTriangular{<:Any,<:AbstractSparseMatrixCSC}) = nnz1(S) -nnz(S::SparseMatrixCSCColumnSubset) = nnz1(S) -nnz1(S) = sum(length.(nzrange.(Ref(S), axes(S, 2)))) +nnz(S::UpperTriangular{<:Any,<:AbstractSparseMatrixCSC}) = nnz1csc(S) +nnz(S::LowerTriangular{<:Any,<:AbstractSparseMatrixCSC}) = nnz1csc(S) +nnz(S::SparseMatrixCSCColumnSubset) = nnz1csc(S) +nnz1csc(S) = sum(length.(nzrange.(Ref(S), axes(S, 2)))) + +nnz(S::AbstractSparseMatrixCSR) = Int(getrowptr(S)[size(S, 1) + 1]) - 1 +nnz(S::ReshapedArray{<:Any,1,<:AbstractSparseMatrixCSR}) = nnz(parent(S)) +nnz(S::AdjOrTrans{<:Any,<:AbstractSparseMatrixCSR}) = nnz(parent(S)) +nnz(S::UpperTriangular{<:Any,<:AbstractSparseMatrixCSR}) = nnz1csr(S) +nnz(S::LowerTriangular{<:Any,<:AbstractSparseMatrixCSR}) = nnz1csr(S) +nnz(S::SparseMatrixCSRRowSubset) = nnz1csr(S) +nnz1csr(S) = sum(length.(nzrange.(Ref(S), axes(S, 1)))) function Base._simple_count(pred, S::AbstractSparseMatrixCSC, init::T) where T init + T(count(pred, nzvalview(S)) + pred(zero(eltype(S)))*(prod(size(S)) - nnz(S))) end +function Base._simple_count(pred, S::AbstractSparseMatrixCSR, init::T) where T + init + T(count(pred, nzvalview(S)) + pred(zero(eltype(S)))*(prod(size(S)) - nnz(S))) +end + """ nonzeros(A) @@ -257,6 +295,9 @@ nonzeros(S::SorF) = getfield(S, :nzval) nonzeros(S::SparseMatrixCSCColumnSubset) = nonzeros(S.parent) nonzeros(S::UpperTriangular{<:Any,<:SparseMatrixCSCUnion}) = nonzeros(S.data) nonzeros(S::LowerTriangular{<:Any,<:SparseMatrixCSCUnion}) = nonzeros(S.data) +nonzeros(S::SparseMatrixCSRRowSubset) = nonzeros(S.parent) +nonzeros(S::UpperTriangular{<:Any,<:SparseMatrixCSRUnion}) = nonzeros(S.data) +nonzeros(S::LowerTriangular{<:Any,<:SparseMatrixCSRUnion}) = nonzeros(S.data) """ rowvals(A::AbstractSparseMatrixCSC) @@ -286,12 +327,26 @@ rowvals(S::SparseMatrixCSCColumnSubset) = rowvals(S.parent) rowvals(S::UpperTriangular{<:Any,<:SparseMatrixCSCUnion}) = rowvals(S.data) rowvals(S::LowerTriangular{<:Any,<:SparseMatrixCSCUnion}) = rowvals(S.data) +""" + colvals(A::AbstractSparseMatrixCSR) + +Return a vector of the column indices of `A`. Any modifications to the returned +vector will mutate `A` as well. Providing access to how the row indices are +stored internally can be useful in conjunction with iterating over structural +nonzero values. See also [`nonzeros`](@ref) and [`nzrange`](@ref). +""" +colvals(S::AbstractSparseMatrixCSR) = getfield(S, :colval) +colvals(S::SparseMatrixCSRRowSubset) = colvals(S.parent) +colvals(S::UpperTriangular{<:Any,<:SparseMatrixCSRUnion}) = colvals(S.data) +colvals(S::LowerTriangular{<:Any,<:SparseMatrixCSRUnion}) = colvals(S.data) + """ nzrange(A::AbstractSparseMatrixCSC, col::Integer) + nzrange(A::AbstractSparseMatrixCSR, row::Integer) Return the range of indices to the structural nonzero values of a sparse matrix -column. In conjunction with [`nonzeros`](@ref) and -[`rowvals`](@ref), this allows for convenient iterating over a sparse matrix : +column/row. In conjunction with [`nonzeros`](@ref) and +[`rowvals`](@ref)/[`colvals`](@ref), this allows for convenient iterating over a sparse matrix, e.g.: A = sparse(I,J,V) rows = rowvals(A) @@ -309,11 +364,15 @@ column. In conjunction with [`nonzeros`](@ref) and Adding or removing nonzero elements to the matrix may invalidate the `nzrange`, one should not mutate the matrix while iterating. """ nzrange(S::AbstractSparseMatrixCSC, col::Integer) = getcolptr(S)[col]:(getcolptr(S)[col+1]-1) +nzrange(S::AbstractSparseMatrixCSR, row::Integer) = getrowptr(S)[row]:(getrowptr(S)[row+1]-1) nzrange(S::SparseMatrixCSCColumnSubset, col::Integer) = nzrange(S.parent, S.indices[2][col]) +nzrange(S::SparseMatrixCSRRowSubset, row::Integer) = nzrange(S.parent, S.indices[1][row]) nzrange(S::UpperTriangular{<:Any,<:SparseMatrixCSCUnion}, i::Integer) = nzrangeup(S.data, i) nzrange(S::LowerTriangular{<:Any,<:SparseMatrixCSCUnion}, i::Integer) = nzrangelo(S.data, i) -const AbstractSparseMatrixCSCInclAdjointAndTranspose = Union{AbstractSparseMatrixCSC,Adjoint{<:Any,<:AbstractSparseMatrixCSC},Transpose{<:Any,<:AbstractSparseMatrixCSC}} +const AbstractSparseMatrixCSCInclAdjointAndTranspose = Union{AbstractSparseMatrixCSC,Adjoint{<:Any,<:AbstractSparseMatrixCSC},Transpose{<:Any,<:AbstractSparseMatrixCSC}} +const AbstractSparseMatrixCSRInclAdjointAndTranspose = Union{AbstractSparseMatrixCSR,Adjoint{<:Any,<:AbstractSparseMatrixCSR},Transpose{<:Any,<:AbstractSparseMatrixCSR}} +const AbstractSparseMatrixCSCOrCSRInclAdjointAndTranspose = Union{AbstractSparseMatrixCSCInclAdjointAndTranspose,AbstractSparseMatrixCSRInclAdjointAndTranspose} function Base.isstored(A::AbstractSparseMatrixCSC, i::Integer, j::Integer) @boundscheck checkbounds(A, i, j) rows = rowvals(A) @@ -322,6 +381,14 @@ function Base.isstored(A::AbstractSparseMatrixCSC, i::Integer, j::Integer) end return false end +function Base.isstored(A::AbstractSparseMatrixCSR, i::Integer, j::Integer) + @boundscheck checkbounds(A, i, j) + cols = colvals(A) + for istored in nzrange(A, i) # could do binary search if the row indices are sorted? + j == cols[istored] && return true + end + return false +end function Base.isstored(A::AdjOrTrans{<:Any,<:AbstractSparseMatrixCSC}, i::Integer, j::Integer) @boundscheck checkbounds(A, i, j) @@ -331,8 +398,16 @@ function Base.isstored(A::AdjOrTrans{<:Any,<:AbstractSparseMatrixCSC}, i::Intege end return false end +function Base.isstored(A::AdjOrTrans{<:Any,<:AbstractSparseMatrixCSR}, i::Integer, j::Integer) + @boundscheck checkbounds(A, i, j) + rows = colvals(parent(A)) + for istored in nzrange(parent(A), j) + i == rows[istored] && return true + end + return false +end -Base.replace_in_print_matrix(A::AbstractSparseMatrixCSCInclAdjointAndTranspose, i::Integer, j::Integer, s::AbstractString) = +Base.replace_in_print_matrix(A::AbstractSparseMatrixCSCOrCSRInclAdjointAndTranspose, i::Integer, j::Integer, s::AbstractString) = Base.isstored(A, i, j) ? s : Base.replace_with_centered_mark(s) function Base.array_summary(io::IO, S::AbstractSparseMatrixCSCInclAdjointAndTranspose, dims::Tuple{Vararg{Base.OneTo}}) @@ -345,8 +420,8 @@ function Base.array_summary(io::IO, S::AbstractSparseMatrixCSCInclAdjointAndTran nothing end -# called by `show(io, MIME("text/plain"), ::AbstractSparseMatrixCSCInclAdjointAndTranspose)` -function Base.print_array(io::IO, S::AbstractSparseMatrixCSCInclAdjointAndTranspose) +# called by `show(io, MIME("text/plain"), ::AbstractSparseMatrixCSCOrCSRInclAdjointAndTranspose)` +function Base.print_array(io::IO, S::AbstractSparseMatrixCSCOrCSRInclAdjointAndTranspose) if max(size(S)...) < 16 Base.print_matrix(io, S) else @@ -395,6 +470,47 @@ function Base.show(io::IO, _S::AbstractSparseMatrixCSCInclAdjointAndTranspose) end end +""" + RowIndices(S::AbstractSparseMatrixCSR) + +Return the row indices of the stored values in `S`. +This is an internal type that is used in displaying sparse matrices, +and is not a part of the public interface. +""" +struct RowIndices{Ti,S<:AbstractSparseMatrixCSR{<:Any,Ti}} <: AbstractVector{Ti} + arr :: S +end + +size(C::RowIndices) = (nnz(C.arr),) +# returns the row index of the n-th non-zero value from the row pointer +@inline function getindex(C::RowIndices, i::Int) + @boundscheck checkbounds(C, i) + rowptr = getrowptr(C.arr) + ind = searchsortedlast(rowptr, i) + eltype(C)(ind) +end + +# always show matrices as `sparse(I, J, K)` +function Base.show(io::IO, _S::AbstractSparseMatrixCSRInclAdjointAndTranspose) + _checkbuffers(_S) + # can't use `findnz`, because that expects all values not to be #undef + S = _S isa Adjoint || _S isa Transpose ? _S.parent : _S + J = rowvals(S) + K = nonzeros(S) + m, n = size(S) + if _S isa Adjoint + print(io, "adjoint(") + elseif _S isa Transpose + print(io, "transpose(") + end + print(io, "sparse(") + show(io, RowIndices(S)) + print(io, ", ", J, ", ", K, ", ", m, ", ", n, ")") + if _S isa Adjoint || _S isa Transpose + print(io, ")") + end +end + const brailleBlocks = UInt16['⠁', '⠂', '⠄', '⡀', '⠈', '⠐', '⠠', '⢀'] function _show_with_braille_patterns(io::IO, S::AbstractSparseMatrixCSCInclAdjointAndTranspose) m, n = size(S) @@ -561,6 +677,8 @@ Base.unaliascopy(S::AbstractSparseMatrixCSC) = typeof(S)(size(S, 1), size(S, 2), copy(S::AbstractSparseMatrixCSC) = SparseMatrixCSC(size(S, 1), size(S, 2), copy(getcolptr(S)), copy(rowvals(S)), copy(nonzeros(S))) +copy(S::MatrixType) where MatrixType <: AbstractSparseMatrixCSR = + MatrixType(size(S, 1), size(S, 2), copy(getrowptr(S)), copy(colvals(S)), copy(nonzeros(S))) copy(S::FixedSparseCSC) = FixedSparseCSC(size(S, 1), size(S, 2), getcolptr(S), rowvals(S), copy(nonzeros(S))) function copyto!(A::AbstractSparseMatrixCSC, B::AbstractSparseMatrixCSC) @@ -615,6 +733,10 @@ copyto!(A::AbstractMatrix, B::AbstractSparseMatrixCSC) = _sparse_copyto!(A, B) # Ambiguity resolution copyto!(A::PermutedDimsArray, B::AbstractSparseMatrixCSC) = _sparse_copyto!(A, B) +copyto!(A::AbstractMatrix, B::AbstractSparseMatrixCSR) = _sparse_copyto!(A, B) +# Ambiguity resolution +copyto!(A::PermutedDimsArray, B::AbstractSparseMatrixCSR) = _sparse_copyto!(A, B) + function _sparse_copyto!(dest::AbstractMatrix, src::AbstractSparseMatrixCSC) (dest === src || isempty(src)) && return dest z = convert(eltype(dest), zero(eltype(src))) # should throw if not possible @@ -634,6 +756,25 @@ function _sparse_copyto!(dest::AbstractMatrix, src::AbstractSparseMatrixCSC) return dest end +function _sparse_copyto!(dest::AbstractMatrix, src::AbstractSparseMatrixCSR) + (dest === src || isempty(src)) && return dest + z = convert(eltype(dest), zero(eltype(src))) # should throw if not possible + isrc = LinearIndices(src) + checkbounds(dest, isrc) + # If src is not dense, zero out the portion of dest spanned by isrc + if widelength(src) > nnz(src) + for i in isrc + @inbounds dest[i] = z + end + end + @inbounds for row in axes(src, 1), ptr in nzrange(src, row) + col = colvals(src)[ptr] + val = nonzeros(src)[ptr] + dest[isrc[row, col]] = val + end + return dest +end + function copyto!(dest::AbstractMatrix, Rdest::CartesianIndices{2}, src::AbstractSparseMatrixCSC{T}, Rsrc::CartesianIndices{2}) where {T} isempty(Rdest) && return dest @@ -693,6 +834,11 @@ function _sparsesimilar(S::AbstractSparseMatrixCSC, ::Type{TvNew}, ::Type{TiNew} newrowval = copyto!(similar(rowvals(S), TiNew), rowvals(S)) return SparseMatrixCSC(size(S, 1), size(S, 2), newcolptr, newrowval, similar(nonzeros(S), TvNew)) end +function _sparsesimilar(S::MatrixType, ::Type{TvNew}, ::Type{TiNew}) where {TvNew,TiNew,MatrixType<:AbstractSparseMatrixCSR} + newrowptr = copyto!(similar(getrowptr(S), TiNew), getrowptr(S)) + newcolval = copyto!(similar(colvals(S), TiNew), colvals(S)) + return MatrixType(size(S, 1), size(S, 2), newcolptr, newrowval, similar(nonzeros(S), TvNew)) +end # parent methods for similar that preserves only storage space (for when new dims are 2-d) _sparsesimilar(S::AbstractSparseMatrixCSC, ::Type{TvNew}, ::Type{TiNew}, dims::Dims{2}) where {TvNew,TiNew} = sizehint!(spzeros(TvNew, TiNew, dims...), length(nonzeros(S))) @@ -742,6 +888,12 @@ function Base.sizehint!(S::SparseMatrixCSC, n::Integer) sizehint!(nonzeros(S), nhint) return S end +function Base.sizehint!(S::AbstractSparseMatrixCSR, n::Integer) + nhint = min(n, widelength(S)) + sizehint!(getcolval(S), nhint) + sizehint!(nonzeros(S), nhint) + return S +end # converting between SparseMatrixCSC types SparseMatrixCSC(S::AbstractSparseMatrixCSC) = copy(S) @@ -1361,6 +1513,12 @@ function halfperm!(X::AbstractSparseMatrixCSC{Tv,Ti}, A::AbstractSparseMatrixCSC _distributevals_halfperm!(X, A, q, f) return X end +function halfperm!(X::AbstractSparseMatrixCSR{Tv,Ti}, A::AbstractSparseMatrixCSR{TvA,Ti}, + q::AbstractVector{<:Integer}, f::F = identity) where {Tv,TvA,Ti,F<:Function} + _computerowptrs_halfperm!(X, A) + _distributevals_halfperm!(X, A, q, f) +return X +end """ Helper method for `halfperm!`. Computes `transpose(A[:,q])`'s column pointers, storing them shifted one position forward in `getcolptr(X)`; `_distributevals_halfperm!` fixes this shift. @@ -1380,6 +1538,22 @@ function _computecolptrs_halfperm!(X::AbstractSparseMatrixCSC{Tv,Ti}, A::Abstrac countsum += overwritten end end +function _computerowptrs_halfperm!(X::AbstractSparseMatrixCSR{Tv,Ti}, A::AbstractSparseMatrixCSR{TvA,Ti}) where {Tv,TvA,Ti} + # Compute `transpose(A[q,:])`'s column counts. Store shifted forward one position in getcolptr(X). + fill!(getrowptr(X), 0) + for k in 1:nnz(A) + getrowptr(X)[colvals(A)[k]+1] += 1 + end + # Compute `transpose(A[q,:])`'s column pointers. Store shifted forward one position in getcolptr(X). + getrowptr(X)[1] = 1 + countsum = 1 + for k in 2:(size(A, 2) + 1) + overwritten = getrowptr(X)[k] + getrowptr(X)[k] = countsum + countsum += overwritten + end + @show getrowptr(X) +end """ Helper method for `halfperm!`. With `transpose(A[:,q])`'s column pointers shifted one position forward in `getcolptr(X)`, computes `map(f, transpose(A[:,q]))` by appropriately @@ -1402,6 +1576,22 @@ respectively. Simultaneously fixes the one-position-forward shift in `getcolptr( end return # kill potential type instability end +@noinline function _distributevals_halfperm!(X::AbstractSparseMatrixCSR{Tv,Ti}, + A::AbstractSparseMatrixCSR{TvA,Ti}, q::AbstractVector{<:Integer}, f::F) where {Tv,TvA,Ti,F<:Function} + resize!(nonzeros(X), nnz(A)) + resize!(colvals(X), nnz(A)) + @inbounds for Xj in 1:size(A, 1) + Ai = q[Xj] + for Ak in nzrange(A, Ai) + Aj = colvals(A)[Ak] + Xk = getrowptr(X)[Aj + 1] + colvals(X)[Xk] = Xj + nonzeros(X)[Xk] = f(nonzeros(A)[Ak]) + getrowptr(X)[Aj + 1] += 1 + end + end + return # kill potential type instability +end """ ftranspose!(X::AbstractSparseMatrixCSC{Tv,Ti}, A::AbstractSparseMatrixCSC{Tv,Ti}, f::Function) where {Tv,Ti} @@ -1433,6 +1623,7 @@ No additional memory is allocated other than resizing the rowval and nzval of `X See `halfperm!` """ transpose!(X::AbstractSparseMatrixCSC{Tv,Ti}, A::AbstractSparseMatrixCSC{Tv,Ti}) where {Tv,Ti} = ftranspose!(X, A, identity) +transpose!(X::AbstractSparseMatrixCSR{Tv,Ti}, A::AbstractSparseMatrixCSR{Tv,Ti}) where {Tv,Ti} = ftranspose!(X, A, identity) """ adjoint!(X::AbstractSparseMatrixCSC{Tv,Ti}, A::AbstractSparseMatrixCSC{Tv,Ti}) where {Tv,Ti} @@ -1454,18 +1645,37 @@ function ftranspose(A::AbstractSparseMatrixCSC{TvA,Ti}, f::Function, eltype::Typ sizehint!(X, nnz(A)) return @if_move_fixed A halfperm!(X, A, 1:size(A, 2), f) end +function ftranspose(A::MatrixType, f::Function, eltype::Type{Tv} = TvA) where {Tv,TvA,Ti,MatrixType<:AbstractSparseMatrixCSR{TvA,Ti}} + X = MatrixType(size(A, 2), size(A, 1), + ones(Ti, size(A, 2)+1), + Vector{Ti}(undef, 0), + Vector{Tv}(undef, 0)) + sizehint!(X, nnz(A)) + return halfperm!(X, A, 1:size(A, 1), f) +end adjoint(A::AbstractSparseMatrixCSC) = Adjoint(A) transpose(A::AbstractSparseMatrixCSC) = Transpose(A) +adjoint(A::AbstractSparseMatrixCSR) = Adjoint(A) +transpose(A::AbstractSparseMatrixCSR) = Transpose(A) Base.copy(A::Adjoint{<:Any,<:AbstractSparseMatrixCSC}) = ftranspose(A.parent, x -> adjoint(copy(x)), eltype(A)) Base.copy(A::Transpose{<:Any,<:AbstractSparseMatrixCSC}) = ftranspose(A.parent, x -> transpose(copy(x)), eltype(A)) +Base.copy(A::Adjoint{<:Any,<:AbstractSparseMatrixCSR}) = + ftranspose(A.parent, x -> adjoint(copy(x)), eltype(A)) +Base.copy(A::Transpose{<:Any,<:AbstractSparseMatrixCSR}) = + ftranspose(A.parent, x -> transpose(copy(x)), eltype(A)) function Base.permutedims(A::AbstractSparseMatrixCSC, (a,b)) (a, b) == (2, 1) && return ftranspose(A, identity) (a, b) == (1, 2) && return copy(A) throw(ArgumentError("no valid permutation of dimensions")) end +function Base.permutedims(A::AbstractSparseMatrixCSR, (a,b)) + (a, b) == (2, 1) && return ftranspose(A, identity) + (a, b) == (1, 2) && return copy(A) + throw(ArgumentError("no valid permutation of dimensions")) +end """ unchecked_noalias_permute!(X::AbstractSparseMatrixCSC{Tv,Ti}, @@ -2737,10 +2947,24 @@ end ((r1 > r2) || (rowvals(A)[r1] != i0)) ? zero(T) : nonzeros(A)[r1] end +@RCI @propagate_inbounds getindex(A::AbstractSparseMatrixCSR, I::Tuple{Integer,Integer}) = getindex(A, I[1], I[2]) + +@RCI @propagate_inbounds function getindex(A::AbstractSparseMatrixCSR{T}, i0::Integer, i1::Integer) where T + @boundscheck checkbounds(A, i0, i1) + c1 = Int(@inbounds getrowptr(A)[i0]) + c2 = Int(@inbounds getrowptr(A)[i0+1]-1) + (c1 > c2) && return zero(T) + c1 = searchsortedfirst(view(colvals(A), c1:c2), i1) + c1 - 1 + ((c1 > c2) || (colvals(A)[c1] != i1)) ? zero(T) : nonzeros(A)[c1] +end + # Colon translation getindex(A::AbstractSparseMatrixCSC, ::Colon, ::Colon) = copy(A) getindex(A::AbstractSparseMatrixCSC, i, ::Colon) = getindex(A, i, 1:size(A, 2)) getindex(A::AbstractSparseMatrixCSC, ::Colon, i) = getindex(A, 1:size(A, 1), i) +getindex(A::AbstractSparseMatrixCSR, ::Colon, ::Colon) = copy(A) +getindex(A::AbstractSparseMatrixCSR, i, ::Colon) = getindex(A, i, 1:size(A, 2)) +getindex(A::AbstractSparseMatrixCSR, ::Colon, i) = getindex(A, 1:size(A, 1), i) function getindex_cols(A::AbstractSparseMatrixCSC{Tv,Ti}, J::AbstractVector) where {Tv,Ti} require_one_based_indexing(A, J) @@ -2779,6 +3003,40 @@ end getindex_traverse_col(::AbstractUnitRange, lo::Integer, hi::Integer) = lo:hi getindex_traverse_col(I::StepRange, lo::Integer, hi::Integer) = step(I) > 0 ? (lo:1:hi) : (hi:-1:lo) +function getindex_rows(A::MatrixType, I::AbstractVector) where {Tv,Ti,MatrixType<:AbstractSparseMatrixCSR{Tv,Ti}} + require_one_based_indexing(A, I) + # for indexing whole columns + (m, n) = size(A) + nI = length(I) + + rowptrA = getrowptr(A); colvalA = colvals(A); nzvalA = nonzeros(A) + + rowptrS = Vector{Ti}(undef, nI+1) + rowptrS[1] = 1 + nnzS = 0 + + @inbounds for i = 1:nI + row = I[i] + 1 <= row <= n || throw(BoundsError()) + nnzS += rowptrA[row+1] - rowptrA[row] + rowptrS[i+1] = nnzS + 1 + end + + colvalS = Vector{Ti}(undef, nnzS) + nzvalS = Vector{Tv}(undef, nnzS) + ptrS = 0 + + @inbounds for i = 1:nI + row = I[i] + for k = rowptrA[row]:rowptrA[row+1]-1 + ptrS += 1 + colvalS[ptrS] = colvalA[k] + nzvalS[ptrS] = nzvalA[k] + end + end + return MatrixType(m, nI, rowptrS, colvalS, nzvalS) +end + function getindex(A::AbstractSparseMatrixCSC{Tv,Ti}, I::AbstractRange, J::AbstractVector) where {Tv,Ti<:Integer} require_one_based_indexing(A, I, J) # Ranges for indexing rows @@ -2827,6 +3085,54 @@ function getindex(A::AbstractSparseMatrixCSC{Tv,Ti}, I::AbstractRange, J::Abstra return @if_move_fixed A SparseMatrixCSC(nI, nJ, colptrS, rowvalS, nzvalS) end +function getindex(A::MatrixType, I::AbstractVector, J::AbstractRange) where {Tv,Ti<:Integer,MatrixType<:AbstractSparseMatrixCSR{Tv,Ti}} + require_one_based_indexing(A, I, J) + # Ranges for indexing rows + (m, n) = size(A) + # whole rows: + if J == 1:n + return getindex_rows(A, J) + end + + nI = length(I) + nI == 0 || (minimum(I) >= 1 && maximum(I) <= m) || throw(BoundsError()) + nJ = length(J) + colptrA = getcolptr(A); rowvalA = rowvals(A); nzvalA = nonzeros(A) + colptrS = Vector{Ti}(undef, nJ+1) + colptrS[1] = 1 + nnzS = 0 + + # Form the structure of the result and compute space + @inbounds for j = 1:nJ + col = J[j] + 1 <= col <= n || throw(BoundsError()) + @simd for k in colptrA[col]:colptrA[col+1]-1 + nnzS += rowvalA[k] in I # `in` is fast for ranges + end + colptrS[j+1] = nnzS+1 + end + + # Populate the values in the result + rowvalS = Vector{Ti}(undef, nnzS) + nzvalS = Vector{Tv}(undef, nnzS) + ptrS = 1 + + @inbounds for j = 1:nJ + col = J[j] + for k = getindex_traverse_col(I, colptrA[col], colptrA[col+1]-1) + rowA = rowvalA[k] + i = rangesearch(I, rowA) + if i > 0 + rowvalS[ptrS] = i + nzvalS[ptrS] = nzvalA[k] + ptrS += 1 + end + end + end + + return MatrixType(nI, nJ, colptrS, rowvalS, nzvalS) +end + function getindex_I_sorted(A::AbstractSparseMatrixCSC{Tv,Ti}, I::AbstractVector, J::AbstractVector) where {Tv,Ti} require_one_based_indexing(A, I, J) # Sorted vectors for indexing rows. diff --git a/test/testgroups b/test/testgroups index a547762c..363d8cf0 100644 --- a/test/testgroups +++ b/test/testgroups @@ -9,4 +9,5 @@ sparsematrix_constructors_indexing sparsematrix_ops sparsevector spqr +transposecscmat umfpack diff --git a/test/transposecscmat.jl b/test/transposecscmat.jl new file mode 100644 index 00000000..2814dad6 --- /dev/null +++ b/test/transposecscmat.jl @@ -0,0 +1,53 @@ +using LinearAlgebra, SparseArrays, Test + +# Test the CSR interface with a transposed SparseMatrixCSC +struct TransposeSparseMatrixCSC{Tv, Ti} <: SparseArrays.AbstractSparseMatrixCSR{Tv, Ti} + A::Transpose{Tv, SparseMatrixCSC{Tv, Ti}} +end + +function TransposeSparseMatrixCSC(A::SparseMatrixCSC) + return TransposeSparseMatrixCSC(Transpose(A)) +end + +function TransposeSparseMatrixCSC{Tv,Ti}(m::Integer, n::Integer, rowptr::Vector{Ti}, colval::Vector{Ti}, nzval::Vector{Tv}) where {Ti,Tv} + A = SparseMatrixCSC(n, m, rowptr, colval, nzval) + return TransposeSparseMatrixCSC(Transpose(A)) +end + +# Dispatches +SparseArrays.getrowptr(A::TransposeSparseMatrixCSC) = SparseArrays.getcolptr(A.A.parent) +SparseArrays.colvals(A::TransposeSparseMatrixCSC) = SparseArrays.rowvals(A.A.parent) +SparseArrays.size(A::TransposeSparseMatrixCSC) = size(A.A) +SparseArrays.nonzeros(A::TransposeSparseMatrixCSC) = nonzeros(A.A.parent) + +@testset "AbstractSparseMatrixCSC interface" begin + A = sprandn(4, 4, 0.5) + AT = TransposeSparseMatrixCSC(A) + # index + @test A[1,1:end] == AT[1:end,1] + @test A[1,1:end] == AT[1:end,1] + @test A[1,1:end] == AT[1:end,1] + @test A[2,2:3] == AT[2:3,2] + @test A[3,4] == AT[4,3] + @test any(A .== AT[1:end,1:end]) + @test A[1:end,1:end] == transpose(AT[1:end,1:end]) + + @test issparse(AT) + @test nnz(A) == nnz(AT) + + ATATT = AT * AT' + @test typeof(ATATT) <: SparseArrays.AbstractSparseMatrixCSR + @test all(ATATT .== A'A) + + ATTAT = AT' * AT + @test typeof(ATTAT) <: SparseArrays.AbstractSparseMatrixCSR + @test all(ATTAT .== A*A') + + A = TransposeSparseMatrixCSC(sprandn(3,2,0.5)) + B1 = TransposeSparseMatrixCSC(sprandn(3,4,0.5)) + B2 = TransposeSparseMatrixCSC(sprandn(3,3,0.5)) + B3 = TransposeSparseMatrixCSC(sprandn(4,3,0.5)) + @test A*B1' == Matrix(A)*Matrix(B1') + @test A*B2 == Matrix(A)*Matrix(B2) + @test A*B3 == Matrix(A)*Matrix(B3) +end