Merge pull request #20009 from Sacha0/okspbcarrays

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

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
 
 

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