/
indexing.jl
97 lines (80 loc) · 5.5 KB
/
indexing.jl
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
A = AxisArray(reshape(1:24, 2,3,4), .1:.1:.2, .1:.1:.3, .1:.1:.4)
D = similar(A)
D[1,1,1,1,1] = 10
@test D[1,1,1,1,1] == D[1] == D.data[1] == 10
# Test slices
@test A == A.data
@test A[:,:,:] == A[Axis{:row}(:)] == A[Axis{:col}(:)] == A[Axis{:page}(:)] == A.data[:,:,:]
# Test UnitRange slices
@test A[1:2,:,:] == A.data[1:2,:,:] == A[Axis{:row}(1:2)] == A[Axis{1}(1:2)] == A[Axis{:row}(ClosedInterval(-Inf,Inf))] == A[[true,true],:,:]
@test @view(A[1:2,:,:]) == A.data[1:2,:,:] == @view(A[Axis{:row}(1:2)]) == @view(A[Axis{1}(1:2)]) == @view(A[Axis{:row}(ClosedInterval(-Inf,Inf))]) == @view(A[[true,true],:,:])
@test A[:,1:2,:] == A.data[:,1:2,:] == A[Axis{:col}(1:2)] == A[Axis{2}(1:2)] == A[Axis{:col}(ClosedInterval(0.0, .25))] == A[:,[true,true,false],:]
@test @view(A[:,1:2,:]) == A.data[:,1:2,:] == @view(A[Axis{:col}(1:2)]) == @view(A[Axis{2}(1:2)]) == @view(A[Axis{:col}(ClosedInterval(0.0, .25))]) == @view(A[:,[true,true,false],:])
@test A[:,:,1:2] == A.data[:,:,1:2] == A[Axis{:page}(1:2)] == A[Axis{3}(1:2)] == A[Axis{:page}(ClosedInterval(-1., .22))] == A[:,:,[true,true,false,false]]
@test @view(A[:,:,1:2]) == @view(A.data[:,:,1:2]) == @view(A[Axis{:page}(1:2)]) == @view(A[Axis{3}(1:2)]) == @view(A[Axis{:page}(ClosedInterval(-1., .22))]) == @view(A[:,:,[true,true,false,false]])
# Test scalar slices
@test A[2,:,:] == A.data[2,:,:] == A[Axis{:row}(2)]
@test A[:,2,:] == A.data[:,2,:] == A[Axis{:col}(2)]
@test A[:,:,2] == A.data[:,:,2] == A[Axis{:page}(2)]
# Test fallback methods
@test A[[1 2; 3 4]] == @view(A[[1 2; 3 4]]) == A.data[[1 2; 3 4]]
@test A[] == A.data[]
# Test axis restrictions
@test A[:,:,:].axes == A.axes
@test A[Axis{:row}(1:2)].axes[1].val == A.axes[1].val[1:2]
@test A[Axis{:row}(1:2)].axes[2:3] == A.axes[2:3]
@test A[Axis{:col}(1:2)].axes[2].val == A.axes[2].val[1:2]
@test A[Axis{:col}(1:2)].axes[[1,3]] == A.axes[[1,3]]
@test A[Axis{:page}(1:2)].axes[3].val == A.axes[3].val[1:2]
@test A[Axis{:page}(1:2)].axes[1:2] == A.axes[1:2]
# Linear indexing across multiple dimensions drops tracking of those dims
@test A[:].axes[1].val == 1:length(A)
@test A[1:2,:].axes[1].val == A.axes[1].val[1:2]
@test A[1:2,:].axes[2].val == 1:Base.trailingsize(A,2)
B = AxisArray(reshape(1:15, 5,3), .1:.1:0.5, [:a, :b, :c])
# Test indexing by Intervals
@test B[ClosedInterval(0.0, 0.5), :] == B[ClosedInterval(0.0, 0.5)] == B[:,:]
@test B[ClosedInterval(0.0, 0.3), :] == B[ClosedInterval(0.0, 0.3)] == B[1:3,:]
@test B[ClosedInterval(0.15, 0.3), :] == B[ClosedInterval(0.15, 0.3)] == B[2:3,:]
@test B[ClosedInterval(0.2, 0.5), :] == B[ClosedInterval(0.2, 0.5)] == B[2:end,:]
@test B[ClosedInterval(0.2, 0.6), :] == B[ClosedInterval(0.2, 0.6)] == B[2:end,:]
@test @view(B[ClosedInterval(0.0, 0.5), :]) == @view(B[ClosedInterval(0.0, 0.5)]) == B[:,:]
@test @view(B[ClosedInterval(0.0, 0.3), :]) == @view(B[ClosedInterval(0.0, 0.3)]) == B[1:3,:]
@test @view(B[ClosedInterval(0.15, 0.3), :]) == @view(B[ClosedInterval(0.15, 0.3)]) == B[2:3,:]
@test @view(B[ClosedInterval(0.2, 0.5), :]) == @view(B[ClosedInterval(0.2, 0.5)]) == B[2:end,:]
@test @view(B[ClosedInterval(0.2, 0.6), :]) == @view(B[ClosedInterval(0.2, 0.6)]) == B[2:end,:]
# Test Categorical indexing
@test B[:, :a] == @view(B[:, :a]) == B[:,1]
@test B[:, :c] == @view(B[:, :c]) == B[:,3]
@test B[:, [:a]] == @view(B[:, [:a]]) == B[:,[1]]
@test B[:, [:a,:c]] == @view(B[:, [:a,:c]]) == B[:,[1,3]]
@test B[Axis{:row}(ClosedInterval(0.15, 0.3))] == @view(B[Axis{:row}(ClosedInterval(0.15, 0.3))]) == B[2:3,:]
A = AxisArray(reshape(1:256, 4,4,4,4), Axis{:d1}(.1:.1:.4), Axis{:d2}(1//10:1//10:4//10), Axis{:d3}(["1","2","3","4"]), Axis{:d4}([:a, :b, :c, :d]))
ax1 = axes(A)[1]
@test A[Axis{:d1}(2)] == A[ax1(2)]
@test A.data[1:2,:,:,:] == A[Axis{:d1}(ClosedInterval(.1,.2))] == A[ClosedInterval(.1,.2),:,:,:] == A[ClosedInterval(.1,.2),:,:,:,1] == A[ClosedInterval(.1,.2)]
@test A.data[:,1:2,:,:] == A[Axis{:d2}(ClosedInterval(1//10,2//10))] == A[:,ClosedInterval(1//10,2//10),:,:] == A[:,ClosedInterval(1//10,2//10),:,:,1] == A[:,ClosedInterval(1//10,2//10)]
@test A.data[:,:,1:2,:] == A[Axis{:d3}(["1","2"])] == A[:,:,["1","2"],:] == A[:,:,["1","2"],:,1] == A[:,:,["1","2"]]
@test A.data[:,:,:,1:2] == A[Axis{:d4}([:a,:b])] == A[:,:,:,[:a,:b]] == A[:,:,:,[:a,:b],1]
A = AxisArray(reshape(1:32, 2, 2, 2, 2, 2), .1:.1:.2, .1:.1:.2, .1:.1:.2, [:a, :b], [:c, :d])
@test A[ClosedInterval(.15, .25), ClosedInterval(.05, .15), ClosedInterval(.15, .25), :a] == A.data[2:2, 1:1, 2:2, 1, :]
@test A[Axis{:dim_5}(2)] == A.data[:, :, :, :, 2]
# Test vectors
v = AxisArray(collect(.1:.1:10.0), .1:.1:10.0)
@test v[Colon()] == v
@test v[:] == v.data[:] == v[Axis{:row}(:)]
@test v[3:8] == v.data[3:8] == v[ClosedInterval(.25,.85)] == v[Axis{:row}(3:8)] == v[Axis{:row}(ClosedInterval(.22,.88))]
# Test repeated intervals
A = AxisArray([1:100 -1:-1:-100], .1:.1:10.0, [:c1, :c2])
@test A[2.0..3.0, :] == A[atindex(-0.5..0.5, 25), :] == [20:30 -20:-1:-30]
@test A[2.0..3.0, [:c1,:c2]] == A[atindex(-0.5..0.5, 25), [:c1, :c2]] == [20:30 -20:-1:-30]
@test A[2.0..3.0, :c1] == A[atindex(-0.5..0.5, 25), :c1] == collect(20:30)
@test A[atindex(-0.5..0.5, 25), :c1] == collect(20:30)
@test A[atindex(-0.5..0.5, [25, 35]), :c1] == [20:30 30:40]
@test_throws BoundsError A[atindex(-0.5..0.5, 5), :c1]
@test_throws BoundsError A[atindex(-0.5..0.5, [5, 15, 25]), :]
# Indexing with CartesianIndex{0}
A = AxisArray(reshape(1:15, 3, 5), :x, :y)
@test A[2,2,CartesianIndex(())] == 5
@test A[2,CartesianIndex(()),2] == 5
@test A[CartesianIndex(()),2,2] == 5