diff --git a/base/sparse/higherorderfns.jl b/base/sparse/higherorderfns.jl index dfd89ff2c3350c..70806e7f5bdb01 100644 --- a/base/sparse/higherorderfns.jl +++ b/base/sparse/higherorderfns.jl @@ -21,8 +21,9 @@ using ..SparseArrays: SparseVector, SparseMatrixCSC, AbstractSparseArray, indtyp # (6) Define general _map_[not]zeropres! capable of handling >2 (input) sparse vectors/matrices. # (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 methods handling combinations of broadcast scalars and sparse vectors/matrices. -# (10) Define methods handling combinations of scalars, sparse vectors/matrices, and structured 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. +# (11) Define (map[!]) methods handling combinations of sparse and structured matrices. # (1) The definitions below provide a common interface to sparse vectors and matrices @@ -927,4 +928,14 @@ promote_containertype(::Type{Tuple}, ::Type{StructuredArray}) = Array @inline _sparsifystructured(A::AbstractSparseArray) = A @inline _sparsifystructured(x) = x + +# (11) map[!] over combinations of sparse and structured matrices +StructuredMatrix = Union{Diagonal,Bidiagonal,Tridiagonal,SymTridiagonal} +SparseOrStructuredMatrix = Union{SparseMatrixCSC,StructuredMatrix} +map{Tf}(f::Tf, A::StructuredMatrix) = _noshapecheck_map(f, _sparsifystructured(A)) +map{Tf,N}(f::Tf, A::SparseOrStructuredMatrix, Bs::Vararg{SparseOrStructuredMatrix,N}) = + (_checksameshape(A, Bs...); _noshapecheck_map(f, _sparsifystructured(A), map(_sparsifystructured, Bs)...)) +map!{Tf,N}(f::Tf, C::SparseMatrixCSC, A::SparseOrStructuredMatrix, Bs::Vararg{SparseOrStructuredMatrix,N}) = + (_checksameshape(C, A, Bs...); _noshapecheck_map!(f, C, _sparsifystructured(A), map(_sparsifystructured, Bs)...)) + end diff --git a/test/sparse/higherorderfns.jl b/test/sparse/higherorderfns.jl index eaf23ced6f4e89..ce19b7d1568ec9 100644 --- a/test/sparse/higherorderfns.jl +++ b/test/sparse/higherorderfns.jl @@ -317,6 +317,34 @@ end end end +@testset "map[!] over combinations of sparse and structured matrices" begin + N, p = 10, 0.4 + A = sprand(N, N, p) + Z, fA = copy(A), Array(A) + D = Diagonal(rand(N)) + B = Bidiagonal(rand(N), rand(N - 1), true) + T = Tridiagonal(rand(N - 1), rand(N), rand(N - 1)) + S = SymTridiagonal(rand(N), rand(N - 1)) + structuredarrays = (D, B, T, S) + fstructuredarrays = map(Array, structuredarrays) + for (X, fX) in zip(structuredarrays, fstructuredarrays) + @test (Q = map(sin, X); Q isa SparseMatrixCSC && Q == sparse(map(sin, fX))) + @test map!(sin, Z, X) == sparse(map(sin, fX)) + @test (Q = map(cos, X); Q isa SparseMatrixCSC && Q == sparse(map(cos, fX))) + @test map!(cos, Z, X) == sparse(map(cos, fX)) + @test (Q = map(+, A, X); Q isa SparseMatrixCSC && Q == sparse(map(+, fA, fX))) + @test map!(+, Z, A, X) == sparse(map(+, fA, fX)) + for (Y, fY) in zip(structuredarrays, fstructuredarrays) + @test (Q = map(+, X, Y); Q isa SparseMatrixCSC && Q == sparse(map(+, fX, fY))) + @test map!(+, Z, X, Y) == sparse(map(+, fX, fY)) + @test (Q = map(*, X, Y); Q isa SparseMatrixCSC && Q == sparse(map(*, fX, fY))) + @test map!(*, Z, X, Y) == sparse(map(*, fX, fY)) + @test (Q = map(+, X, A, Y); Q isa SparseMatrixCSC && Q == sparse(map(+, fX, fA, fY))) + @test map!(+, Z, X, A, Y) == sparse(map(+, fX, fA, fY)) + end + end +end + # Older tests of sparse broadcast, now largely covered by the tests above @testset "assorted tests of sparse broadcast over two input arguments" begin N, p = 10, 0.3