Skip to content

Commit

Permalink
Add more functions implemented with [map]reduce and [map]reducedim (#263
Browse files Browse the repository at this point in the history 
)

* 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
  • Loading branch information
@wsshin @andyferris
wsshin authored and andyferris committed Aug 8, 2017
1 parent c5d7515 commit 7d711ab
Show file tree
Hide file tree
Showing 3 changed files with 142 additions and 60 deletions.
2 changes: 1 addition & 1 deletion src/StaticArrays.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
86 changes: 67 additions & 19 deletions src/mapreduce.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand All @@ -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]...)

Expand All @@ -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

Expand All @@ -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]
Expand All @@ -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
114 changes: 74 additions & 40 deletions test/mapreduce.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 7d711ab

Please sign in to comment.