fast extrema computation on sparse arrays

Author: Clément Fauchereau

stdlib/SparseArrays/src/higherorderfns.jl

@@ -4,7 +4,7 @@ module HigherOrderFns
 
 # This module provides higher order functions specialized for sparse arrays,
 # particularly map[!]/broadcast[!] for SparseVectors and SparseMatrixCSCs at present.
-import Base: map, map!, broadcast, copy, copyto!
+import Base: map, map!, broadcast, copy, copyto!, extrema
 
 using Base: front, tail, to_shape
 using ..SparseArrays: SparseVector, SparseMatrixCSC, AbstractSparseVector, AbstractSparseMatrixCSC,
@@ -29,6 +29,7 @@ using LinearAlgebra
 # (11) Define broadcast[!] methods handling combinations of scalars, sparse vectors/matrices,
 #       structured matrices, and one- and two-dimensional Arrays.
 # (12) Define map[!] methods handling combinations of sparse and structured matrices.
+# (13) Define extrema methods optimized for sparse vectors/matrices.
 
 
 # (0) BroadcastStyle rules and convenience types for dispatch
@@ -1154,4 +1155,27 @@ map(f::Tf, A::SparseOrStructuredMatrix, Bs::Vararg{SparseOrStructuredMatrix,N})
 map!(f::Tf, C::AbstractSparseMatrixCSC, A::SparseOrStructuredMatrix, Bs::Vararg{SparseOrStructuredMatrix,N}) where {Tf,N} =
     (_checksameshape(C, A, Bs...); _noshapecheck_map!(f, C, _sparsifystructured(A), map(_sparsifystructured, Bs)...))
 
+
+# (13) extrema methods optimized for sparse vectors/matrices.
+function extrema_(f, A::SparseVecOrMat)
+    M = length(A)
+    iszero(M) && throw(ArgumentError("Sparse array must have at least one element."))
+    N = nnz(A)
+    iszero(N) && return f(zero(eltype(A))), f(zero(eltype(A)))
+    vmin, vmax = extrema(f, nonzeros(A))
+    if N != M
+        f0 = f(zero(eltype(A)))
+        vmin = min(f0, vmin)
+        vmax = max(f0, vmax)
+    end
+    vmin, vmax
+end
+
+extrema(A::SparseVecOrMat; dims=:) = extrema(identity, A, dims=dims)
+extrema(f, A::SparseVecOrMat; dims=:) = extrema_(f, A, dims)
+extrema_(f, x::SparseVector, ::Colon) = extrema_(f, x)
+extrema_(f, A::AbstractSparseMatrixCSC, ::Colon) = extrema_(f, A)
+extrema_(f, x::SparseVector, dims) = invoke(extrema, Tuple{Any, AbstractArray}, f, x; dims=dims)
+extrema_(f, A::AbstractSparseMatrixCSC, dims) = invoke(extrema, Tuple{Any, AbstractArray}, f, A; dims=dims)
+
 end

stdlib/SparseArrays/test/higherorderfns.jl

@@ -687,4 +687,23 @@ end
     @test SparseMatStyle(Val(3)) == Broadcast.DefaultArrayStyle{3}()
 end
 
+@testset "extrema" begin
+    n = 10
+    A = sprand(n, n, 0.2)
+    B = Array(A)
+    x = sprand(n, 0.2)
+    y = Array(x)
+    f(x) = x^3
+    @test extrema(A) == extrema(B)
+    @test extrema(x) == extrema(y)
+    @test extrema(f, A) == extrema(f, B)
+    @test extrema(f, x) == extrema(f, y)
+    @test extrema(spzeros(n, n)) == (0.0, 0.0)
+    @test extrema(spzeros(n)) == (0.0, 0.0)
+    @test_throws ArgumentError extrema(spzeros(0, 0))
+    @test_throws ArgumentError extrema(spzeros(0))
+    @test extrema(sparse(ones(n, n))) == (1.0, 1.0)
+    @test extrema(sparse(ones(n))) == (1.0, 1.0)
+end
+
 end # module