Skip to content

Commit 29a3563

Browse files
Sacha0Sacha0
authored
Merge pull request #19926 from Sacha0/structuredbc
extend sparse map[!]/broadcast[!] to structured matrices
2 parents 2ded6c4 + 96ae43e commit 29a3563

File tree

3 files changed

+116
-4
lines changed

3 files changed

+116
-4
lines changed

‎base/broadcast.jl

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -199,8 +199,8 @@ arguments to `f` unless it is also listed in the `As`,
199199
as in `broadcast!(f, A, A, B)` to perform `A[:] = broadcast(f, A, B)`.
200200
"""
201201
@inline broadcast!{N}(f, C::AbstractArray, A, Bs::Vararg{Any,N}) =
202-
broadcast_c!(f, containertype(C, A, Bs...), C, A, Bs...)
203-
@inline function broadcast_c!{N}(f, ::Type, C::AbstractArray, A, Bs::Vararg{Any,N})
202+
broadcast_c!(f, containertype(C), containertype(A, Bs...), C, A, Bs...)
203+
@inline function broadcast_c!{N}(f, ::Type, ::Type, C, A, Bs::Vararg{Any,N})
204204
shape = indices(C)
205205
@boundscheck check_broadcast_indices(shape, A, Bs...)
206206
keeps, Idefaults = map_newindexer(shape, A, Bs)

‎base/sparse/higherorderfns.jl

Lines changed: 52 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,9 @@ using ..SparseArrays: SparseVector, SparseMatrixCSC, AbstractSparseArray, indtyp
2121
# (6) Define general _map_[not]zeropres! capable of handling >2 (input) sparse vectors/matrices.
2222
# (7) Define _broadcast_[not]zeropres! specialized for a pair of (input) sparse vectors/matrices.
2323
# (8) Define general _broadcast_[not]zeropres! capable of handling >2 (input) sparse vectors/matrices.
24-
# (9) Define methods handling combinations of broadcast scalars and sparse vectors/matrices.
24+
# (9) Define (broadcast[!]) methods handling combinations of broadcast scalars and sparse vectors/matrices.
25+
# (10) Define (broadcast[!]) methods handling combinations of scalars, sparse vectors/matrices, and structured matrices.
26+
# (11) Define (map[!]) methods handling combinations of sparse and structured matrices.
2527

2628

2729
# (1) The definitions below provide a common interface to sparse vectors and matrices
@@ -850,7 +852,7 @@ promote_containertype(::Type{Tuple}, ::Type{AbstractSparseArray}) = Array
850852
promote_containertype(::Type{AbstractSparseArray}, ::Type{Array}) = Array
851853
promote_containertype(::Type{AbstractSparseArray}, ::Type{Tuple}) = Array
852854

853-
# broadcast[!] entry points for combinations of sparse arrays and other types
855+
# broadcast[!] entry points for combinations of sparse arrays and other (scalar) types
854856
@inline function broadcast_c{N}(f, ::Type{AbstractSparseArray}, mixedargs::Vararg{Any,N})
855857
parevalf, passedargstup = capturescalars(f, mixedargs)
856858
return broadcast(parevalf, passedargstup...)
@@ -885,4 +887,52 @@ end
885887
broadcast{Tf,T}(f::Tf, ::Type{T}, A::SparseMatrixCSC) = broadcast(y -> f(T, y), A)
886888
broadcast{Tf,T}(f::Tf, A::SparseMatrixCSC, ::Type{T}) = broadcast(x -> f(x, T), A)
887889

890+
891+
# (10) broadcast[!] over combinations of scalars, sparse vectors/matrices, and structured matrices
892+
893+
# structured array container type promotion
894+
immutable StructuredArray end
895+
_containertype{T<:Diagonal}(::Type{T}) = StructuredArray
896+
_containertype{T<:Bidiagonal}(::Type{T}) = StructuredArray
897+
_containertype{T<:Tridiagonal}(::Type{T}) = StructuredArray
898+
_containertype{T<:SymTridiagonal}(::Type{T}) = StructuredArray
899+
promote_containertype(::Type{StructuredArray}, ::Type{StructuredArray}) = StructuredArray
900+
# combinations involving sparse arrays continue in the structured array funnel
901+
promote_containertype(::Type{StructuredArray}, ::Type{AbstractSparseArray}) = StructuredArray
902+
promote_containertype(::Type{AbstractSparseArray}, ::Type{StructuredArray}) = StructuredArray
903+
# combinations involving scalars continue in the structured array funnel
904+
promote_containertype(::Type{StructuredArray}, ::Type{Any}) = StructuredArray
905+
promote_containertype(::Type{Any}, ::Type{StructuredArray}) = StructuredArray
906+
# combinations involving arrays divert to the generic array code
907+
promote_containertype(::Type{StructuredArray}, ::Type{Array}) = Array
908+
promote_containertype(::Type{Array}, ::Type{StructuredArray}) = Array
909+
# combinations involving tuples divert to the generic array code
910+
promote_containertype(::Type{StructuredArray}, ::Type{Tuple}) = Array
911+
promote_containertype(::Type{Tuple}, ::Type{StructuredArray}) = Array
912+
913+
# for combinations involving sparse/structured arrays and scalars only,
914+
# promote all structured arguments to sparse and then rebroadcast
915+
@inline broadcast_c{N}(f, ::Type{StructuredArray}, As::Vararg{Any,N}) =
916+
broadcast(f, map(_sparsifystructured, As)...)
917+
@inline broadcast_c!{N}(f, ::Type{AbstractSparseArray}, ::Type{StructuredArray}, C, B, As::Vararg{Any,N}) =
918+
broadcast!(f, C, _sparsifystructured(B), map(_sparsifystructured, As)...)
919+
@inline broadcast_c!{N}(f, CT::Type, ::Type{StructuredArray}, C, B, As::Vararg{Any,N}) =
920+
broadcast_c!(f, CT, Array, C, B, As...)
921+
@inline _sparsifystructured(S::SymTridiagonal) = SparseMatrixCSC(S)
922+
@inline _sparsifystructured(T::Tridiagonal) = SparseMatrixCSC(T)
923+
@inline _sparsifystructured(B::Bidiagonal) = SparseMatrixCSC(B)
924+
@inline _sparsifystructured(D::Diagonal) = SparseMatrixCSC(D)
925+
@inline _sparsifystructured(A::AbstractSparseArray) = A
926+
@inline _sparsifystructured(x) = x
927+
928+
929+
# (11) map[!] over combinations of sparse and structured matrices
930+
StructuredMatrix = Union{Diagonal,Bidiagonal,Tridiagonal,SymTridiagonal}
931+
SparseOrStructuredMatrix = Union{SparseMatrixCSC,StructuredMatrix}
932+
map{Tf}(f::Tf, A::StructuredMatrix) = _noshapecheck_map(f, _sparsifystructured(A))
933+
map{Tf,N}(f::Tf, A::SparseOrStructuredMatrix, Bs::Vararg{SparseOrStructuredMatrix,N}) =
934+
(_checksameshape(A, Bs...); _noshapecheck_map(f, _sparsifystructured(A), map(_sparsifystructured, Bs)...))
935+
map!{Tf,N}(f::Tf, C::SparseMatrixCSC, A::SparseOrStructuredMatrix, Bs::Vararg{SparseOrStructuredMatrix,N}) =
936+
(_checksameshape(C, A, Bs...); _noshapecheck_map!(f, C, _sparsifystructured(A), map(_sparsifystructured, Bs)...))
937+
888938
end

‎test/sparse/higherorderfns.jl

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -285,6 +285,68 @@ end
285285
end
286286
end
287287

288+
@testset "broadcast[!] over combinations of scalars, structured matrices, and sparse vectors/matrices" begin
289+
N, p = 10, 0.4
290+
s = rand()
291+
V = sprand(N, p)
292+
A = sprand(N, N, p)
293+
Z = copy(A)
294+
sparsearrays = (V, A)
295+
fV, fA = map(Array, sparsearrays)
296+
D = Diagonal(rand(N))
297+
B = Bidiagonal(rand(N), rand(N - 1), true)
298+
T = Tridiagonal(rand(N - 1), rand(N), rand(N - 1))
299+
S = SymTridiagonal(rand(N), rand(N - 1))
300+
structuredarrays = (D, B, T, S)
301+
fstructuredarrays = map(Array, structuredarrays)
302+
for (X, fX) in zip(structuredarrays, fstructuredarrays)
303+
@test (Q = broadcast(sin, X); Q isa SparseMatrixCSC && Q == sparse(broadcast(sin, fX)))
304+
@test broadcast!(sin, Z, X) == sparse(broadcast(sin, fX))
305+
@test (Q = broadcast(cos, X); Q isa SparseMatrixCSC && Q == sparse(broadcast(cos, fX)))
306+
@test broadcast!(cos, Z, X) == sparse(broadcast(cos, fX))
307+
@test (Q = broadcast(*, s, X); Q isa SparseMatrixCSC && Q == sparse(broadcast(*, s, fX)))
308+
@test broadcast!(*, Z, s, X) == sparse(broadcast(*, s, fX))
309+
@test (Q = broadcast(+, V, A, X); Q isa SparseMatrixCSC && Q == sparse(broadcast(+, fV, fA, fX)))
310+
@test broadcast!(+, Z, V, A, X) == sparse(broadcast(+, fV, fA, fX))
311+
@test (Q = broadcast(*, s, V, A, X); Q isa SparseMatrixCSC && Q == sparse(broadcast(*, s, fV, fA, fX)))
312+
@test broadcast!(*, Z, s, V, A, X) == sparse(broadcast(*, s, fV, fA, fX))
313+
for (Y, fY) in zip(structuredarrays, fstructuredarrays)
314+
@test (Q = broadcast(+, X, Y); Q isa SparseMatrixCSC && Q == sparse(broadcast(+, fX, fY)))
315+
@test broadcast!(+, Z, X, Y) == sparse(broadcast(+, fX, fY))
316+
@test (Q = broadcast(*, X, Y); Q isa SparseMatrixCSC && Q == sparse(broadcast(*, fX, fY)))
317+
@test broadcast!(*, Z, X, Y) == sparse(broadcast(*, fX, fY))
318+
end
319+
end
320+
end
321+
322+
@testset "map[!] over combinations of sparse and structured matrices" begin
323+
N, p = 10, 0.4
324+
A = sprand(N, N, p)
325+
Z, fA = copy(A), Array(A)
326+
D = Diagonal(rand(N))
327+
B = Bidiagonal(rand(N), rand(N - 1), true)
328+
T = Tridiagonal(rand(N - 1), rand(N), rand(N - 1))
329+
S = SymTridiagonal(rand(N), rand(N - 1))
330+
structuredarrays = (D, B, T, S)
331+
fstructuredarrays = map(Array, structuredarrays)
332+
for (X, fX) in zip(structuredarrays, fstructuredarrays)
333+
@test (Q = map(sin, X); Q isa SparseMatrixCSC && Q == sparse(map(sin, fX)))
334+
@test map!(sin, Z, X) == sparse(map(sin, fX))
335+
@test (Q = map(cos, X); Q isa SparseMatrixCSC && Q == sparse(map(cos, fX)))
336+
@test map!(cos, Z, X) == sparse(map(cos, fX))
337+
@test (Q = map(+, A, X); Q isa SparseMatrixCSC && Q == sparse(map(+, fA, fX)))
338+
@test map!(+, Z, A, X) == sparse(map(+, fA, fX))
339+
for (Y, fY) in zip(structuredarrays, fstructuredarrays)
340+
@test (Q = map(+, X, Y); Q isa SparseMatrixCSC && Q == sparse(map(+, fX, fY)))
341+
@test map!(+, Z, X, Y) == sparse(map(+, fX, fY))
342+
@test (Q = map(*, X, Y); Q isa SparseMatrixCSC && Q == sparse(map(*, fX, fY)))
343+
@test map!(*, Z, X, Y) == sparse(map(*, fX, fY))
344+
@test (Q = map(+, X, A, Y); Q isa SparseMatrixCSC && Q == sparse(map(+, fX, fA, fY)))
345+
@test map!(+, Z, X, A, Y) == sparse(map(+, fX, fA, fY))
346+
end
347+
end
348+
end
349+
288350
# Older tests of sparse broadcast, now largely covered by the tests above
289351
@testset "assorted tests of sparse broadcast over two input arguments" begin
290352
N, p = 10, 0.3

0 commit comments

Comments
 (0)