Extend sparse broadcast to one- and two-dimensional Arrays and test.

base/sparse/higherorderfns.jl

@@ -22,7 +22,8 @@ using ..SparseArrays: SparseVector, SparseMatrixCSC, AbstractSparseArray, indtyp
 # (7) Define _broadcast_[not]zeropres! specialized for a pair of (input) sparse vectors/matrices.
 # (8) Define general _broadcast_[not]zeropres! capable of handling >2 (input) sparse vectors/matrices.
 # (9) Define (broadcast[!]) methods handling combinations of broadcast scalars and sparse vectors/matrices.
-# (10) Define (broadcast[!]) methods handling combinations of scalars, sparse vectors/matrices, and structured matrices.
+# (10) Define (broadcast[!]) methods handling combinations of scalars, sparse vectors/matrices,
+#       structured matrices, and one- and two-dimensional Arrays.
 # (11) Define (map[!]) methods handling combinations of sparse and structured matrices.
 
 
@@ -133,7 +134,8 @@ _maxnnzfrom(shape::NTuple{2}, A::SparseVector) = nnz(A) * div(shape[1], A.n) * s
 _maxnnzfrom(shape::NTuple{2}, A::SparseMatrixCSC) = nnz(A) * div(shape[1], A.m) * div(shape[2], A.n)
 @inline _maxnnzfrom_each(shape, ::Tuple{}) = ()
 @inline _maxnnzfrom_each(shape, As) = (_maxnnzfrom(shape, first(As)), _maxnnzfrom_each(shape, tail(As))...)
-@inline _unchecked_maxnnzbcres(shape, As) = min(_densennz(shape), sum(_maxnnzfrom_each(shape, As)))
+@inline _unchecked_maxnnzbcres(shape, As::Tuple) = min(_densennz(shape), sum(_maxnnzfrom_each(shape, As)))
+@inline _unchecked_maxnnzbcres(shape, As...) = _unchecked_maxnnzbcres(shape, As)
 @inline _checked_maxnnzbcres(shape::NTuple{1}, As...) = shape[1] != 0 ? _unchecked_maxnnzbcres(shape, As) : 0
 @inline _checked_maxnnzbcres(shape::NTuple{2}, As...) = shape[1] != 0 && shape[2] != 0 ? _unchecked_maxnnzbcres(shape, As) : 0
 @inline function _allocres(shape::NTuple{1}, indextype, entrytype, maxnnz)
@@ -846,10 +848,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
@@ -888,40 +889,146 @@ 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)
 
 
-# (10) broadcast[!] over combinations of scalars, sparse vectors/matrices, and structured matrices
-
-# 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}) =
+# (10) 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, then
+# promote arguments to sparse as appropriate and rebroadcast.
+
+
+# first (Broadcast containertype) dispatch layer's promotion logic
+
+immutable PromoteToSparse end
+
+# broadcast containertype definitions for structured matrices
+_containertype{T<:Diagonal}(::Type{T}) = PromoteToSparse
+_containertype{T<:Bidiagonal}(::Type{T}) = PromoteToSparse
+_containertype{T<:Tridiagonal}(::Type{T}) = PromoteToSparse
+_containertype{T<:SymTridiagonal}(::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 ecollections 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
+# NOTE: Array in the broadcast containertype promotion mechanism can mean a number
+# of things (see "Issue..." just below). As such, below we have to implement a
+# containertype promotion mechanism for sparse broadcast that allows us to disambiguate
+# Array. Ultimately, disambiguating Array (-> separate AbstractArray and Array) in
+# Base.Broadcast is likely wise (see "Issue..." just below).
+#
+# Issue: `Array` has multiple meanings in the broadcast containertype promotion mechanism:
+# (1) `Array`, in the sense of "this object is an Array (or this is a collection of Arrays)";
+# (2) `AbstractArray`, in the sense of "this object is an AbstractArray (or this is a collection
+#   of AbstractArrays) but where that AbstractArray is (/a subset of that collection of
+#   AbstractArrays are) not of a specific AbstractArray subtype for which special
+#   handling is defined;
+# (3) `AbstractArray`, in the sense of "this is a collection of objects that need be funneled
+#   to the generic AbstractArray broadcast code".
+# Presently we conflate these three meanings in `Array`. Conflating meanings (2) and (3)
+# might be fine, but conflating meaning (1) with the other two prevents separating objects
+# that are Arrays (from e.g. collections of Arrays and Tuples) for special handling, as is
+# necessary e.g. to handle Arrays in sparse broadcast.
+#
+# We should probably disambiguate these meanings in Broadcast. One approach to doing so
+# would be to replace `Array` with `AbstractArray` in the existing containertype promotion
+# mechanism, and then separately define containertype promotion methods for Array as we do
+# for Tuple.
+#
+# The tricky question is what to do about the promote_containertype(ct, ::Type{Array}) = Array
+# (after the change suggested above, promote_containertype(ct, ::Type{AbstractArray})) = AbstractArray)
+# methods. These methods are ambiguity magnets, and there is a tendency to write similar
+# such methods whenever extending broadcast to a new type, which ultimately results
+# in having to write a combinatorial explosion of ambiguity-killing methods.
+#
+# Perhaps the following makes a reasonable model: don't define methods like
+# promote_containertype(ct, ::Type{Array}) = Array, and discourage definition of such
+# methods for new types. Instead (1) define promote_containertype(cta, ctb) = AbstractArray as
+# primary fallback, such that any type pair without clearly defined behavior gets funneled
+# to the generic AbstractArray broadcast code, and (2) encourage writing only explicit, tight
+# promote_containertype definitions.
+#
+# The above should improve the extensibility and maintainability of Broadcast.
+
+
+# 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 methods to override the immediately preceding method
+_spcontainertype{T<:Diagonal}(::Type{T}) = PromoteToSparse
+_spcontainertype{T<:Bidiagonal}(::Type{T}) = PromoteToSparse
+_spcontainertype{T<:Tridiagonal}(::Type{T}) = PromoteToSparse
+_spcontainertype{T<:SymTridiagonal}(::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 spbroadcast_c!{N}(f, ::Type{AbstractSparseArray}, ::Type{AbstractArray}, C, B, As::Vararg{Any,N}) =
+    broadcast_c!(f, Array, 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 _sparsifystructured(M::Matrix) = SparseMatrixCSC(M)
+@inline _sparsifystructured(V::Vector) = SparseVector(V)
 @inline _sparsifystructured(A::AbstractSparseArray) = A
 @inline _sparsifystructured(x) = x
 

test/sparse/higherorderfns.jl

@@ -285,7 +285,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)
@@ -317,6 +317,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