Skip to content

Commit

Permalink
Refactoring: Add AbstractSparseMatrixCSC (#33039)
Browse files Browse the repository at this point in the history
* Add AbstractSparseMatrixCSC

* Replace ::SparseMatrixCSC -> ::AbstractSparseMatrixCSC if reasonable

* Replace <:SparseMatrixCSC -> <:AbstractSparseMatrixCSC

* Use AbstractSparseMatrixCSC in more dispatches
  • Loading branch information
tkf authored and ViralBShah committed Aug 24, 2019
1 parent d1979e3 commit da9685b
Show file tree
Hide file tree
Showing 6 changed files with 359 additions and 352 deletions.
7 changes: 7 additions & 0 deletions stdlib/SparseArrays/src/abstractsparse.jl
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,13 @@ of type `Tv` and index type `Ti`. Alias for `AbstractSparseArray{Tv,Ti,2}`.
"""
const AbstractSparseMatrix{Tv,Ti} = AbstractSparseArray{Tv,Ti,2}

"""
AbstractSparseMatrixCSC{Tv,Ti<:Integer} <: AbstractSparseMatrix{Tv,Ti}
Supertype for matrix with compressed sparse column (CSC).
"""
abstract type AbstractSparseMatrixCSC{Tv,Ti<:Integer} <: AbstractSparseMatrix{Tv,Ti} end

"""
issparse(S)
Expand Down
44 changes: 22 additions & 22 deletions stdlib/SparseArrays/src/higherorderfns.jl
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ module HigherOrderFns
import Base: map, map!, broadcast, copy, copyto!

using Base: front, tail, to_shape
using ..SparseArrays: SparseVector, SparseMatrixCSC, AbstractSparseVector,
using ..SparseArrays: SparseVector, SparseMatrixCSC, AbstractSparseVector, AbstractSparseMatrixCSC,
AbstractSparseMatrix, AbstractSparseArray, indtype, nnz, nzrange,
SparseVectorUnion, AdjOrTransSparseVectorUnion, nonzeroinds, nonzeros, rowvals, getcolptr
using Base.Broadcast: BroadcastStyle, Broadcasted, flatten
Expand All @@ -33,13 +33,13 @@ using LinearAlgebra

# (0) BroadcastStyle rules and convenience types for dispatch

SparseVecOrMat = Union{SparseVector,SparseMatrixCSC}
SparseVecOrMat = Union{SparseVector,AbstractSparseMatrixCSC}

# broadcast container type promotion for combinations of sparse arrays and other types
struct SparseVecStyle <: Broadcast.AbstractArrayStyle{1} end
struct SparseMatStyle <: Broadcast.AbstractArrayStyle{2} end
Broadcast.BroadcastStyle(::Type{<:SparseVector}) = SparseVecStyle()
Broadcast.BroadcastStyle(::Type{<:SparseMatrixCSC}) = SparseMatStyle()
Broadcast.BroadcastStyle(::Type{<:AbstractSparseMatrixCSC}) = SparseMatStyle()
const SPVM = Union{SparseVecStyle,SparseMatStyle}

# SparseVecStyle handles 0-1 dimensions, SparseMatStyle 0-2 dimensions.
Expand Down Expand Up @@ -109,23 +109,23 @@ const SpBroadcasted2{Style<:SPVM,Axes,F,Args<:Tuple{SparseVecOrMat,SparseVecOrMa
# as n-by-one sparse matrices which, though technically incorrect, is how broacast[!] views
# sparse vectors in practice.
@inline numrows(A::SparseVector) = length(A)
@inline numrows(A::SparseMatrixCSC) = size(A, 1)
@inline numrows(A::AbstractSparseMatrixCSC) = size(A, 1)
@inline numcols(A::SparseVector) = 1
@inline numcols(A::SparseMatrixCSC) = size(A, 2)
@inline numcols(A::AbstractSparseMatrixCSC) = size(A, 2)
# numrows and numcols respectively yield size(A, 1) and size(A, 2), but avoid a branch
@inline columns(A::SparseVector) = 1
@inline columns(A::SparseMatrixCSC) = 1:size(A, 2)
@inline columns(A::AbstractSparseMatrixCSC) = 1:size(A, 2)
@inline colrange(A::SparseVector, j) = 1:length(nonzeroinds(A))
@inline colrange(A::SparseMatrixCSC, j) = nzrange(A, j)
@inline colrange(A::AbstractSparseMatrixCSC, j) = nzrange(A, j)
@inline colstartind(A::SparseVector, j) = one(indtype(A))
@inline colboundind(A::SparseVector, j) = convert(indtype(A), length(nonzeroinds(A)) + 1)
@inline colstartind(A::SparseMatrixCSC, j) = getcolptr(A)[j]
@inline colboundind(A::SparseMatrixCSC, j) = getcolptr(A)[j + 1]
@inline colstartind(A::AbstractSparseMatrixCSC, j) = getcolptr(A)[j]
@inline colboundind(A::AbstractSparseMatrixCSC, j) = getcolptr(A)[j + 1]
@inline storedinds(A::SparseVector) = nonzeroinds(A)
@inline storedinds(A::SparseMatrixCSC) = rowvals(A)
@inline storedinds(A::AbstractSparseMatrixCSC) = rowvals(A)
@inline storedvals(A::SparseVecOrMat) = nonzeros(A)
@inline setcolptr!(A::SparseVector, j, val) = val
@inline setcolptr!(A::SparseMatrixCSC, j, val) = getcolptr(A)[j] = val
@inline setcolptr!(A::AbstractSparseMatrixCSC, j, val) = getcolptr(A)[j] = val
function trimstorage!(A::SparseVecOrMat, maxstored)
resize!(storedinds(A), maxstored)
resize!(storedvals(A), maxstored)
Expand All @@ -140,12 +140,12 @@ end

# (2) map[!] entry points
map(f::Tf, A::SparseVector) where {Tf} = _noshapecheck_map(f, A)
map(f::Tf, A::SparseMatrixCSC) where {Tf} = _noshapecheck_map(f, A)
map(f::Tf, A::SparseMatrixCSC, Bs::Vararg{SparseMatrixCSC,N}) where {Tf,N} =
map(f::Tf, A::AbstractSparseMatrixCSC) where {Tf} = _noshapecheck_map(f, A)
map(f::Tf, A::AbstractSparseMatrixCSC, Bs::Vararg{SparseMatrixCSC,N}) where {Tf,N} =
(_checksameshape(A, Bs...); _noshapecheck_map(f, A, Bs...))
map(f::Tf, A::SparseVecOrMat, Bs::Vararg{SparseVecOrMat,N}) where {Tf,N} =
(_checksameshape(A, Bs...); _noshapecheck_map(f, A, Bs...))
map!(f::Tf, C::SparseMatrixCSC, A::SparseMatrixCSC, Bs::Vararg{SparseMatrixCSC,N}) where {Tf,N} =
map!(f::Tf, C::AbstractSparseMatrixCSC, A::AbstractSparseMatrixCSC, Bs::Vararg{SparseMatrixCSC,N}) where {Tf,N} =
(_checksameshape(C, A, Bs...); _noshapecheck_map!(f, C, A, Bs...))
map!(f::Tf, C::SparseVecOrMat, A::SparseVecOrMat, Bs::Vararg{SparseVecOrMat,N}) where {Tf,N} =
(_checksameshape(C, A, Bs...); _noshapecheck_map!(f, C, A, Bs...))
Expand Down Expand Up @@ -216,7 +216,7 @@ end
@inline _densennz(shape::NTuple{2}) = shape[1] * shape[2]
_maxnnzfrom(shape::NTuple{1}, A::SparseVector) = nnz(A) * div(shape[1], length(A))
_maxnnzfrom(shape::NTuple{2}, A::SparseVector) = nnz(A) * div(shape[1], length(A)) * shape[2]
_maxnnzfrom(shape::NTuple{2}, A::SparseMatrixCSC) = nnz(A) * div(shape[1], size(A, 1)) * div(shape[2], size(A, 2))
_maxnnzfrom(shape::NTuple{2}, A::AbstractSparseMatrixCSC) = nnz(A) * div(shape[1], size(A, 1)) * div(shape[2], size(A, 2))
@inline _maxnnzfrom_each(shape, ::Tuple{}) = ()
@inline _maxnnzfrom_each(shape, As) = (_maxnnzfrom(shape, first(As)), _maxnnzfrom_each(shape, tail(As))...)
@inline _unchecked_maxnnzbcres(shape, As::Tuple) = min(_densennz(shape), sum(_maxnnzfrom_each(shape, As)))
Expand Down Expand Up @@ -277,13 +277,13 @@ function _map_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, A::SparseVecOrMa
end
# helper functions for these methods and some of those below
@inline _densecoloffsets(A::SparseVector) = 0
@inline _densecoloffsets(A::SparseMatrixCSC) = 0:size(A, 1):(size(A, 1)*(size(A, 2) - 1))
@inline _densecoloffsets(A::AbstractSparseMatrixCSC) = 0:size(A, 1):(size(A, 1)*(size(A, 2) - 1))
function _densestructure!(A::SparseVector)
expandstorage!(A, length(A))
copyto!(nonzeroinds(A), 1:length(A))
return A
end
function _densestructure!(A::SparseMatrixCSC)
function _densestructure!(A::AbstractSparseMatrixCSC)
nnzA = size(A, 1) * size(A, 2)
expandstorage!(A, nnzA)
copyto!(getcolptr(A), 1:size(A, 1):(nnzA + 1))
Expand Down Expand Up @@ -812,7 +812,7 @@ function _broadcast_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, A::SparseV
return C
end
_finishempty!(C::SparseVector) = C
_finishempty!(C::SparseMatrixCSC) = (fill!(getcolptr(C), 1); C)
_finishempty!(C::AbstractSparseMatrixCSC) = (fill!(getcolptr(C), 1); C)

# special case - vector outer product
_copy(f::typeof(*), x::SparseVectorUnion, y::AdjOrTransSparseVectorUnion) = _outer(x, y)
Expand Down Expand Up @@ -1012,7 +1012,7 @@ end
end

_copy(f, args::SparseVector...) = _shapecheckbc(f, args...)
_copy(f, args::SparseMatrixCSC...) = _shapecheckbc(f, args...)
_copy(f, args::AbstractSparseMatrixCSC...) = _shapecheckbc(f, args...)
_copy(f, args::SparseVecOrMat...) = _diffshape_broadcast(f, args...)
# Otherwise, we incorporate scalars into the function and re-dispatch
function _copy(f, args...)
Expand Down Expand Up @@ -1112,8 +1112,8 @@ end
end

# NOTE: The following two method definitions work around #19096.
broadcast(f::Tf, ::Type{T}, A::SparseMatrixCSC) where {Tf,T} = broadcast(y -> f(T, y), A)
broadcast(f::Tf, A::SparseMatrixCSC, ::Type{T}) where {Tf,T} = broadcast(x -> f(x, T), A)
broadcast(f::Tf, ::Type{T}, A::AbstractSparseMatrixCSC) where {Tf,T} = broadcast(y -> f(T, y), A)
broadcast(f::Tf, A::AbstractSparseMatrixCSC, ::Type{T}) where {Tf,T} = broadcast(x -> f(x, T), A)


# (11) broadcast[!] over combinations of scalars, sparse vectors/matrices, structured matrices,
Expand Down Expand Up @@ -1151,7 +1151,7 @@ _sparsifystructured(x) = x
SparseOrStructuredMatrix = Union{SparseMatrixCSC,LinearAlgebra.StructuredMatrix}
map(f::Tf, A::SparseOrStructuredMatrix, Bs::Vararg{SparseOrStructuredMatrix,N}) where {Tf,N} =
(_checksameshape(A, Bs...); _noshapecheck_map(f, _sparsifystructured(A), map(_sparsifystructured, Bs)...))
map!(f::Tf, C::SparseMatrixCSC, A::SparseOrStructuredMatrix, Bs::Vararg{SparseOrStructuredMatrix,N}) where {Tf,N} =
map!(f::Tf, C::AbstractSparseMatrixCSC, A::SparseOrStructuredMatrix, Bs::Vararg{SparseOrStructuredMatrix,N}) where {Tf,N} =
(_checksameshape(C, A, Bs...); _noshapecheck_map!(f, C, _sparsifystructured(A), map(_sparsifystructured, Bs)...))

end
Loading

0 comments on commit da9685b

Please sign in to comment.