map and mapreduce

src/fillbroadcast.jl

@@ -2,6 +2,58 @@
 
 map(f::Function, r::AbstractFill) = Fill(f(getindex_value(r)), axes(r))
 
+function map(f::Function, vs::AbstractFill{<:Any,1}...)
+    stop = mapreduce(length, min, vs)
+    val = f(map(getindex_value, vs)...)
+    Fill(val, stop)
+end
+
+function map(f::Function, rs::AbstractFill...)
+    if _maplinear(rs...)
+        map(f, map(vec, rs)...)
+    else
+        val = f(map(getindex_value, rs)...)
+        Fill(val, axes(first(rs)))
+    end
+end
+
+function _maplinear(rs...) # tries to match Base's behaviour, could perhaps hook in more deeply
+    if any(ndims(r)==1 for r in rs)
+        return true
+    else
+        r1 = axes(first(rs))
+        for r in rs
+            axes(r) == r1 || throw(DimensionMismatch(
+            "dimensions must match: a has dims $r1, b has dims $(axes(r))"))
+        end
+        return false
+    end
+end
+
+### mapreduce
+
+mapreduce(f, op, A::AbstractFill; kw...) = reduce(op, map(f, A); kw...)
+
+function mapreduce(f, op, A::AbstractFill, B::AbstractFill; kw...)
+    val(_...) = f(getindex_value(A), getindex_value(B))
+    reduce(op, map(val, A, B); kw...)
+end
+
+# These are particularly useful because mapreduce(*, +, A, B; dims) is slow in Base,
+# but can be re-written as some mapreduce(g, +, C; dims) which is fast.
+
+function mapreduce(f, op, A::AbstractFill, B::AbstractArray, Cs::AbstractArray...; kw...)
+    g(b, cs...) = f(getindex_value(A), b, cs...)
+    mapreduce(g, op, B, Cs...; kw...)
+end
+function mapreduce(f, op, A::AbstractArray, B::AbstractFill, Cs::AbstractArray...; kw...)
+    h(a, cs...) = f(a, getindex_value(B), cs...)
+    mapreduce(h, op, A, Cs...; kw...)
+end
+function mapreduce(f, op, A::AbstractFill, B::AbstractFill, Cs::AbstractArray...; kw...)
+    gh(cs...) = f(getindex_value(A), getindex_value(B), cs...)
+    mapreduce(gh, op, A, Cs...; kw...)
+end
 
 ### Unary broadcasting
 

test/runtests.jl

@@ -758,15 +758,46 @@ end
 end
 
 @testset "map" begin
-    x = Ones(5)
-    @test map(exp,x) === Fill(exp(1.0),5)
-    @test map(isone,x) === Fill(true,5)
+    x1 = Ones(5)
+    @test map(exp,x1) === Fill(exp(1.0),5)
+    @test map(isone,x1) === Fill(true,5)
 
-    x = Zeros(5)
-    @test map(exp,x) === exp.(x)
+    x0 = Zeros(5)
+    @test map(exp,x0) === exp.(x0)
 
-    x = Fill(2,5,3)
-    @test map(exp,x) === Fill(exp(2),5,3)
+    x2 = Fill(2,5,3)
+    @test map(exp,x2) === Fill(exp(2),5,3)
+
+    @test map(+, x1, x2) === Fill(3.0, 5)
+    @test map(+, x2, x2) === x2 .+ x2
+    @test_throws DimensionMismatch map(+, x2', x2)
+end
+
+@testset "mapreduce" begin
+    x = rand(3, 4)
+    y = fill(1.0, 3, 4)
+    Y = Fill(1.0, 3, 4)
+    O = Ones(3, 4)
+
+    @test_broken mapreduce(exp, +, Y) == mapreduce(exp, +, y)
+
+    # Two arrays
+    @test mapreduce(*, +, x, Y) == mapreduce(*, +, x, y)
+    @test mapreduce(*, +, Y, x) == mapreduce(*, +, y, x)
+    @test mapreduce(*, +, x, O) == mapreduce(*, +, x, y)
+    @test mapreduce(*, +, Y, O) == mapreduce(*, +, y, y)
+
+    @test mapreduce(*, +, x, Y, dims=1, init=5.0) == mapreduce(*, +, x, y, dims=1, init=5.0)
+    @test mapreduce(*, +, Y, x, dims=1, init=5.0) == mapreduce(*, +, y, x, dims=1, init=5.0)
+    @test mapreduce(*, +, x, O, dims=1, init=5.0) == mapreduce(*, +, x, y, dims=1, init=5.0)
+    @test mapreduce(*, +, Y, O, dims=1, init=5.0) == mapreduce(*, +, y, y, dims=1, init=5.0)
+
+    # More than two
+    @test mapreduce(*, +, x, Y, x) == mapreduce(*, +, x, y, x)
+    @test mapreduce(*, +, Y, x, x) == mapreduce(*, +, y, x, x)
+    @test mapreduce(*, +, x, O, Y) == mapreduce(*, +, x, y, y)
+    @test mapreduce(*, +, Y, O, Y) == mapreduce(*, +, y, y, y)
+    @test mapreduce(*, +, Y, O, Y, x) == mapreduce(*, +, y, y, y, x)
 end
 
 @testset "Offset indexing" begin