Skip to content

Sparse matrix: fix fast implementation of findnext and findprev for cartesian coordinates #32007

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
May 23, 2019
Merged

Sparse matrix: fix fast implementation of findnext and findprev for cartesian coordinates #32007

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
310a47a
Revert "sparse findnext findprev hash performance improved (#31354)"
tkluck May 11, 2019
c45eff0
_sparse_findnextnz: work with CartesianIndex
tkluck May 11, 2019
54802d1
sparse findXXX: cherry-pick test cases from #31354
KlausC May 11, 2019
5b73a97
sparse findprev/next: skip zeros if the predicate is false on zero(el…
tkluck May 11, 2019
9c3a15b
_sparse_findnextnz: search m.colptr using searchsorted for performance
tkluck May 18, 2019
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 17 additions & 7 deletions stdlib/SparseArrays/src/abstractsparse.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,21 +63,31 @@ end

# The following two methods should be overloaded by concrete types to avoid
# allocating the I = findall(...)
_sparse_findnextnz(v::AbstractSparseArray, i::Integer) = (I = findall(!iszero, v); n = searchsortedfirst(I, i); n<=length(I) ? I[n] : nothing)
_sparse_findprevnz(v::AbstractSparseArray, i::Integer) = (I = findall(!iszero, v); n = searchsortedlast(I, i); !iszero(n) ? I[n] : nothing)

function findnext(f::typeof(!iszero), v::AbstractSparseArray, i::Integer)
_sparse_findnextnz(v::AbstractSparseArray, i) = (I = findall(!iszero, v); n = searchsortedfirst(I, i); n<=length(I) ? I[n] : nothing)
_sparse_findprevnz(v::AbstractSparseArray, i) = (I = findall(!iszero, v); n = searchsortedlast(I, i); !iszero(n) ? I[n] : nothing)

function findnext(f::Function, v::AbstractSparseArray, i)
# short-circuit the case f == !iszero because that avoids
# allocating e.g. zero(BigInt) for the f(zero(...)) test.
if nnz(v) == length(v) || (f != (!iszero) && f(zero(eltype(v))))
return invoke(findnext, Tuple{Function,Any,Any}, f, v, i)
end
j = _sparse_findnextnz(v, i)
while j !== nothing && !f(v[j])
j = _sparse_findnextnz(v, j+1)
j = _sparse_findnextnz(v, nextind(v, j))
end
return j
end

function findprev(f::typeof(!iszero), v::AbstractSparseArray, i::Integer)
function findprev(f::Function, v::AbstractSparseArray, i)
# short-circuit the case f == !iszero because that avoids
# allocating e.g. zero(BigInt) for the f(zero(...)) test.
if nnz(v) == length(v) || (f != (!iszero) && f(zero(eltype(v))))
return invoke(findprev, Tuple{Function,Any,Any}, f, v, i)
end
j = _sparse_findprevnz(v, i)
while j !== nothing && !f(v[j])
j = _sparse_findprevnz(v, j-1)
j = _sparse_findprevnz(v, prevind(v, j))
end
return j
end
Expand Down
105 changes: 16 additions & 89 deletions stdlib/SparseArrays/src/sparsematrix.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1312,36 +1312,34 @@ function findnz(S::SparseMatrixCSC{Tv,Ti}) where {Tv,Ti}
return (I, J, V)
end

function _sparse_findnextnz(m::SparseMatrixCSC, i::Integer)
if i > length(m)
return nothing
end
row, col = Tuple(CartesianIndices(m)[i])
function _sparse_findnextnz(m::SparseMatrixCSC, ij::CartesianIndex{2})
row, col = Tuple(ij)
col > m.n && return nothing

lo, hi = m.colptr[col], m.colptr[col+1]
n = searchsortedfirst(m.rowval, row, lo, hi-1, Base.Order.Forward)
if lo <= n <= hi-1
return LinearIndices(m)[m.rowval[n], col]
return CartesianIndex(m.rowval[n], col)
end
nextcol = findnext(c->(c>hi), m.colptr, col+1)
nextcol === nothing && return nothing
nextcol = searchsortedfirst(m.colptr, hi + 1, col + 1, length(m.colptr), Base.Order.Forward)
nextcol > length(m.colptr) && return nothing
nextlo = m.colptr[nextcol-1]
return LinearIndices(m)[m.rowval[nextlo], nextcol-1]
return CartesianIndex(m.rowval[nextlo], nextcol - 1)
end

function _sparse_findprevnz(m::SparseMatrixCSC, i::Integer)
if iszero(i)
return nothing
end
row, col = Tuple(CartesianIndices(m)[i])
function _sparse_findprevnz(m::SparseMatrixCSC, ij::CartesianIndex{2})
row, col = Tuple(ij)
iszero(col) && return nothing

lo, hi = m.colptr[col], m.colptr[col+1]
n = searchsortedlast(m.rowval, row, lo, hi-1, Base.Order.Forward)
if lo <= n <= hi-1
return LinearIndices(m)[m.rowval[n], col]
return CartesianIndex(m.rowval[n], col)
end
prevcol = findprev(c->(c<lo), m.colptr, col-1)
prevcol === nothing && return nothing
prevcol = searchsortedlast(m.colptr, lo - 1, 1, col - 1, Base.Order.Forward)
prevcol < 1 && return nothing
prevhi = m.colptr[prevcol+1]
return LinearIndices(m)[m.rowval[prevhi-1], prevcol]
return CartesianIndex(m.rowval[prevhi-1], prevcol)
end


Expand All @@ -1361,77 +1359,6 @@ function sparse_sortedlinearindices!(I::Vector{Ti}, V::Vector, m::Int, n::Int) w
return SparseMatrixCSC(m, n, colptr, I, V)
end

# findfirst/next/prev/last
function _idxfirstnz(A::SparseMatrixCSC, ij::CartesianIndex{2})
nzr = nzrange(A, ij[2])
searchk = searchsortedfirst(A.rowval, ij[1], first(nzr), last(nzr), Forward)
return _idxnextnz(A, searchk)
end

function _idxlastnz(A::SparseMatrixCSC, ij::CartesianIndex{2})
nzr = nzrange(A, ij[2])
searchk = searchsortedlast(A.rowval, ij[1], first(nzr), last(nzr), Forward)
return _idxprevnz(A, searchk)
end

function _idxnextnz(A::SparseMatrixCSC, idx::Integer)
nnza = nnz(A)
nzval = nonzeros(A)
z = zero(eltype(A))
while idx <= nnza
nzv = nzval[idx]
!isequal(nzv, z) && return idx, nzv
idx += 1
end
return zero(idx), z
end

function _idxprevnz(A::SparseMatrixCSC, idx::Integer)
nzval = nonzeros(A)
z = zero(eltype(A))
while idx > 0
nzv = nzval[idx]
!isequal(nzv, z) && return idx, nzv
idx -= 1
end
return zero(idx), z
end

function _idx_to_cartesian(A::SparseMatrixCSC, idx::Integer)
rowval = rowvals(A)
i = rowval[idx]
j = searchsortedlast(A.colptr, idx, 1, size(A, 2), Base.Order.Forward)
return CartesianIndex(i, j)
end

function Base.findnext(pred::Function, A::SparseMatrixCSC, ij::CartesianIndex{2})
if nnz(A) == length(A) || pred(zero(eltype(A)))
return invoke(findnext, Tuple{Function,Any,Any}, pred, A, ij)
end
idx, nzv = _idxfirstnz(A, ij)
while idx > 0
if pred(nzv)
return _idx_to_cartesian(A, idx)
end
idx, nzv = _idxnextnz(A, idx + 1)
end
return nothing
end

function Base.findprev(pred::Function, A::SparseMatrixCSC, ij::CartesianIndex{2})
if nnz(A) == length(A) || pred(zero(eltype(A)))
return invoke(findprev, Tuple{Function,Any,Any}, pred, A, ij)
end
idx, nzv = _idxlastnz(A, ij)
while idx > 0
if pred(nzv)
return _idx_to_cartesian(A, idx)
end
idx, nzv = _idxprevnz(A, idx - 1)
end
return nothing
end

"""
sprand([rng],[type],m,[n],p::AbstractFloat,[rfn])

Expand Down