Skip to content

Commit 5f03cc4

Browse files
author
Michael Abbott
committed
move findindex etc to lookup.jl, untouched
1 parent 827907e commit 5f03cc4

File tree

3 files changed

+134
-131
lines changed

3 files changed

+134
-131
lines changed

‎src/AxisKeys.jl‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@ module AxisKeys
33
include("struct.jl")
44
export KeyedArray, axiskeys
55

6+
include("lookup.jl")
7+
68
include("names.jl")
79
export NamedDimsArray, dimnames
810

‎src/lookup.jl‎

Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,132 @@
1+
"""
2+
(A::KeyedArray)("a", 2.0, :γ) == A[1, 2, 3]
3+
A(:γ) == view(A, :, :, 3)
4+
5+
`KeyedArray`s are callable, and this behaves much like indexing,
6+
except that it searches for the given keys in `axiskeys(A)`,
7+
instead of `axes(A)` for indices.
8+
9+
A single key may be used to indicate a slice, provided that its type
10+
only matches the eltype of one `axiskeys(A,d)`.
11+
You can also slice explicitly with `A("a", :, :)`, both of these return a `view`.
12+
13+
An extra trailing colon (when all other indices are fixed) will return
14+
a zero-dimensional `view`. This allows setting one value by
15+
writing `A("a", 2.0, :γ, :) .= 100`.
16+
17+
Also accepts functions like `A(<=(2.0))` and selectors,
18+
see `Nearest` and `Index`.
19+
"""
20+
@inline @propagate_inbounds (A::KeyedArray)(args...) = getkey(A, args...)
21+
22+
@inline function getkey(A, args...)
23+
if length(args) == ndims(A)
24+
inds_raw = map(findindex, args, axiskeys(A))
25+
inds = Base.to_indices(A, inds_raw)
26+
@boundscheck checkbounds(A, inds...)
27+
return @inbounds get_or_view(A, inds...)
28+
29+
elseif length(args) > ndims(A) && all(args[ndims(A)+1:end] .== (:)) # trailing colons
30+
args_nd = args[1:ndims(A)]
31+
inds_raw = map(findindex, args_nd, axiskeys(A))
32+
inds = Base.to_indices(A, inds_raw)
33+
@boundscheck checkbounds(A, inds...)
34+
if inds isa NTuple{<:Any, Int}
35+
return @inbounds view(keyless(A), inds...) # zero-dim view of underlying
36+
else
37+
return @inbounds get_or_view(A, inds...)
38+
end
39+
40+
elseif length(args)==1
41+
arg = first(args)
42+
d = inferdim(arg, axiskeys(A))
43+
i = findindex(arg, axiskeys(A,d))
44+
inds = ntuple(n -> n==d ? i : (:), ndims(A))
45+
@boundscheck checkbounds(A, inds...)
46+
return @inbounds view(A, inds...)
47+
48+
end
49+
50+
if length(args) != ndims(A)
51+
throw(ArgumentError(string("wrong number of keys: got ", length(args),
52+
" arguments, expected ndims(A) = ", ndims(A)," and perhaps a trailing colon.")))
53+
else
54+
throw(ArgumentError("can't understand what to do with $args, sorry"))
55+
end
56+
end
57+
58+
@propagate_inbounds get_or_view(A, inds::Integer...) = getindex(A, inds...)
59+
@propagate_inbounds get_or_view(A, inds...) = view(A, inds...)
60+
61+
@propagate_inbounds function setkey!(A, val, args...)
62+
length(args) == ndims(A) || error("wrong number of keys")
63+
inds = map((v,r) -> findindex(v,r), args, axiskeys(A))
64+
setindex!(A, val, inds...)
65+
end
66+
67+
68+
"""
69+
findindex(key, vec)
70+
71+
This is usually `findfirst(isequal(key), vec)`, and will error if it finds `nothing`.
72+
But it also understands `findindex(:, vec) = (:)`,
73+
and `findindex(array, vec) = vcat((findindex(x, vec) for x in array)...)`.
74+
75+
It also understands functions `findindex(<(4), vec) = findall(x -> x<4, vec)`,
76+
and selectors like `Nearest(key)` and `Interval(lo,hi)`.
77+
"""
78+
@inline function findindex(a, r::AbstractArray)
79+
i = findfirst(isequal(a), r)
80+
i === nothing && throw(ArgumentError("could not find key $(repr(a)) in vector $r"))
81+
i
82+
end
83+
84+
findindex(a::Colon, r::AbstractArray) = Colon()
85+
86+
findindex(a::Union{AbstractArray, Base.Generator}, r::AbstractArray) =
87+
reduce(vcat, [findindex(x, r) for x in a])
88+
89+
findindex(f::Function, r::AbstractArray) = findall(f, r)
90+
91+
# It's possible this should be a method of to_indices or one of its friends?
92+
# https://docs.julialang.org/en/v1/base/arrays/#Base.to_indices
93+
94+
95+
"""
96+
inferdim(key, axiskeys::Tuple)
97+
98+
When you call `A(key)` for `ndims(A) > 1`, this returns which `d` you meant,
99+
if unambigous, by comparing types & gradually widening.
100+
"""
101+
@generated inferdim(arg, tup) = _inferdim(arg, map(eltype, Tuple(tup.parameters)))
102+
103+
function _inferdim(argT, types, subtypes=())
104+
types == subtypes && throw(ArgumentError("key of type $argT doesn't match any dimensions"))
105+
106+
# First look for direct match
107+
ds = findall(T -> argT <: T, types)
108+
109+
if length(ds) == 1
110+
return first(ds)
111+
elseif length(ds) >= 2
112+
throw(ArgumentError("key of type $argT is ambiguous, matches dimensions $(Tuple(ds))"))
113+
end
114+
115+
# If no direct match, look for a container whose eltype matches:
116+
if argT <: Selector || argT <: AbstractArray || argT <: Interval
117+
ds = findall(T -> eltype(argT) <: T, types)
118+
elseif argT <: Base.Fix2 # Base.Fix2{typeof(==),Int64}
119+
ds = findall(T -> argT.parameters[2] <: T, types)
120+
end
121+
122+
if length(ds) == 1
123+
return first(ds)
124+
elseif length(ds) >= 2
125+
throw(ArgumentError("key of type $argT is ambiguous, matches dimensions $(Tuple(ds))"))
126+
end
127+
128+
# Otherwise, widen the key types and try again.
129+
# This will recurse until types stop changing.
130+
supers = map(T -> supertype(T) == Any ? T : supertype(T), types)
131+
return _inferdim(argT, supers, types)
132+
end

‎src/struct.jl‎

Lines changed: 0 additions & 131 deletions
Original file line numberDiff line numberDiff line change
@@ -131,134 +131,3 @@ end
131131
@inbounds setindex!(parent(A), val, I...)
132132
val
133133
end
134-
135-
"""
136-
(A::KeyedArray)("a", 2.0, :γ) == A[1, 2, 3]
137-
A(:γ) == view(A, :, :, 3)
138-
139-
`KeyedArray`s are callable, and this behaves much like indexing,
140-
except that it searches for the given keys in `axiskeys(A)`,
141-
instead of `axes(A)` for indices.
142-
143-
A single key may be used to indicate a slice, provided that its type
144-
only matches the eltype of one `axiskeys(A,d)`.
145-
You can also slice explicitly with `A("a", :, :)`, both of these return a `view`.
146-
147-
An extra trailing colon (when all other indices are fixed) will return
148-
a zero-dimensional `view`. This allows setting one value by
149-
writing `A("a", 2.0, :γ, :) .= 100`.
150-
151-
Also accepts functions like `A(<=(2.0))` and selectors,
152-
see `Nearest` and `Index`.
153-
"""
154-
@inline @propagate_inbounds (A::KeyedArray)(args...) = getkey(A, args...)
155-
156-
@inline function getkey(A, args...)
157-
if length(args) == ndims(A)
158-
inds_raw = map(findindex, args, axiskeys(A))
159-
inds = Base.to_indices(A, inds_raw)
160-
@boundscheck checkbounds(A, inds...)
161-
return @inbounds get_or_view(A, inds...)
162-
163-
elseif length(args) > ndims(A) && all(args[ndims(A)+1:end] .== (:)) # trailing colons
164-
args_nd = args[1:ndims(A)]
165-
inds_raw = map(findindex, args_nd, axiskeys(A))
166-
inds = Base.to_indices(A, inds_raw)
167-
@boundscheck checkbounds(A, inds...)
168-
if inds isa NTuple{<:Any, Int}
169-
return @inbounds view(keyless(A), inds...) # zero-dim view of underlying
170-
else
171-
return @inbounds get_or_view(A, inds...)
172-
end
173-
174-
elseif length(args)==1
175-
arg = first(args)
176-
d = inferdim(arg, axiskeys(A))
177-
i = findindex(arg, axiskeys(A,d))
178-
inds = ntuple(n -> n==d ? i : (:), ndims(A))
179-
@boundscheck checkbounds(A, inds...)
180-
return @inbounds view(A, inds...)
181-
182-
end
183-
184-
if length(args) != ndims(A)
185-
throw(ArgumentError(string("wrong number of keys: got ", length(args),
186-
" arguments, expected ndims(A) = ", ndims(A)," and perhaps a trailing colon.")))
187-
else
188-
throw(ArgumentError("can't understand what to do with $args, sorry"))
189-
end
190-
end
191-
192-
@propagate_inbounds get_or_view(A, inds::Integer...) = getindex(A, inds...)
193-
@propagate_inbounds get_or_view(A, inds...) = view(A, inds...)
194-
195-
@propagate_inbounds function setkey!(A, val, args...)
196-
length(args) == ndims(A) || error("wrong number of keys")
197-
inds = map((v,r) -> findindex(v,r), args, axiskeys(A))
198-
setindex!(A, val, inds...)
199-
end
200-
201-
"""
202-
inferdim(key, axiskeys::Tuple)
203-
204-
When you call `A(key)` for `ndims(A) > 1`, this returns which `d` you meant,
205-
if unambigous, by comparing types & gradually widening.
206-
"""
207-
@generated inferdim(arg, tup) = _inferdim(arg, map(eltype, Tuple(tup.parameters)))
208-
209-
function _inferdim(argT, types, subtypes=())
210-
types == subtypes && throw(ArgumentError("key of type $argT doesn't match any dimensions"))
211-
212-
# First look for direct match
213-
ds = findall(T -> argT <: T, types)
214-
215-
if length(ds) == 1
216-
return first(ds)
217-
elseif length(ds) >= 2
218-
throw(ArgumentError("key of type $argT is ambiguous, matches dimensions $(Tuple(ds))"))
219-
end
220-
221-
# If no direct match, look for a container whose eltype matches:
222-
if argT <: Selector || argT <: AbstractArray || argT <: Interval
223-
ds = findall(T -> eltype(argT) <: T, types)
224-
elseif argT <: Base.Fix2 # Base.Fix2{typeof(==),Int64}
225-
ds = findall(T -> argT.parameters[2] <: T, types)
226-
end
227-
228-
if length(ds) == 1
229-
return first(ds)
230-
elseif length(ds) >= 2
231-
throw(ArgumentError("key of type $argT is ambiguous, matches dimensions $(Tuple(ds))"))
232-
end
233-
234-
# Otherwise, widen the key types and try again.
235-
# This will recurse until types stop changing.
236-
supers = map(T -> supertype(T) == Any ? T : supertype(T), types)
237-
return _inferdim(argT, supers, types)
238-
end
239-
240-
"""
241-
findindex(key, vec)
242-
243-
This is usually `findfirst(isequal(key), vec)`, and will error if it finds `nothing`.
244-
But it also understands `findindex(:, vec) = (:)`,
245-
and `findindex(array, vec) = vcat((findindex(x, vec) for x in array)...)`.
246-
247-
It also understands functions `findindex(<(4), vec) = findall(x -> x<4, vec)`,
248-
and selectors like `Nearest(key)` and `Interval(lo,hi)`.
249-
"""
250-
@inline function findindex(a, r::AbstractArray)
251-
i = findfirst(isequal(a), r)
252-
i === nothing && throw(ArgumentError("could not find key $(repr(a)) in vector $r"))
253-
i
254-
end
255-
256-
findindex(a::Colon, r::AbstractArray) = Colon()
257-
258-
findindex(a::Union{AbstractArray, Base.Generator}, r::AbstractArray) =
259-
reduce(vcat, [findindex(x, r) for x in a])
260-
261-
findindex(f::Function, r::AbstractArray) = findall(f, r)
262-
263-
# It's possible this should be a method of to_indices or one of its friends?
264-
# https://docs.julialang.org/en/v1/base/arrays/#Base.to_indices

0 commit comments

Comments
 (0)