Add more functions implemented with [map]reduce and [map]reducedim (#263)

* Use StaticArrays' reduce and mapreduce for iszero and count(f, a)

* Add more functions implemented with [map]reduce and [map]reducedim

* Handle possibility of eltype change in mapreducedim and diff

Author: wsshin

src/StaticArrays.jl

@@ -10,7 +10,7 @@ import Base: getindex, setindex!, size, similar, vec, show,
              hcat, vcat, ones, zeros, eye, one, cross, vecdot, reshape, fill,
              fill!, det, logdet, inv, eig, eigvals, expm, logm, sqrtm, lyap, trace, kron, diag, vecnorm, norm, dot, diagm, diag,
              lu, svd, svdvals, svdfact, factorize, ishermitian, issymmetric, isposdef,
-             sum, diff, prod, count, any, all, minimum,
+             iszero, sum, diff, prod, count, any, all, minimum,
              maximum, extrema, mean, copy, rand, randn, randexp, rand!, randn!,
              randexp!, normalize, normalize!, read, read!, write
 

src/mapreduce.jl

@@ -101,9 +101,11 @@ end
     _mapreducedim(f, op, Size(a), a, Val{D}, v0)
 end
 
-@generated function _mapreducedim(f, op, ::Size{S}, a::StaticArray, ::Type{Val{D}}) where {S, D}
+@generated function _mapreducedim(f, op, ::Size{S}, a::StaticArray, ::Type{Val{D}}) where {S,D}
     N = length(S)
     Snew = ([n==D ? 1 : S[n] for n = 1:N]...)
+    T0 = eltype(a)
+    T = :((T1 = Base.promote_op(f, $T0); Base.promote_op(op, T1, T1)))
 
     exprs = Array{Expr}(Snew)
     itr = [1:n for n ∈ Snew]
@@ -118,14 +120,13 @@ end
         exprs[i...] = expr
     end
 
-    # TODO element type might change
     return quote
         @_inline_meta
-        @inbounds return similar_type(a, Size($Snew))(tuple($(exprs...)))
+        @inbounds return similar_type(a, $T, Size($Snew))(tuple($(exprs...)))
     end
 end
 
-@generated function _mapreducedim(f, op, ::Size{S}, a::StaticArray, ::Type{Val{D}}, v0) where {S, D}
+@generated function _mapreducedim(f, op, ::Size{S}, a::StaticArray, ::Type{Val{D}}, v0::T) where {S,D,T}
     N = length(S)
     Snew = ([n==D ? 1 : S[n] for n = 1:N]...)
 
@@ -142,10 +143,9 @@ end
         exprs[i...] = expr
     end
 
-    # TODO element type might change
     return quote
         @_inline_meta
-        @inbounds return similar_type(a, Size($Snew))(tuple($(exprs...)))
+        @inbounds return similar_type(a, T, Size($Snew))(tuple($(exprs...)))
     end
 end
 
@@ -160,33 +160,82 @@ end
 ## reducedim ##
 ###############
 
-@inline reducedim(op, a::StaticArray, ::Val{D}) where {D} = mapreducedim(identity, op, a, Val{D})
-@inline reducedim(op, a::StaticArray, ::Val{D}, v0) where {D} = mapreducedim(identity, op, a, Val{D}, v0)
+@inline reducedim(op, a::StaticArray, ::Type{Val{D}}) where {D} = mapreducedim(identity, op, a, Val{D})
+@inline reducedim(op, a::StaticArray, ::Type{Val{D}}, v0) where {D} = mapreducedim(identity, op, a, Val{D}, v0)
 
 #######################
 ## related functions ##
 #######################
 
 # These are all similar in Base but not @inline'd
-@inline sum(a::StaticArray{<:Any, T}) where {T} = reduce(+, zero(T), a)
-@inline sum(f::Base.Callable, a::StaticArray) = mapreduce(f, +, a)
-@inline prod(a::StaticArray{<:Any, T}) where {T} = reduce(*, one(T), a)
-@inline count(a::StaticArray{<:Any, Bool}) = reduce(+, 0, a)
-@inline all(a::StaticArray{<:Any, Bool}) = reduce(&, true, a)  # non-branching versions
-@inline any(a::StaticArray{<:Any, Bool}) = reduce(|, false, a) # (benchmarking needed)
+#
+# Implementation notes:
+#
+# 1. When providing an initial value v0, note that its location is different in reduce and
+# reducedim: v0 comes earlier than collection in reduce, whereas it is the last argument in
+# reducedim.  The same difference exists between mapreduce and mapreducedim.
+#
+# 2. mapreduce and mapreducedim usually do not take initial value v0, because we don't
+# always know the return type of an arbitrary mapping function f.  (We usually want to use
+# some initial value such as one(T) or zero(T) as v0, where T is the return type of f, but
+# if users provide type-unstable f, its return type cannot be known.)  Therefore, mapped
+# versions of the functions implemented below usually require the collection to have at
+# least two entries.
+#
+# 3. Exceptions are the ones that require Boolean mapping functions.  For example, f in
+# all and any must return Bool, so we know the appropriate v0 is true and false,
+# respectively.  Therefore, all(f, ...) and any(f, ...) are implemented by mapreduce(f, ...)
+# with an initial value v0 = true and false.
+@inline iszero(a::StaticArray{<:Any,T}) where {T} = reduce((x,y) -> x && (y==zero(T)), true, a)
+
+@inline sum(a::StaticArray{<:Any,T}) where {T} = reduce(+, zero(T), a)
+@inline sum(f::Function, a::StaticArray) = mapreduce(f, +, a)
+@inline sum(a::StaticArray{<:Any,T}, ::Type{Val{D}}) where {T,D} = reducedim(+, a, Val{D}, zero(T))
+@inline sum(f::Function, a::StaticArray, ::Type{Val{D}}) where D = mapreducedim(f, +, a, Val{D})
+
+@inline prod(a::StaticArray{<:Any,T}) where {T} = reduce(*, one(T), a)
+@inline prod(f::Function, a::StaticArray{<:Any,T}) where {T} = mapreduce(f, *, a)
+@inline prod(a::StaticArray{<:Any,T}, ::Type{Val{D}}) where {T,D} = reducedim(*, a, Val{D}, one(T))
+@inline prod(f::Function, a::StaticArray{<:Any,T}, ::Type{Val{D}}) where {T,D} = mapreducedim(f, *, a, Val{D})
+
+@inline count(a::StaticArray{<:Any,Bool}) = reduce(+, 0, a)
+@inline count(f::Function, a::StaticArray) = mapreduce(x->f(x)::Bool, +, 0, a)
+@inline count(a::StaticArray{<:Any,Bool}, ::Type{Val{D}}) where {D} = reducedim(+, a, Val{D}, 0)
+@inline count(f::Function, a::StaticArray, ::Type{Val{D}}) where {D} = mapreducedim(x->f(x)::Bool, +, a, Val{D}, 0)
+
+@inline all(a::StaticArray{<:Any,Bool}) = reduce(&, true, a)  # non-branching versions
+@inline all(f::Function, a::StaticArray) = mapreduce(x->f(x)::Bool, &, true, a)
+@inline all(a::StaticArray{<:Any,Bool}, ::Type{Val{D}}) where {D} = reducedim(&, a, Val{D}, true)
+@inline all(f::Function, a::StaticArray, ::Type{Val{D}}) where {D} = mapreducedim(x->f(x)::Bool, &, a, Val{D}, true)
+
+@inline any(a::StaticArray{<:Any,Bool}) = reduce(|, false, a) # (benchmarking needed)
+@inline any(f::Function, a::StaticArray) = mapreduce(x->f(x)::Bool, |, false, a) # (benchmarking needed)
+@inline any(a::StaticArray{<:Any,Bool}, ::Type{Val{D}}) where {D} = reducedim(|, a, Val{D}, false)
+@inline any(f::Function, a::StaticArray, ::Type{Val{D}}) where {D} = mapreducedim(x->f(x)::Bool, |, a, Val{D}, false)
+
 @inline mean(a::StaticArray) = sum(a) / length(a)
+@inline mean(f::Function, a::StaticArray) = sum(f, a) / length(a)
+@inline mean(a::StaticArray, ::Type{Val{D}}) where {D} = sum(a, Val{D}) / size(a, D)
+@inline mean(f::Function, a::StaticArray, ::Type{Val{D}}) where {D} = sum(f, a, Val{D}) / size(a, D)
+
 @inline minimum(a::StaticArray) = reduce(min, a) # base has mapreduce(idenity, scalarmin, a)
+@inline minimum(f::Function, a::StaticArray) = mapreduce(f, min, a)
+@inline minimum(a::StaticArray, ::Type{Val{D}}) where {D} = reducedim(min, a, Val{D})
+@inline minimum(f::Function, a::StaticArray, ::Type{Val{D}}) where {D} = mapreducedim(f, min, a, Val{D})
+
 @inline maximum(a::StaticArray) = reduce(max, a) # base has mapreduce(idenity, scalarmax, a)
-@inline minimum(a::StaticArray, dim::Type{Val{D}}) where {D} = reducedim(min, a, dim)
-@inline maximum(a::StaticArray, dim::Type{Val{D}}) where {D} = reducedim(max, a, dim)
+@inline maximum(f::Function, a::StaticArray) = mapreduce(f, max, a)
+@inline maximum(a::StaticArray, ::Type{Val{D}}) where {D} = reducedim(max, a, Val{D})
+@inline maximum(f::Function, a::StaticArray, ::Type{Val{D}}) where {D} = mapreducedim(f, max, a, Val{D})
 
 # Diff is slightly different
 @inline diff(a::StaticArray) = diff(a, Val{1})
 @inline diff(a::StaticArray, ::Type{Val{D}}) where {D} = _diff(Size(a), a, Val{D})
 
-@generated function _diff(::Size{S}, a::StaticArray, ::Type{Val{D}}) where {S, D}
+@generated function _diff(::Size{S}, a::StaticArray, ::Type{Val{D}}) where {S,D}
     N = length(S)
     Snew = ([n==D ? S[n]-1 : S[n] for n = 1:N]...)
+    T = Base.promote_op(-, eltype(a), eltype(a))
 
     exprs = Array{Expr}(Snew)
     itr = [1:n for n = Snew]
@@ -197,9 +246,8 @@ end
         exprs[i1...] = :(a[$(i2...)] - a[$(i1...)])
     end
 
-    # TODO element type might change
     return quote
         @_inline_meta
-        @inbounds return similar_type(a, Size($Snew))(tuple($(exprs...)))
+        @inbounds return similar_type(a, $T, Size($Snew))(tuple($(exprs...)))
     end
 end

test/mapreduce.jl

@@ -23,48 +23,82 @@
         @test mv3 == @MVector [7, 9, 11, 13]
     end
 
-    @testset "reduce" begin
-        v1 = @SVector [2,4,6,8]
-        @test reduce(+, v1) === 20
-        @test reduce(+, 0, v1) === 20
-        @test sum(v1) === 20
-        @test sum(abs2, v1) === 120
-        @test prod(v1) === 384
-        @test mean(v1) === 5.
-        @test maximum(v1) === 8
-        @test minimum(v1) === 2
-        vb = @SVector [true, false, true, false]
-        @test count(vb) === 2
-        @test any(vb)
-    end
-
-    @testset "reduce in dim" begin
-        a = @SArray rand(4,3,2)
-        @test maximum(a, Val{1}) == maximum(a, 1)
-        @test maximum(a, Val{2}) == maximum(a, 2)
-        @test maximum(a, Val{3}) == maximum(a, 3)
-        @test minimum(a, Val{1}) == minimum(a, 1)
-        @test minimum(a, Val{2}) == minimum(a, 2)
-        @test minimum(a, Val{3}) == minimum(a, 3)
-        @test diff(a) == diff(a, Val{1}) == a[2:end,:,:] - a[1:end-1,:,:]
-        @test diff(a, Val{2}) == a[:,2:end,:] - a[:,1:end-1,:]
-        @test diff(a, Val{3}) == a[:,:,2:end] - a[:,:,1:end-1]
-
-        a = @SArray rand(4,3)  # as of Julia v0.5, diff() for regular Array is defined only for vectors and matrices
-        @test diff(a) == diff(a, Val{1}) == diff(a, 1)
-        @test diff(a, Val{2}) == diff(a, 2)
-
-        @test reducedim(max, a, Val{1}, -1.) == reducedim(max, a, 1, -1.)
-        @test reducedim(max, a, Val{2}, -1.) == reducedim(max, a, 2, -1.)
+    @testset "[map]reduce and [map]reducedim" begin
+        a = rand(4,3); sa = SMatrix{4,3}(a); (I,J) = size(a)
+        v1 = [2,4,6,8]; sv1 = SVector{4}(v1)
+        v2 = [4,3,2,1]; sv2 = SVector{4}(v2)
+        @test reduce(+, sv1) === reduce(+, v1)
+        @test reduce(+, 0, sv1) === reduce(+, 0, v1)
+        @test reducedim(max, sa, Val{1}, -1.) === SMatrix{1,J}(reducedim(max, a, 1, -1.))
+        @test reducedim(max, sa, Val{2}, -1.) === SMatrix{I,1}(reducedim(max, a, 2, -1.))
+        @test mapreduce(-, +, sv1) === mapreduce(-, +, v1)
+        @test mapreduce(-, +, 0, sv1) === mapreduce(-, +, 0, v1)
+        @test mapreduce(*, +, sv1, sv2) === 40
+        @test mapreduce(*, +, 0, sv1, sv2) === 40
+        @test mapreducedim(x->x^2, max, sa, Val{1}, -1.) == SMatrix{1,J}(mapreducedim(x->x^2, max, a, 1, -1.))
+        @test mapreducedim(x->x^2, max, sa, Val{2}, -1.) == SMatrix{I,1}(mapreducedim(x->x^2, max, a, 2, -1.))
     end
 
-    @testset "mapreduce" begin
-        v1 = @SVector [2,4,6,8]
-        v2 = @SVector [4,3,2,1]
-        @test mapreduce(-, +, v1) === -20
-        @test mapreduce(-, +, 0, v1) === -20
-        @test mapreduce(*, +, v1, v2) === 40
-        @test mapreduce(*, +, 0, v1, v2) === 40
+    @testset "implemented by [map]reduce and [map]reducedim" begin
+        I, J, K = 2, 2, 2
+        OSArray = SArray{Tuple{I,J,K}}  # original
+        RSArray1 = SArray{Tuple{1,J,K}}  # reduced in dimension 1
+        RSArray2 = SArray{Tuple{I,1,K}}  # reduced in dimension 2
+        RSArray3 = SArray{Tuple{I,J,1}}  # reduced in dimension 3
+        a = randn(I,J,K); sa = OSArray(a)
+        b = rand(Bool,I,J,K); sb = OSArray(b)
+        z = zeros(I,J,K); sz = OSArray(z)
+
+        @test iszero(sz) == iszero(z)
+
+        @test sum(sa) === sum(a)
+        @test sum(abs2, sa) === sum(abs2, a)
+        @test sum(sa, Val{2}) === RSArray2(sum(a, 2))
+        @test sum(abs2, sa, Val{2}) === RSArray2(sum(abs2, a, 2))
+
+        @test prod(sa) === prod(a)
+        @test prod(abs2, sa) === prod(abs2, a)
+        @test prod(sa, Val{2}) === RSArray2(prod(a, 2))
+        @test prod(abs2, sa, Val{2}) === RSArray2(prod(abs2, a, 2))
+
+        @test count(sb) === count(b)
+        @test count(x->x>0, sa) === count(x->x>0, a)
+        @test count(sb, Val{2}) === RSArray2(reshape([count(b[i,:,k]) for i = 1:I, k = 1:K], (I,1,K)))
+        @test count(x->x>0, sa, Val{2}) === RSArray2(reshape([count(x->x>0, a[i,:,k]) for i = 1:I, k = 1:K], (I,1,K)))
+
+        @test all(sb) === all(b)
+        @test all(x->x>0, sa) === all(x->x>0, a)
+        @test all(sb, Val{2}) === RSArray2(all(b, 2))
+        @test all(x->x>0, sa, Val{2}) === RSArray2(all(x->x>0, a, 2))
+
+        @test any(sb) === any(b)
+        @test any(x->x>0, sa) === any(x->x>0, a)
+        @test any(sb, Val{2}) === RSArray2(any(b, 2))
+        @test any(x->x>0, sa, Val{2}) === RSArray2(any(x->x>0, a, 2))
+
+        @test mean(sa) === mean(a)
+        @test mean(abs2, sa) === mean(abs2, a)
+        @test mean(sa, Val{2}) === RSArray2(mean(a, 2))
+        @test mean(abs2, sa, Val{2}) === RSArray2(mean(abs2.(a), 2))
+
+        @test minimum(sa) === minimum(a)
+        @test minimum(abs2, sa) === minimum(abs2, a)
+        @test minimum(sa, Val{2}) === RSArray2(minimum(a, 2))
+        @test minimum(abs2, sa, Val{2}) === RSArray2(minimum(abs2, a, 2))
+
+        @test maximum(sa) === maximum(a)
+        @test maximum(abs2, sa) === maximum(abs2, a)
+        @test maximum(sa, Val{2}) === RSArray2(maximum(a, 2))
+        @test maximum(abs2, sa, Val{2}) === RSArray2(maximum(abs2, a, 2))
+
+        @test diff(sa, Val{1}) === RSArray1(a[2:end,:,:] - a[1:end-1,:,:])
+        @test diff(sa, Val{2}) === RSArray2(a[:,2:end,:] - a[:,1:end-1,:])
+        @test diff(sa, Val{3}) === RSArray3(a[:,:,2:end] - a[:,:,1:end-1])
+
+        # as of Julia v0.6, diff() for regular Array is defined only for vectors and matrices
+        m = randn(4,3); sm = SMatrix{4,3}(m)
+        @test diff(sm, Val{1}) == diff(m, 1) == diff(sm) == diff(m)
+        @test diff(sm, Val{2}) == diff(m, 2)
     end
 
     @testset "broadcast and broadcast!" begin