[WIP] generic sparse indices

NEWS.md

@@ -34,6 +34,7 @@ New library functions
 * `isunordered(x)` returns true if `x` is value that is normally unordered, such as `NaN` or `missing`.
 * New macro `Base.@invokelatest f(args...; kwargs...)` provides a convenient way to call `Base.invokelatest(f, args...; kwargs...)` ([#37971])
 * New macro `Base.@invoke f(arg1::T1, arg2::T2; kwargs...)` provides an easier syntax to call `invoke(f, Tuple{T1,T2}, arg1, arg2; kwargs...)` ([#38438])
+* New function `Base.eachstoredindex` returns an index iterator over the structural non-zero indices of a container ([#40103])
 
 New library features
 --------------------

base/abstractarray.jl

@@ -323,6 +323,16 @@ function _all_match_first(f::F, inds, A, B...) where F<:Function
 end
 _all_match_first(f::F, inds) where F<:Function = true
 
+"""
+    eachstoredindex(A)
+
+Returns an iterable over the indices of `A` where the values are structurally non-zero.
+It falls back to `eachindex(A)` and can be redefined by array types.
+`eachstoredindex(A)` is not guaranteed to return the same shape as `eachindex(A)`.
+"""
+eachstoredindex(A) = eachindex(A)
+
+
 # keys with an IndexStyle
 keys(s::IndexStyle, A::AbstractArray, B::AbstractArray...) = eachindex(s, A, B...)
 

stdlib/LinearAlgebra/src/diagonal.jl

@@ -111,6 +111,8 @@ end
 
 parent(D::Diagonal) = D.diag
 
+Base.eachstoredindex(D::Diagonal) = diagind(D)
+
 ishermitian(D::Diagonal{<:Real}) = true
 ishermitian(D::Diagonal{<:Number}) = isreal(D.diag)
 ishermitian(D::Diagonal) = all(ishermitian, D.diag)

stdlib/SparseArrays/src/sparsematrix.jl

@@ -1515,7 +1515,7 @@ function findnz(S::AbstractSparseMatrixCSC{Tv,Ti}) where {Tv,Ti}
     V = Vector{Tv}(undef, numnz)
 
     count = 1
-    @inbounds for col = 1 : size(S, 2), k = getcolptr(S)[col] : (getcolptr(S)[col+1]-1)
+    @inbounds for col in 1:size(S, 2), k in nzrange(S, col)
         I[count] = rowvals(S)[k]
         J[count] = col
         V[count] = nonzeros(S)[k]
@@ -1525,6 +1525,33 @@ function findnz(S::AbstractSparseMatrixCSC{Tv,Ti}) where {Tv,Ti}
     return (I, J, V)
 end
 
+function Base.eachstoredindex(S::AbstractSparseMatrixCSC{Tv,Ti}) where {Tv,Ti}
+    numnz = nnz(S)
+    indices = Vector{CartesianIndex{2}}(undef, numnz)
+    count = 1
+    @inbounds for col in 1:size(S, 2), k in nzrange(S, col)
+        indices[count] = CartesianIndex(rowvals(S)[k], col)
+        count += 1
+    end
+    return indices
+end
+
+function Base.eachstoredindex(S::Union{Adjoint{T, ST}, Transpose{T, ST}}) where {T, ST <: AbstractSparseMatrixCSC{T}}
+    P = parent(S)
+    numnz = nnz(P)
+    indices = Vector{CartesianIndex{2}}(undef, numnz)
+    count = 1
+    @inbounds for col = 1 : size(P, 2), k = getcolptr(P)[col] : (getcolptr(P)[col+1]-1)
+        indices[count] = CartesianIndex(col, rowvals(P)[k])
+        count += 1
+    end
+    return indices
+end
+
+function Base.eachstoredindex(S::Union{Symmetric{T, ST}, LinearAlgebra.AbstractTriangular{T, ST}, UpperHessenberg{T, ST}, Hermitian{T, ST}}) where {T, ST <: AbstractSparseMatrixCSC{T}}
+    return Base.eachstoredindex(parent(S))
+end
+
 function _sparse_findnextnz(m::AbstractSparseMatrixCSC, ij::CartesianIndex{2})
     row, col = Tuple(ij)
     col > size(m, 2) && return nothing

stdlib/SparseArrays/src/sparsevector.jl

@@ -77,6 +77,9 @@ function nonzeroinds(x::SparseColumnView)
     @inbounds y = view(rowvals(A), nzrange(A, colidx))
     return y
 end
+
+Base.eachstoredindex(x::SparseVector) = getfield(x, :nzind)
+
 nonzeroinds(x::SparseVectorView) = nonzeroinds(parent(x))
 
 rowvals(x::SparseVectorUnion) = nonzeroinds(x)

stdlib/SparseArrays/test/sparse.jl

@@ -1833,7 +1833,9 @@ 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 nonzeroinds(V) == [1,3]
+    @test Base.eachstoredindex(V) == [1,3]
     @test nonzeros(V) == [0.0,0.0]
+    @test V[Base.eachstoredindex(V)] == nonzeros(V)
 end
 
 @testset "expandptr" begin
@@ -2875,6 +2877,14 @@ end
         @test SparseMatrixCSC(at(wr(A))) == Matrix(at(wr(B)))
     end
 
+    @testset "eachstoredindex($(wr))" for wr in (UpperTriangular, LowerTriangular,
+                                                 UnitUpperTriangular, UnitLowerTriangular,
+                                                 Hermitian, (Hermitian, :L), Symmetric, (Symmetric, :L), Transpose, Adjoint)
+        S = dowrap(wr, A)
+        sum(S[Base.eachstoredindex(S)]) == sum(S)
+        sum(S[Base.eachstoredindex(S)]) == sum(Matrix(S))
+    end
+
     @test sparse([1,2,3,4,5]') == SparseMatrixCSC([1 2 3 4 5])
     @test sparse(UpperTriangular(A')) == UpperTriangular(B')
     @test sparse(Adjoint(UpperTriangular(A'))) == Adjoint(UpperTriangular(B'))

stdlib/SparseArrays/test/sparsevector.jl

@@ -31,7 +31,9 @@ x1_full[SparseArrays.nonzeroinds(spv_x1)] = nonzeros(spv_x1)
     @test count(!iszero, x) == 3
     @test nnz(x) == 3
     @test SparseArrays.nonzeroinds(x) == [2, 5, 6]
+    @test Base.eachstoredindex(x) == [2, 5, 6]
     @test nonzeros(x) == [1.25, -0.75, 3.5]
+    @test nonzeros(x) == x[Base.eachstoredindex(x)]
     @test count(SparseVector(8, [2, 5, 6], [true,false,true])) == 2
 end