Skip to content

Commit 6f279a4

Browse files
johnnychen94johnnychen94
committed
[WIP] rework OffsetArrays constructor
1 parent dc867d5 commit 6f279a4

File tree

2 files changed

+120
-139
lines changed

2 files changed

+120
-139
lines changed

‎src/OffsetArrays.jl

Lines changed: 76 additions & 139 deletions
Original file line numberDiff line numberDiff line change
@@ -10,62 +10,14 @@ end
1010
export OffsetArray, OffsetMatrix, OffsetVector
1111

1212
include("axes.jl")
13+
include("utils.jl")
1314

14-
## OffsetArray
15-
struct OffsetArray{T,N,AA<:AbstractArray} <: AbstractArray{T,N}
16-
parent::AA
17-
offsets::NTuple{N,Int}
18-
function OffsetArray{T, N, AA}(parent::AA, offsets::NTuple{N, Int}) where {T, N, AA<:AbstractArray}
19-
overflow_check.(axes(parent), offsets)
20-
new{T, N, AA}(parent, offsets)
21-
end
22-
end
23-
OffsetVector{T,AA<:AbstractArray} = OffsetArray{T,1,AA}
24-
OffsetMatrix{T,AA<:AbstractArray} = OffsetArray{T,2,AA}
25-
26-
function overflow_check(r, offset::T) where T
27-
# This gives some performance boost https://github.com/JuliaLang/julia/issues/33273
28-
throw_upper_overflow_error() = throw(ArgumentError("Boundary overflow detected: offset $offset should be equal or less than $(typemax(T) - last(r))"))
29-
throw_lower_overflow_error() = throw(ArgumentError("Boundary overflow detected: offset $offset should be equal or greater than $(typemin(T) - first(r))"))
30-
31-
if offset > 0 && last(r) > typemax(T) - offset
32-
throw_upper_overflow_error()
33-
elseif offset < 0 && first(r) < typemin(T) - offset
34-
throw_lower_overflow_error()
35-
end
36-
end
37-
38-
## OffsetArray constructors
39-
40-
offset(axparent::AbstractUnitRange, ax::AbstractUnitRange) = first(ax) - first(axparent)
41-
offset(axparent::AbstractUnitRange, ax::Integer) = 1 - first(axparent)
42-
43-
function OffsetArray(A::AbstractArray{T,N}, offsets::NTuple{N,Int}) where {T,N}
44-
OffsetArray{T,N,typeof(A)}(A, offsets)
45-
end
46-
OffsetArray(A::AbstractArray{T,0}, offsets::Tuple{}) where T =
47-
OffsetArray{T,0,typeof(A)}(A, ())
48-
49-
OffsetArray(A::AbstractArray{T,N}, offsets::Vararg{Int,N}) where {T,N} =
50-
OffsetArray(A, offsets)
51-
OffsetArray(A::AbstractArray{T,0}) where {T} = OffsetArray(A, ())
52-
15+
# Techniquely we know the length of CartesianIndices
16+
const OffsetAxisKnownLength = Union{Integer, UnitRange, Base.OneTo, IdentityUnitRange, IdOffsetRange}
17+
const OffsetAxis = Union{OffsetAxisKnownLength, CartesianIndices, Colon}
5318
const ArrayInitializer = Union{UndefInitializer, Missing, Nothing}
54-
OffsetArray{T,N}(init::ArrayInitializer, inds::Indices{N}) where {T,N} =
55-
OffsetArray(Array{T,N}(init, map(indexlength, inds)), map(indexoffset, inds))
56-
OffsetArray{T}(init::ArrayInitializer, inds::Indices{N}) where {T,N} = OffsetArray{T,N}(init, inds)
57-
OffsetArray{T,N}(init::ArrayInitializer, inds::Vararg{AbstractUnitRange,N}) where {T,N} = OffsetArray{T,N}(init, inds)
58-
OffsetArray{T}(init::ArrayInitializer, inds::Vararg{AbstractUnitRange,N}) where {T,N} = OffsetArray{T,N}(init, inds)
59-
60-
# OffsetVector constructors
61-
OffsetVector(A::AbstractVector, offset) = OffsetArray(A, offset)
62-
OffsetVector{T}(init::ArrayInitializer, inds::AbstractUnitRange) where {T} = OffsetArray{T}(init, inds)
63-
64-
# OffsetMatrix constructors
65-
OffsetMatrix(A::AbstractMatrix, offset1, offset2) = OffsetArray(A, offset1, offset2)
66-
OffsetMatrix(A::AbstractMatrix, I::CartesianIndices{2}) = OffsetArray(A, I)
67-
OffsetMatrix{T}(init::ArrayInitializer, inds1::AbstractUnitRange, inds2::AbstractUnitRange) where {T} = OffsetArray{T}(init, inds1, inds2)
6819

20+
## OffsetArray
6921
"""
7022
OffsetArray(A, indices...)
7123
@@ -116,84 +68,81 @@ ERROR: [...]
11668
```
11769
11870
"""
119-
function OffsetArray(A::AbstractArray{T,N}, inds::NTuple{N,AbstractUnitRange}) where {T,N}
120-
axparent = axes(A)
121-
lA = map(length, axparent)
122-
lI = map(length, inds)
123-
lA == lI || throw(DimensionMismatch("supplied axes do not agree with the size of the array (got size $lA for the array and $lI for the indices"))
124-
OffsetArray(A, map(offset, axparent, inds))
71+
struct OffsetArray{T,N,AA<:AbstractArray} <: AbstractArray{T,N}
72+
parent::AA
73+
offsets::NTuple{N,Int}
74+
function OffsetArray{T, N, AA}(parent::AA, offsets::NTuple{N, Int}) where {T, N, AA<:AbstractArray}
75+
overflow_check.(axes(parent), offsets)
76+
new{T, N, AA}(parent, offsets)
77+
end
12578
end
126-
OffsetArray(A::AbstractArray{T,N}, inds::Vararg{AbstractUnitRange,N}) where {T,N} =
127-
OffsetArray(A, inds)
79+
const OffsetVector{T,AA<:AbstractArray} = OffsetArray{T,1,AA}
80+
const OffsetMatrix{T,AA<:AbstractArray} = OffsetArray{T,2,AA}
12881

129-
uncolonindices(A::AbstractArray{<:Any,N}, inds::NTuple{N,Any}) where {N} = uncolonindices(axes(A), inds)
130-
uncolonindices(ax::Tuple, inds::Tuple) = (first(inds), uncolonindices(tail(ax), tail(inds))...)
131-
uncolonindices(ax::Tuple, inds::Tuple{Colon, Vararg{Any}}) = (first(ax), uncolonindices(tail(ax), tail(inds))...)
132-
uncolonindices(::Tuple{}, ::Tuple{}) = ()
82+
function overflow_check(r, offset::T) where T
83+
# This gives some performance boost https://github.com/JuliaLang/julia/issues/33273
84+
throw_upper_overflow_error() = throw(ArgumentError("Boundary overflow detected: offset $offset should be equal or less than $(typemax(T) - last(r))"))
85+
throw_lower_overflow_error() = throw(ArgumentError("Boundary overflow detected: offset $offset should be equal or greater than $(typemin(T) - first(r))"))
13386

134-
function OffsetArray(A::AbstractArray{T,N}, inds::NTuple{N,Union{AbstractUnitRange, CartesianIndices{1}, Colon}}) where {T,N}
135-
OffsetArray(A, uncolonindices(A, inds))
136-
end
137-
OffsetArray(A::AbstractArray{T,N}, inds::Vararg{Union{AbstractUnitRange, CartesianIndices{1}, Colon},N}) where {T,N} =
138-
OffsetArray(A, uncolonindices(A, inds))
139-
140-
# Specify offsets using CartesianIndices (issue #71)
141-
# Support a mix of AbstractUnitRanges and CartesianIndices{1}
142-
# Extract the range r from CartesianIndices((r,))
143-
function stripCartesianIndices(inds::Tuple{CartesianIndices{1},Vararg{Any}})
144-
I = first(inds)
145-
Ir = convert(Tuple{AbstractUnitRange{Int}}, I) |> first
146-
(Ir, stripCartesianIndices(tail(inds))...)
147-
end
148-
stripCartesianIndices(inds::Tuple)= (first(inds), stripCartesianIndices(tail(inds))...)
149-
stripCartesianIndices(::Tuple{}) = ()
150-
151-
OffsetArray(A::AbstractArray{<:Any,N}, inds::NTuple{N,Union{CartesianIndices{1}, AbstractUnitRange}}) where {N} =
152-
OffsetArray(A, stripCartesianIndices(inds))
153-
OffsetArray(A::AbstractArray{<:Any,N}, inds::Vararg{Union{CartesianIndices{1}, AbstractUnitRange},N}) where {N} =
154-
OffsetArray(A, inds)
155-
156-
# Support an arbitrary CartesianIndices alongside colons and ranges
157-
# The total number of indices should equal ndims(arr)
158-
# We split the CartesianIndices{N} into N CartesianIndices{1} indices to facilitate dispatch
159-
splitCartesianIndices(c::CartesianIndices{0}) = ()
160-
function splitCartesianIndices(c::CartesianIndices)
161-
c1, ct = Base.IteratorsMD.split(c, Val(1))
162-
(c1, splitCartesianIndices(ct)...)
163-
end
164-
function splitCartesianIndices(t::Tuple{CartesianIndices, Vararg{Any}})
165-
(splitCartesianIndices(first(t))..., splitCartesianIndices(tail(t))...)
166-
end
167-
function splitCartesianIndices(t::Tuple)
168-
(first(t), splitCartesianIndices(tail(t))...)
87+
if offset > 0 && last(r) > typemax(T) - offset
88+
throw_upper_overflow_error()
89+
elseif offset < 0 && first(r) < typemin(T) - offset
90+
throw_lower_overflow_error()
91+
end
16992
end
170-
splitCartesianIndices(::Tuple{}) = ()
93+
## OffsetArray constructors
17194

172-
function OffsetArray(A::AbstractArray, inds::Tuple{Vararg{Union{AbstractUnitRange, CartesianIndices, Colon}}})
173-
OffsetArray(A, splitCartesianIndices(inds))
174-
end
175-
function OffsetArray(A::AbstractArray, inds::Vararg{Union{AbstractUnitRange, CartesianIndices, Colon}})
176-
OffsetArray(A, inds)
95+
for FT in (:OffsetArray, :OffsetVector, :OffsetMatrix)
96+
# The only route out to inner constructor
97+
@eval function $FT(A::AbstractArray{T}, offsets::NTuple{N, Integer}) where {T, N}
98+
ndims(A) == N || throw(DimensionMismatch("Array dimensions should equal to number of offsets"))
99+
OffsetArray{T, ndims(A), typeof(A)}(A, offsets)
100+
end
101+
# nested OffsetArrays
102+
@eval function $FT(A::OffsetArray{T}, offsets::NTuple{N, Integer}) where {T,N}
103+
$FT(parent(A), A.offsets .+ offsets)
104+
end
105+
# convert ranges to offsets
106+
@eval function $FT(A::AbstractArray{T}, inds::NTuple{N,OffsetAxisKnownLength}) where {T,N}
107+
axparent = axes(A)
108+
lA = map(length, axparent)
109+
lI = map(length, inds)
110+
lA == lI || throw(DimensionMismatch("supplied axes do not agree with the size of the array (got size $lA for the array and $lI for the indices"))
111+
$FT(A, map(_offset, axparent, inds))
112+
end
113+
# lower CartesianIndices and Colon
114+
@eval function $FT(A::AbstractArray{T}, inds::NTuple{N, OffsetAxis}) where {T, N}
115+
indsN = _uncolonindices(A, _stripCartesianIndices(_splitCartesianIndices(inds)))
116+
$FT(A, indsN)
117+
end
118+
@eval function $FT(A::AbstractArray{T}, inds::Vararg{OffsetAxis,N}) where {T, N}
119+
$FT(A, _uncolonindices(A, _stripCartesianIndices(_splitCartesianIndices(inds))))
120+
end
121+
@eval $FT(A::AbstractArray, inds::CartesianIndices) = $FT(A, inds.indices)
177122
end
178123

179-
# Add methods to initialize OffsetArrays using CartesianIndices (issue #71)
180-
function OffsetArray{T,N}(init::ArrayInitializer, inds::Tuple{Vararg{Union{AbstractUnitRange, CartesianIndices}}}) where {T,N}
181-
indsN = stripCartesianIndices(splitCartesianIndices(inds))
182-
OffsetArray{T,N}(init, indsN)
124+
# array initialization
125+
OffsetArray{T,N}(init::ArrayInitializer, inds::NTuple{N, OffsetAxisKnownLength}) where {T,N} =
126+
OffsetArray(Array{T,N}(init, map(_indexlength, inds)), map(_indexoffset, inds))
127+
OffsetArray{T,N}(init::ArrayInitializer, inds::Vararg{OffsetAxisKnownLength,N}) where {T,N} = OffsetArray{T,N}(init, inds)
128+
function OffsetArray{T, N}(init::ArrayInitializer, inds::NTuple{N, Union{OffsetAxisKnownLength, CartesianIndices}}) where {T, N}
129+
OffsetArray{T, N}(init, _stripCartesianIndices(_splitCartesianIndices(inds)))
183130
end
184-
function OffsetArray{T}(init::ArrayInitializer, inds::Tuple{Vararg{Union{AbstractUnitRange, CartesianIndices}}}) where {T}
185-
indsN = stripCartesianIndices(splitCartesianIndices(inds))
186-
OffsetArray{T}(init, indsN)
131+
function OffsetArray{T, N}(init::ArrayInitializer, inds::Vararg{Union{OffsetAxisKnownLength, CartesianIndices}}) where {T, N}
132+
OffsetArray{T, N}(init, inds)
187133
end
188-
OffsetArray{T,N}(init::ArrayInitializer, inds::Vararg{Union{AbstractUnitRange, CartesianIndices}}) where {T,N} = OffsetArray{T,N}(init, inds)
189-
OffsetArray{T}(init::ArrayInitializer, inds::Vararg{Union{AbstractUnitRange, CartesianIndices}}) where {T} = OffsetArray{T}(init, inds)
134+
OffsetArray{T, N}(init::ArrayInitializer, inds::CartesianIndices{N}) where {T,N} = OffsetArray{T, N}(init, inds.indices)
190135

191-
# avoid a level of indirection when nesting OffsetArrays
192-
function OffsetArray(A::OffsetArray, offsets::NTuple{N,Int}) where {N}
193-
OffsetArray(parent(A), offsets .+ A.offsets)
136+
OffsetArray{T}(init::ArrayInitializer, inds::NTuple{N, OffsetAxisKnownLength}) where {T,N} = OffsetArray{T,N}(init, inds)
137+
OffsetArray{T}(init::ArrayInitializer, inds::Vararg{OffsetAxisKnownLength,N}) where {T,N} = OffsetArray{T,N}(init, inds)
138+
function OffsetArray{T}(init::ArrayInitializer, inds::NTuple{N, Union{OffsetAxisKnownLength, CartesianIndices}}) where {T, N}
139+
indsN = _stripCartesianIndices(_splitCartesianIndices(inds)) # CartesianIndices might contain multiple dimensions
140+
OffsetArray{T, length(indsN)}(init, _stripCartesianIndices(_splitCartesianIndices(inds)))
141+
end
142+
function OffsetArray{T}(init::ArrayInitializer, inds::Vararg{Union{OffsetAxisKnownLength, CartesianIndices}, N}) where {T, N}
143+
OffsetArray{T}(init, inds)
194144
end
195-
OffsetArray(A::OffsetArray{T,0}, inds::Tuple{}) where {T} = OffsetArray(parent(A), ())
196-
# OffsetArray(A::OffsetArray{T,N}, inds::Tuple{}) where {T,N} = error("this should never be called")
145+
OffsetArray{T}(init::ArrayInitializer, inds::CartesianIndices{N}) where {T,N} = OffsetArray{T, N}(init, inds.indices)
197146

198147
Base.IndexStyle(::Type{OA}) where {OA<:OffsetArray} = IndexStyle(parenttype(OA))
199148
parenttype(::Type{OffsetArray{T,N,AA}}) where {T,N,AA} = AA
@@ -211,27 +160,24 @@ Base.eachindex(::IndexLinear, A::OffsetVector) = axes(A, 1)
211160
@inline Base.axes(A::OffsetArray, d) = d <= ndims(A) ? IdOffsetRange(axes(parent(A), d), A.offsets[d]) : IdOffsetRange(axes(parent(A), d))
212161
@inline Base.axes1(A::OffsetArray{T,0}) where {T} = IdOffsetRange(axes(parent(A), 1)) # we only need to specialize this one
213162

214-
const OffsetAxisKnownLength = Union{Integer, UnitRange, Base.OneTo, IdentityUnitRange, IdOffsetRange}
215-
216163
Base.similar(A::OffsetArray, ::Type{T}, dims::Dims) where T =
217164
similar(parent(A), T, dims)
218165
function Base.similar(A::AbstractArray, ::Type{T}, inds::Tuple{OffsetAxisKnownLength,Vararg{OffsetAxisKnownLength}}) where T
219-
B = similar(A, T, map(indexlength, inds))
220-
return OffsetArray(B, map(offset, axes(B), inds))
166+
B = similar(A, T, map(_indexlength, inds))
167+
return OffsetArray(B, map(_offset, axes(B), inds))
221168
end
222169

223170
# reshape accepts a single colon
224-
const OffsetAxis = Union{OffsetAxisKnownLength, Colon}
225171
Base.reshape(A::AbstractArray, inds::OffsetAxis...) = reshape(A, inds)
226172
function Base.reshape(A::AbstractArray, inds::Tuple{OffsetAxis,Vararg{OffsetAxis}})
227-
AR = reshape(A, map(indexlength, inds))
228-
return OffsetArray(AR, map(offset, axes(AR), inds))
173+
AR = reshape(A, map(_indexlength, inds))
174+
return OffsetArray(AR, map(_offset, axes(AR), inds))
229175
end
230176

231177
# Reshaping OffsetArrays can "pop" the original OffsetArray wrapper and return
232178
# an OffsetArray(reshape(...)) instead of an OffsetArray(reshape(OffsetArray(...)))
233179
Base.reshape(A::OffsetArray, inds::Tuple{OffsetAxis,Vararg{OffsetAxis}}) =
234-
OffsetArray(reshape(parent(A), map(indexlength, inds)), map(indexoffset, inds))
180+
OffsetArray(reshape(parent(A), map(_indexlength, inds)), map(_indexoffset, inds))
235181
# And for non-offset axes, we can just return a reshape of the parent directly
236182
Base.reshape(A::OffsetArray, inds::Tuple{Union{Integer,Base.OneTo},Vararg{Union{Integer,Base.OneTo}}}) = reshape(parent(A), inds)
237183
Base.reshape(A::OffsetArray, inds::Dims) = reshape(parent(A), inds)
@@ -241,8 +187,8 @@ Base.reshape(A::OffsetArray, inds::Union{Int,Colon}...) = reshape(parent(A), ind
241187
Base.reshape(A::OffsetArray, inds::Tuple{Vararg{Union{Int,Colon}}}) = reshape(parent(A), inds)
242188

243189
function Base.similar(::Type{T}, shape::Tuple{OffsetAxis,Vararg{OffsetAxis}}) where {T<:AbstractArray}
244-
P = T(undef, map(indexlength, shape))
245-
OffsetArray(P, map(offset, axes(P), shape))
190+
P = T(undef, map(_indexlength, shape))
191+
OffsetArray(P, map(_offset, axes(P), shape))
246192
end
247193

248194
Base.fill(v, inds::NTuple{N, Union{Integer, AbstractUnitRange}}) where {N} =
@@ -339,15 +285,6 @@ Base.pop!(A::OffsetVector) = pop!(A.parent)
339285
Base.append!(A::OffsetVector, items) = (append!(A.parent, items); A)
340286
Base.empty!(A::OffsetVector) = (empty!(A.parent); A)
341287

342-
### Low-level utilities ###
343-
344-
indexoffset(r::AbstractRange) = first(r) - 1
345-
indexoffset(i::Integer) = 0
346-
indexoffset(i::Colon) = 0
347-
indexlength(r::AbstractRange) = length(r)
348-
indexlength(i::Integer) = i
349-
indexlength(i::Colon) = Colon()
350-
351288
# These functions keep the summary compact
352289
function Base.inds2string(inds::Tuple{Vararg{Union{IdOffsetRange, IdentityUnitRange{<:IdOffsetRange}}}})
353290
Base.inds2string(map(UnitRange, inds))

‎src/utils.jl

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
### Low-level utilities ###
2+
3+
_indexoffset(r::AbstractRange) = first(r) - 1
4+
_indexoffset(i::Integer) = 0
5+
_indexoffset(i::Colon) = 0
6+
_indexlength(r::AbstractRange) = length(r)
7+
_indexlength(i::Integer) = i
8+
_indexlength(i::Colon) = Colon()
9+
10+
_offset(axparent::AbstractUnitRange, ax::AbstractUnitRange) = first(ax) - first(axparent)
11+
_offset(axparent::AbstractUnitRange, ax::CartesianIndices) = _offset(axparent, first(ax.indices))
12+
_offset(axparent::AbstractUnitRange, ax::Integer) = 1 - first(axparent)
13+
14+
_uncolonindices(A::AbstractArray{<:Any,N}, inds::NTuple{N,Any}) where {N} = _uncolonindices(axes(A), inds)
15+
_uncolonindices(ax::Tuple, inds::Tuple) = (first(inds), _uncolonindices(tail(ax), tail(inds))...)
16+
_uncolonindices(ax::Tuple, inds::Tuple{Colon, Vararg{Any}}) = (first(ax), _uncolonindices(tail(ax), tail(inds))...)
17+
_uncolonindices(::Tuple{}, ::Tuple{}) = ()
18+
19+
# Specify offsets using CartesianIndices (issue #71)
20+
# Support a mix of AbstractUnitRanges and CartesianIndices{1}
21+
# Extract the range r from CartesianIndices((r,))
22+
function _stripCartesianIndices(inds::Tuple{CartesianIndices{1},Vararg{Any}})
23+
I = first(inds)
24+
Ir = convert(Tuple{AbstractUnitRange{Int}}, I) |> first
25+
(Ir, _stripCartesianIndices(tail(inds))...)
26+
end
27+
_stripCartesianIndices(inds::Tuple)= (first(inds), _stripCartesianIndices(tail(inds))...)
28+
_stripCartesianIndices(::Tuple{}) = ()
29+
30+
# Support an arbitrary CartesianIndices alongside colons and ranges
31+
# The total number of indices should equal ndims(arr)
32+
# We split the CartesianIndices{N} into N CartesianIndices{1} indices to facilitate dispatch
33+
_splitCartesianIndices(c::CartesianIndices{0}) = ()
34+
function _splitCartesianIndices(c::CartesianIndices)
35+
c1, ct = Base.IteratorsMD.split(c, Val(1))
36+
(c1, _splitCartesianIndices(ct)...)
37+
end
38+
function _splitCartesianIndices(t::Tuple{CartesianIndices, Vararg{Any}})
39+
(_splitCartesianIndices(first(t))..., _splitCartesianIndices(tail(t))...)
40+
end
41+
function _splitCartesianIndices(t::Tuple)
42+
(first(t), _splitCartesianIndices(tail(t))...)
43+
end
44+
_splitCartesianIndices(::Tuple{}) = ()

0 commit comments

Comments
 (0)