diff --git a/stdlib/SparseArrays/src/SparseArrays.jl b/stdlib/SparseArrays/src/SparseArrays.jl index f43ff255e003b..55f112048ee05 100644 --- a/stdlib/SparseArrays/src/SparseArrays.jl +++ b/stdlib/SparseArrays/src/SparseArrays.jl @@ -36,8 +36,11 @@ export AbstractSparseArray, AbstractSparseMatrix, AbstractSparseVector, issparse, nonzeros, nzrange, rowvals, sparse, sparsevec, spdiagm, sprand, sprandn, spzeros, nnz, permute, findnz +export unwrap, iswrsparse + include("abstractsparse.jl") include("sparsematrix.jl") +include("sparseconvert.jl") include("sparsevector.jl") include("higherorderfns.jl") include("linalg.jl") diff --git a/stdlib/SparseArrays/src/sparseconvert.jl b/stdlib/SparseArrays/src/sparseconvert.jl new file mode 100644 index 0000000000000..bdadc3a4729ff --- /dev/null +++ b/stdlib/SparseArrays/src/sparseconvert.jl @@ -0,0 +1,288 @@ +# This file is a part of Julia. License is MIT: https://julialang.org/license + +import LinearAlgebra: AbstractTriangular + +""" + SparseMatrixCSCSymmHerm + +`Symmetric` or `Hermitian` of a `SparseMatrixCSC` or `SparseMatrixCSCView`. +""" +const SparseMatrixCSCSymmHerm{Tv,Ti} = Union{Symmetric{Tv,<:SparseMatrixCSCUnion{Tv,Ti}}, + Hermitian{Tv,<:SparseMatrixCSCUnion{Tv,Ti}}} + +const AbstractTriangularSparse{Tv,Ti} = AbstractTriangular{Tv,<:SparseMatrixCSCUnion{Tv,Ti}} + +# converting Symmetric/Hermitian/AbstractTriangular/SubArray of SparseMatrixCSC +# and Transpose/Adjoint of AbstractTriangular of SparseMatrixCSC to SparseMatrixCSC +for wr in (Symmetric, Hermitian, Transpose, Adjoint, + UpperTriangular, LowerTriangular, UnitUpperTriangular, UnitLowerTriangular, + SubArray) + + @eval SparseMatrixCSC(A::$wr) = _sparsem(A) + @eval SparseMatrixCSC{Tv}(A::$wr{Tv}) where Tv = _sparsem(A) + @eval SparseMatrixCSC{Tv}(A::$wr) where Tv = SparseMatrixCSC{Tv}(_sparsem(A)) + @eval SparseMatrixCSC{Tv,Ti}(A::$wr) where {Tv,Ti} = SparseMatrixCSC{Tv,Ti}(_sparsem(A)) +end + +""" + iswrsparse(::S) + iswrsparse(::Type{S}) + +Returns `true` if type `S` is backed by a sparse array, and `false` otherwise. +""" +iswrsparse(::T) where T<:AbstractArray = iswrsparse(T) +iswrsparse(::Type) = false +iswrsparse(::Type{T}) where T<:AbstractSparseArray = true + +""" + depth(::Type{S}) + +Returns 0 for unwrapped S, and nesting depth for wrapped (nested) abstract arrays. +""" +depth(::T) where T = depth(T) +depth(::Type{T}) where T<:AbstractArray = 0 + +for wr in (Symmetric, Hermitian, + LowerTriangular, UnitLowerTriangular, UpperTriangular, UnitUpperTriangular, + Transpose, Adjoint, SubArray, + Diagonal, Bidiagonal, Tridiagonal, SymTridiagonal) + + pl = wr === SubArray ? :($wr{<:Any,<:Any,T}) : :($wr{<:Any,T}) + @eval iswrsparse(::Type{<:$pl}) where T = iswrsparse(T) + @eval depth(::Type{<:$pl}) where T = depth(T) + 1 +end + +# convert parent and re-wrap in same wrapper +_sparsewrap(A::Symmetric) = Symmetric(_sparsem(parent(A)), A.uplo == 'U' ? :U : :L) +_sparsewrap(A::Hermitian) = Hermitian(_sparsem(parent(A)), A.uplo == 'U' ? :U : :L) +_sparsewrap(A::SubArray) = SubArray(_sparsem(parent(A)), A.indices) +for ty in ( LowerTriangular, UnitLowerTriangular, + UpperTriangular, UnitUpperTriangular, + Transpose, Adjoint) + + @eval _sparsewrap(A::$ty) = $ty(_sparsem(parent(A))) +end +function _sparsewrap(A::Union{Diagonal,Bidiagonal,Tridiagonal,SymTridiagonal}) + dropzeros!(sparse(A)) +end + +""" + unwrap(A::AbstractMatrix) + +In case A is a wrapper type (`SubArray, Symmetric, Adjoint, SubArray, Triangular, Tridiagonal`, etc.) +convert to `Matrix` or `SparseMatrixCSC`, depending on final storage type of A. +For other types return A itself. +""" +unwrap(A::Any) = A +unwrap(A::AbstractMatrix) = iswrsparse(A) ? convert(SparseMatrixCSC, A) : convert(Array, A) + +import Base.copy +copy(A::SubArray{T,2}) where T = getindex(unwrap(parent(A)), A.indices...) + +# For pure sparse matrices and vectors return A. +# For wrapped sparse matrices or vectors convert to SparseMatrixCSC. +# Handle nested wrappers properly. +# Use abstract matrix fallback if A is not sparse. +function _sparsem(@nospecialize A::AbstractArray{Tv}) where Tv + if iswrsparse(A) + if depth(A) >= 1 + _sparsem(_sparsewrap(A)) + else + A + end + else + # explicitly call abstract matrix fallback using getindex(A,...) + invoke(SparseMatrixCSC{Tv,Int}, Tuple{AbstractMatrix}, A) + end +end + +_sparsem(A::AbstractSparseMatrix) = A +_sparsem(A::AbstractSparseVector) = A + +# Transpose/Adjoint of sparse vector (returning sparse matrix) +function _sparsem(A::Union{Transpose{<:Any,<:AbstractSparseVector},Adjoint{<:Any,<:AbstractSparseVector}}) + B = parent(A) + n = length(B) + Ti = eltype(B.nzind) + fadj = A isa Transpose ? transpose : adjoint + colptr = fill!(Vector{Ti}(undef, n + 1), 0) + colptr[1] = 1 + colptr[B.nzind .+ 1] .= 1 + cumsum!(colptr, colptr) + rowval = fill!(similar(B.nzind), 1) + nzval = fadj.(nonzeros(B)) + SparseMatrixCSC(1, n, colptr, rowval, nzval) +end + +function _sparsem(A::Union{Transpose{<:Any,<:SparseMatrixCSC},Adjoint{<:Any,<:SparseMatrixCSC}}) + ftranspose(parent(A), A isa Transpose ? transpose : adjoint) +end + +# Symmetric/Hermitian of sparse matrix +_sparsem(A::SparseMatrixCSCSymmHerm) = _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) +# view of sparse matrix +_sparsem(S::SubArray{<:Any,2,<:SparseMatrixCSC}) = getindex(S.parent,S.indices...) + +# 4 cases: (Symmetric|Hermitian) variants (:U|:L) +function _sparsem(fnzrange::Function, sA::SparseMatrixCSCSymmHerm{Tv}) where {Tv} + A = sA.data + rowval = rowvals(A) + nzval = nonzeros(A) + m, n = size(A) + Ti = eltype(rowval) + fadj = sA isa Symmetric ? transpose : adjoint + newcolptr = Vector{Ti}(undef, n+1) + diagmap = fadj == transpose ? identity : real + + newcolptr[1] = 1 + colrange = fnzrange === nzrangeup ? (1:n) : (n:-1:1) + @inbounds for j = colrange + r = fnzrange(A, j); r1 = r.start; r2 = r.stop + newcolptr[j+1] = r2 - r1 + 1 + for k = r1:r2 + row = rowval[k] + if row != j + newcolptr[row+1] += 1 + end + end + end + cumsum!(newcolptr, newcolptr) + nz = newcolptr[n+1] - 1 + newrowval = Vector{Ti}(undef, nz) + newnzval = Vector{Tv}(undef, nz) + @inbounds for j = 1:n + newk = newcolptr[j] + for k = fnzrange(A, j) + i = rowval[k] + nzv = nzval[k] + if i != j + newrowval[newk] = i + newnzval[newk] = nzv + newk += 1 + ni = newcolptr[i] + newrowval[ni] = j + newnzval[ni] = fadj(nzv) + newcolptr[i] = ni + 1 + else + newrowval[newk] = i + newnzval[newk] = diagmap(nzv) + newk += 1 + end + end + newcolptr[j] = newk + end + _sparse_gen(m, n, newcolptr, newrowval, newnzval) +end + +# 2 cases: Unit(Upper|Lower)Triangular{Tv,SparseMatrixCSC} +function _sparsem(A::AbstractTriangularSparse{Tv}) where Tv + S = A.data + rowval = rowvals(S) + nzval = nonzeros(S) + m, n = size(S) + Ti = eltype(rowval) + fnzrange = A isa Union{UpperTriangular,UnitUpperTriangular} ? nzrangeup : nzrangelo + unit = A isa Union{UnitUpperTriangular,UnitLowerTriangular} + newcolptr = Vector{Ti}(undef, n+1) + newrowval = Vector{Ti}(undef, nnz(S)) + newnzval = Vector{Tv}(undef, nnz(S)) + newcolptr[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 = rowval[k] + if i != j || i == j && !unit + newrowval[newk] = i + newnzval[newk] = nzval[k] + newk += 1 + end + end + if unit + uplo && (newkk = newk) + newrowval[newkk] = j + newnzval[newkk] = one(Tv) + newk += uplo + end + newcolptr[j+1] = newk + end + nz = newcolptr[n+1] - 1 + resize!(newrowval, nz) + resize!(newnzval, nz) + SparseMatrixCSC(m, n, newcolptr, newrowval, newnzval) +end + +# 8 cases: (Transpose|Adjoint){Tv,[Unit](Upper|Lower)Triangular} +function _sparsem(taA::Union{Transpose{Tv,<:AbstractTriangularSparse}, + Adjoint{Tv,<:AbstractTriangularSparse}}) where {Tv} + + sA = taA.parent + A = sA.data + rowval = rowvals(A) + nzval = nonzeros(A) + m, n = size(A) + Ti = eltype(rowval) + 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} + + newcolptr = Vector{Ti}(undef, n+1) + fill!(newcolptr, 1unit) + newcolptr[1] = 1 + @inbounds for j = 1:n + for k = fnzrange(A, j) + i = rowval[k] + if i != j || i == j && !unit + newcolptr[i+1] += 1 + end + end + end + cumsum!(newcolptr, newcolptr) + nz = newcolptr[n+1] - 1 + newrowval = Vector{Ti}(undef, nz) + newnzval = Vector{Tv}(undef, nz) + + @inbounds for j = 1:n + if !uplo && unit + ni = newcolptr[j] + newrowval[ni] = j + newnzval[ni] = fadj(one(Tv)) + newcolptr[j] = ni + 1 + end + for k = fnzrange(A, j) + i = rowval[k] + nzv = nzval[k] + if i != j || i == j && !unit + ni = newcolptr[i] + newrowval[ni] = j + newnzval[ni] = fadj(nzv) + newcolptr[i] = ni + 1 + end + end + if uplo && unit + ni = newcolptr[j] + newrowval[ni] = j + newnzval[ni] = fadj(one(Tv)) + newcolptr[j] = ni + 1 + end + end + _sparse_gen(n, m, newcolptr, newrowval, newnzval) +end + +function _sparse_gen(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) +end + diff --git a/stdlib/SparseArrays/test/sparse.jl b/stdlib/SparseArrays/test/sparse.jl index 062c7b8870715..afadd1540e984 100644 --- a/stdlib/SparseArrays/test/sparse.jl +++ b/stdlib/SparseArrays/test/sparse.jl @@ -2510,6 +2510,42 @@ end end end +@testset "wrappers of sparse" begin + m = n = 10 + A = spzeros(ComplexF64, m, n) + A[:,1] = 1:m + A[:,2] = [1 3 0 0 0 0 0 0 0 0]' + A[:,3] = [2 4 0 0 0 0 0 0 0 0]' + A[:,4] = [0 0 0 0 5 3 0 0 0 0]' + A[:,5] = [0 0 0 0 6 2 0 0 0 0]' + A[:,6] = [0 0 0 0 7 4 0 0 0 0]' + A[:,7:n] = rand(ComplexF64, m, n-6) + B = Matrix(A) + dowrap(wr, A) = wr(A) + dowrap(wr::Tuple, A) = (wr[1])(A, wr[2:end]...) + + @testset "sparse($wr(A))" for wr in ( + Symmetric, (Symmetric, :L), Hermitian, (Hermitian, :L), + Transpose, Adjoint, + UpperTriangular, LowerTriangular, + UnitUpperTriangular, UnitLowerTriangular, + (view, 3:6, 2:5)) + + @test SparseMatrixCSC(dowrap(wr, A)) == Matrix(dowrap(wr, B)) + end + + @testset "sparse($at($wr))" for at = (Transpose, Adjoint), wr = + (UpperTriangular, LowerTriangular, + UnitUpperTriangular, UnitLowerTriangular) + + @test SparseMatrixCSC(at(wr(A))) == Matrix(at(wr(B))) + end + + @test sparse([1,2,3,4,5]') == SparseMatrixCSC([1 2 3 4 5]) + @test sparse(UpperTriangular(A')) == UpperTriangular(B') + @test sparse(Adjoint(UpperTriangular(A'))) == Adjoint(UpperTriangular(B')) +end + @testset "Ti cannot store all potential values #31024" begin @test_throws ArgumentError SparseMatrixCSC(128, 1, [Int8(1), Int8(1)], Int8[], Int[]) @test_throws ArgumentError SparseMatrixCSC(12, 12, [Int8(1), Int8(1)], Int8[], Int[])