Reduce to flatten{N} only

Author: iamed2

src/AxisArrays.jl

@@ -3,6 +3,7 @@ __precompile__()
 module AxisArrays
 
 using Base: tail
+import Base.Iterators: repeated
 using RangeArrays, IntervalSets
 using IterTools
 using Compat

src/categoricalvector.jl

@@ -42,8 +42,12 @@ A[(:a,:x), :]
 A[(:a,:x,:x), :]
 ```
 """
-immutable CategoricalVector{T} <: AbstractVector{T}
-    data::AbstractVector{T}
+immutable CategoricalVector{T, A<:AbstractVector{T}} <: AbstractVector{T}
+    data::A
+end
+
+function CategoricalVector(data::AbstractVector{T}) where T
+    CategoricalVector{T, typeof(data)}(data)
 end
 
 Base.getindex(v::CategoricalVector, idx::Int) = v.data[idx]
@@ -54,16 +58,16 @@ Base.size(v::CategoricalVector) = size(v.data)
 Base.size(v::CategoricalVector, i) = size(v.data, i)
 Base.indices(v::CategoricalVector) = indices(v.data)
 
-axistrait{T}(::Type{CategoricalVector{T}}) = Categorical
+axistrait(::Type{CategoricalVector{T,A}}) where {T,A} = Categorical
 checkaxis(::CategoricalVector) = nothing
 
 
 ## Add some special indexing for CategoricalVector{Tuple}'s to achieve something like
 ## Panda's hierarchical indexing
 
-axisindexes{T<:Tuple,S}(ax::Axis{S,CategoricalVector{T}}, idx) = axisindexes(ax, (idx,))
+axisindexes{T<:Tuple,S,A}(ax::Axis{S,CategoricalVector{T,A}}, idx) = axisindexes(ax, (idx,))
 
-function axisindexes{T<:Tuple,S}(ax::Axis{S,CategoricalVector{T}}, idx::Tuple)
+function axisindexes{T<:Tuple,S,A}(ax::Axis{S,CategoricalVector{T,A}}, idx::Tuple)
     collect(filter(ax_idx->_tuple_matches(ax.val[ax_idx], idx), indices(ax.val)...))
 end
 
@@ -77,5 +81,5 @@ function _tuple_matches(element::Tuple, idx::Tuple)
     return true
 end
 
-axisindexes{T<:Tuple,S}(ax::Axis{S,CategoricalVector{T}}, idx::AbstractArray) =
+axisindexes{T<:Tuple,S,A}(ax::Axis{S,CategoricalVector{T,A}}, idx::AbstractArray) =
     vcat([axisindexes(ax, i) for i in idx]...)

src/combine.jl

@@ -140,88 +140,122 @@ function Base.join{T,N,D,Ax}(As::AxisArray{T,N,D,Ax}...; fillvalue::T=zero(T),
 
 end #join
 
-function greatest_common_axis(As::AxisArray...)
-    length(As) == 1 && return ndims(first(As))
+function _flatten_array_axes(array_name, array_axes...)
+    ((array_name, (idx isa Tuple ? idx : (idx,))...) for idx in product((Ax.val for Ax in array_axes)...))
+end
 
-    for (i, zip_axes) in enumerate(zip(axes.(As)...))
-        if !all(ax -> ax == zip_axes[1], zip_axes[2:end])
-            return i - 1
-        end
+function _flatten_axes(array_names, array_axes)
+    collect(Iterators.flatten(map(array_names, array_axes) do tup_name, tup_array_axes
+        _flatten_array_axes(tup_name, tup_array_axes...)
+    end))
+end
+
+function _splitall{N}(::Type{Val{N}}, As...)
+    tuple((Base.IteratorsMD.split(A, Val{N}) for A in As)...)
+end
+
+function _reshapeall{N}(::Type{Val{N}}, As...)
+    tuple((reshape(A, Val{N}) for A in As)...)
+end
+
+function _check_common_axes(common_axis_tuple)
+    if !all(axisname(first(common_axis_tuple)) .=== axisname.(common_axis_tuple[2:end]))
+        throw(ArgumentError("Leading common axes must have the same name in each array"))
     end
 
-    return minimum(map(ndims, As))
+    return nothing
 end
 
-function flatten_array_axes(array_name, array_axes)
-    map(zip(repeated(array_name), product(map(Ax->Ax.val, array_axes)...))) do tup
-        tup_name, tup_idx = tup
-        return (tup_name, tup_idx...)
+function _flat_axis_eltype(LType, trailing_axes)
+    eltypes = map(trailing_axes) do array_trailing_axes
+        Tuple{LType, eltype.(array_trailing_axes)...}
     end
+
+    return typejoin(eltypes...)
 end
 
-function flatten_axes(array_names, array_axes)
-    collect(chain(map(flatten_array_axes, array_names, array_axes)...))
+function flatten{N, NA}(::Type{Val{N}}, As::Vararg{AxisArray, NA})
+    flatten(Val{N}, ntuple(identity, Val{NA}), As...)
 end
 
 """
     flatten(As::AxisArray...) -> AxisArray
-    flatten(last_dim::Integer, As::AxisArray...) -> AxisArray
+    flatten(last_dim::Type{Val{N}}, As::AxisArray...) -> AxisArray
+    flatten(last_dim::Type{Val{N}}, labels::Tuple, As::AxisArray...) -> AxisArray
 
-Concatenates AxisArrays with equal leading axes into a single AxisArray.
+Concatenates AxisArrays with N equal leading axes into a single AxisArray.
 All additional axes in any of the arrays are flattened into a single additional
 CategoricalVector{Tuple} axis.
 
 ### Arguments
 
-* `last_dim::Integer`: (optional) the greatest common dimension to share between all input
-                       arrays. The remaining axes are flattened. If this argument is not
-                       provided, the greatest common axis found among the input arrays is
-                       used. All preceeding axes must also be common to each input array, at
-                       the same dimension. Values from 0 up to one more than the minimum
-                       number of dimensions across all input arrays are allowed.
+* `::Type{Val{N}}`:   the greatest common dimension to share between all input
+                      arrays. The remaining axes are flattened. All N axes must be common
+                      to each input array, at the same dimension. Values from 0 up to the
+                      minimum number of dimensions across all input arrays are allowed.
+* `labels::Tuple`:    (optional) a label for each AxisArray in As which is used in the flat
+                      axis
 * `As::AxisArray...`: AxisArrays to be flattened together.
 """
-function flatten(As::AxisArray...; kwargs...)
-    gca = greatest_common_axis(As...)
-
-    return _flatten(gca, As...; kwargs...)
-end
-
-function flatten(last_dim::Integer, As::AxisArray...; kwargs...)
-    last_dim >= 0 || throw(ArgumentError("last_dim must be at least 0"))
-
-    if last_dim > minimum(map(ndims, As))
-        throw(ArgumentError(
-            "There must be at least $last_dim (last_dim) axes in each argument"
-        ))
+@generated function flatten{N, AN, LType}(::Type{Val{N}}, labels::NTuple{AN, LType}, As::Vararg{AxisArray, AN})
+    if N < 0
+        throw(ArgumentError("flatten dimension N must be at least 0"))
     end
 
-    if last_dim > greatest_common_axis(As...)
+    if N > minimum(ndims.(As))
         throw(ArgumentError(
-            "The first $last_dim axes don't all match across all arguments"
+            """
+            flatten dimension N must not be greater than the maximum number of dimensions
+            across all input arrays
+            """
         ))
     end
 
-    return _flatten(last_dim, As...; kwargs...)
-end
+    flat_dim = Val{N + 1}
+    flat_dim_int = Int(N) + 1
 
-function _flatten(
-    last_dim::Integer,
-    As::AxisArray...;
-    array_names=1:length(As),
-    axis_name=nothing,
-)
-    common_axes = axes(As[1])[1:last_dim]
-
-    if axis_name === nothing
-        axis_name = _defaultdimname(last_dim + 1)
-    elseif !isa(axis_name, Symbol)
-        throw(ArgumentError("axis_name must be a Symbol"))
-    end
+    common_axes, trailing_axes = zip(_splitall(Val{N}, axisparams.(As)...)...)
+
+    foreach(_check_common_axes, zip(common_axes...))
+
+    new_common_axes = first(common_axes)
+    flat_axis_eltype = _flat_axis_eltype(LType, trailing_axes)
+    flat_axis_type = CategoricalVector{flat_axis_eltype, Vector{flat_axis_eltype}}
+
+    new_axes_type = Tuple{new_common_axes..., Axis{:flat, flat_axis_type}}
+    new_eltype = Base.promote_eltype(As...)
 
-    new_data = cat(last_dim + 1, (view(A.data, repeated(:, last_dim + 1)...) for A in As)...)
-    new_axis = flatten_axes(array_names, map(A -> axes(A)[last_dim+1:end], As))
+    quote
+        common_axes, trailing_axes = zip(_splitall(Val{N}, axes.(As)...)...)
 
-    # TODO: Consider creating a SortedVector axis when all flattened axes are Dimensional
-    return AxisArray(new_data, common_axes..., CategoricalVector(new_axis))
+        for common_axis_tuple in zip(common_axes...)
+            if !isempty(common_axis_tuple)
+                for common_axis in common_axis_tuple[2:end]
+                    if !all(axisvalues(common_axis) .== axisvalues(common_axis_tuple[1]))
+                        throw(ArgumentError(
+                            """
+                            Leading common axes must be identical across
+                            all input arrays"""
+                        ))
+                    end
+                end
+            end
+        end
+
+        array_data = cat($flat_dim, _reshapeall($flat_dim, As...)...)
+
+        axis_array_type = AxisArray{
+            $new_eltype,
+            $flat_dim_int,
+            Array{$new_eltype, $flat_dim_int},
+            $new_axes_type
+        }
+
+        new_axes = (
+            first(common_axes)...,
+            Axis{:flat, $flat_axis_type}($flat_axis_type(_flatten_axes(labels, trailing_axes))),
+        )
+
+        return axis_array_type(array_data, new_axes)
+    end
 end

src/core.jl

@@ -508,6 +508,15 @@ end
 axes(A::AbstractArray) = default_axes(A)
 axes(A::AbstractArray, dim::Int) = default_axes(A)[dim]
 
+"""
+    axisparams(::AxisArray) -> Vararg{::Type{Axis}}
+    axisparams(::Type{AxisArray}) -> Vararg{::Type{Axis}}
+
+Returns the axis parameters for an AxisArray.
+"""
+axisparams{T,N,D,Ax}(::AxisArray{T,N,D,Ax}) = (Ax.parameters...)
+axisparams{T,N,D,Ax}(::Type{AxisArray{T,N,D,Ax}}) = (Ax.parameters...)
+
 ### Axis traits ###
 @compat abstract type AxisTrait end
 immutable Dimensional <: AxisTrait end

test/combine.jl

@@ -52,15 +52,21 @@ ABdata[3:6,3:6,:,2] = Bdata
 A1 = AxisArray(A1data, Axis{:X}(1:2), Axis{:Y}(1:2))
 A2 = AxisArray(reshape(A2data, size(A2data)..., 1), Axis{:X}(1:2), Axis{:Y}(1:2), Axis{:Z}([:foo]))
 
-@test flatten(A1, A2; array_names=[:A1, :A2]) == AxisArray(cat(3, A1data, A2data), Axis{:X}(1:2), Axis{:Y}(1:2), Axis{:page}(CategoricalVector([(:A1,), (:A2, :foo)])))
-@test flatten(A1; array_names=[:foo]) == AxisArray(reshape(A1, 2, 2, 1), Axis{:X}(1:2), Axis{:Y}(1:2), Axis{:page}(CategoricalVector([(:foo,)])))
-@test flatten(A1; array_names=[:a], axis_name=:ax) == AxisArray(reshape(A1.data, size(A1)..., 1), axes(A1)..., Axis{:ax}(CategoricalVector([(:a,)])))
+@test @inferred(flatten(Val{2}, A1, A2)) == AxisArray(cat(3, A1data, A2data), Axis{:X}(1:2), Axis{:Y}(1:2), Axis{:flat}(CategoricalVector([(1,), (2, :foo)])))
+@test @inferred(flatten(Val{2}, A1)) == AxisArray(reshape(A1, 2, 2, 1), Axis{:X}(1:2), Axis{:Y}(1:2), Axis{:flat}(CategoricalVector([(1,)])))
+@test @inferred(flatten(Val{2}, A1)) == AxisArray(reshape(A1.data, size(A1)..., 1), axes(A1)..., Axis{:flat}(CategoricalVector([(1,)])))
 
-@test_throws ArgumentError flatten(-1, A1)
-@test_throws ArgumentError flatten(10, A1)
+@test @inferred(flatten(Val{2}, (:A1, :A2), A1, A2)) == AxisArray(cat(3, A1data, A2data), Axis{:X}(1:2), Axis{:Y}(1:2), Axis{:flat}(CategoricalVector([(:A1,), (:A2, :foo)])))
+@test @inferred(flatten(Val{2}, (:foo,), A1)) == AxisArray(reshape(A1, 2, 2, 1), Axis{:X}(1:2), Axis{:Y}(1:2), Axis{:flat}(CategoricalVector([(:foo,)])))
+@test @inferred(flatten(Val{2}, (:a,), A1)) == AxisArray(reshape(A1.data, size(A1)..., 1), axes(A1)..., Axis{:flat}(CategoricalVector([(:a,)])))
+
+@test @inferred(flatten(Val{0}, A1)) == AxisArray(vec(A1data), Axis{:flat}(CategoricalVector(collect(IterTools.product((1,), axisvalues(A1)...)))))
+@test @inferred(flatten(Val{1}, A1)) == AxisArray(A1data, Axis{:row}(1:2), Axis{:flat}(CategoricalVector(collect(IterTools.product((1,), axisvalues(A1)[2])))))
+
+@test_throws ArgumentError flatten(Val{-1}, A1)
+@test_throws ArgumentError flatten(Val{10}, A1)
 
 A1ᵀ = transpose(A1)
-@test flatten(A1, A1ᵀ) == flatten(0, A1, A1ᵀ)
-@test_throws ArgumentError flatten(-1, A1, A1ᵀ)
-@test_throws ArgumentError flatten(1, A1, A1ᵀ)
-@test_throws ArgumentError flatten(10, A1, A1ᵀ)
+@test_throws ArgumentError flatten(Val{-1}, A1, A1ᵀ)
+@test_throws ArgumentError flatten(Val{1}, A1, A1ᵀ)
+@test_throws ArgumentError flatten(Val{10}, A1, A1ᵀ)

test/runtests.jl

@@ -1,5 +1,6 @@
 using AxisArrays
 using Base.Test
+import IterTools
 
 @testset "AxisArrays" begin
     # during this time there was an ambiguity in base with checkbounds_linear_indices