JuliaArrays / AxisArrays.jl

Performant arrays where each dimension can have a named axis with values
http://JuliaArrays.github.io/AxisArrays.jl/latest/
Other
200 stars 41 forks source link

support permutedims and reshape #186

Open schlichtanders opened 4 years ago

schlichtanders commented 4 years ago

one common thing is to using broadcasting on named vectors, for which one is put into another dimension.

permutedims(a) as well as reshape(a, 1, :) can be used to bring a column vector a to row format. reshape(a, 1, 1, :) would create a 3rd-dim vector.

It would be great if these operations preserve the indexing.

using AxisArrays
AxisArray(1:3, ["A", "B", "C"])["A"] # works and returns 1
permutedims(AxisArray(1:3, ["A", "B", "C"]))["A"]

gives an error

ArgumentError: invalid index: "A" of type String

Stacktrace:
 [1] to_index(::String) at ./indices.jl:297
 [2] to_index(::Base.ReshapedArray{Int64,2,AxisArray{Int64,1,UnitRange{Int64},Tuple{Axis{:row,Array{String,1}}}},Tuple{}}, ::String) at ./indices.jl:274
 [3] to_indices at ./indices.jl:325 [inlined]
 [4] to_indices at ./indices.jl:322 [inlined]
 [5] getindex(::Base.ReshapedArray{Int64,2,AxisArray{Int64,1,UnitRange{Int64},Tuple{Axis{:row,Array{String,1}}}},Tuple{}}, ::String) at ./abstractarray.jl:980
 [6] top-level scope at In[182]:1

The ultimate goal would be that the following works and preserves indexes as would be intuitive:

a = AxisArray(1:3, ["A", "B", "C"])
tuple.(a, permutedims(a))
# currently returns a plain standard Array
3×3 Array{Tuple{Int64,Int64},2}:
 (1, 1)  (1, 2)  (1, 3)
 (2, 1)  (2, 2)  (2, 3)
 (3, 1)  (3, 2)  (3, 3)

and similar for using reshape instead of permutedims

nickrobinson251 commented 3 years ago

perhaps we can keep the AxisArray wrapper on the outside, like we do with adjoint and transpose?

https://github.com/JuliaArrays/AxisArrays.jl/blob/9b91d546b28d96cd980e0a86d9c860c3689881d7/src/core.jl#L411-L412

e.g. define

Base.permutedims(A::AxisArray{T,1}) where {T} = AxisArray(permutedims(A.data), Axis{:permutedim}(Base.OneTo(1)), A.axes[1])

then we'd be able to do e.g.

julia> permutedims(AxisArray(1:3, ["A", "B", "C"]))[Axis{:row}("A")]
1-dimensional AxisArray{Int64,1,...} with axes:
    :permutedim, Base.OneTo(1)
And data, a 1-element Vector{Int64}:
 1