diag of SparseMatrixCSC should always return SparseVector

JuliaLang · Aug 14, 2017 · d4abafb · d4abafb
1 parent 3f86c1a
commit d4abafb
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 3 deletions.
diff --git a/base/sparse/linalg.jl b/base/sparse/linalg.jl
@@ -879,7 +879,7 @@ for f in (:\, :Ac_ldiv_B, :At_ldiv_B)
             if m == n
                 if istril(A)
                     if istriu(A)
-                        return ($f)(Diagonal(A), B)
+                        return ($f)(Diagonal(Vector(diag(A))), B)
                     else
                         return ($f)(LowerTriangular(A), B)
                     end

diff --git a/base/sparse/sparsematrix.jl b/base/sparse/sparsematrix.jl
@@ -3412,8 +3412,6 @@ function trace(A::SparseMatrixCSC{Tv}) where Tv
     s
 end
 
-diag(A::SparseMatrixCSC{Tv}) where {Tv} = Tv[d for d in SpDiagIterator(A)]
-
 function diagm(v::SparseMatrixCSC{Tv,Ti}) where {Tv,Ti}
     if size(v,1) != 1 && size(v,2) != 1
         throw(DimensionMismatch("input should be nx1 or 1xn"))

diff --git a/test/sparse/sparse.jl b/test/sparse/sparse.jl
@@ -1323,6 +1323,16 @@ end
     @test diagm(sparse(ones(5,1))) == speye(5)
 end
 
+@testset "diag" begin
+    for T in (Float64, Complex128)
+        S = sprand(T, 5, 5, 0.5)
+        A = Matrix(S)
+        @test diag(S)::SparseVector{T,Int}     == diag(A)
+        @test diag(S, 1)::SparseVector{T,Int}  == diag(A, 1)
+        @test diag(S, -1)::SparseVector{T,Int} == diag(A, -1)
+    end
+end
+
 @testset "expandptr" begin
     A = speye(5)
     @test Base.SparseArrays.expandptr(A.colptr) == collect(1:5)