-
-
Notifications
You must be signed in to change notification settings - Fork 16
/
indexing.jl
107 lines (86 loc) · 3.92 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
98
99
100
101
102
103
104
105
106
107
nodeselect(ns, i, I...) = ns[i][I...]
nodechild(ns, i, j) = ns[i].nodes[j]
function bisect_search(a, i)
first(searchsorted(a, i))
end
Base.IndexStyle(::Type{<:AbstractMultiScaleArray}) = IndexLinear()
@inline function getindex(m::AbstractMultiScaleArray, i::Int)
idx = bisect_search(m.end_idxs, i)
idx > 1 && (i -= m.end_idxs[idx - 1]) # also works with values
(isempty(m.values) || idx < length(m.end_idxs)) ? nodeselect(m.nodes, idx, i) :
m.values[i]
end
@inline function setindex!(m::AbstractMultiScaleArray, nodes, i::Int)
idx = bisect_search(m.end_idxs, i) # +1 for 1-based indexing
idx > 1 && (i -= m.end_idxs[idx - 1])
if isempty(m.values) || idx < length(m.end_idxs)
m.nodes[idx][i] = nodes
else
m.values[i] = nodes
end
end
@inline getindex(m::AbstractMultiScaleArrayLeaf, i::Int) = m.values[i]
@inline getindex(m::AbstractMultiScaleArrayLeaf, i::Int...) = m.values[i[1]]
@inline function getindex(m::AbstractMultiScaleArray, i, I...)
if isempty(m.values) || i < length(m.end_idxs)
length(I) == 1 ? nodechild(m.nodes, i, I[1]) : nodeselect(m.nodes, i, I...)
else
m.values[I...]
end
end
@inline getindex(m::AbstractMultiScaleArrayLeaf, i...) = m.values[i[1]]
@inline getindex(m::AbstractMultiScaleArrayLeaf, i::CartesianIndex{1}) = m.values[i[1]]
@inline setindex!(m::AbstractMultiScaleArrayLeaf, values, i::Int) = (m.values[i] = values)
@inline function setindex!(m::AbstractMultiScaleArrayLeaf, values, i::Int...)
(m.values[i[1]] = values)
end
@inline function setindex!(m::AbstractMultiScaleArray, values, i::CartesianIndex{1})
(m[i[1]] = values)
end
@inline function setindex!(m::AbstractMultiScaleArray, values, i, I::Int...)
if isempty(m.values) || i < length(m.end_idxs)
mx = m.nodes[i]
length(I) == 1 ? (mx[I[1]] = values) : (mx[I...] = values)
else
m.values[I...] = values
end
end
@inline getindex(m::AbstractMultiScaleArray, ::Colon) = [m[i] for i in 1:length(m)]
@inline getindex(m::AbstractMultiScaleArrayLeaf, ::Colon) = m.values
@inline getindex(m::AbstractMultiScaleArray, i::CartesianIndex{1}) = m[i[1]] # (i, )
eachindex(m::AbstractMultiScaleArray) = 1:length(m)
endof(m::AbstractMultiScaleArray) = length(m)
Base.eltype(::Type{<:AbstractMultiScaleArray{B}}) where {B <: Number} = B
# For now, default to Float64 (like v0.5 version did) if not specified
function Base.eltype(::Type{T}) where {T <: AbstractMultiScaleArray}
eltype(fieldtype(T, :values)) === Any && return Float64
error("Invalid AbstractMultiScaleArray, type of values not specified")
end
Base.eltype(::T) where {T <: AbstractMultiScaleArray} = eltype(T)
#broadcast_getindex(m::AbstractMultiScaleArrayLeaf, i::Int) = (println("here");m[i])
#broadcast_getindex(m::AbstractMultiScaleArray, i::Int) = (println("here");m[i])
#broadcast_getindex(m::AbstractMultiScaleArray, i::Int...) = m[i]
getindices(m::AbstractMultiScaleArrayHead) = 1:length(m)
function getindices(m::AbstractMultiScaleArrayHead, i::Int)
((i > 1) ? (m.end_idxs[i - 1] + 1) : 1):m.end_idxs[i]
end
function getindices(m::AbstractMultiScaleArrayHead, i, I::Int...)
getindices(m.nodes[i], ((i > 1) ? (m.end_idxs[i - 1] + 1) : 1), m.end_idxs[i], I...)
end
function getindices(m::AbstractMultiScaleArray, bot_idx, top_idx, i::Int)
(bot_idx + ((i > 1) ? m.end_idxs[i - 1] : 0)):(top_idx + m.end_idxs[i] - length(m))
end
function getindices(m::AbstractMultiScaleArray, bot_idx, top_idx, i, I::Int...)
getindices(m.nodes[i],
bot_idx + ((i > 1) ? m.end_idxs[i - 1] : 0),
top_idx + m.end_idxs[i] - length(m),
I...)
end
function getindices(m::AbstractMultiScaleArrayLeaf, bot_idx, top_idx, i::Int)
i > length(m) && error("Final index is larger than length of leaf")
top_idx -= length(m) - i
top_idx:top_idx
end
function getindices(m::AbstractMultiScaleArrayLeaf, bot_idx, top_idx, I::Int...)
error("Indexes past the bottom.")
end