More AxisArray constructors

- also in the background @defdim is broken up into smaller macros

Author: Tokazama

Project.toml

@@ -1,7 +1,7 @@
 name = "AxisIndices"
 uuid = "f52c9ee2-1b1c-4fd8-8546-6350938c7f11"
 authors = ["Tokazama <[email protected]>"]
-version = "0.4.4"
+version = "0.4.5"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"

src/Arrays/AbstractAxisArray.jl

@@ -325,3 +325,10 @@ Base.vcat(A::AbstractAxisArray{T,N}) where {T,N} = A
 
 Base.cat(A::AbstractAxisArray{T,N}; dims) where {T,N} = A
 
+function Base.convert(::Type{T}, A::AbstractArray) where {T<:AbstractAxisArray}
+    if A isa T
+        return A
+    else
+        return T(A)
+    end
+end

src/Arrays/AxisArray.jl

@@ -225,17 +225,15 @@ julia> size(AxisArray{Int,2}(undef, (["a", "b"], [:one, :two])))
 (2, 2)
 
 """
-function AxisArray{T,N}(x::AbstractArray{T2,N}, axis_keys::Tuple, check_length::Bool=true) where {T,T2,N}
-    return AxisArray{T,N}(map(T, x), axis_keys, check_length)
+function AxisArray{T,N}(A::AbstractArray{T2,N}, axis_keys::Tuple, check_length::Bool=true) where {T,T2,N}
+    return AxisArray{T,N}(copyto!(Array{T}(undef, size(A)), A), axis_keys, check_length)
 end
 
 function AxisArray{T,N}(init::ArrayInitializer, args...) where {T,N}
     return AxisArray{T,N}(init, args)
 end
 
-function AxisArray{T,N}(x::AbstractArray, args...) where {T,N}
-    return AxisArray{T,N}(x, args)
-end
+AxisArray{T,N}(x::AbstractArray, args...) where {T,N} = AxisArray{T,N}(x, args)
 
 function AxisArray{T,N}(init::ArrayInitializer, axs::Tuple{Vararg{Any,N}}) where {T,N}
     return AxisArray{T,N}(init, map(to_axis, axs))
@@ -246,14 +244,27 @@ function AxisArray{T,N}(init::ArrayInitializer, axs::AbstractAxes{N}) where {T,N
     return AxisArray{T,N,typeof(p),typeof(axs)}(p, axs)
 end
 
-function AxisArray{T,N}(
-    x::AbstractArray{T,N},
-    axs::Tuple{Vararg{Any,N2}},
+function AxisArray{T,N}(x::AbstractArray{T,N}, axs::Tuple, check_length::Bool=true) where {T,N}
+    return AxisArray{T,N,typeof(x)}(x, axs, check_length)
+end
+
+###
+### AxisArray{T,N,P}
+###
+AxisArray{T,N,P}(A::AbstractArray, args...) where {T,N,P} = AxisArray{T,N,P}(A, args)
+
+function AxisArray{T,N,P}(
+    x::AbstractArray,
+    axs::Tuple,
     check_length::Bool=true
-) where {T,N,N2}
+) where {T,N,P}
+
+    return AxisArray{T,N,P}(convert(P, x), axs, check_length)
+end
 
+function AxisArray{T,N,P}(x::P, axs::Tuple, check_length::Bool=true) where {T,N,P<:AbstractArray{T,N}}
     axs = to_axes((), axs, axes(x), check_length, Staticness(x))
-    return AxisArray{T,N,typeof(x),typeof(axs)}(x, axs)
+    return AxisArray{T,N,P,typeof(axs)}(x, axs)
 end
 
 ###
@@ -301,3 +312,8 @@ function Base.reshape(A::AbstractArray, shp::Tuple{<:AbstractAxis,Vararg{<:Abstr
     return AxisArray{eltype(p),ndims(p),typeof(p),typeof(axs)}(p, axs)
 end
 
+AxisArray{T,N,P}(A::AxisArray{T,N,P}) where {T,N,P} = A
+
+function AxisArray{T,N,P}(A::AxisArray) where {T,N,P}
+    return AxisArray{T,N,P}(convert(P, parent(A)), axes(A))
+end

src/Arrays/OffsetArray.jl

@@ -37,7 +37,28 @@ end
 
 OffsetArray{T,N}(A, inds::Vararg) where {T,N} = OffsetArray{T,N}(A, inds)
 
-@inline function OffsetArray{T,N}(A::AbstractArray{T,N}, inds::NTuple{M,Any}) where {T,N,M}
+function OffsetArray{T,N}(A::AbstractArray{T,N}, inds::Tuple) where {T,N}
+    return OffsetArray{T,N,typeof(A)}(A, inds)
+end
+
+function OffsetArray{T,N}(A::AbstractArray{T2,N}, inds::Tuple) where {T,T2,N}
+    return OffsetArray{T,N}(copyto!(Array{T}(undef, size(A)), A), inds)
+end
+
+
+function OffsetArray{T,N,P}(A::AbstractArray, inds::NTuple{M,Any}) where {T,N,P<:AbstractArray{T,N},M}
+    return OffsetArray{T,N,P}(convert(P, A))
+end
+
+OffsetArray{T,N,P}(A::OffsetArray{T,N,P}) where {T,N,P} = A
+
+function OffsetArray{T,N,P}(A::OffsetArray) where {T,N,P}
+    p = convert(P, parent(A))
+    axs = map(assign_indices, axes(A), axes(p))
+    return AxisArray{T,N,P,typeof(axs)}(p, axs)
+end
+
+function OffsetArray{T,N,P}(A::P, inds::NTuple{M,Any}) where {T,N,P<:AbstractArray{T,N},M}
     S = Staticness(A)
     if N === M
         axs = map((x, y) -> OffsetAxis(as_staticness(S, x), as_staticness(S, y)), inds, axes(A))

src/Interface/names.jl

@@ -61,6 +61,94 @@ end
 # So `condition` also has to be pure, which shouldn't be hard because it should basically
 # just be comparing symbols
 
+macro def_naxis(name, name_dim)
+    nname = Symbol(:n, name)
+    nname_doc = """
+        $nname(x) -> Int
+
+    Returns the size along the dimension corresponding to the $name.
+    """
+
+    esc(quote
+        @doc $nname_doc
+        @inline $nname(x) = Base.size(x, $name_dim(x))
+    end)
+end
+
+
+macro def_axis_keys(name, name_dims)
+    name_keys = Symbol(name, :_keys)
+    name_keys_doc = """
+        $name_keys(x)
+
+    Returns the keys corresponding to the $name axis
+    """
+    esc(quote
+        @doc $name_keys_doc
+        @inline $name_keys(x) = keys(axes(x, $name_dims(x)))
+    end)
+end
+
+
+
+macro def_axis_indices(name, name_dim)
+    name_indices = Symbol(name, :_indices)
+    name_indices_doc = """
+        $name_indices(x)
+
+    Returns the indices corresponding to the $name axis
+    """
+    esc(quote
+        @doc $name_indices_doc
+        @inline $name_indices(x) = indices(axes(x, $name_dim(x)))
+    end)
+end
+
+
+# TODO I'm not sure this is the best name for this one
+macro def_axis_type(name, name_dim)
+    name_type = Symbol(name, :_axis_type)
+    name_type_doc = """
+        $name_type(x)
+
+    Returns the key type corresponding to the $name axis.
+    """
+
+    esc(quote
+        @doc $name_type_doc
+        @inline $name_type(x) = keytype(axes(x, $name_dim(x)))
+    end)
+end
+
+macro def_selectdim(name, name_dim)
+    name_selectdim = Symbol(:select_, name, :dim)
+    name_selectdim_doc = """
+        $name_selectdim(x, i)
+
+    Return a view of all the data of `x` where the index for the $name dimension equals `i`.
+    """
+
+    esc(quote
+        @doc $name_selectdim_doc
+        @inline $name_selectdim(x, i) = selectdim(x, $name_dim(x), i)
+
+    end)
+end
+
+macro def_eachslice(name, name_dim)
+    each_name = Symbol(:each_, name)
+    each_name_doc = """
+        $each_name(x)
+
+    Create a generator that iterates over the $name dimensions `A`, returning views that select
+    all the data from the other dimensions in `A`.
+    """
+    esc(quote
+        @doc $each_name_doc
+        @inline $each_name(x) = eachslice(x, dims=$name_dim(x))
+    end)
+end
+
 """
     @defdim name condition
 
@@ -92,7 +180,16 @@ julia> @defdim time is_time
     `@defdim` should be considered experimental and subject to change
 
 """
-macro defdim(name, condition)
+macro defdim(
+    name,
+    condition,
+    def_naxis::Bool=true,
+    def_axis_keys::Bool=true,
+    def_axis_indices::Bool=true,
+    def_axis_type::Bool=true,
+    def_selectdim::Bool=true,
+    def_eachslice::Bool=true,
+)
 
     dim_noerror_name = Symbol(:dim_noerror_, name)
 
@@ -103,13 +200,6 @@ macro defdim(name, condition)
     Returns the dimension corresponding to $name.
     """
 
-    nname = Symbol(:n, name)
-    nname_doc = """
-        $nname(x) -> Int
-
-    Returns the size along the dimension corresponding to the $name.
-    """
-
     has_name_dim = Symbol(:has_, name, :dim)
     has_name_dim_doc = """
         $has_name_dim(x) -> Bool
@@ -130,42 +220,6 @@ macro defdim(name, condition)
     Returns an `AxisIterator` along the $name axis.
     """
 
-    name_indices = Symbol(name, :_indices)
-    name_indices_doc = """
-        $name_indices(x)
-
-    Returns the indices corresponding to the $name axis
-    """
-
-    name_keys = Symbol(name, :_keys)
-    name_keys_doc = """
-        $name_keys(x)
-
-    Returns the keys corresponding to the $name axis
-    """
-
-    name_type = Symbol(name, :_axis_type)
-    name_type_doc = """
-        $name_type(x)
-
-    Returns the key type corresponding to the $name axis.
-    """
-
-    name_selectdim = Symbol(:select_, name, :dim)
-    name_selectdim_doc = """
-        $name_selectdim(x, i)
-
-    Return a view of all the data of `x` where the index for the $name dimension equals `i`.
-    """
-
-    each_name = Symbol(:each_, name)
-    each_name_doc = """
-        $each_name(x)
-
-    Create a generator that iterates over the $name dimensions `A`, returning views that select
-    all the data from the other dimensions in `A`.
-    """
-
     err_msg = "Method $(Symbol(condition)) is not true for any dimensions of "
 
     esc(quote
@@ -186,9 +240,6 @@ macro defdim(name, condition)
             end
         end
 
-        @doc $nname_doc
-        @inline $nname(x) = Base.size(x, $name_dim(x))
-
         @doc $has_name_dim_doc
         @inline $has_name_dim(x) = !($dim_noerror_name(dimnames(x)) === 0)
 
@@ -198,22 +249,30 @@ macro defdim(name, condition)
         @doc $name_axis_itr
         @inline $name_axis(x, sz; kwargs...) = AxisIterator(axes(x, $name_dim(x)), sz; kwargs...)
 
-        @doc $name_keys_doc
-        @inline $name_keys(x) = keys($name_axis(x))
+        if $def_naxis
+            Interface.@def_naxis($name, $name_dim)
+        end
 
-        @doc $name_indices_doc
-        @inline $name_indices(x) = values($name_axis(x))
+        if $def_axis_keys
+            Interface.@def_axis_keys($name, $name_dim)
+        end
 
-        @doc $name_type_doc
-        @inline $name_type(x) = keytype($name_axis(x))
+        if $def_axis_indices
+            Interface.@def_axis_indices($name, $name_dim)
+        end
 
-        @doc $name_selectdim_doc
-        @inline $name_selectdim(x, i) = selectdim(x, $name_dim(x), i)
+        if $def_axis_type
+            Interface.@def_axis_type($name, $name_dim)
+        end
 
-        @doc $each_name_doc
-        @inline $each_name(x) = eachslice(x, dims=$name_dim(x))
+        if $def_selectdim
+            Interface.@def_selectdim($name, $name_dim)
+        end
+
+        if $def_eachslice
+            Interface.@def_eachslice($name, $name_dim)
+        end
 
         nothing
     end)
 end
-

test/Arrays/Arrays.jl

@@ -1,5 +1,6 @@
 
 include("vectors.jl")
+include("AxisArray.jl")
 
 
 @testset "permuteddimsview" begin