Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Reshape with an offset parent array #245

Merged
merged 10 commits into from
Jul 15, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 24 additions & 7 deletions src/OffsetArrays.jl
Original file line number Diff line number Diff line change
Expand Up @@ -340,22 +340,39 @@ _similar_axes_or_length(AT, ax::I, ::I) where {I} = similar(AT, map(_indexlength
# reshape accepts a single colon
Base.reshape(A::AbstractArray, inds::OffsetAxis...) = reshape(A, inds)
function Base.reshape(A::AbstractArray, inds::Tuple{OffsetAxis,Vararg{OffsetAxis}})
AR = reshape(A, map(_indexlength, inds))
AR = reshape(no_offset_view(A), map(_indexlength, inds))
O = OffsetArray(AR, map(_offset, axes(AR), inds))
return _popreshape(O, axes(AR), _filterreshapeinds(inds))
end

# Reshaping OffsetArrays can "pop" the original OffsetArray wrapper and return
# an OffsetArray(reshape(...)) instead of an OffsetArray(reshape(OffsetArray(...)))
# Short-circuit for AbstractVectors if the axes are compatible to get around the Base restriction
# to 1-based vectors
function _reshape(A::AbstractVector, inds::Tuple{OffsetAxis})
@noinline throw_dimerr(ind::Integer) = throw(
DimensionMismatch("parent has $(size(A,1)) elements, which is incompatible with length $ind"))
@noinline throw_dimerr(ind) = throw(
DimensionMismatch("parent has $(size(A,1)) elements, which is incompatible with indices $ind"))
_checksize(first(inds), size(A,1)) || throw_dimerr(first(inds))
A
end
_reshape(A, inds) = _reshape2(A, inds)
_reshape2(A, inds) = reshape(A, inds)
# avoid a stackoverflow by relegating to the parent if no_offset_view returns an offsetarray
_reshape2(A::OffsetArray, inds) = reshape(parent(A), inds)
_reshape_nov(A, inds) = _reshape(no_offset_view(A), inds)

Base.reshape(A::OffsetArray, inds::Tuple{OffsetAxis,Vararg{OffsetAxis}}) =
OffsetArray(reshape(parent(A), map(_indexlength, inds)), map(_indexoffset, inds))
OffsetArray(_reshape(parent(A), inds), map(_toaxis, inds))
# And for non-offset axes, we can just return a reshape of the parent directly
Base.reshape(A::OffsetArray, inds::Tuple{Union{Integer,Base.OneTo},Vararg{Union{Integer,Base.OneTo}}}) = reshape(parent(A), inds)
Base.reshape(A::OffsetArray, inds::Dims) = reshape(parent(A), inds)
Base.reshape(A::OffsetArray, ::Colon) = reshape(parent(A), Colon())
Base.reshape(A::OffsetArray, inds::Tuple{Union{Integer,Base.OneTo},Vararg{Union{Integer,Base.OneTo}}}) = _reshape_nov(A, inds)
Base.reshape(A::OffsetArray, inds::Dims) = _reshape_nov(A, inds)
Base.reshape(A::OffsetVector, ::Colon) = A
Base.reshape(A::OffsetArray, inds::Union{Int,Colon}...) = reshape(parent(A), inds)
Base.reshape(A::OffsetArray, inds::Tuple{Vararg{Union{Int,Colon}}}) = reshape(parent(A), inds)
Base.reshape(A::OffsetVector, ::Tuple{Colon}) = A
Base.reshape(A::OffsetArray, ::Colon) = reshape(A, (Colon(),))
Base.reshape(A::OffsetArray, inds::Union{Int,Colon}...) = reshape(A, inds)
Base.reshape(A::OffsetArray, inds::Tuple{Vararg{Union{Int,Colon}}}) = _reshape_nov(A, inds)

# permutedims in Base does not preserve axes, and can not be fixed in a non-breaking way
# This is a stopgap solution
Expand Down
9 changes: 8 additions & 1 deletion src/utils.jl
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,18 @@

_indexoffset(r::AbstractRange) = first(r) - 1
_indexoffset(i::Integer) = 0
_indexoffset(i::Colon) = 0
_indexlength(r::AbstractRange) = length(r)
_indexlength(i::Integer) = Int(i)
_indexlength(i::Colon) = Colon()

# utility methods used in reshape
# we don't use _indexlength in this to avoid converting the arguments to Int
_checksize(ind::Integer, s) = ind == s
_checksize(ind::AbstractUnitRange, s) = length(ind) == s

_toaxis(i::Integer) = Base.OneTo(i)
_toaxis(i) = i

_strip_IdOffsetRange(r::IdOffsetRange) = parent(r)
_strip_IdOffsetRange(r) = r

Expand Down
5 changes: 5 additions & 0 deletions test/customranges.jl
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,10 @@ for Z in [:ZeroBasedRange, :ZeroBasedUnitRange]
@boundscheck checkbounds(A, r)
OffsetArrays._indexedby(A[r.a], axes(r))
end

@eval Base.reshape(z::$Z, inds::Tuple{}) = reshape(parent(z), inds)
@eval Base.reshape(z::$Z, inds::Tuple{Int, Vararg{Int}}) = reshape(parent(z), inds)
@eval Base.reshape(z::$Z, inds::Tuple{Union{Int, AbstractUnitRange{<:Integer}}, Vararg{Union{Int, AbstractUnitRange{<:Integer}}}}) = reshape(parent(z), inds)
end

# A basic range that does not have specialized vector indexing methods defined
Expand All @@ -75,6 +79,7 @@ struct CustomRange{T,A<:AbstractRange{T}} <: AbstractRange{T}
end
Base.parent(r::CustomRange) = r.a
Base.size(r::CustomRange) = size(parent(r))
Base.length(r::CustomRange) = length(parent(r))
Base.axes(r::CustomRange) = axes(parent(r))
Base.first(r::CustomRange) = first(parent(r))
Base.last(r::CustomRange) = last(parent(r))
Expand Down
43 changes: 43 additions & 0 deletions test/runtests.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1756,6 +1756,49 @@ end
@test axes(R) == (1:2, 1:3)
R = reshape(zeros(6,1), 1:2, :)
@test axes(R) == (1:2, 1:3)

r = OffsetArray(ZeroBasedRange(3:4), 1);
@test reshape(r, 2) == 3:4
@test reshape(r, (2,)) == 3:4
@test reshape(r, :) == 3:4
@test reshape(r, (:,)) == 3:4

# getindex for a reshaped array that wraps an offset array is broken on 1.0
if VERSION >= v"1.1"
@test reshape(r, (2,:,4:4)) == OffsetArray(reshape(3:4, 2, 1, 1), 1:2, 1:1, 4:4)
end

# reshape works even if the parent doesn't have 1-based indices
# this works even if the parent doesn't support the reshape
r = OffsetArray(IdentityUnitRange(0:1), -1)
@test reshape(r, 2) == 0:1
@test reshape(r, (2,)) == 0:1
@test reshape(r, :) == OffsetArray(0:1, -1:0)
@test reshape(r, (:,)) == OffsetArray(0:1, -1:0)

@test reshape(ones(2:3, 4:5), (2, :)) == ones(2,2)

# more than one colon is not allowed
@test_throws Exception reshape(ones(3:4, 4:5, 1:2), :, :, 2)
@test_throws Exception reshape(ones(3:4, 4:5, 1:2), :, 2, :)

A = OffsetArray(rand(4, 4), -1, -1);
B = reshape(A, (2, :))
@test axes(B, 1) == 1:2
@test axes(B, 2) == 1:8

# some more exotic vector types
r = OffsetVector(CustomRange(ZeroBasedRange(0:2)), -2)
r2 = reshape(r, :)
@test r2 == r
r2 = reshape(r, 3)
@test axes(r2, 1) == 1:3
@test r2 == no_offset_view(r)
@test_throws Exception reshape(r, length(r) + 1)
@test_throws Exception reshape(r, 1:length(r) + 1)
rp = parent(r)
@test axes(reshape(rp, 4:6), 1) == 4:6
@test axes(reshape(r, (3,1))) == (1:3, 1:1)
end

@testset "permutedims" begin
Expand Down