Create

Project.toml

@@ -1,6 +1,6 @@
 name = "NamedDims"
 uuid = "356022a1-0364-5f58-8944-0da4b18d706f"
-authors = ["Lyndon White <lyndon.white@invenialabs.co.uk>"]
+authors = ["Invenia Technical Computing Corporation"]
 version = "0.1.0"
 
 [extras]

src/NamedDims.jl

@@ -1,5 +1,9 @@
 module NamedDims
 
-greet() = print("Hello World!")
+export NamedDimsArray, name2dim, dim_names
+
+include("wrapper_array.jl")
+include("name2dim.jl")
+include("base_functions.jl")
 
 end # module

src/base_functions.jl

@@ -0,0 +1,5 @@
+
+
+function Base.sum(f::Base.Callable, a::NamedDimsArray; dims)
+    return sum(f, parent(a); dims=name2dim(A, dims))
+end

src/name2dim.jl

@@ -0,0 +1,47 @@
+
+function NamedNotFoundException(name, cands)
+    return ArgumentError(
+        "No dimensioned called $name exists. Dimension names: $(join(cands, ", "))"
+    )
+end
+
+###################################
+
+# We change the name into a `Val` because that will trigger constant propergation
+# and allow this to resolve at compile time when the `name` is a constant
+name2dim(namemap, name) = _name2dim(namemap, Val(name))
+
+# `_name2dim` returns 0 for names not found, because throwing an error would make it inpure
+# And then it couldn't run at compile time.
+
+# Generic case
+#TODO: for speed this could be made into a generated function, so it happens at compile time
+function _name2dim(::Type{T}, ::Val{name}) where {T<:Tuple, name}
+    return something(findfirst(isequal(name), T.parameters), 0)
+end
+
+## Hand roll out special cases to help this optimize to run at compile time
+_name2dim(::Type{Tuple{A}}, ::Val{A}) where A = 1
+_name2dim(::Type{Tuple{A}}, ::Val{name}) where {A,name} = 0
+
+_name2dim(::Type{Tuple{A,B}}, ::Val{A}) where {A,B} = 1
+_name2dim(::Type{Tuple{A,B}}, ::Val{B}) where {A,B} = 2
+_name2dim(::Type{Tuple{A,B}}, ::Val{name}) where {A,B,name} = 0
+
+_name2dim(::Type{Tuple{A,B,C}}, ::Val{A}) where {A,B,C} = 1
+_name2dim(::Type{Tuple{A,B,C}}, ::Val{B}) where {A,B,C} = 2
+_name2dim(::Type{Tuple{A,B,C}}, ::Val{C}) where {A,B,C} = 3
+_name2dim(::Type{Tuple{A,B,C}}, ::Val{name}) where {A,B,C,name} = 0
+
+_name2dim(::Type{Tuple{A,B,C,D}}, ::Val{A}) where {A,B,C,D} = 1
+_name2dim(::Type{Tuple{A,B,C,D}}, ::Val{B}) where {A,B,C,D} = 2
+_name2dim(::Type{Tuple{A,B,C,D}}, ::Val{C}) where {A,B,C,D} = 3
+_name2dim(::Type{Tuple{A,B,C,D}}, ::Val{D}) where {A,B,C,D} = 4
+_name2dim(::Type{Tuple{A,B,C,D}}, ::Val{name}) where {A,B,C,D, name} = 0
+
+_name2dim(::Type{Tuple{A,B,C,D,E}}, ::Val{A}) where {A,B,C,D,E} = 1
+_name2dim(::Type{Tuple{A,B,C,D,E}}, ::Val{B}) where {A,B,C,D,E} = 2
+_name2dim(::Type{Tuple{A,B,C,D,E}}, ::Val{C}) where {A,B,C,D,E} = 3
+_name2dim(::Type{Tuple{A,B,C,D,E}}, ::Val{D}) where {A,B,C,D,E} = 4
+_name2dim(::Type{Tuple{A,B,C,D,E}}, ::Val{E}) where {A,B,C,D,E} = 5
+_name2dim(::Type{Tuple{A,B,C,D,E}}, ::Val{name}) where {A,B,C,D,E, name} = 0

src/wrapper_array.jl

@@ -0,0 +1,81 @@
+
+struct NamedDimsArray{L<:Tuple, T, N, A<:AbstractArray{T,N}} <: AbstractArray{T,N}
+    data::A
+end
+
+function NamedDimsArray(orig::AbstractArray{T,N}, names) where {T, N}
+    if length(names) != N
+        throw(ArgumentError("A $N dimentional array, need $N dimension names. Got: $names"))
+    end
+    names_tt = Tuple{names...}
+    return NamedDimsArray{names_tt, T, N, typeof(orig)}(orig)
+end
+
+
+Base.parent(x::NamedDimsArray) = x.data
+
+
+"""
+    dim_names(A)
+
+Returns a tuple of containing the names of all the dimensions of the array `A`.
+"""
+dim_names(::Type{<:NamedDimsArray{L}}) where L = Tuple(L.parameters)
+dim_names(x::T) where T = dim_names(T)
+
+
+name2dim(a::NamedDimsArray{L}, name) where L = name2dim(L, name)
+
+
+
+#############################
+# AbstractArray Interface
+# https://docs.julialang.org/en/v1/manual/interfaces/index.html#man-interface-array-1
+
+Base.size(a::NamedDimsArray) = size(parent(a))
+Base.getindex(A::NamedDimsArray, inds...) = getindex(parent(A), inds...)
+Base.setindex(A::NamedDimsArray, value, inds...) = setindex(parent(A), value, inds...)
+
+
+###############################
+# kwargs indexing
+
+"""
+    order_named_inds(A, named_inds...)
+
+Returns the indices that have the names and values given by `named_inds`
+sorted into the order expected for the dimension of the array `A`.
+If any dimensions of `A` are not present in the named_inds,
+then they are given the value `:`, for slicing
+
+For example:
+```
+A = NamedDimArray(rand(4,4), (:x,, :y))
+order_named_inds(A; y=10, x=13) == (13,10)
+order_named_inds(A; x=2, y=1:3) == (2, 1:3)
+order_named_inds(A; y=5) == (:, 5)
+```
+
+This provides the core indexed lookup for `getindex` and `setindex` on the Array `A`
+"""
+function order_named_inds(A; named_inds...)
+    keys(named_inds) ⊆ dim_names(A) || throw(
+        DimensionMismatch("Expected $(dim_names(A)), got $(keys(named_inds))")
+    )
+
+
+    inds = map(dim_names(A)) do name
+        get(named_inds, name, :)  # default to slicing
+    end
+end
+
+function Base.getindex(A::NamedDimsArray; named_inds...)
+    inds = order_named_inds(A; named_inds...)
+    return getindex(parent(A), inds...)
+end
+
+
+function Base.setindex(A::NamedDimsArray, value; named_inds...)
+    inds = order_named_inds(A; named_inds...)
+    return setindex(parent(A), value, inds...)
+end

test/base_functions.jl

@@ -0,0 +1,11 @@
+using NamedDims
+using Test
+
+@testset "sum" begin
+    nda = NamedDimsArray([10 20; 30 40], (:x, :y))
+
+    @test sum(nda) == 100 == sum(identity, nda)
+
+    @test sum(nda; dims=:x) == [40 60]
+    @test sum(nda; dims=1) == [40 60]
+end

test/name2dim.jl

@@ -0,0 +1,17 @@
+using NamedDims
+using Test
+
+
+@testset "name2dim" begin
+    @testset "small case, that hits unrolled method" begin
+        @test name2dim(Tuple{:x, :y}, :x)==1
+        @test name2dim(Tuple{:x, :y}, :y)==2
+        @test name2dim(Tuple{:x, :y}, :z)==0  # not found
+    end
+    @testset "large case that hits generic fallback" begin
+        @test name2dim(Tuple{:x, :y, :a, :b, :c, :d}, :x)==1
+        @test name2dim(Tuple{:x, :y, :a, :b, :c, :d}, :a)==3
+        @test name2dim(Tuple{:x, :y, :a, :b, :c, :d}, :d)==6
+        @test name2dim(Tuple{:x, :y, :a, :b, :c, :d}, :z)==0 # not found
+    end
+end

test/runtests.jl

@@ -3,4 +3,8 @@ using Test
 
 @testset "NamedDims.jl" begin
     # Write your own tests here.
+
+    include("name2dim.jl")
+    include("wrapper_array.jl")
+    include("base_functions.jl")
 end

test/wrapper_array.jl

@@ -0,0 +1,25 @@
+using NamedDims
+using Test
+
+
+@testset "get the parent array that was wrapped" begin
+    orig = [1 2; 3 4]
+    @test parent(NamedDimsArray(orig, (:x, :y))) === orig
+end
+
+
+@testset "get the named array that was wrapped" begin
+    @test dim_names(NamedDimsArray([10 20; 30 40], (:x, :y))) === (:x, :y)
+end
+
+
+@testset "getindex" begin
+    nda = NamedDimsArray([10 20; 30 40], (:x, :y))
+
+    @test nda[x=1, y=1] == nda[y=1, x=1] == nda[1, 1] == 10
+    @test nda[y=end,x=end] == nda[end, end] == 40
+
+    # Missing dims become slices
+    @test nda[y=1] == nda[y=1, x=:] == nda[:, 1] == [10; 30]
+
+end