Skip to content

Commit

Permalink
use to_indices to simplify constructor dispatch
Browse files Browse the repository at this point in the history
to_indices converts colons to ranges. This is followed by converting CartesianIndices to ranges.
Fix indexing IdOffsetRange using offset ranges, ensuring that indexing is idempotent
  • Loading branch information
jishnub committed Sep 25, 2020
1 parent 2ed1598 commit 9f26eb1
Show file tree
Hide file tree
Showing 8 changed files with 320 additions and 39 deletions.
3 changes: 2 additions & 1 deletion Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,9 @@ Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
CatIndices = "aafaddc9-749c-510e-ac4f-586e18779b91"
DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab"
Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
EllipsisNotation = "da5c29d0-fa7d-589e-88eb-ea29b0a81949"
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

[targets]
test = ["Aqua", "CatIndices", "DelimitedFiles", "Documenter", "Test", "LinearAlgebra"]
test = ["Aqua", "CatIndices", "DelimitedFiles", "Documenter", "Test", "LinearAlgebra", "EllipsisNotation"]
2 changes: 1 addition & 1 deletion docs/src/index.md
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@ Base.require_one_based_indexing(OA)
```

[`OffsetArrays.Origin`](@ref) can be convenient if you want to directly specify the origin of the output
OffsetArray, it will automatically compute the needed offsets. For example:
OffsetArray, it will automatically compute the corresponding offsets. For example:

```@repl index
OffsetArray(A, OffsetArrays.Origin(-1, -1))
Expand Down
25 changes: 25 additions & 0 deletions docs/src/internals.md
Original file line number Diff line number Diff line change
Expand Up @@ -143,3 +143,28 @@ OffsetArrays.IdOffsetRange(-3:3)
julia> Ao[ax, 0][1] == Ao[ax[1], 0]
true
```

## Using custom axis types

While a wide variety of `AbstractUnitRange`s provided by `Base` may be used as indices to construct an `OffsetArray`, at times it might be convenient to define custom types. The `OffsetArray` constructor accepts any type that may be converted to an `AbstractUnitRange`. This proceeds through a two-step process. Let's assume that the constructor called is `OffsetArray(A, indices)`.

1. In the first step, the constructor calls `to_indices(A, axes(A), indices)` to lower `indices` to a `Tuple` of `AbstractUnitRange`s. This step converts --- among other things --- `Colon`s to axis ranges. Custom types may extend `Base.to_indices(A, axes(A), indices)` with the desired conversion of `indices` to `Tuple{Vararg{AbstractUnitRange{Int}}}` if this is feasible.

2. In the second step, the result of the previous step is passed to `OffsetArrays._toAbsUnitRanges`. This step is only necessary if the previous step didn't return a `Tuple` of `AbstractUnitRange`s. This step allows an additional customization option: a type may be converted either to a single `AbstractUnitRange{Int}`, or to a `Tuple` of them. A custom type might specify which of these two behaviours is desired by extending [`OffsetArrays.AxisConversionStyle`](@ref). An example of a type that is acted upon at this stage is `CartesianIndices`, which is converted to a `Tuple` of `AbstractUnitRange`s.

For example, here is a custom type that leads to zero-based indexing:

```jldoctest; setup = :(using OffsetArrays)
julia> struct ZeroBasedIndexing end
julia> Base.to_indices(A, inds, ::Tuple{ZeroBasedIndexing}) = map(x -> 0:length(x)-1, inds)
julia> a = zeros(3, 3);
julia> oa = OffsetArray(a, ZeroBasedIndexing());
julia> axes(oa)
(OffsetArrays.IdOffsetRange(0:2), OffsetArrays.IdOffsetRange(0:2))
```

Note that zero-based indexing may also be achieved using [`OffsetArrays.Origin`](@ref).
1 change: 1 addition & 0 deletions docs/src/reference.md
Original file line number Diff line number Diff line change
Expand Up @@ -7,4 +7,5 @@ OffsetMatrix
OffsetArrays.Origin
OffsetArrays.IdOffsetRange
OffsetArrays.no_offset_view
OffsetArrays.AxisConversionStyle
```
88 changes: 60 additions & 28 deletions src/OffsetArrays.jl
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ include("origin.jl")

# Technically we know the length of CartesianIndices but we need to convert it first, so here we
# don't put it in OffsetAxisKnownLength.
const OffsetAxisKnownLength = Union{Integer, AbstractUnitRange, IdOffsetRange}
const OffsetAxisKnownLength = Union{Integer, AbstractUnitRange}
const OffsetAxis = Union{OffsetAxisKnownLength, CartesianIndices, Colon}
const ArrayInitializer = Union{UndefInitializer, Missing, Nothing}

Expand Down Expand Up @@ -51,7 +51,7 @@ julia> OffsetArray(reshape(1:6, 2, 3), 0:1, -1:1)
1 3 5
2 4 6
julia> OffsetArray(reshape(1:6, 2, 3), :, -1:1) # : as a placeholder means no offset is applied at this dimension
julia> OffsetArray(reshape(1:6, 2, 3), :, -1:1) # : as a placeholder to indicate that no offset is to be applied to this dimension
2×3 OffsetArray(reshape(::UnitRange{$Int}, 2, 3), 1:2, -1:1) with eltype $Int with indices 1:2×-1:1:
1 3 5
2 4 6
Expand Down Expand Up @@ -123,54 +123,86 @@ function overflow_check(r, offset::T) where T
throw_lower_overflow_error()
end
end
## OffsetArray constructors

# The only route out to inner constructor
function OffsetArray(A::AbstractArray{T,N}, offsets::NTuple{N,Integer}) where {T,N}
OffsetArray{T, ndims(A), typeof(A)}(A, offsets)
end
# Nested OffsetArrays
function OffsetArray(A::OffsetArray{<:Any,N}, offsets::NTuple{N, Integer}) where {N}
OffsetArray(parent(A), A.offsets .+ offsets)
end

for (FT, ND) in ((:OffsetVector, :1), (:OffsetMatrix, :2))
@eval function $FT(A::AbstractArray{<:Any,$ND}, offsets::NTuple{$ND,Integer})
OffsetArray{eltype(A), $ND, typeof(A)}(A, offsets)
end
@eval function $FT(A::OffsetArray{<:Any,$ND}, offsets::NTuple{$ND, Integer})
$FT(parent(A), A.offsets .+ offsets)
end
FTstr = string(FT)
@eval function $FT(A::AbstractArray{<:Any,N}, offsets::NTuple{N,Integer}) where {N}
throw(ArgumentError($FTstr*" requires a $(string($ND))D array"))
end
# Ambiguity resolution
@eval function $FT(A::AbstractArray{<:Any,0}, offsets::NTuple{0,Integer})
throw(ArgumentError($FTstr*" requires a $(string($ND))D array"))
end
end

## OffsetArray constructors
for FT in (:OffsetArray, :OffsetVector, :OffsetMatrix)
# The only route out to inner constructor
@eval function $FT(A::AbstractArray{T, N}, offsets::NTuple{N, Integer}) where {T, N}
ndims(A) == N || throw(DimensionMismatch("The number of offsets $(N) should equal ndims(A) = $(ndims(A))"))
OffsetArray{T, ndims(A), typeof(A)}(A, offsets)
@eval function $FT(A::AbstractArray, offsets::Tuple{Vararg{Integer}})
throw(DimensionMismatch("offsets $offsets are not compatible with a $(ndims(A))D array"))
end
@eval function $FT(A::AbstractArray, inds::Tuple{})
throw(DimensionMismatch("indices $inds are not compatible with a $(ndims(A))D array"))
end
@eval function $FT(A::AbstractArray, inds::Indices)
throw(DimensionMismatch("indices $inds are not compatible with a $(ndims(A))D array"))
end

# In general, indices get converted to AbstractUnitRanges.
# CartesianIndices{N} get converted to N ranges
@eval function $FT(A::AbstractArray, inds::Tuple)
inds2 = to_indices(A, axes(A), inds)
OffsetArray(A, _toAbstractUnitRanges(inds2))
end
# nested OffsetArrays
@eval $FT(A::OffsetArray{T, N}, offsets::NTuple{N, Integer}) where {T,N} = $FT(parent(A), A.offsets .+ offsets)

@eval $FT(A::AbstractArray, inds::Vararg) = OffsetArray(A, inds)

# convert ranges to offsets
@eval function $FT(A::AbstractArray{T}, inds::NTuple{N,OffsetAxisKnownLength}) where {T,N}
@eval function $FT(A::AbstractArray{<:Any,N}, inds::Indices{N}) where {N}
axparent = axes(A)
lA = map(length, axparent)
lI = map(length, inds)
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"))
$FT(A, map(_offset, axparent, inds))
end
# lower CartesianIndices and Colon
@eval function $FT(A::AbstractArray{T}, inds::NTuple{N, OffsetAxis}) where {T, N}
indsN = _uncolonindices(A, _expandCartesianIndices(inds))
$FT(A, indsN)
end
@eval $FT(A::AbstractArray{T}, inds::Vararg{OffsetAxis,N}) where {T, N} = $FT(A, inds)

@eval $FT(A::AbstractArray, origin::Origin) = OffsetArray(A, origin(A))
end

# Ambiguity resolution
OffsetArray(A::AbstractArray{<:Any,0}, ::Tuple{}) = OffsetArray{eltype(A),0,typeof(A)}(A, ())
OffsetArray(A::OffsetArray{<:Any,0}, ::Tuple{}) = typeof(A)(A.parent, ())

# array initialization
function OffsetArray{T,N}(init::ArrayInitializer, inds::NTuple{N, OffsetAxisKnownLength}) where {T,N}
function OffsetArray{T,N}(init::ArrayInitializer, inds::Tuple{Vararg{OffsetAxisKnownLength}}) where {T,N}
length(inds) == N || throw(DimensionMismatch("indices $inds are not compatible with a $(N)D array"))
AA = Array{T,N}(init, map(_indexlength, inds))
OffsetArray{T, N, typeof(AA)}(AA, map(_indexoffset, inds))
end
function OffsetArray{T, N}(init::ArrayInitializer, inds::NTuple{NT, Union{OffsetAxisKnownLength, CartesianIndices}}) where {T, N, NT}
# NT is probably not the actual dimension of the array; CartesianIndices might contain multiple dimensions
indsN = _expandCartesianIndices(inds)
length(indsN) == N || throw(DimensionMismatch("The number of offsets $(length(indsN)) should equal ndims(A) = $N"))
OffsetArray{T, N}(init, indsN)
function OffsetArray{T, N}(init::ArrayInitializer, inds::Tuple) where {T, N}
OffsetArray{T, N}(init, _toAbstractUnitRanges(inds))
end
OffsetArray{T,N}(init::ArrayInitializer, inds::Union{OffsetAxisKnownLength, CartesianIndices}...) where {T,N} = OffsetArray{T,N}(init, inds)
OffsetArray{T,N}(init::ArrayInitializer, inds::Vararg) where {T,N} = OffsetArray{T,N}(init, inds)

OffsetArray{T}(init::ArrayInitializer, inds::NTuple{N, OffsetAxisKnownLength}) where {T,N} = OffsetArray{T,N}(init, inds)
function OffsetArray{T}(init::ArrayInitializer, inds::NTuple{N, Union{OffsetAxisKnownLength, CartesianIndices}}) where {T, N}
# N is probably not the actual dimension of the array; CartesianIndices might contain multiple dimensions
indsN = _expandCartesianIndices(inds)
OffsetArray{T, length(indsN)}(init, indsN)
function OffsetArray{T}(init::ArrayInitializer, inds::Tuple) where {T}
OffsetArray{T}(init, _toAbstractUnitRanges(inds))
end
OffsetArray{T}(init::ArrayInitializer, inds::Union{OffsetAxisKnownLength, CartesianIndices}...) where {T} = OffsetArray{T}(init, inds)
OffsetArray{T}(init::ArrayInitializer, inds::Vararg) where {T} = OffsetArray{T}(init, inds)

Base.IndexStyle(::Type{OA}) where {OA<:OffsetArray} = IndexStyle(parenttype(OA))
parenttype(::Type{OffsetArray{T,N,AA}}) where {T,N,AA} = AA
Expand Down
8 changes: 7 additions & 1 deletion src/axes.jl
Original file line number Diff line number Diff line change
Expand Up @@ -132,6 +132,9 @@ offset_coerce(::Type{I}, r::AbstractUnitRange) where I<:AbstractUnitRange{T} whe
Base.reduced_index(i::IdOffsetRange) = typeof(i)(first(i):first(i))
# Workaround for #92 on Julia < 1.4
Base.reduced_index(i::IdentityUnitRange{<:IdOffsetRange}) = typeof(i)(first(i):first(i))
for f in [:firstindex, :lastindex]
@eval Base.$f(r::IdOffsetRange) = $f(r.parent) .+ r.offset
end

@inline function Base.iterate(r::IdOffsetRange)
ret = iterate(r.parent)
Expand All @@ -151,9 +154,12 @@ end
@propagate_inbounds function Base.getindex(r::IdOffsetRange, s::AbstractUnitRange{<:Integer})
return r.parent[s .- r.offset] .+ r.offset
end
@propagate_inbounds function Base.getindex(r::IdOffsetRange, s::IdOffsetRange)
@propagate_inbounds function Base.getindex(r::IdOffsetRange, s::IdentityUnitRange)
return IdOffsetRange(r.parent[s .- r.offset], r.offset)
end
@propagate_inbounds function Base.getindex(r::IdOffsetRange, s::IdOffsetRange)
return IdOffsetRange(r.parent[s.parent .+ (s.offset - r.offset)] .+ (r.offset - s.offset), s.offset)
end

# offset-preserve broadcasting
Broadcast.broadcasted(::Base.Broadcast.DefaultArrayStyle{1}, ::typeof(-), r::IdOffsetRange{T}, x::Integer) where T =
Expand Down
61 changes: 53 additions & 8 deletions src/utils.jl
Original file line number Diff line number Diff line change
Expand Up @@ -8,14 +8,59 @@ _indexlength(i::Integer) = i
_indexlength(i::Colon) = Colon()

_offset(axparent::AbstractUnitRange, ax::AbstractUnitRange) = first(ax) - first(axparent)
_offset(axparent::AbstractUnitRange, ax::CartesianIndices) = _offset(axparent, first(ax.indices))
_offset(axparent::AbstractUnitRange, ax::Integer) = 1 - first(axparent)

_uncolonindices(A::AbstractArray{<:Any,N}, inds::NTuple{N,Any}) where {N} = _uncolonindices(axes(A), inds)
_uncolonindices(ax::Tuple, inds::Tuple) = (first(inds), _uncolonindices(tail(ax), tail(inds))...)
_uncolonindices(ax::Tuple, inds::Tuple{Colon, Vararg{Any}}) = (first(ax), _uncolonindices(tail(ax), tail(inds))...)
_uncolonindices(::Tuple{}, ::Tuple{}) = ()
"""
OffsetArrays.AxisConversionStyle(typeof(indices))
_expandCartesianIndices(inds::Tuple{<:CartesianIndices, Vararg{Any}}) = (convert(Tuple{Vararg{AbstractUnitRange{Int}}}, inds[1])..., _expandCartesianIndices(Base.tail(inds))...)
_expandCartesianIndices(inds::Tuple{Any,Vararg{Any}}) = (inds[1], _expandCartesianIndices(Base.tail(inds))...)
_expandCartesianIndices(::Tuple{}) = ()
`AxisConversionStyle` declares if `indices` should be converted to a single `AbstractUnitRange{Int}`
or to a `Tuple{Vararg{AbstractUnitRange{Int}}}` while flattening custom types into indices.
This method is called after `to_indices(A::Array, axes(A), indices)` to provide
further information in case `to_indices` does not return a `Tuple` of `AbstractUnitRange{Int}`.
Custom index types should extend `AxisConversionStyle` and return either `OffsetArray.SingleRange()`,
which is the default, or `OffsetArray.TupleOfRanges()`. In the former case, the type `T` should
define `Base.convert(::Type{AbstractUnitRange{Int}}, ::T)`, whereas in the latter it should define
`Base.convert(::Type{Tuple{Vararg{AbstractUnitRange{Int}}}}, ::T)`.
An example of the latter is `CartesianIndices`, which is converted to a `Tuple` of
`AbstractUnitRange{Int}` while flattening the indices.
# Example
```jldoctest; setup=:(using OffsetArrays)
julia> struct NTupleOfUnitRanges{N}
x ::NTuple{N, UnitRange{Int}}
end
julia> Base.to_indices(A, inds, t::Tuple{NTupleOfUnitRanges{N}}) where {N} = t;
julia> OffsetArrays.AxisConversionStyle(::Type{NTupleOfUnitRanges{N}}) where {N} = OffsetArrays.TupleOfRanges();
julia> Base.convert(::Type{Tuple{Vararg{AbstractUnitRange{Int}}}}, t::NTupleOfUnitRanges) = t.x;
julia> a = zeros(3, 3);
julia> inds = NTupleOfUnitRanges((3:5, 2:4));
julia> oa = OffsetArray(a, inds);
julia> axes(oa, 1) == 3:5
true
julia> axes(oa, 2) == 2:4
true
```
"""
abstract type AxisConversionStyle end
struct SingleRange <: AxisConversionStyle end
struct TupleOfRanges <: AxisConversionStyle end

AxisConversionStyle(::Type) = SingleRange()
AxisConversionStyle(::Type{<:CartesianIndices}) = TupleOfRanges()

_convertfirst(t::Tuple) = _convertfirst(AxisConversionStyle(typeof(first(t))), first(t))
_convertfirst(::SingleRange, x) = (convert(AbstractUnitRange{Int}, x),)
_convertfirst(::TupleOfRanges, x) = convert(Tuple{Vararg{AbstractUnitRange{Int}}}, x)

_toAbstractUnitRanges(t::Tuple) = (_convertfirst(t)..., _toAbstractUnitRanges(tail(t))...)
_toAbstractUnitRanges(::Tuple{}) = ()
Loading

0 comments on commit 9f26eb1

Please sign in to comment.