Skip to content

Commit 5d859e3

Browse files
johnnychen94johnnychen94
committed
[WIP] rework OffsetArrays constructor
1 parent 93e60d3 commit 5d859e3

File tree

2 files changed

+117
-135
lines changed

2 files changed

+117
-135
lines changed

‎src/OffsetArrays.jl

Lines changed: 73 additions & 135 deletions
Original file line numberDiff line numberDiff line change
@@ -10,57 +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-
if offset > 0 && last(r) > typemax(T) - offset
28-
throw(ArgumentError("Boundary overflow detected: offset $offset should be equal or less than $(typemax(T) - last(r))"))
29-
elseif offset < 0 && first(r) < typemin(T) - offset
30-
throw(ArgumentError("Boundary overflow detected: offset $offset should be equal or greater than $(typemin(T) - first(r))"))
31-
end
32-
end
33-
## OffsetArray constructors
34-
35-
offset(axparent::AbstractUnitRange, ax::AbstractUnitRange) = first(ax) - first(axparent)
36-
offset(axparent::AbstractUnitRange, ax::Integer) = 1 - first(axparent)
37-
38-
function OffsetArray(A::AbstractArray{T,N}, offsets::NTuple{N,Int}) where {T,N}
39-
OffsetArray{T,N,typeof(A)}(A, offsets)
40-
end
41-
OffsetArray(A::AbstractArray{T,0}, offsets::Tuple{}) where T =
42-
OffsetArray{T,0,typeof(A)}(A, ())
43-
44-
OffsetArray(A::AbstractArray{T,N}, offsets::Vararg{Int,N}) where {T,N} =
45-
OffsetArray(A, offsets)
46-
OffsetArray(A::AbstractArray{T,0}) where {T} = OffsetArray(A, ())
47-
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}
4818
const ArrayInitializer = Union{UndefInitializer, Missing, Nothing}
49-
OffsetArray{T,N}(init::ArrayInitializer, inds::Indices{N}) where {T,N} =
50-
OffsetArray(Array{T,N}(init, map(indexlength, inds)), map(indexoffset, inds))
51-
OffsetArray{T}(init::ArrayInitializer, inds::Indices{N}) where {T,N} = OffsetArray{T,N}(init, inds)
52-
OffsetArray{T,N}(init::ArrayInitializer, inds::Vararg{AbstractUnitRange,N}) where {T,N} = OffsetArray{T,N}(init, inds)
53-
OffsetArray{T}(init::ArrayInitializer, inds::Vararg{AbstractUnitRange,N}) where {T,N} = OffsetArray{T,N}(init, inds)
54-
55-
# OffsetVector constructors
56-
OffsetVector(A::AbstractVector, offset) = OffsetArray(A, offset)
57-
OffsetVector{T}(init::ArrayInitializer, inds::AbstractUnitRange) where {T} = OffsetArray{T}(init, inds)
58-
59-
# OffsetMatrix constructors
60-
OffsetMatrix(A::AbstractMatrix, offset1, offset2) = OffsetArray(A, offset1, offset2)
61-
OffsetMatrix(A::AbstractMatrix, I::CartesianIndices{2}) = OffsetArray(A, I)
62-
OffsetMatrix{T}(init::ArrayInitializer, inds1::AbstractUnitRange, inds2::AbstractUnitRange) where {T} = OffsetArray{T}(init, inds1, inds2)
6319

20+
## OffsetArray
6421
"""
6522
OffsetArray(A, indices...)
6623
@@ -111,84 +68,77 @@ ERROR: [...]
11168
```
11269
11370
"""
114-
function OffsetArray(A::AbstractArray{T,N}, inds::NTuple{N,AbstractUnitRange}) where {T,N}
115-
axparent = axes(A)
116-
lA = map(length, axparent)
117-
lI = map(length, inds)
118-
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"))
119-
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
12078
end
121-
OffsetArray(A::AbstractArray{T,N}, inds::Vararg{AbstractUnitRange,N}) where {T,N} =
122-
OffsetArray(A, inds)
123-
124-
uncolonindices(A::AbstractArray{<:Any,N}, inds::NTuple{N,Any}) where {N} = uncolonindices(axes(A), inds)
125-
uncolonindices(ax::Tuple, inds::Tuple) = (first(inds), uncolonindices(tail(ax), tail(inds))...)
126-
uncolonindices(ax::Tuple, inds::Tuple{Colon, Vararg{Any}}) = (first(ax), uncolonindices(tail(ax), tail(inds))...)
127-
uncolonindices(::Tuple{}, ::Tuple{}) = ()
79+
const OffsetVector{T,AA<:AbstractArray} = OffsetArray{T,1,AA}
80+
const OffsetMatrix{T,AA<:AbstractArray} = OffsetArray{T,2,AA}
12881

129-
function OffsetArray(A::AbstractArray{T,N}, inds::NTuple{N,Union{AbstractUnitRange, CartesianIndices{1}, Colon}}) where {T,N}
130-
OffsetArray(A, uncolonindices(A, inds))
131-
end
132-
OffsetArray(A::AbstractArray{T,N}, inds::Vararg{Union{AbstractUnitRange, CartesianIndices{1}, Colon},N}) where {T,N} =
133-
OffsetArray(A, uncolonindices(A, inds))
134-
135-
# Specify offsets using CartesianIndices (issue #71)
136-
# Support a mix of AbstractUnitRanges and CartesianIndices{1}
137-
# Extract the range r from CartesianIndices((r,))
138-
function stripCartesianIndices(inds::Tuple{CartesianIndices{1},Vararg{Any}})
139-
I = first(inds)
140-
Ir = convert(Tuple{AbstractUnitRange{Int}}, I) |> first
141-
(Ir, stripCartesianIndices(tail(inds))...)
142-
end
143-
stripCartesianIndices(inds::Tuple)= (first(inds), stripCartesianIndices(tail(inds))...)
144-
stripCartesianIndices(::Tuple{}) = ()
145-
146-
OffsetArray(A::AbstractArray{<:Any,N}, inds::NTuple{N,Union{CartesianIndices{1}, AbstractUnitRange}}) where {N} =
147-
OffsetArray(A, stripCartesianIndices(inds))
148-
OffsetArray(A::AbstractArray{<:Any,N}, inds::Vararg{Union{CartesianIndices{1}, AbstractUnitRange},N}) where {N} =
149-
OffsetArray(A, inds)
150-
151-
# Support an arbitrary CartesianIndices alongside colons and ranges
152-
# The total number of indices should equal ndims(arr)
153-
# We split the CartesianIndices{N} into N CartesianIndices{1} indices to facilitate dispatch
154-
splitCartesianIndices(c::CartesianIndices{0}) = ()
155-
function splitCartesianIndices(c::CartesianIndices)
156-
c1, ct = Base.IteratorsMD.split(c, Val(1))
157-
(c1, splitCartesianIndices(ct)...)
158-
end
159-
function splitCartesianIndices(t::Tuple{CartesianIndices, Vararg{Any}})
160-
(splitCartesianIndices(first(t))..., splitCartesianIndices(tail(t))...)
161-
end
162-
function splitCartesianIndices(t::Tuple)
163-
(first(t), splitCartesianIndices(tail(t))...)
82+
function overflow_check(r, offset::T) where T
83+
if offset > 0 && last(r) > typemax(T) - offset
84+
throw(ArgumentError("Boundary overflow detected: offset $offset should be equal or less than $(typemax(T) - last(r))"))
85+
elseif offset < 0 && first(r) < typemin(T) - offset
86+
throw(ArgumentError("Boundary overflow detected: offset $offset should be equal or greater than $(typemin(T) - first(r))"))
87+
end
16488
end
165-
splitCartesianIndices(::Tuple{}) = ()
89+
## OffsetArray constructors
16690

167-
function OffsetArray(A::AbstractArray, inds::Tuple{Vararg{Union{AbstractUnitRange, CartesianIndices, Colon}}})
168-
OffsetArray(A, splitCartesianIndices(inds))
169-
end
170-
function OffsetArray(A::AbstractArray, inds::Vararg{Union{AbstractUnitRange, CartesianIndices, Colon}})
171-
OffsetArray(A, inds)
91+
for FT in (:OffsetArray, :OffsetVector, :OffsetMatrix)
92+
# The only route out to inner constructor
93+
@eval function $FT(A::AbstractArray{T}, offsets::NTuple{N, Integer}) where {T, N}
94+
ndims(A) == N || throw(DimensionMismatch("Array dimensions should equal to number of offsets"))
95+
OffsetArray{T, ndims(A), typeof(A)}(A, offsets)
96+
end
97+
# nested OffsetArrays
98+
@eval function $FT(A::OffsetArray{T}, offsets::NTuple{N, Integer}) where {T,N}
99+
$FT(parent(A), A.offsets .+ offsets)
100+
end
101+
# convert ranges to offsets
102+
@eval function $FT(A::AbstractArray{T}, inds::NTuple{N,OffsetAxisKnownLength}) where {T,N}
103+
axparent = axes(A)
104+
lA = map(length, axparent)
105+
lI = map(length, inds)
106+
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"))
107+
$FT(A, map(_offset, axparent, inds))
108+
end
109+
# lower CartesianIndices and Colon
110+
@eval function $FT(A::AbstractArray{T}, inds::NTuple{N, OffsetAxis}) where {T, N}
111+
indsN = _uncolonindices(A, _stripCartesianIndices(_splitCartesianIndices(inds)))
112+
$FT(A, indsN)
113+
end
114+
@eval function $FT(A::AbstractArray{T}, inds::Vararg{OffsetAxis,N}) where {T, N}
115+
$FT(A, _uncolonindices(A, _stripCartesianIndices(_splitCartesianIndices(inds))))
116+
end
117+
@eval $FT(A::AbstractArray, inds::CartesianIndices) = $FT(A, inds.indices)
172118
end
173119

174-
# Add methods to initialize OffsetArrays using CartesianIndices (issue #71)
175-
function OffsetArray{T,N}(init::ArrayInitializer, inds::Tuple{Vararg{Union{AbstractUnitRange, CartesianIndices}}}) where {T,N}
176-
indsN = stripCartesianIndices(splitCartesianIndices(inds))
177-
OffsetArray{T,N}(init, indsN)
120+
# array initialization
121+
OffsetArray{T,N}(init::ArrayInitializer, inds::NTuple{N, OffsetAxisKnownLength}) where {T,N} =
122+
OffsetArray(Array{T,N}(init, map(_indexlength, inds)), map(_indexoffset, inds))
123+
OffsetArray{T,N}(init::ArrayInitializer, inds::Vararg{OffsetAxisKnownLength,N}) where {T,N} = OffsetArray{T,N}(init, inds)
124+
function OffsetArray{T, N}(init::ArrayInitializer, inds::NTuple{N, Union{OffsetAxisKnownLength, CartesianIndices}}) where {T, N}
125+
OffsetArray{T, N}(init, _stripCartesianIndices(_splitCartesianIndices(inds)))
178126
end
179-
function OffsetArray{T}(init::ArrayInitializer, inds::Tuple{Vararg{Union{AbstractUnitRange, CartesianIndices}}}) where {T}
180-
indsN = stripCartesianIndices(splitCartesianIndices(inds))
181-
OffsetArray{T}(init, indsN)
127+
function OffsetArray{T, N}(init::ArrayInitializer, inds::Vararg{Union{OffsetAxisKnownLength, CartesianIndices}}) where {T, N}
128+
OffsetArray{T, N}(init, inds)
182129
end
183-
OffsetArray{T,N}(init::ArrayInitializer, inds::Vararg{Union{AbstractUnitRange, CartesianIndices}}) where {T,N} = OffsetArray{T,N}(init, inds)
184-
OffsetArray{T}(init::ArrayInitializer, inds::Vararg{Union{AbstractUnitRange, CartesianIndices}}) where {T} = OffsetArray{T}(init, inds)
130+
OffsetArray{T, N}(init::ArrayInitializer, inds::CartesianIndices{N}) where {T,N} = OffsetArray{T, N}(init, inds.indices)
185131

186-
# avoid a level of indirection when nesting OffsetArrays
187-
function OffsetArray(A::OffsetArray, offsets::NTuple{N,Int}) where {N}
188-
OffsetArray(parent(A), offsets .+ A.offsets)
132+
OffsetArray{T}(init::ArrayInitializer, inds::NTuple{N, OffsetAxisKnownLength}) where {T,N} = OffsetArray{T,N}(init, inds)
133+
OffsetArray{T}(init::ArrayInitializer, inds::Vararg{OffsetAxisKnownLength,N}) where {T,N} = OffsetArray{T,N}(init, inds)
134+
function OffsetArray{T}(init::ArrayInitializer, inds::NTuple{N, Union{OffsetAxisKnownLength, CartesianIndices}}) where {T, N}
135+
indsN = _stripCartesianIndices(_splitCartesianIndices(inds)) # CartesianIndices might contain multiple dimensions
136+
OffsetArray{T, length(indsN)}(init, _stripCartesianIndices(_splitCartesianIndices(inds)))
189137
end
190-
OffsetArray(A::OffsetArray{T,0}, inds::Tuple{}) where {T} = OffsetArray(parent(A), ())
191-
# OffsetArray(A::OffsetArray{T,N}, inds::Tuple{}) where {T,N} = error("this should never be called")
138+
function OffsetArray{T}(init::ArrayInitializer, inds::Vararg{Union{OffsetAxisKnownLength, CartesianIndices}, N}) where {T, N}
139+
OffsetArray{T}(init, inds)
140+
end
141+
OffsetArray{T}(init::ArrayInitializer, inds::CartesianIndices{N}) where {T,N} = OffsetArray{T, N}(init, inds.indices)
192142

193143
Base.IndexStyle(::Type{OA}) where {OA<:OffsetArray} = IndexStyle(parenttype(OA))
194144
parenttype(::Type{OffsetArray{T,N,AA}}) where {T,N,AA} = AA
@@ -206,27 +156,24 @@ Base.eachindex(::IndexLinear, A::OffsetVector) = axes(A, 1)
206156
@inline Base.axes(A::OffsetArray, d) = d <= ndims(A) ? IdOffsetRange(axes(parent(A), d), A.offsets[d]) : IdOffsetRange(axes(parent(A), d))
207157
@inline Base.axes1(A::OffsetArray{T,0}) where {T} = IdOffsetRange(axes(parent(A), 1)) # we only need to specialize this one
208158

209-
const OffsetAxisKnownLength = Union{Integer, UnitRange, Base.OneTo, IdentityUnitRange, IdOffsetRange}
210-
211159
Base.similar(A::OffsetArray, ::Type{T}, dims::Dims) where T =
212160
similar(parent(A), T, dims)
213161
function Base.similar(A::AbstractArray, ::Type{T}, inds::Tuple{OffsetAxisKnownLength,Vararg{OffsetAxisKnownLength}}) where T
214-
B = similar(A, T, map(indexlength, inds))
215-
return OffsetArray(B, map(offset, axes(B), inds))
162+
B = similar(A, T, map(_indexlength, inds))
163+
return OffsetArray(B, map(_offset, axes(B), inds))
216164
end
217165

218166
# reshape accepts a single colon
219-
const OffsetAxis = Union{OffsetAxisKnownLength, Colon}
220167
Base.reshape(A::AbstractArray, inds::OffsetAxis...) = reshape(A, inds)
221168
function Base.reshape(A::AbstractArray, inds::Tuple{OffsetAxis,Vararg{OffsetAxis}})
222-
AR = reshape(A, map(indexlength, inds))
223-
return OffsetArray(AR, map(offset, axes(AR), inds))
169+
AR = reshape(A, map(_indexlength, inds))
170+
return OffsetArray(AR, map(_offset, axes(AR), inds))
224171
end
225172

226173
# Reshaping OffsetArrays can "pop" the original OffsetArray wrapper and return
227174
# an OffsetArray(reshape(...)) instead of an OffsetArray(reshape(OffsetArray(...)))
228175
Base.reshape(A::OffsetArray, inds::Tuple{OffsetAxis,Vararg{OffsetAxis}}) =
229-
OffsetArray(reshape(parent(A), map(indexlength, inds)), map(indexoffset, inds))
176+
OffsetArray(reshape(parent(A), map(_indexlength, inds)), map(_indexoffset, inds))
230177
# And for non-offset axes, we can just return a reshape of the parent directly
231178
Base.reshape(A::OffsetArray, inds::Tuple{Union{Integer,Base.OneTo},Vararg{Union{Integer,Base.OneTo}}}) = reshape(parent(A), inds)
232179
Base.reshape(A::OffsetArray, inds::Dims) = reshape(parent(A), inds)
@@ -235,8 +182,8 @@ Base.reshape(A::OffsetArray, inds::Union{Int,Colon}...) = reshape(parent(A), ind
235182
Base.reshape(A::OffsetArray, inds::Tuple{Vararg{Union{Int,Colon}}}) = reshape(parent(A), inds)
236183

237184
function Base.similar(::Type{T}, shape::Tuple{OffsetAxis,Vararg{OffsetAxis}}) where {T<:AbstractArray}
238-
P = T(undef, map(indexlength, shape))
239-
OffsetArray(P, map(offset, axes(P), shape))
185+
P = T(undef, map(_indexlength, shape))
186+
OffsetArray(P, map(_offset, axes(P), shape))
240187
end
241188

242189
Base.fill(v, inds::NTuple{N, Union{Integer, AbstractUnitRange}}) where {N} =
@@ -333,15 +280,6 @@ Base.pop!(A::OffsetVector) = pop!(A.parent)
333280
Base.append!(A::OffsetVector, items) = (append!(A.parent, items); A)
334281
Base.empty!(A::OffsetVector) = (empty!(A.parent); A)
335282

336-
### Low-level utilities ###
337-
338-
indexoffset(r::AbstractRange) = first(r) - 1
339-
indexoffset(i::Integer) = 0
340-
indexoffset(i::Colon) = 0
341-
indexlength(r::AbstractRange) = length(r)
342-
indexlength(i::Integer) = i
343-
indexlength(i::Colon) = Colon()
344-
345283
# These functions keep the summary compact
346284
function Base.inds2string(inds::Tuple{Vararg{Union{IdOffsetRange, IdentityUnitRange{<:IdOffsetRange}}}})
347285
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)