Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

extend sparse broadcast to one- and two-dimensional Arrays, better version #20009

Merged
merged 1 commit into from
Jan 25, 2017
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
124 changes: 90 additions & 34 deletions base/sparse/higherorderfns.jl
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,8 @@ using ..SparseArrays: SparseVector, SparseMatrixCSC, AbstractSparseArray, indtyp
# (8) Define _broadcast_[not]zeropres! specialized for a pair of (input) sparse vectors/matrices.
# (9) Define general _broadcast_[not]zeropres! capable of handling >2 (input) sparse vectors/matrices.
# (10) Define (broadcast[!]) methods handling combinations of broadcast scalars and sparse vectors/matrices.
# (11) Define (broadcast[!]) methods handling combinations of scalars, sparse vectors/matrices, and structured matrices.
# (11) Define (broadcast[!]) methods handling combinations of scalars, sparse vectors/matrices,
# structured matrices, and one- and two-dimensional Arrays.
# (12) Define (map[!]) methods handling combinations of sparse and structured matrices.


Expand Down Expand Up @@ -944,10 +945,9 @@ _containertype{T<:SparseVecOrMat}(::Type{T}) = AbstractSparseArray
# combinations of sparse arrays with broadcast scalars should yield sparse arrays
promote_containertype(::Type{Any}, ::Type{AbstractSparseArray}) = AbstractSparseArray
promote_containertype(::Type{AbstractSparseArray}, ::Type{Any}) = AbstractSparseArray
# combinations of sparse arrays with anything else should fall back to generic dense broadcast
promote_containertype(::Type{Array}, ::Type{AbstractSparseArray}) = Array
# combinations of sparse arrays with tuples should divert to the generic AbstractArray broadcast code
# (we handle combinations involving dense vectors/matrices below)
promote_containertype(::Type{Tuple}, ::Type{AbstractSparseArray}) = Array
promote_containertype(::Type{AbstractSparseArray}, ::Type{Array}) = Array
promote_containertype(::Type{AbstractSparseArray}, ::Type{Tuple}) = Array

# broadcast[!] entry points for combinations of sparse arrays and other (scalar) types
Expand Down Expand Up @@ -986,41 +986,97 @@ broadcast{Tf,T}(f::Tf, ::Type{T}, A::SparseMatrixCSC) = broadcast(y -> f(T, y),
broadcast{Tf,T}(f::Tf, A::SparseMatrixCSC, ::Type{T}) = broadcast(x -> f(x, T), A)


# (11) broadcast[!] over combinations of scalars, sparse vectors/matrices, and structured matrices
# (11) broadcast[!] over combinations of scalars, sparse vectors/matrices, structured matrices,
# and one- and two-dimensional Arrays (via promotion of structured matrices and Arrays)
#
# for combinations involving only scalars, sparse arrays, structured matrices, and dense
# vectors/matrices, promote all structured matrices and dense vectors/matrices to sparse
# and rebroadcast. otherwise, divert to generic AbstractArray broadcast code.
#
# this requires three steps: segregate combinations to promote to sparse via Broadcast's
# containertype promotion and dispatch layer (broadcast_c[!], containertype,
# promote_containertype), separate ambiguous cases from the preceding dispatch
# layer in sparse broadcast's internal containertype promotion and dispatch layer
# (spbroadcast_c[!], spcontainertype, promote_spcontainertype), and then promote
# arguments to sparse as appropriate and rebroadcast.

# structured array container type promotion
immutable StructuredArray end
_containertype{T<:Diagonal}(::Type{T}) = StructuredArray
_containertype{T<:Bidiagonal}(::Type{T}) = StructuredArray
_containertype{T<:Tridiagonal}(::Type{T}) = StructuredArray
_containertype{T<:SymTridiagonal}(::Type{T}) = StructuredArray
promote_containertype(::Type{StructuredArray}, ::Type{StructuredArray}) = StructuredArray
# combinations involving sparse arrays continue in the structured array funnel
promote_containertype(::Type{StructuredArray}, ::Type{AbstractSparseArray}) = StructuredArray
promote_containertype(::Type{AbstractSparseArray}, ::Type{StructuredArray}) = StructuredArray
# combinations involving scalars continue in the structured array funnel
promote_containertype(::Type{StructuredArray}, ::Type{Any}) = StructuredArray
promote_containertype(::Type{Any}, ::Type{StructuredArray}) = StructuredArray
# combinations involving arrays divert to the generic array code
promote_containertype(::Type{StructuredArray}, ::Type{Array}) = Array
promote_containertype(::Type{Array}, ::Type{StructuredArray}) = Array
# combinations involving tuples divert to the generic array code
promote_containertype(::Type{StructuredArray}, ::Type{Tuple}) = Array
promote_containertype(::Type{Tuple}, ::Type{StructuredArray}) = Array

# for combinations involving sparse/structured arrays and scalars only,
# promote all structured arguments to sparse and then rebroadcast
@inline broadcast_c{N}(f, ::Type{StructuredArray}, As::Vararg{Any,N}) =
# first (Broadcast containertype) dispatch layer's promotion logic
immutable PromoteToSparse end

# broadcast containertype definitions for structured matrices
StructuredMatrix = Union{Diagonal,Bidiagonal,Tridiagonal,SymTridiagonal}
_containertype{T<:StructuredMatrix}(::Type{T}) = PromoteToSparse

# combinations explicitly involving Tuples and PromoteToSparse collections
# divert to the generic AbstractArray broadcast code
promote_containertype(::Type{PromoteToSparse}, ::Type{Tuple}) = Array
promote_containertype(::Type{Tuple}, ::Type{PromoteToSparse}) = Array
# combinations involving scalars and PromoteToSparse collections continue in the promote-to-sparse funnel
promote_containertype(::Type{PromoteToSparse}, ::Type{Any}) = PromoteToSparse
promote_containertype(::Type{Any}, ::Type{PromoteToSparse}) = PromoteToSparse
# combinations involving sparse arrays and PromoteToSparse collections continue in the promote-to-sparse funnel
promote_containertype(::Type{PromoteToSparse}, ::Type{AbstractSparseArray}) = PromoteToSparse
promote_containertype(::Type{AbstractSparseArray}, ::Type{PromoteToSparse}) = PromoteToSparse
# combinations involving Arrays and PromoteToSparse collections continue in the promote-to-sparse funnel
promote_containertype(::Type{PromoteToSparse}, ::Type{Array}) = PromoteToSparse
promote_containertype(::Type{Array}, ::Type{PromoteToSparse}) = PromoteToSparse
# combinations involving Arrays and sparse arrays continue in the promote-to-sparse funnel
promote_containertype(::Type{AbstractSparseArray}, ::Type{Array}) = PromoteToSparse
promote_containertype(::Type{Array}, ::Type{AbstractSparseArray}) = PromoteToSparse

# second (internal sparse broadcast containertype) dispatch layer's promotion logic
# mostly just disambiguates Array from the main containertype promotion mechanism
# AbstractArray serves as a marker to shunt to the generic AbstractArray broadcast code
_spcontainertype(x) = _containertype(x)
_spcontainertype{T<:Vector}(::Type{T}) = Vector
_spcontainertype{T<:Matrix}(::Type{T}) = Matrix
_spcontainertype{T<:Ref}(::Type{T}) = AbstractArray
_spcontainertype{T<:AbstractArray}(::Type{T}) = AbstractArray
# need the following two methods to override the immediately preceding method
_spcontainertype{T<:StructuredMatrix}(::Type{T}) = PromoteToSparse
_spcontainertype{T<:SparseVecOrMat}(::Type{T}) = AbstractSparseArray
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should this be only concrete SparseVecOrMat, or cover the entire corresponding abstract types?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

At present the outer (i.e. Broadcast's) container type promotion mechanism only funnels SparseVector and SparseMatrixCSC to sparse broadcast[!], so the inner (sparse broadcast[!]'s) container type promotion mechanism only need handle SparseVector and SparseMatrixCSC. Not capturing other <:AbstractSparseArray is intentional: sparse broadcast[!] isn't built to handle <:AbstractSparseArray that are not SparseVector or SparseMatrixCSC, there being no established interface for such <:AbstractSparseArrays.

If we want sparse broadcast[!] to handle combinations involving non-SparseVecOrMat <:AbstractSparseArray at this time, I can add a commit or open a separate PR promoting non-SparseVecOrMat <:AbstractSparseArray to SparseVector or SparseMatrixCSC as appropriate (as with structured matrices). Thoughts? Thanks!

Copy link
Contributor

@tkelman tkelman Jan 25, 2017

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should think through how the AbstractSparse types are used a bit, maybe not right here right now though. Anything that's assuming the exact CSC or concrete SparseVector data structures, fields, storage order etc should only be defined for those concrete types. But any other behavior we think should generically apply for all sparse arrays, I think it's worth widening some signatures and letting the definers of any concrete subtypes override more specific behaviors where they want.

I guess in this particular case the tradeoff would be between getting method errors vs dispatching to more generic dense behavior? The latter's likely better for now until we work out a better extensibility mechanism for how to get broadcast on custom array types to do interesting things.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I guess in this particular case the tradeoff would be between getting method errors vs dispatching to more generic dense behavior? The latter's likely better for now until we work out a better extensibility mechanism for how to get broadcast on custom array types to do interesting things.

Yes, and agreed. (Though for the promotion approach to work, definition of convert([SparseVector|SparseMatrixCSC], UnknownSparseVectorOrMatrixType) alone should suffice.) Best!

spcontainertype(x) = _spcontainertype(typeof(x))
spcontainertype(ct1, ct2) = promote_spcontainertype(spcontainertype(ct1), spcontainertype(ct2))
@inline spcontainertype(ct1, ct2, cts...) = promote_spcontainertype(spcontainertype(ct1), spcontainertype(ct2, cts...))

promote_spcontainertype{T}(::Type{T}, ::Type{T}) = T
# combinations involving AbstractArrays and/or Tuples divert to the generic AbstractArray broadcast code
DivertToAbsArrayBC = Union{Type{AbstractArray},Type{Tuple}}
promote_spcontainertype{T<:DivertToAbsArrayBC}(::T, ct) = AbstractArray
promote_spcontainertype{T<:DivertToAbsArrayBC}(ct, ::T) = AbstractArray
promote_spcontainertype{S<:DivertToAbsArrayBC,T<:DivertToAbsArrayBC}(::S, ::T) = AbstractArray
# combinations involving scalars, sparse arrays, structured matrices (PromoteToSparse),
# dense vectors/matrices, and PromoteToSparse collections continue in the promote-to-sparse funnel
FunnelToSparseBC = Union{Type{Any},Type{Vector},Type{Matrix},Type{PromoteToSparse},Type{AbstractSparseArray}}
promote_spcontainertype{S<:FunnelToSparseBC,T<:FunnelToSparseBC}(::S, ::T) = PromoteToSparse


# first (Broadcast containertype) dispatch layer
# (broadcast_c[!], containertype, promote_containertype)
@inline broadcast_c{N}(f, ::Type{PromoteToSparse}, As::Vararg{Any,N}) =
spbroadcast_c(f, spcontainertype(As...), As...)
@inline broadcast_c!{N}(f, ::Type{AbstractSparseArray}, ::Type{PromoteToSparse}, C, B, As::Vararg{Any,N}) =
spbroadcast_c!(f, AbstractSparseArray, spcontainertype(B, As...), C, B, As...)
# where destination C is not an AbstractSparseArray, divert to generic AbstractArray broadcast code
@inline broadcast_c!{N}(f, CT::Type, ::Type{PromoteToSparse}, C, B, As::Vararg{Any,N}) =
broadcast_c!(f, CT, Array, C, B, As...)

# second (internal sparse broadcast containertype) dispatch layer
# (spbroadcast_c[!], spcontainertype, promote_spcontainertype)
@inline spbroadcast_c{N}(f, ::Type{PromoteToSparse}, As::Vararg{Any,N}) =
broadcast(f, map(_sparsifystructured, As)...)
@inline broadcast_c!{N}(f, ::Type{AbstractSparseArray}, ::Type{StructuredArray}, C, B, As::Vararg{Any,N}) =
@inline spbroadcast_c{N}(f, ::Type{AbstractArray}, As::Vararg{Any,N}) =
broadcast_c(f, Array, As...)
@inline spbroadcast_c!{N}(f, ::Type{AbstractSparseArray}, ::Type{PromoteToSparse}, C, B, As::Vararg{Any,N}) =
broadcast!(f, C, _sparsifystructured(B), map(_sparsifystructured, As)...)
@inline broadcast_c!{N}(f, CT::Type, ::Type{StructuredArray}, C, B, As::Vararg{Any,N}) =
broadcast_c!(f, CT, Array, C, B, As...)
@inline _sparsifystructured(S::SymTridiagonal) = SparseMatrixCSC(S)
@inline _sparsifystructured(T::Tridiagonal) = SparseMatrixCSC(T)
@inline _sparsifystructured(B::Bidiagonal) = SparseMatrixCSC(B)
@inline _sparsifystructured(D::Diagonal) = SparseMatrixCSC(D)
@inline spbroadcast_c!{N}(f, ::Type{AbstractSparseArray}, ::Type{AbstractArray}, C, B, As::Vararg{Any,N}) =
broadcast_c!(f, Array, Array, C, B, As...)

@inline _sparsifystructured(A::AbstractSparseArray) = A
@inline _sparsifystructured(M::StructuredMatrix) = SparseMatrixCSC(M)
@inline _sparsifystructured(M::Matrix) = SparseMatrixCSC(M)
@inline _sparsifystructured(V::Vector) = SparseVector(V)
@inline _sparsifystructured(x) = x


Expand Down
18 changes: 17 additions & 1 deletion test/sparse/higherorderfns.jl
Original file line number Diff line number Diff line change
Expand Up @@ -355,7 +355,7 @@ end
end
end

@testset "broadcast[!] over combinations of scalars, structured matrices, and sparse vectors/matrices" begin
@testset "broadcast[!] over combinations of scalars, sparse arrays, structured matrices, and dense vectors/matrices" begin
N, p = 10, 0.4
s = rand()
V = sprand(N, p)
Expand Down Expand Up @@ -387,6 +387,22 @@ end
@test broadcast!(*, Z, X, Y) == sparse(broadcast(*, fX, fY))
end
end
C = Array(sprand(N, 0.4))
M = Array(sprand(N, N, 0.4))
densearrays = (C, M)
fD, fB = Array(D), Array(B)
for X in densearrays
@test (Q = broadcast(+, D, X); Q isa SparseMatrixCSC && Q == sparse(broadcast(+, fD, X)))
@test broadcast!(+, Z, D, X) == sparse(broadcast(+, fD, X))
@test (Q = broadcast(*, s, B, X); Q isa SparseMatrixCSC && Q == sparse(broadcast(*, s, fB, X)))
@test broadcast!(*, Z, s, B, X) == sparse(broadcast(*, s, fB, X))
@test (Q = broadcast(+, V, B, X); Q isa SparseMatrixCSC && Q == sparse(broadcast(+, fV, fB, X)))
@test broadcast!(+, Z, V, B, X) == sparse(broadcast(+, fV, fB, X))
@test (Q = broadcast(+, V, A, X); Q isa SparseMatrixCSC && Q == sparse(broadcast(+, fV, fA, X)))
@test broadcast!(+, Z, V, A, X) == sparse(broadcast(+, fV, fA, X))
@test (Q = broadcast(*, s, V, A, X); Q isa SparseMatrixCSC && Q == sparse(broadcast(*, s, fV, fA, X)))
@test broadcast!(*, Z, s, V, A, X) == sparse(broadcast(*, s, fV, fA, X))
end
end

@testset "map[!] over combinations of sparse and structured matrices" begin
Expand Down