diag of SparseMatrixCSC should always return SparseVector (#23261)

* diag of SparseMatrixCSC should always return SparseVector

* remove SpDiagIterator

Author: fredrikekre

NEWS.md

@@ -312,6 +312,8 @@ Deprecated or removed
   * `Base.cpad` has been removed; use an appropriate combination of `rpad` and `lpad`
     instead ([#23187]).
 
+  * `Base.SparseArrays.SpDiagIterator` has been removed ([#23261]).
+
 Command-line option changes
 ---------------------------
 

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

base/sparse/sparsematrix.jl

@@ -3380,40 +3380,39 @@ function expandptr(V::Vector{<:Integer})
     res
 end
 
-## diag and related using an iterator
 
-mutable struct SpDiagIterator{Tv,Ti}
-    A::SparseMatrixCSC{Tv,Ti}
-    n::Int
-end
-SpDiagIterator(A::SparseMatrixCSC) = SpDiagIterator(A,minimum(size(A)))
-
-length(d::SpDiagIterator) = d.n
-start(d::SpDiagIterator) = 1
-done(d::SpDiagIterator, j) = j > d.n
-
-function next(d::SpDiagIterator{Tv}, j) where Tv
-    A = d.A
-    r1 = Int(A.colptr[j])
-    r2 = Int(A.colptr[j+1]-1)
-    (r1 > r2) && (return (zero(Tv), j+1))
-    r1 = searchsortedfirst(A.rowval, j, r1, r2, Forward)
-    (((r1 > r2) || (A.rowval[r1] != j)) ? zero(Tv) : A.nzval[r1], j+1)
+function diag(A::SparseMatrixCSC{Tv,Ti}, d::Integer=0) where {Tv,Ti}
+    m, n = size(A)
+    k = Int(d)
+    if !(-m <= k <= n)
+        throw(ArgumentError("requested diagonal, $k, out of bounds in matrix of size ($m, $n)"))
+    end
+    l = k < 0 ? min(m+k,n) : min(n-k,m)
+    r, c = k <= 0 ? (-k, 0) : (0, k) # start row/col -1
+    ind = Vector{Ti}()
+    val = Vector{Tv}()
+    for i in 1:l
+        r += 1; c += 1
+        r1 = Int(A.colptr[c])
+        r2 = Int(A.colptr[c+1]-1)
+        r1 > r2 && continue
+        r1 = searchsortedfirst(A.rowval, r, r1, r2, Forward)
+        ((r1 > r2) || (A.rowval[r1] != r)) && continue
+        push!(ind, i)
+        push!(val, A.nzval[r1])
+    end
+    return SparseVector{Tv,Ti}(l, ind, val)
 end
 
 function trace(A::SparseMatrixCSC{Tv}) where Tv
-    if size(A,1) != size(A,2)
-        throw(DimensionMismatch("expected square matrix"))
-    end
+    n = checksquare(A)
     s = zero(Tv)
-    for d in SpDiagIterator(A)
-        s += d
+    for i in 1:n
+        s += A[i,i]
     end
-    s
+    return 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"))

test/sparse/sparse.jl

@@ -1323,6 +1323,27 @@ end
     @test diagm(sparse(ones(5,1))) == speye(5)
 end
 
+@testset "diag" begin
+    for T in (Float64, Complex128)
+        S1 = sprand(T,  5,  5, 0.5)
+        S2 = sprand(T, 10,  5, 0.5)
+        S3 = sprand(T,  5, 10, 0.5)
+        for S in (S1, S2, S3)
+            A = Matrix(S)
+            @test diag(S)::SparseVector{T,Int} == diag(A)
+            for k in -size(S,1):size(S,2)
+                @test diag(S, k)::SparseVector{T,Int} == diag(A, k)
+            end
+            @test_throws ArgumentError diag(S, -size(S,1)-1)
+            @test_throws ArgumentError diag(S,  size(S,2)+1)
+        end
+    end
+    # test that stored zeros are still stored zeros in the diagonal
+    S = sparse([1,3],[1,3],[0.0,0.0]); V = diag(S)
+    @test V.nzind == [1,3]
+    @test V.nzval == [0.0,0.0]
+end
+
 @testset "expandptr" begin
     A = speye(5)
     @test Base.SparseArrays.expandptr(A.colptr) == collect(1:5)