Merge pull request #71 from JuliaArrays/gs/map

map and map!
JuliaArrays · Mar 24, 2017 · c56dc40 · c56dc40
2 parents 4a7742b + b8d3ed6
commit c56dc40
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 1 deletion.
diff --git a/src/combine.jl b/src/combine.jl
@@ -9,7 +9,7 @@ function equalvalued(X::NTuple)
     return allequal
 end #equalvalued
 
-sizes{T<:AxisArray}(As::T...) = tuple(zip(map(size, As)...)...)
+sizes{T<:AxisArray}(As::T...) = tuple(zip(map(a -> map(length, indices(a)), As)...)...)
 matchingdims{N,T<:AxisArray}(As::NTuple{N,T}) = all(equalvalued, sizes(As...))
 matchingdimsexcept{N,T<:AxisArray}(As::NTuple{N,T}, n::Int) = all(equalvalued, sizes(As[[1:n-1; n+1:end]]...))
 

diff --git a/src/core.jl b/src/core.jl
@@ -290,6 +290,23 @@ Base.ctranspose{T}(A::AxisArray{T,2}) = AxisArray(ctranspose(A.data), A.axes[2],
 Base.transpose{T}(A::AxisArray{T,1})  = AxisArray(transpose(A.data), Axis{:transpose}(Base.OneTo(1)), A.axes[1])
 Base.ctranspose{T}(A::AxisArray{T,1}) = AxisArray(ctranspose(A.data), Axis{:transpose}(Base.OneTo(1)), A.axes[1])
 
+Base.map!{F}(f::F, A::AxisArray) = (map!(f, A.data); A)
+Base.map(f, A::AxisArray) = AxisArray(map(f, A.data), A.axes...)
+
+function Base.map!{F,T,N,D,Ax<:Tuple{Vararg{Axis}}}(f::F, dest::AxisArray{T,N,D,Ax},
+                                                  As::AxisArray{T,N,D,Ax}...)
+    matchingdims((dest, As...)) || error("All axes must be identically-valued")
+    data = map(a -> a.data, As)
+    map!(f, dest.data, data...)
+    return dest
+end
+
+function Base.map{T,N,D,Ax<:Tuple{Vararg{Axis}}}(f, As::AxisArray{T,N,D,Ax}...)
+    matchingdims(As) || error("All axes must be identically-valued")
+    data = map(a -> a.data, As)
+    return AxisArray(map(f, data...), As[1].axes...)
+end
+
 permutation(to::Union{AbstractVector{Int},Tuple{Int,Vararg{Int}}}, from::Symbols) = to
 
 """

diff --git a/test/core.jl b/test/core.jl
@@ -202,3 +202,26 @@ A[0] = 12
 A = AxisArray(OffsetArrays.OffsetArray(rand(4,5), -1:2, 5:9), :x, :y)
 @test indices(A) == (-1:2, 5:9)
 @test linearindices(A) == 1:20
+
+@test AxisArrays.matchingdims((A, A))
+
+f1(x) = x < 0
+A2 = map(f1, A)
+@test isa(A2, AxisArray)
+@test A2.axes == A.axes
+@test A2.data == map(f1, A.data)
+
+map!(~, A2)
+@test isa(A2, AxisArray)
+@test A2.axes == A.axes
+@test A2.data == ~map(f1, A).data
+
+A2 = map(+, A, A)
+@test isa(A2, AxisArray)
+@test A2.axes == A.axes
+@test A2.data == A.data .+ A.data
+
+map!(*, A2, A, A)
+@test isa(A2, AxisArray)
+@test A2.axes == A.axes
+@test A2.data == A.data .* A.data