Extend sparse map[!] to structured matrices by promoting structured matrices to sparse.

Author: Sacha0

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

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