Make sure reductions benefit from sparsity

src/sparsematrix.jl

@@ -209,8 +209,8 @@ nnz(S::LowerTriangular{<:Any,<:AbstractSparseMatrixCSC}) = nnz1(S)
 nnz(S::SparseMatrixCSCView) = nnz1(S)
 nnz1(S) = sum(length.(nzrange.(Ref(S), axes(S, 2))))
 
-function count(pred, S::AbstractSparseMatrixCSC)
-    count(pred, nzvalview(S)) + pred(zero(eltype(S)))*(prod(size(S)) - nnz(S))
+function Base._simple_count(pred, S::AbstractSparseMatrixCSC, init::T) where T
+    init + T(count(pred, nzvalview(S)) + pred(zero(eltype(S)))*(prod(size(S)) - nnz(S)))
 end
 
 """
@@ -304,7 +304,7 @@ function Base.isstored(A::AbstractSparseMatrixCSC, i::Integer, j::Integer)
     return false
 end
 
-function Base.isstored(A::Union{Adjoint{<:Any,<:AbstractSparseMatrixCSC},Transpose{<:Any,<:AbstractSparseMatrixCSC}}, i::Integer, j::Integer)
+function Base.isstored(A::AdjOrTrans{<:Any,<:AbstractSparseMatrixCSC}, i::Integer, j::Integer)
     @boundscheck checkbounds(A, i, j)
     cols = rowvals(parent(A))
     for istored in nzrange(parent(A), i)
@@ -2116,8 +2116,7 @@ nzeq(A::Transpose{<:Any,<:AbstractSparseMatrixCSCInclAdjointAndTranspose},
 # the case where the RHS is both adjoint and transposed, i.e. where it
 # is in CSC format again.)
 function ==(A::AbstractSparseMatrixCSC,
-            B::Union{Adjoint{<:Any,<:AbstractSparseMatrixCSCInclAdjointAndTranspose},
-                     Transpose{<:Any,<:AbstractSparseMatrixCSCInclAdjointAndTranspose}})
+            B::AdjOrTrans{<:Any,<:AbstractSparseMatrixCSCInclAdjointAndTranspose})
     # Different sizes are always different
     size(A) ≠ size(B) && return false
     # Compare nonzero elements
@@ -2182,15 +2181,22 @@ _mapreducezeros(f, op::Union{typeof(min),typeof(max)}, ::Type{T}, nzeros::Intege
 _mapreducezeros(f::Base.ExtremaMap, op::typeof(Base._extrema_rf), ::Type{T}, nzeros::Integer, v0) where {T} =
     nzeros == 0 ? v0 : op(v0, f(zero(T)))
 
-function Base._mapreduce(f, op::typeof(*), ::Base.IndexCartesian, A::AbstractSparseMatrixCSC{T}) where T
-    nzeros = widelength(A)-nnz(A)
+# Specialized mapreduce for any and all
+Base._any(f, A::AbstractSparseMatrixCSC, ::Colon) =
+    Base._mapreduce(f, |, IndexCartesian(), A)
+Base._all(f, A::AbstractSparseMatrixCSC, ::Colon) =
+    Base._mapreduce(f, &, IndexCartesian(), A)
+
+function Base._mapreduce(f, op::Union{typeof(Base.mul_prod),typeof(*)}, ::Base.IndexCartesian, A::AbstractSparseMatrixCSC{T}) where T
+    nnzA = nnz(A)
+    nzeros = widelength(A) - nnzA
     if nzeros == 0
         # No zeros, so don't compute f(0) since it might throw
         Base._mapreduce(f, op, nzvalview(A))
     else
         v = f(zero(T))^(nzeros)
-        # Bail out early if initial reduction value is zero
-        v == zero(T) ? v : v*Base._mapreduce(f, op, nzvalview(A))
+        # Bail out early if initial reduction value is zero or if there are no stored elements
+        (v == zero(T) || nnzA == 0) ? v : v*Base._mapreduce(f, op, nzvalview(A))
     end
 end
 

src/sparsevector.jl

@@ -92,7 +92,10 @@ const SVorFSV{Tv,Ti} = Union{SparseVector{Tv,Ti},FixedSparseVector{Tv,Ti}}
 
 length(x::SVorFSV)   = getfield(x, :n)
 size(x::SVorFSV)     = (getfield(x, :n),)
-count(f, x::AbstractCompressedVector) = count(f, nonzeros(x)) + f(zero(eltype(x)))*(length(x) - nnz(x))
+
+function Base._simple_count(f, x::AbstractCompressedVector, init::T) where T
+    init + T(count(f, nonzeros(x)) + f(zero(eltype(x)))*(length(x) - nnz(x)))
+end
 
 # implement the nnz - nzrange - nonzeros - rowvals interface for sparse vectors
 
@@ -1019,6 +1022,8 @@ function Vector(x::AbstractSparseVector{Tv}) where Tv
 end
 Array(x::AbstractSparseVector) = Vector(x)
 
+Base.iszero(x::AbstractSparseVector) = iszero(nonzeros(x))
+
 ### Array manipulation
 
 vec(x::AbstractSparseVector) = x
@@ -1487,6 +1492,9 @@ function Base._mapreduce(f, op, ::IndexCartesian, A::SparseVectorUnion{T}) where
     _mapreducezeros(f, op, T, rest, ini)
 end
 
+Base._any(f, A::SparseVectorUnion, ::Colon) = Base._mapreduce(f, |, IndexCartesian(), A)
+Base._all(f, A::SparseVectorUnion, ::Colon) = Base._mapreduce(f, &, IndexCartesian(), A)
+
 function Base.mapreducedim!(f, op, R::AbstractVector, A::SparseVectorUnion)
     # dim1 reduction could be safely replaced with a mapreduce
     if length(R) == 1

test/sparsematrix_constructors_indexing.jl

@@ -445,12 +445,18 @@ end
 
 @testset "setindex" begin
     a = spzeros(Int, 10, 10)
-    @test count(!iszero, a) == 0
+    @test count(!iszero, a) == count((!iszero).(a)) == 0
+    @test count(!iszero, a') == count((!iszero).(a')) == 0
+    @test count(!iszero, transpose(a)) == count(transpose((!iszero).(a))) == 0
     a[1,:] .= 1
-    @test count(!iszero, a) == 10
+    @test count(!iszero, a) == count((!iszero).(a)) == 10
+    @test count(!iszero, a, init=2) == count((!iszero).(a), init=2) == 12
+    @test count(!iszero, a, init=Int128(2))::Int128 == 12
+    @test count(!iszero, a') == count(((!iszero).(a))') == 10
+    @test count(!iszero, transpose(a)) == count(transpose((!iszero).(a))) == 10
     @test a[1,:] == sparse(fill(1,10))
     a[:,2] .= 2
-    @test count(!iszero, a) == 19
+    @test count(!iszero, a) == count((!iszero).(a)) == 19
     @test a[:,2] == sparse(fill(2,10))
     b = copy(a)
 

test/sparsematrix_ops.jl

@@ -174,7 +174,7 @@ dA = Array(sA)
     pA = sparse(rand(3, 7))
     p28227 = sparse(Real[0 0.5])
 
-    for arr in (se33, sA, pA, p28227)
+    for arr in (se33, sA, pA, p28227, spzeros(3, 3))
         for f in (sum, prod, minimum, maximum)
             farr = Array(arr)
             @test f(arr) ≈ f(farr)
@@ -183,6 +183,11 @@ dA = Array(sA)
             @test f(arr, dims=(1, 2)) ≈ [f(farr)]
             @test isequal(f(arr, dims=3), f(farr, dims=3))
         end
+        for f in (+, *, min, max)
+            farr = Array(arr)
+            @test mapreduce(identity, f, arr) ≈ mapreduce(identity, f, farr)
+            @test mapreduce(x -> x + 1, f, arr) ≈ mapreduce(x -> x + 1, f, farr)
+        end
     end
 
     for f in (sum, prod, minimum, maximum)
@@ -199,6 +204,34 @@ dA = Array(sA)
         # @test f(x->sqrt(x-1), pA .+ 1, dims=3) ≈ f(pA)
     end
 
+    @testset "logical reductions" begin
+        v = spzeros(Bool, 5, 2)
+        @test !any(v)
+        @test !all(v)
+        @test iszero(v)
+        @test count(v) == 0
+        v = SparseMatrixCSC(5, 2, [1, 2, 2], [1], [false])
+        @test !any(v)
+        @test !all(v)
+        @test iszero(v)
+        @test count(v) == 0
+        v = SparseMatrixCSC(5, 2, [1, 2, 2], [1], [true])
+        @test any(v)
+        @test !all(v)
+        @test !iszero(v)
+        @test count(v) == 1
+        v[2,1] = true
+        @test any(v)
+        @test !all(v)
+        @test !iszero(v)
+        @test count(v) == 2
+        v .= true
+        @test any(v)
+        @test all(v)
+        @test !iszero(v)
+        @test count(v) == length(v)
+    end
+
     @testset "empty cases" begin
         errchecker(str) = occursin("reducing over an empty collection is not allowed", str) ||
                           occursin("collection slices must be non-empty", str)

test/sparsevector.jl

@@ -33,6 +33,7 @@ x1_full[SparseArrays.nonzeroinds(spv_x1)] = nonzeros(spv_x1)
     @test SparseArrays.nonzeroinds(x) == [2, 5, 6]
     @test nonzeros(x) == [1.25, -0.75, 3.5]
     @test count(SparseVector(8, [2, 5, 6], [true,false,true])) == 2
+    @test count(SparseVector(8, [2, 5, 6], [true,false,true]), init=Int16(2))::Int16 == 4
     y = SparseVector(8, Int128[4], [5])
     @test y isa SparseVector{Int,Int128}
     @test @inferred size(y) == (@inferred(length(y))::Int128,)
@@ -931,6 +932,27 @@ end
         v[3] = 2
         @test argmax(v) == 3
     end
+
+    let
+        v = spzeros(Bool, 5)
+        @test !any(v)
+        @test !all(v)
+        @test iszero(v)
+        @test count(v) == 0
+        v = SparseVector(5, [1], [false])
+        @test !any(v)
+        @test !all(v)
+        @test iszero(v)
+        @test count(v) == 0
+        v[2] = true
+        @test any(v)
+        @test !all(v)
+        @test count(v) == 1
+        v .= true
+        @test any(v)
+        @test all(v)
+        @test count(v) == length(v)
+    end
 end
 
 ### linalg