diff --git a/src/sparsematrix.jl b/src/sparsematrix.jl index f2f8673b..f5ac4ad5 100644 --- a/src/sparsematrix.jl +++ b/src/sparsematrix.jl @@ -1777,6 +1777,64 @@ function permute(A::AbstractSparseMatrixCSC{Tv,Ti}, p::AbstractVector{<:Integer} unchecked_noalias_permute!(X, A, p, q, C) end +## Sorting + +#sorting TODO: integrate with `Base.Sort.IEEEFloatOptimization`'s partitioning by zero +searchsortedfirst_discard_keywords(v::AbstractVector, x; lt=isless, by=identity, + rev::Union{Bool,Nothing}=nothing, order::Base.Order.Ordering=Forward, kws...) = + searchsortedfirst(v, x, Base.Order.ord(lt,by,rev,order)) + +""" +Sort the stored entries of each column of `A` in place, rewriting the row indices so that +the values sorting before `zero(eltype(A))` end up at the top of their column and the +remaining values at the bottom, with the structural zeros in between. `nnz(A)` and the +column pointers are left untouched. +""" +function _sortcolumns!(A::AbstractSparseMatrixCSC; kws...) + require_one_based_indexing(A) + rows = rowvals(A) + vals = nonzeros(A) + m = size(A, 1) + z = zero(eltype(A)) + for j in axes(A, 2) + r = nzrange(A, j) + isempty(r) && continue + col = view(vals, r) + sort!(col; kws...) + # `i-1` stored values sort before the structural zeros and `length(r)-i+1` after + i = searchsortedfirst_discard_keywords(col, z; kws...) + k = first(r) + @inbounds for t in 1:i-1 + rows[k] = t + k += 1 + end + @inbounds for t in (m - length(r) + i):m + rows[k] = t + k += 1 + end + end + return A +end + +function Base.sort!(A::AbstractSparseMatrixCSC; dims::Integer, kws...) + if dims == 1 + _sortcolumns!(A; kws...) + elseif dims == 2 + # the rows of `A` are the columns of `transpose(A)`, which is cheap to form and + # cheap to transpose back once its columns are sorted + At = ftranspose(A, identity) + _sortcolumns!(At; kws...) + transpose!(A, At) + else + throw(ArgumentError(lazy"dimension out of range, got dims = $dims, expected 1 or 2")) + end + return A +end + +# the generic `Base.sort` for matrices goes through `permutedims`/`reshape` and does not +# return a `SparseMatrixCSC` for `dims = 1` +Base.sort(A::AbstractSparseMatrixCSC; kws...) = sort!(copy(A); kws...) + ## fkeep! and children tril!, triu!, droptol!, dropzeros[!] function _fkeep!(f::F, A::AbstractSparseMatrixCSC) where F<:Function diff --git a/src/sparsevector.jl b/src/sparsevector.jl index 7bb56f6e..84b72054 100644 --- a/src/sparsevector.jl +++ b/src/sparsevector.jl @@ -2272,10 +2272,8 @@ function _densifystarttolastnz!(x::SparseVector) x end -#sorting TODO: integrate with `Base.Sort.IEEEFloatOptimization`'s partitioning by zero -searchsortedfirst_discard_keywords(v::AbstractVector, x; lt=isless, by=identity, - rev::Union{Bool,Nothing}=nothing, order::Base.Order.Ordering=Forward, kws...) = - searchsortedfirst(v,x,Base.Order.ord(lt,by,rev,order)) +# `searchsortedfirst_discard_keywords` is defined alongside the sparse matrix sorting +# methods in sparsematrix.jl function sort!(x::AbstractCompressedVector; kws...) nz = nonzeros(x) sort!(nz; kws...) diff --git a/test/sparsematrix_ops.jl b/test/sparsematrix_ops.jl index a0d941a3..c2639e6d 100644 --- a/test/sparsematrix_ops.jl +++ b/test/sparsematrix_ops.jl @@ -759,4 +759,97 @@ end @test isdiag(S) end +@testset "sort/sort! of a sparse matrix" begin + # `sort` of a dense 0-dimension-along-`dims` matrix errors in Base, so those sizes are + # compared against the input itself rather than against a dense reference + @testset "size = ($m, $n), density = $d" for (m, n) in ((6, 5), (1, 1), (0, 3), (3, 0), + (1, 9), (9, 1), (20, 13)), + d in (0.0, 0.05, 0.3, 1.0) + A = sprand(m, n, d) + M = Matrix(A) + for dims in (1, 2), kws in ((;), (; rev=true), (; by=abs), (; by=x -> -x), + (; lt=(x, y) -> isless(y, x)), + (; alg=Base.DEFAULT_STABLE)) + expected = (m == 0 || n == 0) ? M : sort(M; dims, kws...) + B = copy(A) + @test sort!(B; dims, kws...) === B + @test B isa SparseMatrixCSC + @test Matrix(B) == expected + # sorting only moves the stored entries around + @test nnz(B) == nnz(A) + S = sort(A; dims, kws...) + @test S isa SparseMatrixCSC + @test Matrix(S) == expected + @test A == sparse(M) # `sort` leaves its argument alone + end + end + + @testset "index type $Ti" for Ti in (Int32, Int64) + A = SparseMatrixCSC{Float64,Ti}(sprand(11, 7, 0.4)) + for dims in (1, 2) + @test sort(A; dims) isa SparseMatrixCSC{Float64,Ti} + @test Matrix(sort(A; dims)) == sort(Matrix(A); dims) + end + end + + @testset "keyword arguments" begin + A = sprand(50, 50, 0.1) + # `scratch` is forwarded to the underlying `sort!` and ignored by the search for + # where the structural zeros belong (see #335) + @test Matrix(sort!(copy(A); dims=1, scratch=Vector{Float64}(undef, 50))) == + sort(Matrix(A); dims=1) + @test_throws MethodError sort!(copy(A); dims=1, banana=:blue) + @test_throws ArgumentError sort!(copy(A); dims=3) + @test_throws ArgumentError sort!(copy(A); dims=0) + @test_throws UndefKeywordError sort!(copy(A)) + end + + @testset "empty and zero-size matrices" begin + # `Base.sort` on a zero-size *dense* matrix throws `ArgumentError: step cannot be + # zero`, so there is no dense reference to compare against here; the sparse methods + # just return the (empty) matrix unchanged + @testset "size = ($m, $n)" for (m, n) in ((0, 3), (3, 0), (0, 0)) + A = spzeros(m, n) + for dims in (1, 2) + B = copy(A) + @test sort!(B; dims) === B + @test size(B) == (m, n) + @test nnz(B) == 0 + @test B == A + S = sort(A; dims) + @test S isa SparseMatrixCSC{Float64,Int} + @test size(S) == (m, n) + @test nnz(S) == 0 + end + end + + # structurally empty, but not zero-size: here dense does give a reference + @testset "all structural zeros, size = ($m, $n)" for (m, n) in ((1, 1), (5, 4)) + A = spzeros(m, n) + for dims in (1, 2) + B = sort!(copy(A); dims) + @test Matrix(B) == sort(Matrix(A); dims) + @test nnz(B) == 0 + @test getcolptr(B) == getcolptr(A) + end + end + + # a single column/row that is entirely structural next to a populated one + A = SparseMatrixCSC(4, 3, [1, 1, 5, 5], [1, 2, 3, 4], [1.0, -2.0, 0.0, 3.0]) + for dims in (1, 2) + @test Matrix(sort(A; dims)) == sort(Matrix(A); dims) + @test nnz(sort(A; dims)) == nnz(A) + end + end + + @testset "stored zeros" begin + # column 1 stores an explicit zero next to structural zeros + A = SparseMatrixCSC(4, 2, [1, 3, 4], [1, 3, 2], [0.0, -1.0, 2.0]) + for dims in (1, 2) + @test Matrix(sort(A; dims)) == sort(Matrix(A); dims) + @test nnz(sort(A; dims)) == nnz(A) + end + end +end + end # module