diff --git a/base/sparse/higherorderfns.jl b/base/sparse/higherorderfns.jl index f5d04de81766d..acffdd0048efb 100644 --- a/base/sparse/higherorderfns.jl +++ b/base/sparse/higherorderfns.jl @@ -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. @@ -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 @@ -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 +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 diff --git a/test/sparse/higherorderfns.jl b/test/sparse/higherorderfns.jl index 5558d5dc33d53..a8589f205ade6 100644 --- a/test/sparse/higherorderfns.jl +++ b/test/sparse/higherorderfns.jl @@ -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) @@ -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