Merge pull request #86 from wsshin/reducer-with-dim

Implement versions of maximum(), minimum(), diff() with Val{dim} as second argument

Author: andyferris

src/StaticArrays.jl

@@ -3,12 +3,12 @@ __precompile__()
 module StaticArrays
 
 import Base: @pure, @propagate_inbounds, getindex, setindex!, size, similar,
-             length, convert, promote_op, map, map!, reduce, mapreduce,
-             broadcast, broadcast!, conj, transpose, ctranspose, hcat, vcat,
-             ones, zeros, eye, one, cross, vecdot, reshape, fill, fill!, det,
-             inv, eig, eigvals, trace, vecnorm, norm, dot, diagm, sum, prod,
-             count, any, all, sumabs, sumabs2, minimum, maximum, extrema, mean,
-             copy
+             length, convert, promote_op, map, map!, reduce, reducedim,
+             mapreduce, broadcast, broadcast!, conj, transpose, ctranspose,
+             hcat, vcat, ones, zeros, eye, one, cross, vecdot, reshape, fill,
+             fill!, det, inv, eig, eigvals, trace, vecnorm, norm, dot, diagm,
+             sum, diff, prod, count, any, all, sumabs, sumabs2, minimum,
+             maximum, extrema, mean, copy
 
 export StaticScalar, StaticArray, StaticVector, StaticMatrix
 export Scalar, SArray, SVector, SMatrix

src/mapreduce.jl

@@ -90,6 +90,36 @@ end
     end
 end
 
+@generated function reducedim{D}(op, a::StaticArray, ::Type{Val{D}})
+    S = size(a)
+    if S[D] == 1
+        return :(return a)
+    else
+        N = ndims(a)
+        Snew = ([n==D ? 1 : S[n] for n = 1:N]...)
+        newtype = similar_type(a, Snew)
+
+        exprs = Array{Expr}(Snew)
+        itr = [1:n for n = Snew]
+        for i = Base.product(itr...)
+            ik = copy([i...])
+            ik[D] = 2
+            expr = :(op(a[$(i...)], a[$(ik...)]))
+            for k = 3:S[D]
+                ik[D] = k
+                expr = :(op($expr, a[$(ik...)]))
+            end
+
+            exprs[i...] = expr
+        end
+
+        return quote
+            $(Expr(:meta,:inline))
+            @inbounds return $(Expr(:call, newtype, Expr(:tuple, exprs...)))
+        end
+    end
+end
+
 # These are all similar in Base but not @inline'd
 @inline sum{T}(a::StaticArray{T}) = reduce(+, zero(T), a)
 @inline prod{T}(a::StaticArray{T}) = reduce(*, one(T), a)
@@ -101,8 +131,29 @@ end
 @inline sumabs2{T}(a::StaticArray{T}) = mapreduce(abs2, +, zero(T), a)
 @inline minimum(a::StaticArray) = reduce(min, a) # base has mapreduce(idenity, scalarmin, a)
 @inline maximum(a::StaticArray) = reduce(max, a) # base has mapreduce(idenity, scalarmax, a)
+@inline minimum{D}(a::StaticArray, dim::Type{Val{D}}) = reducedim(min, a, dim)
+@inline maximum{D}(a::StaticArray, dim::Type{Val{D}}) = reducedim(max, a, dim)
+
+@generated function diff{D}(a::StaticArray, ::Type{Val{D}}=Val{1})
+    S = size(a)
+    N = ndims(a)
+    Snew = ([n==D ? S[n]-1:S[n] for n = 1:N]...)
+    newtype = similar_type(a, Snew)
+
+    exprs = Array{Expr}(Snew)
+    itr = [1:n for n = Snew]
+
+    for i1 = Base.product(itr...)
+        i2 = copy([i1...])
+        i2[D] = i1[D] + 1
+        exprs[i1...] = :(a[$(i2...)] - a[$(i1...)])
+    end
 
-
+    return quote
+        $(Expr(:meta,:inline))
+        @inbounds return $(Expr(:call, newtype, Expr(:tuple, exprs...)))
+    end
+end
 
 ###############
 ## mapreduce ##

test/mapreduce.jl

@@ -24,6 +24,23 @@
         @test prod(v1) === 384
     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)
+    end
+
     @testset "mapreduce" begin
         v1 = @SVector [2,4,6,8]
         v2 = @SVector [4,3,2,1]