AxisIndices.jl
Here are some reasons you should try AxisIndices
- Flexible design for customizing multidimensional indexing behavior
- It's fast. StaticRanges are used to speed up indexing ranges. If something is slow, please create a detailed issue.
- Works with Julia's standard library (in progress). The end goal of AxisIndices is to fully integrate with the standard library wherever possible. If you can find a relevant method that isn't supported in
BaseorStatisticsthen it's likely an oversight, so make an issue.LinearAlgebra,MappedArrays,OffsetArrays, andNamedDimsalso have some form of support.
Note that in the Julia REPL an AxisArray prints as follows, which may not be apparent in the online documentation.
A Simple AxisArray
The simplest form of an AxisArray just wraps a standard array.
julia> using AxisIndices
julia> x = reshape(1:8, 2, 4);
julia> ax = AxisArray(x)
2×4 AxisArray(reshape(::UnitRange{Int64}, 2, 4)
• axes:
1 = 1:2
2 = 1:4
)
1 2 3 4
1 1 3 5 7
2 2 4 6 8
Simply wrapping x allows us to use functions to access its elements.
When using functions as indexing arguments, the axis corresponding to each argument is ultimately filtered by the function.
julia> ax[:, >(2)]
2×2 AxisArray(::Matrix{Int64}
• axes:
1 = 1:2
2 = 1:2
)
1 2
1 5 7
2 6 8
julia> ax[:, >(2)] == ax[:,filter(>(2), axes(ax, 2))] == ax[:, 3:4]
trueThis can be particularly helpful when indexing arguments for large arrays would otherwise require combining two or more non-continuous sets of indices. For example, if we wanted to get every element except for those at one index along the second dimension you would need to do something like:
julia> not_index = 2; # the index we don't want to include
julia> axis = axes(x, 2); # the axis that we want to refer to
julia> inds_before = firstindex(axis):(not_index - 1); # all of the indices before `not_index`
julia> inds_after = (not_index + 1):lastindex(axis); # all of the indices after `not_index`
julia> x[:, vcat(inds_before, inds_after)]
2×3 Matrix{Int64}:
1 5 7
2 6 8
Using an AxisArray, this only requires one line of code
julia> ax[:, !=(2)]
2×3 AxisArray(::Matrix{Int64}
• axes:
1 = 1:2
2 = 1:3
)
1 2 3
1 1 5 7
2 2 6 8 We can using ChainedFixes to combine multiple functions.
julia> using ChainedFixes
julia> ax[:, or(<(2), >(3))] # == ax[:, [1, 4]]
2×2 AxisArray(::Matrix{Int64}
• axes:
1 = 1:2
2 = 1:2
)
1 2
1 1 7
2 2 8
julia> ax[:, and(>(1), <(4))]
2×2 AxisArray(::Matrix{Int64}
• axes:
1 = 1:2
2 = 1:2
)
1 2
1 3 5
2 4 6
Although these examples are simple and could be done by hand (i.e. without producing the indices programmatically), arrays that are larger or have unknown indices are more easily managed.
An AxisArray with Keys
All arguments after the array passed to AxisArray are applied to corresponding axes.
We can bind a set of keys to an axis when constructing an AxisArray by providing them in the corresponding axis argument position.
Whenever altering an axis we need to provide an argument for each dimension.
In the following example we pass nothing, (.1:.1:.4)s, which means the first axis won't be changed and the second axis will have (.1:.1:.4)s bound to it.
julia> import Unitful: s
julia> ax = AxisArray(x, nothing, (.1:.1:.4)s)
2×4 AxisArray(reshape(::UnitRange{Int64}, 2, 4)
• axes:
1 = 1:2
2 = (0.1:0.1:0.4) s
)
0.1 s 0.2 s 0.3 s 0.4 s
1 1 3 5 7
2 2 4 6 8
We can still use functions to access these elements
julia> ax[:, <(0.3s)]
2×2 AxisArray(::Matrix{Int64}
• axes:
1 = 1:2
2 = (0.1:0.1:0.2) s
)
0.1 s 0.2 s
1 1 3
2 2 4
This also allows us to use keys as indexing arguments...
julia> ax[1, 0.1s]
1...or as intervals.
julia> ax[:, 0.1s..0.3s]
2×3 AxisArray(::Matrix{Int64}
• axes:
1 = 1:2
2 = (0.1:0.1:0.3) s
)
0.1 s 0.2 s 0.3 s
1 1 3 5
2 2 4 6
Offset Indexing
Indices don't have to start at one if we don't want them to.
julia> ax = AxisArray(x, 2:3, 2:5)
2×4 AxisArray(reshape(::UnitRange{Int64}, 2, 4)
• axes:
1 = 2:3
2 = 2:5
)
2 3 4 5
2 1 3 5 7
3 2 4 6 8
julia> ax[:,2]
2-element AxisArray(::Vector{Int64}
• axes:
1 = 2:3
)
1
2 1
3 2
If you don't know the length of each axis beforehand you can use offset.
The argument passed to offset specifies how much the standard indices should be adjusted by.
To start indexing at 2 we need to offset one-based indexing by +1.
julia> ax = AxisArray(x, offset(1), offset(1))
2×4 AxisArray(reshape(::UnitRange{Int64}, 2, 4)
• axes:
1 = 2:3
2 = 2:5
)
2 3 4 5
2 1 3 5 7
3 2 4 6 8
We can also center each axis.
julia> ax = AxisArray(x, center, center)
2×4 AxisArray(reshape(::UnitRange{Int64}, 2, 4)
• axes:
1 = -1:0
2 = -2:1
)
-2 -1 0 1
-1 1 3 5 7
0 2 4 6 8
The default origin of each centered axis is zero, but we can choose any origin.
julia> ax = AxisArray(reshape(1:9, 3, 3), center(10), center(10))
3×3 AxisArray(reshape(::UnitRange{Int64}, 3, 3)
• axes:
1 = 9:11
2 = 9:11
)
9 10 11
9 1 4 7
10 2 5 8
11 3 6 9
Static Sizing
Sometimes we know the size of the arrays we'll be working with beforehand.
This can be encoded in the axis using ArrayInterface.StaticInt.
julia> using ArrayInterface
julia> import ArrayInterface: StaticInt
julia> ax = AxisArray{Int}(
undef, # initialize empty array
StaticInt(1):StaticInt(2), # first axis with known size of two
StaticInt(1):StaticInt(2) # second axis with known size of two
);
julia> ArrayInterface.known_length(typeof(ax)) # size is known at compile time
4
julia> ax[1:2, 1:2] .= x[1:2, 1:2]; # underlying type is mutable `Array`, so we can assign new values
julia> ax
2×2 AxisArray(::Matrix{Int64}
• axes:
1 = 1:2
2 = 1:2
)
1 2
1 1 3
2 2 4
Encoding Types as Axes
If each element along a particular axis corresponds to a field of a type then we can encode that information in the axis.
julia> ax = AxisArray(reshape(1:4, 2, 2), StructAxis{ComplexF64}(), [:a, :b])
2×2 AxisArray(reshape(::UnitRange{Int64}, 2, 2)
• axes:
1 = [:re, :im]
2 = [:a, :b]
)
:a :b
:re 1 3
:im 2 4
We can then create a lazy mapping of that type across views of the array.
julia> axview = struct_view(ax)
2-element AxisArray(mappedarray(ComplexF64, view(reshape(::UnitRange{Int64}, 2, 2), 1, :), view(reshape(::UnitRange{Int64}, 2, 2), 2, :))
• axes:
1 = [:a, :b]
)
1
:a 1.0 + 2.0im
:b 3.0 + 4.0im
julia> axview[:b]
3.0 + 4.0im
Attaching Metadata
Using the Metadata package, metadata can be added to an AxisArray.
julia> using Metadata
julia> mx = attach_metadata(AxisArray(x))
2×4 attach_metadata(AxisArray(reshape(::UnitRange{Int64}, 2, 4)
• axes:
1 = 1:2
2 = 1:4
), ::Dict{Symbol, Any}
• metadata:
)
1 2 3 4
1 1 3 5 7
2 2 4 6 8
julia> mx.m1 = 1;
julia> mx.m1
1
Metadata can also be attached to an axis.
julia> m = (a = 1, b = 2);
julia> ax = AxisArray(x, nothing, attach_metadata(m));
julia> metadata(ax, dim=2)
(a = 1, b = 2)
Padded Axes
We can also pad axes in various ways.
julia> x = [:a, :b, :c, :d];
julia> AxisArray(x, circular_pad(first_pad=2, last_pad=2))
8-element AxisArray(::Vector{Symbol}
• axes:
1 = -1:6
)
1
-1 :c
0 :d
1 :a
2 :b
3 :c
4 :d
5 :a
6 :b
julia> AxisArray(x, replicate_pad(first_pad=2, last_pad=2))
8-element AxisArray(::Vector{Symbol}
• axes:
1 = -1:6
)
1
-1 :a
0 :a
1 :a
2 :b
3 :c
4 :d
5 :d
6 :d
julia> AxisArray(x, symmetric_pad(first_pad=2, last_pad=2))
8-element AxisArray(::Vector{Symbol}
• axes:
1 = -1:6
)
1
-1 :c
0 :b
1 :a
2 :b
3 :c
4 :d
5 :c
6 :b
julia> AxisArray(x, reflect_pad(first_pad=2, last_pad=2))
8-element AxisArray(::Vector{Symbol}
• axes:
1 = -1:6
)
1
-1 :b
0 :a
1 :a
2 :b
3 :c
4 :d
5 :d
6 :c
julia> AxisArray(3:4, zero_pad(sym_pad=2))
6-element AxisArray(::UnitRange{Int64}
• axes:
1 = -1:4
)
1
-1 0
0 0
1 3
2 4
3 0
4 0
julia> AxisArray(3:4, one_pad(sym_pad=2))
6-element AxisArray(::UnitRange{Int64}
• axes:
1 = -1:4
)
1
-1 1
0 1
1 3
2 4
3 1
4 1
Named Axes
Names can be attached to each dimension/axis using NamedAxisArray.
julia> nax = NamedAxisArray(reshape(1:4, 2, 2), x = [:a, :b], y = ["c", "d"])
2×2 NamedDimsArray(AxisArray(reshape(::UnitRange{Int64}, 2, 2)
• axes:
x = [:a, :b]
y = ["c", "d"]
))
"c" "d"
:a 1 3
:b 2 4