Skip to content

Commit c66f3b2

Browse files
Sacha0Sacha0
committed
Extend sparse broadcast to one- and two-dimensional Arrays and test.
1 parent 6f9fe0b commit c66f3b2

File tree

2 files changed

+107
-35
lines changed

2 files changed

+107
-35
lines changed

‎base/sparse/higherorderfns.jl

+90-34
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,8 @@ using ..SparseArrays: SparseVector, SparseMatrixCSC, AbstractSparseArray, indtyp
2323
# (8) Define _broadcast_[not]zeropres! specialized for a pair of (input) sparse vectors/matrices.
2424
# (9) Define general _broadcast_[not]zeropres! capable of handling >2 (input) sparse vectors/matrices.
2525
# (10) Define (broadcast[!]) methods handling combinations of broadcast scalars and sparse vectors/matrices.
26-
# (11) Define (broadcast[!]) methods handling combinations of scalars, sparse vectors/matrices, and structured matrices.
26+
# (11) Define (broadcast[!]) methods handling combinations of scalars, sparse vectors/matrices,
27+
# structured matrices, and one- and two-dimensional Arrays.
2728
# (12) Define (map[!]) methods handling combinations of sparse and structured matrices.
2829

2930

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

953953
# 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),
986986
broadcast{Tf,T}(f::Tf, A::SparseMatrixCSC, ::Type{T}) = broadcast(x -> f(x, T), A)
987987

988988

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

991-
# structured array container type promotion
992-
immutable StructuredArray end
993-
_containertype{T<:Diagonal}(::Type{T}) = StructuredArray
994-
_containertype{T<:Bidiagonal}(::Type{T}) = StructuredArray
995-
_containertype{T<:Tridiagonal}(::Type{T}) = StructuredArray
996-
_containertype{T<:SymTridiagonal}(::Type{T}) = StructuredArray
997-
promote_containertype(::Type{StructuredArray}, ::Type{StructuredArray}) = StructuredArray
998-
# combinations involving sparse arrays continue in the structured array funnel
999-
promote_containertype(::Type{StructuredArray}, ::Type{AbstractSparseArray}) = StructuredArray
1000-
promote_containertype(::Type{AbstractSparseArray}, ::Type{StructuredArray}) = StructuredArray
1001-
# combinations involving scalars continue in the structured array funnel
1002-
promote_containertype(::Type{StructuredArray}, ::Type{Any}) = StructuredArray
1003-
promote_containertype(::Type{Any}, ::Type{StructuredArray}) = StructuredArray
1004-
# combinations involving arrays divert to the generic array code
1005-
promote_containertype(::Type{StructuredArray}, ::Type{Array}) = Array
1006-
promote_containertype(::Type{Array}, ::Type{StructuredArray}) = Array
1007-
# combinations involving tuples divert to the generic array code
1008-
promote_containertype(::Type{StructuredArray}, ::Type{Tuple}) = Array
1009-
promote_containertype(::Type{Tuple}, ::Type{StructuredArray}) = Array
10101003

1011-
# for combinations involving sparse/structured arrays and scalars only,
1012-
# promote all structured arguments to sparse and then rebroadcast
1013-
@inline broadcast_c{N}(f, ::Type{StructuredArray}, As::Vararg{Any,N}) =
1004+
# first (Broadcast containertype) dispatch layer's promotion logic
1005+
immutable PromoteToSparse end
1006+
1007+
# broadcast containertype definitions for structured matrices
1008+
StructuredMatrix = Union{Diagonal,Bidiagonal,Tridiagonal,SymTridiagonal}
1009+
_containertype{T<:StructuredMatrix}(::Type{T}) = PromoteToSparse
1010+
1011+
# combinations explicitly involving Tuples and PromoteToSparse collections
1012+
# divert to the generic AbstractArray broadcast code
1013+
promote_containertype(::Type{PromoteToSparse}, ::Type{Tuple}) = Array
1014+
promote_containertype(::Type{Tuple}, ::Type{PromoteToSparse}) = Array
1015+
# combinations involving scalars and PromoteToSparse collections continue in the promote-to-sparse funnel
1016+
promote_containertype(::Type{PromoteToSparse}, ::Type{Any}) = PromoteToSparse
1017+
promote_containertype(::Type{Any}, ::Type{PromoteToSparse}) = PromoteToSparse
1018+
# combinations involving sparse arrays and PromoteToSparse collections continue in the promote-to-sparse funnel
1019+
promote_containertype(::Type{PromoteToSparse}, ::Type{AbstractSparseArray}) = PromoteToSparse
1020+
promote_containertype(::Type{AbstractSparseArray}, ::Type{PromoteToSparse}) = PromoteToSparse
1021+
# combinations involving Arrays and PromoteToSparse collections continue in the promote-to-sparse funnel
1022+
promote_containertype(::Type{PromoteToSparse}, ::Type{Array}) = PromoteToSparse
1023+
promote_containertype(::Type{Array}, ::Type{PromoteToSparse}) = PromoteToSparse
1024+
# combinations involving Arrays and sparse arrays continue in the promote-to-sparse funnel
1025+
promote_containertype(::Type{AbstractSparseArray}, ::Type{Array}) = PromoteToSparse
1026+
promote_containertype(::Type{Array}, ::Type{AbstractSparseArray}) = PromoteToSparse
1027+
1028+
# second (internal sparse broadcast containertype) dispatch layer's promotion logic
1029+
# mostly just disambiguates Array from the main containertype promotion mechanism
1030+
# AbstractArray serves as a marker to shunt to the generic AbstractArray broadcast code
1031+
_spcontainertype(x) = _containertype(x)
1032+
_spcontainertype{T<:Vector}(::Type{T}) = Vector
1033+
_spcontainertype{T<:Matrix}(::Type{T}) = Matrix
1034+
_spcontainertype{T<:Ref}(::Type{T}) = AbstractArray
1035+
_spcontainertype{T<:AbstractArray}(::Type{T}) = AbstractArray
1036+
# need the following two methods to override the immediately preceding method
1037+
_spcontainertype{T<:StructuredMatrix}(::Type{T}) = PromoteToSparse
1038+
_spcontainertype{T<:SparseVecOrMat}(::Type{T}) = AbstractSparseArray
1039+
spcontainertype(x) = _spcontainertype(typeof(x))
1040+
spcontainertype(ct1, ct2) = promote_spcontainertype(spcontainertype(ct1), spcontainertype(ct2))
1041+
@inline spcontainertype(ct1, ct2, cts...) = promote_spcontainertype(spcontainertype(ct1), spcontainertype(ct2, cts...))
1042+
1043+
promote_spcontainertype{T}(::Type{T}, ::Type{T}) = T
1044+
# combinations involving AbstractArrays and/or Tuples divert to the generic AbstractArray broadcast code
1045+
DivertToAbsArrayBC = Union{Type{AbstractArray},Type{Tuple}}
1046+
promote_spcontainertype{T<:DivertToAbsArrayBC}(::T, ct) = AbstractArray
1047+
promote_spcontainertype{T<:DivertToAbsArrayBC}(ct, ::T) = AbstractArray
1048+
promote_spcontainertype{S<:DivertToAbsArrayBC,T<:DivertToAbsArrayBC}(::S, ::T) = AbstractArray
1049+
# combinations involving scalars, sparse arrays, structured matrices (PromoteToSparse),
1050+
# dense vectors/matrices, and PromoteToSparse collections continue in the promote-to-sparse funnel
1051+
FunnelToSparseBC = Union{Type{Any},Type{Vector},Type{Matrix},Type{PromoteToSparse},Type{AbstractSparseArray}}
1052+
promote_spcontainertype{S<:FunnelToSparseBC,T<:FunnelToSparseBC}(::S, ::T) = PromoteToSparse
1053+
1054+
1055+
# first (Broadcast containertype) dispatch layer
1056+
# (broadcast_c[!], containertype, promote_containertype)
1057+
@inline broadcast_c{N}(f, ::Type{PromoteToSparse}, As::Vararg{Any,N}) =
1058+
spbroadcast_c(f, spcontainertype(As...), As...)
1059+
@inline broadcast_c!{N}(f, ::Type{AbstractSparseArray}, ::Type{PromoteToSparse}, C, B, As::Vararg{Any,N}) =
1060+
spbroadcast_c!(f, AbstractSparseArray, spcontainertype(B, As...), C, B, As...)
1061+
# where destination C is not an AbstractSparseArray, divert to generic AbstractArray broadcast code
1062+
@inline broadcast_c!{N}(f, CT::Type, ::Type{PromoteToSparse}, C, B, As::Vararg{Any,N}) =
1063+
broadcast_c!(f, CT, Array, C, B, As...)
1064+
1065+
# second (internal sparse broadcast containertype) dispatch layer
1066+
# (spbroadcast_c[!], spcontainertype, promote_spcontainertype)
1067+
@inline spbroadcast_c{N}(f, ::Type{PromoteToSparse}, As::Vararg{Any,N}) =
10141068
broadcast(f, map(_sparsifystructured, As)...)
1015-
@inline broadcast_c!{N}(f, ::Type{AbstractSparseArray}, ::Type{StructuredArray}, C, B, As::Vararg{Any,N}) =
1069+
@inline spbroadcast_c{N}(f, ::Type{AbstractArray}, As::Vararg{Any,N}) =
1070+
broadcast_c(f, Array, As...)
1071+
@inline spbroadcast_c!{N}(f, ::Type{AbstractSparseArray}, ::Type{PromoteToSparse}, C, B, As::Vararg{Any,N}) =
10161072
broadcast!(f, C, _sparsifystructured(B), map(_sparsifystructured, As)...)
1017-
@inline broadcast_c!{N}(f, CT::Type, ::Type{StructuredArray}, C, B, As::Vararg{Any,N}) =
1018-
broadcast_c!(f, CT, Array, C, B, As...)
1019-
@inline _sparsifystructured(S::SymTridiagonal) = SparseMatrixCSC(S)
1020-
@inline _sparsifystructured(T::Tridiagonal) = SparseMatrixCSC(T)
1021-
@inline _sparsifystructured(B::Bidiagonal) = SparseMatrixCSC(B)
1022-
@inline _sparsifystructured(D::Diagonal) = SparseMatrixCSC(D)
1073+
@inline spbroadcast_c!{N}(f, ::Type{AbstractSparseArray}, ::Type{AbstractArray}, C, B, As::Vararg{Any,N}) =
1074+
broadcast_c!(f, Array, Array, C, B, As...)
1075+
10231076
@inline _sparsifystructured(A::AbstractSparseArray) = A
1077+
@inline _sparsifystructured(M::StructuredMatrix) = SparseMatrixCSC(M)
1078+
@inline _sparsifystructured(M::Matrix) = SparseMatrixCSC(M)
1079+
@inline _sparsifystructured(V::Vector) = SparseVector(V)
10241080
@inline _sparsifystructured(x) = x
10251081

10261082

‎test/sparse/higherorderfns.jl

+17-1
Original file line numberDiff line numberDiff line change
@@ -355,7 +355,7 @@ end
355355
end
356356
end
357357

358-
@testset "broadcast[!] over combinations of scalars, structured matrices, and sparse vectors/matrices" begin
358+
@testset "broadcast[!] over combinations of scalars, sparse arrays, structured matrices, and dense vectors/matrices" begin
359359
N, p = 10, 0.4
360360
s = rand()
361361
V = sprand(N, p)
@@ -387,6 +387,22 @@ end
387387
@test broadcast!(*, Z, X, Y) == sparse(broadcast(*, fX, fY))
388388
end
389389
end
390+
C = Array(sprand(N, 0.4))
391+
M = Array(sprand(N, N, 0.4))
392+
densearrays = (C, M)
393+
fD, fB = Array(D), Array(B)
394+
for X in densearrays
395+
@test (Q = broadcast(+, D, X); Q isa SparseMatrixCSC && Q == sparse(broadcast(+, fD, X)))
396+
@test broadcast!(+, Z, D, X) == sparse(broadcast(+, fD, X))
397+
@test (Q = broadcast(*, s, B, X); Q isa SparseMatrixCSC && Q == sparse(broadcast(*, s, fB, X)))
398+
@test broadcast!(*, Z, s, B, X) == sparse(broadcast(*, s, fB, X))
399+
@test (Q = broadcast(+, V, B, X); Q isa SparseMatrixCSC && Q == sparse(broadcast(+, fV, fB, X)))
400+
@test broadcast!(+, Z, V, B, X) == sparse(broadcast(+, fV, fB, X))
401+
@test (Q = broadcast(+, V, A, X); Q isa SparseMatrixCSC && Q == sparse(broadcast(+, fV, fA, X)))
402+
@test broadcast!(+, Z, V, A, X) == sparse(broadcast(+, fV, fA, X))
403+
@test (Q = broadcast(*, s, V, A, X); Q isa SparseMatrixCSC && Q == sparse(broadcast(*, s, fV, fA, X)))
404+
@test broadcast!(*, Z, s, V, A, X) == sparse(broadcast(*, s, fV, fA, X))
405+
end
390406
end
391407

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

0 commit comments

Comments
 (0)