Merge pull request #11514 from JuliaLang/sjk/fix-round-digits

Fix round(::AbstractArray, digits::Int) etc.

Author: simonster

base/floatfuncs.jl

@@ -77,6 +77,15 @@ for f in (:trunc,:floor,:ceil,:round)
         function ($f){T}(::Type{T}, x::AbstractArray)
             reshape([ ($f)(T, x[i])::T for i in eachindex(x) ], size(x))
         end
+        function ($f){R}(x::AbstractArray{R,1}, digits::Integer, base::Integer=10)
+            [ ($f)(x[i], digits, base) for i = 1:length(x) ]
+        end
+        function ($f){R}(x::AbstractArray{R,2}, digits::Integer, base::Integer=10)
+            [ ($f)(x[i,j], digits, base) for i = 1:size(x,1), j = 1:size(x,2) ]
+        end
+        function ($f)(x::AbstractArray, digits::Integer, base::Integer=10)
+            reshape([ ($f)(x[i], digits, base) for i in eachindex(x) ], size(x))
+        end
     end
 end
 
@@ -135,7 +144,7 @@ end
 
 for f in (:round, :ceil, :floor, :trunc)
     @eval begin
-        function ($f)(x, digits::Integer, base::Integer=10)
+        function ($f)(x::Real, digits::Integer, base::Integer=10)
             x = float(x)
             og = convert(eltype(x),base)^digits
             r = ($f)(x * og) / og
@@ -144,15 +153,9 @@ for f in (:round, :ceil, :floor, :trunc)
                 if digits > 0
                     return x
                 elseif x > 0
-                    if ceil == $f
-                        return convert(eltype(x), Inf)
-                    end
-                    return zero(x)
+                    return $(:ceil == f ? :(convert(eltype(x), Inf)) : :(zero(x)))
                 elseif x < 0
-                    if floor == $f
-                        return -convert(eltype(x), Inf)
-                    end
-                    return -zero(x)
+                    return $(:floor == f ? :(-convert(eltype(x), Inf)) : :(-zero(x)))
                 else
                     return x
                 end

test/numbers.jl

@@ -1860,6 +1860,16 @@ x = 0.0
 @test approx_eq(round(pi,2,2), 3.25)
 @test approx_eq(round(pi,3,2), 3.125)
 @test approx_eq(round(pi,3,5), 3.144)
+# vectorized trunc/round/floor/ceil with digits/base argument
+a = rand(2, 2, 2)
+for f in (trunc, round, floor, ceil)
+    @test f(a[:, 1, 1], 2) == map(x->f(x, 2), a[:, 1, 1])
+    @test f(a[:, :, 1], 2) == map(x->f(x, 2), a[:, :, 1])
+    @test f(a, 9, 2) == map(x->f(x, 9, 2), a)
+    @test f(a[:, 1, 1], 9, 2) == map(x->f(x, 9, 2), a[:, 1, 1])
+    @test f(a[:, :, 1], 9, 2) == map(x->f(x, 9, 2), a[:, :, 1])
+    @test f(a, 9, 2) == map(x->f(x, 9, 2), a)
+ end
 # significant digits (would be nice to have a smart vectorized
 # version of signif)
 @test approx_eq(signif(123.456,1), 100.)