Add efficient O(nnz) isdiag for sparse matrices

src/SparseArrays.jl

@@ -17,7 +17,7 @@ using LinearAlgebra: AdjOrTrans, AdjointFactorization, TransposeFactorization, m
 import Base: +, -, *, \, /, ==, zero
 import Base: Matrix, Vector
 import LinearAlgebra: mul!, ldiv!, rdiv!, cholesky, adjoint!, diag, eigen, dot,
-    issymmetric, istril, istriu, lu, tr, transpose!, tril!, triu!, isbanded,
+    issymmetric, istril, istriu, lu, tr, transpose!, tril!, triu!, isbanded, isdiag,
     cond, diagm, factorize, ishermitian, norm, opnorm, lmul!, rmul!, tril, triu,
     matprod_dest, generic_matvecmul!, generic_matmatmul!, copytrito!
 

src/sparsematrix.jl

@@ -4200,6 +4200,21 @@ function istril(A::AbstractSparseMatrixCSC, k::Integer=0)
     return true
 end
 
+function isdiag(A::AbstractSparseMatrixCSC)
+    m, n = size(A)
+    colptr = getcolptr(A)
+    rowval = rowvals(A)
+    nzval = nonzeros(A)
+    @inbounds for col in 1:n
+        for k in colptr[col]:(colptr[col + 1] - 1)
+            if rowval[k] != col && _isnotzero(nzval[k])
+                return false
+            end
+        end
+    end
+    return true
+end
+
 _nnz(v::AbstractSparseVector) = nnz(v)
 _nnz(v::AbstractVector) = length(v)
 

test/sparsematrix_ops.jl

@@ -639,4 +639,38 @@ end
     end
 end
 
+@testset "isdiag" begin
+    # Diagonal matrices should return true
+    @test isdiag(sparse(Diagonal(1:4)))
+    @test isdiag(sparse(Diagonal([1.0, 2.0, 3.0])))
+    @test isdiag(spzeros(5, 5))  # Empty matrix is diagonal
+
+    # Non-diagonal matrices should return false
+    @test !isdiag(sparse(Tridiagonal(1:3, 1:4, 1:3)))
+    @test !isdiag(sparse(Bidiagonal(1:4, 1:3, :U)))
+    @test !isdiag(sparse(Bidiagonal(1:4, 1:3, :L)))
+    @test !isdiag(sparse([1 2; 3 4]))
+
+    # Non-square diagonal matrices should return true (consistent with generic isdiag)
+    @test isdiag(sparse([1 0 0; 0 2 0]))  # 2x3 diagonal
+    @test isdiag(sparse([1 0; 0 2; 0 0]))  # 3x2 diagonal
+    @test isdiag(spzeros(3, 5))  # Empty non-square matrix is diagonal
+    @test isdiag(spzeros(5, 3))
+
+    # Non-square non-diagonal matrices should return false
+    @test !isdiag(sparse([1 1 0; 0 2 0]))  # Off-diagonal element
+    @test !isdiag(sparse([1 0; 0 2; 1 0]))  # Off-diagonal element
+
+    # Consistency with dense isdiag
+    for T in Any[Diagonal(1:4), Tridiagonal(1:3, 1:4, 1:3),
+                 Bidiagonal(1:4, 1:3, :U), diagm(-1=>1:3, 1=>1:3)]
+        S = sparse(T)
+        @test isdiag(S) == isdiag(T)
+    end
+
+    # Explicit zeros on off-diagonal should still be diagonal
+    S = sparse([1, 2, 1], [1, 2, 2], [1.0, 2.0, 0.0])
+    @test isdiag(S)
+end
+
 end # module