1517 lines
52 KiB
Julia
1517 lines
52 KiB
Julia
# This file is a part of Julia. License is MIT: https://julialang.org/license
|
||
|
||
### Multidimensional iterators
|
||
module IteratorsMD
|
||
import Base: eltype, length, size, start, done, next, last, in, getindex,
|
||
setindex!, IndexStyle, min, max, zero, one, isless, eachindex,
|
||
ndims, iteratorsize, convert
|
||
|
||
importall ..Base.Operators
|
||
import Base: simd_outer_range, simd_inner_length, simd_index
|
||
using Base: IndexLinear, IndexCartesian, AbstractCartesianIndex, fill_to_length, tail
|
||
|
||
export CartesianIndex, CartesianRange
|
||
|
||
"""
|
||
CartesianIndex(i, j, k...) -> I
|
||
CartesianIndex((i, j, k...)) -> I
|
||
|
||
Create a multidimensional index `I`, which can be used for
|
||
indexing a multidimensional array `A`. In particular, `A[I]` is
|
||
equivalent to `A[i,j,k...]`. One can freely mix integer and
|
||
`CartesianIndex` indices; for example, `A[Ipre, i, Ipost]` (where
|
||
`Ipre` and `Ipost` are `CartesianIndex` indices and `i` is an
|
||
`Int`) can be a useful expression when writing algorithms that
|
||
work along a single dimension of an array of arbitrary
|
||
dimensionality.
|
||
|
||
A `CartesianIndex` is sometimes produced by [`eachindex`](@ref), and
|
||
always when iterating with an explicit [`CartesianRange`](@ref).
|
||
|
||
# Example
|
||
|
||
```jldoctest
|
||
julia> A = reshape(collect(1:16), (2, 2, 2, 2))
|
||
2×2×2×2 Array{Int64,4}:
|
||
[:, :, 1, 1] =
|
||
1 3
|
||
2 4
|
||
|
||
[:, :, 2, 1] =
|
||
5 7
|
||
6 8
|
||
|
||
[:, :, 1, 2] =
|
||
9 11
|
||
10 12
|
||
|
||
[:, :, 2, 2] =
|
||
13 15
|
||
14 16
|
||
|
||
julia> A[CartesianIndex((1, 1, 1, 1))]
|
||
1
|
||
|
||
julia> A[CartesianIndex((1, 1, 1, 2))]
|
||
9
|
||
|
||
julia> A[CartesianIndex((1, 1, 2, 1))]
|
||
5
|
||
```
|
||
"""
|
||
struct CartesianIndex{N} <: AbstractCartesianIndex{N}
|
||
I::NTuple{N,Int}
|
||
CartesianIndex{N}(index::NTuple{N,Integer}) where {N} = new(index)
|
||
end
|
||
|
||
CartesianIndex(index::NTuple{N,Integer}) where {N} = CartesianIndex{N}(index)
|
||
CartesianIndex(index::Integer...) = CartesianIndex(index)
|
||
CartesianIndex{N}(index::Vararg{Integer,N}) where {N} = CartesianIndex{N}(index)
|
||
# Allow passing tuples smaller than N
|
||
CartesianIndex{N}(index::Tuple) where {N} = CartesianIndex{N}(fill_to_length(index, 1, Val{N}))
|
||
CartesianIndex{N}(index::Integer...) where {N} = CartesianIndex{N}(index)
|
||
CartesianIndex{N}() where {N} = CartesianIndex{N}(())
|
||
# Un-nest passed CartesianIndexes
|
||
CartesianIndex(index::Union{Integer, CartesianIndex}...) = CartesianIndex(flatten(index))
|
||
flatten(I::Tuple{}) = I
|
||
flatten(I::Tuple{Any}) = I
|
||
flatten(I::Tuple{<:CartesianIndex}) = I[1].I
|
||
@inline flatten(I) = _flatten(I...)
|
||
@inline _flatten() = ()
|
||
@inline _flatten(i, I...) = (i, _flatten(I...)...)
|
||
@inline _flatten(i::CartesianIndex, I...) = (i.I..., _flatten(I...)...)
|
||
CartesianIndex(index::Tuple{Vararg{Union{Integer, CartesianIndex}}}) = CartesianIndex(index...)
|
||
|
||
# length
|
||
length(::CartesianIndex{N}) where {N} = N
|
||
length(::Type{CartesianIndex{N}}) where {N} = N
|
||
|
||
# indexing
|
||
getindex(index::CartesianIndex, i::Integer) = index.I[i]
|
||
|
||
# zeros and ones
|
||
zero(::CartesianIndex{N}) where {N} = zero(CartesianIndex{N})
|
||
zero(::Type{CartesianIndex{N}}) where {N} = CartesianIndex(ntuple(x -> 0, Val{N}))
|
||
one(::CartesianIndex{N}) where {N} = one(CartesianIndex{N})
|
||
one(::Type{CartesianIndex{N}}) where {N} = CartesianIndex(ntuple(x -> 1, Val{N}))
|
||
|
||
# arithmetic, min/max
|
||
@inline (-)(index::CartesianIndex{N}) where {N} =
|
||
CartesianIndex{N}(map(-, index.I))
|
||
@inline (+)(index1::CartesianIndex{N}, index2::CartesianIndex{N}) where {N} =
|
||
CartesianIndex{N}(map(+, index1.I, index2.I))
|
||
@inline (-)(index1::CartesianIndex{N}, index2::CartesianIndex{N}) where {N} =
|
||
CartesianIndex{N}(map(-, index1.I, index2.I))
|
||
@inline min(index1::CartesianIndex{N}, index2::CartesianIndex{N}) where {N} =
|
||
CartesianIndex{N}(map(min, index1.I, index2.I))
|
||
@inline max(index1::CartesianIndex{N}, index2::CartesianIndex{N}) where {N} =
|
||
CartesianIndex{N}(map(max, index1.I, index2.I))
|
||
|
||
@inline (+)(i::Integer, index::CartesianIndex) = index+i
|
||
@inline (+)(index::CartesianIndex{N}, i::Integer) where {N} = CartesianIndex{N}(map(x->x+i, index.I))
|
||
@inline (-)(index::CartesianIndex{N}, i::Integer) where {N} = CartesianIndex{N}(map(x->x-i, index.I))
|
||
@inline (-)(i::Integer, index::CartesianIndex{N}) where {N} = CartesianIndex{N}(map(x->i-x, index.I))
|
||
@inline (*)(a::Integer, index::CartesianIndex{N}) where {N} = CartesianIndex{N}(map(x->a*x, index.I))
|
||
@inline (*)(index::CartesianIndex, a::Integer) = *(a,index)
|
||
|
||
# comparison
|
||
@inline isless(I1::CartesianIndex{N}, I2::CartesianIndex{N}) where {N} = _isless(0, I1.I, I2.I)
|
||
@inline function _isless(ret, I1::NTuple{N,Int}, I2::NTuple{N,Int}) where N
|
||
newret = ifelse(ret==0, icmp(I1[N], I2[N]), ret)
|
||
_isless(newret, Base.front(I1), Base.front(I2))
|
||
end
|
||
_isless(ret, ::Tuple{}, ::Tuple{}) = ifelse(ret==1, true, false)
|
||
icmp(a, b) = ifelse(isless(a,b), 1, ifelse(a==b, 0, -1))
|
||
|
||
# hashing
|
||
const cartindexhash_seed = UInt == UInt64 ? 0xd60ca92f8284b8b0 : 0xf2ea7c2e
|
||
function Base.hash(ci::CartesianIndex, h::UInt)
|
||
h += cartindexhash_seed
|
||
for i in ci.I
|
||
h = hash(i, h)
|
||
end
|
||
return h
|
||
end
|
||
|
||
# Iteration
|
||
"""
|
||
CartesianRange(Istart::CartesianIndex, Istop::CartesianIndex) -> R
|
||
CartesianRange(sz::Dims) -> R
|
||
CartesianRange(istart:istop, jstart:jstop, ...) -> R
|
||
|
||
Define a region `R` spanning a multidimensional rectangular range
|
||
of integer indices. These are most commonly encountered in the
|
||
context of iteration, where `for I in R ... end` will return
|
||
[`CartesianIndex`](@ref) indices `I` equivalent to the nested loops
|
||
|
||
for j = jstart:jstop
|
||
for i = istart:istop
|
||
...
|
||
end
|
||
end
|
||
|
||
Consequently these can be useful for writing algorithms that
|
||
work in arbitrary dimensions.
|
||
|
||
```jldoctest
|
||
julia> foreach(println, CartesianRange((2, 2, 2)))
|
||
CartesianIndex{3}((1, 1, 1))
|
||
CartesianIndex{3}((2, 1, 1))
|
||
CartesianIndex{3}((1, 2, 1))
|
||
CartesianIndex{3}((2, 2, 1))
|
||
CartesianIndex{3}((1, 1, 2))
|
||
CartesianIndex{3}((2, 1, 2))
|
||
CartesianIndex{3}((1, 2, 2))
|
||
CartesianIndex{3}((2, 2, 2))
|
||
```
|
||
"""
|
||
struct CartesianRange{I<:CartesianIndex}
|
||
start::I
|
||
stop::I
|
||
end
|
||
|
||
CartesianRange(index::CartesianIndex) = CartesianRange(one(index), index)
|
||
CartesianRange(::Tuple{}) = CartesianRange{CartesianIndex{0}}(CartesianIndex{0}(()),CartesianIndex{0}(()))
|
||
CartesianRange{N}(sz::NTuple{N,Int}) = CartesianRange(CartesianIndex(sz))
|
||
CartesianRange{N}(rngs::NTuple{N,Union{Integer,AbstractUnitRange}}) =
|
||
CartesianRange(CartesianIndex(map(first, rngs)), CartesianIndex(map(last, rngs)))
|
||
|
||
convert(::Type{NTuple{N,UnitRange{Int}}}, R::CartesianRange{CartesianIndex{N}}) where {N} =
|
||
map((f,l)->f:l, first(R).I, last(R).I)
|
||
convert(::Type{NTuple{N,UnitRange}}, R::CartesianRange) where {N} =
|
||
convert(NTuple{N,UnitRange{Int}}, R)
|
||
convert(::Type{Tuple{Vararg{UnitRange{Int}}}}, R::CartesianRange{CartesianIndex{N}}) where {N} =
|
||
convert(NTuple{N,UnitRange{Int}}, R)
|
||
convert(::Type{Tuple{Vararg{UnitRange}}}, R::CartesianRange) =
|
||
convert(Tuple{Vararg{UnitRange{Int}}}, R)
|
||
|
||
ndims(R::CartesianRange) = length(R.start)
|
||
ndims(::Type{CartesianRange{I}}) where {I<:CartesianIndex} = length(I)
|
||
|
||
eachindex(::IndexCartesian, A::AbstractArray) = CartesianRange(indices(A))
|
||
|
||
@inline eachindex(::IndexCartesian, A::AbstractArray, B::AbstractArray...) =
|
||
CartesianRange(maxsize((), A, B...))
|
||
maxsize(sz) = sz
|
||
@inline maxsize(sz, A, B...) = maxsize(maxt(sz, size(A)), B...)
|
||
@inline maxt(a::Tuple{}, b::Tuple{}) = ()
|
||
@inline maxt(a::Tuple{}, b::Tuple) = b
|
||
@inline maxt(a::Tuple, b::Tuple{}) = a
|
||
@inline maxt(a::Tuple, b::Tuple) = (max(a[1], b[1]), maxt(tail(a), tail(b))...)
|
||
|
||
eltype(::Type{CartesianRange{I}}) where {I} = I
|
||
iteratorsize(::Type{<:CartesianRange}) = Base.HasShape()
|
||
|
||
@inline function start(iter::CartesianRange{<:CartesianIndex})
|
||
if any(map(>, iter.start.I, iter.stop.I))
|
||
return iter.stop+1
|
||
end
|
||
iter.start
|
||
end
|
||
@inline function next(iter::CartesianRange{I}, state) where I<:CartesianIndex
|
||
state, I(inc(state.I, iter.start.I, iter.stop.I))
|
||
end
|
||
# increment & carry
|
||
@inline inc(::Tuple{}, ::Tuple{}, ::Tuple{}) = ()
|
||
@inline inc(state::Tuple{Int}, start::Tuple{Int}, stop::Tuple{Int}) = (state[1]+1,)
|
||
@inline function inc(state, start, stop)
|
||
if state[1] < stop[1]
|
||
return (state[1]+1,tail(state)...)
|
||
end
|
||
newtail = inc(tail(state), tail(start), tail(stop))
|
||
(start[1], newtail...)
|
||
end
|
||
@inline done(iter::CartesianRange{<:CartesianIndex}, state) = state.I[end] > iter.stop.I[end]
|
||
|
||
# 0-d cartesian ranges are special-cased to iterate once and only once
|
||
start(iter::CartesianRange{<:CartesianIndex{0}}) = false
|
||
next(iter::CartesianRange{<:CartesianIndex{0}}, state) = iter.start, true
|
||
done(iter::CartesianRange{<:CartesianIndex{0}}, state) = state
|
||
|
||
size(iter::CartesianRange{<:CartesianIndex}) = map(dimlength, iter.start.I, iter.stop.I)
|
||
dimlength(start, stop) = stop-start+1
|
||
|
||
length(iter::CartesianRange) = prod(size(iter))
|
||
|
||
last(iter::CartesianRange) = iter.stop
|
||
|
||
@inline function in(i::I, r::CartesianRange{I}) where I<:CartesianIndex
|
||
_in(true, i.I, r.start.I, r.stop.I)
|
||
end
|
||
_in(b, ::Tuple{}, ::Tuple{}, ::Tuple{}) = b
|
||
@inline _in(b, i, start, stop) = _in(b & (start[1] <= i[1] <= stop[1]), tail(i), tail(start), tail(stop))
|
||
|
||
simd_outer_range(iter::CartesianRange{CartesianIndex{0}}) = iter
|
||
function simd_outer_range(iter::CartesianRange)
|
||
start = CartesianIndex(tail(iter.start.I))
|
||
stop = CartesianIndex(tail(iter.stop.I))
|
||
CartesianRange(start, stop)
|
||
end
|
||
|
||
simd_inner_length(iter::CartesianRange{<:CartesianIndex{0}}, ::CartesianIndex) = 1
|
||
simd_inner_length(iter::CartesianRange, I::CartesianIndex) = iter.stop[1]-iter.start[1]+1
|
||
|
||
simd_index(iter::CartesianRange{<:CartesianIndex{0}}, ::CartesianIndex, I1::Int) = iter.start
|
||
@inline function simd_index(iter::CartesianRange, Ilast::CartesianIndex, I1::Int)
|
||
CartesianIndex((I1+iter.start[1], Ilast.I...))
|
||
end
|
||
|
||
# Split out the first N elements of a tuple
|
||
@inline split(t, V::Type{<:Val}) = _split((), t, V)
|
||
@inline _split(tN, trest, V) = _split((tN..., trest[1]), tail(trest), V)
|
||
# exit either when we've exhausted the input tuple or when tN has length N
|
||
@inline _split(tN::NTuple{N,Any}, ::Tuple{}, ::Type{Val{N}}) where {N} = tN, () # ambig.
|
||
@inline _split(tN, ::Tuple{}, ::Type{Val{N}}) where {N} = tN, ()
|
||
@inline _split(tN::NTuple{N,Any}, trest, ::Type{Val{N}}) where {N} = tN, trest
|
||
|
||
@inline function split(I::CartesianIndex, V::Type{<:Val})
|
||
i, j = split(I.I, V)
|
||
CartesianIndex(i), CartesianIndex(j)
|
||
end
|
||
function split(R::CartesianRange, V::Type{<:Val})
|
||
istart, jstart = split(first(R), V)
|
||
istop, jstop = split(last(R), V)
|
||
CartesianRange(istart, istop), CartesianRange(jstart, jstop)
|
||
end
|
||
end # IteratorsMD
|
||
|
||
|
||
using .IteratorsMD
|
||
|
||
## Bounds-checking with CartesianIndex
|
||
# Disallow linear indexing with CartesianIndex
|
||
function checkbounds(::Type{Bool}, A::AbstractArray, i::Union{CartesianIndex, AbstractArray{<:CartesianIndex}})
|
||
@_inline_meta
|
||
checkbounds_indices(Bool, indices(A), (i,))
|
||
end
|
||
|
||
@inline checkbounds_indices(::Type{Bool}, ::Tuple{}, I::Tuple{CartesianIndex,Vararg{Any}}) =
|
||
checkbounds_indices(Bool, (), (I[1].I..., tail(I)...))
|
||
@inline checkbounds_indices(::Type{Bool}, IA::Tuple{Any}, I::Tuple{CartesianIndex,Vararg{Any}}) =
|
||
checkbounds_indices(Bool, IA, (I[1].I..., tail(I)...))
|
||
@inline checkbounds_indices(::Type{Bool}, IA::Tuple, I::Tuple{CartesianIndex,Vararg{Any}}) =
|
||
checkbounds_indices(Bool, IA, (I[1].I..., tail(I)...))
|
||
|
||
# Indexing into Array with mixtures of Integers and CartesianIndices is
|
||
# extremely performance-sensitive. While the abstract fallbacks support this,
|
||
# codegen has extra support for SIMDification that sub2ind doesn't (yet) support
|
||
@propagate_inbounds getindex(A::Array, i1::Union{Integer, CartesianIndex}, I::Union{Integer, CartesianIndex}...) =
|
||
A[to_indices(A, (i1, I...))...]
|
||
@propagate_inbounds setindex!(A::Array, v, i1::Union{Integer, CartesianIndex}, I::Union{Integer, CartesianIndex}...) =
|
||
(A[to_indices(A, (i1, I...))...] = v; A)
|
||
|
||
# Support indexing with an array of CartesianIndex{N}s
|
||
# Here we try to consume N of the indices (if there are that many available)
|
||
# The first two simply handle ambiguities
|
||
@inline function checkbounds_indices(::Type{Bool}, ::Tuple{},
|
||
I::Tuple{AbstractArray{CartesianIndex{N}},Vararg{Any}}) where N
|
||
checkindex(Bool, (), I[1]) & checkbounds_indices(Bool, (), tail(I))
|
||
end
|
||
@inline function checkbounds_indices(::Type{Bool}, IA::Tuple{Any},
|
||
I::Tuple{AbstractArray{CartesianIndex{0}},Vararg{Any}})
|
||
checkbounds_indices(Bool, IA, tail(I))
|
||
end
|
||
@inline function checkbounds_indices(::Type{Bool}, IA::Tuple{Any},
|
||
I::Tuple{AbstractArray{CartesianIndex{N}},Vararg{Any}}) where N
|
||
checkindex(Bool, IA, I[1]) & checkbounds_indices(Bool, (), tail(I))
|
||
end
|
||
@inline function checkbounds_indices(::Type{Bool}, IA::Tuple,
|
||
I::Tuple{AbstractArray{CartesianIndex{N}},Vararg{Any}}) where N
|
||
IA1, IArest = IteratorsMD.split(IA, Val{N})
|
||
checkindex(Bool, IA1, I[1]) & checkbounds_indices(Bool, IArest, tail(I))
|
||
end
|
||
|
||
function checkindex(::Type{Bool}, inds::Tuple, I::AbstractArray{<:CartesianIndex})
|
||
b = true
|
||
for i in I
|
||
b &= checkbounds_indices(Bool, inds, (i,))
|
||
end
|
||
b
|
||
end
|
||
|
||
# combined count of all indices, including CartesianIndex and
|
||
# AbstractArray{CartesianIndex}
|
||
# rather than returning N, it returns an NTuple{N,Bool} so the result is inferrable
|
||
@inline index_ndims(i1, I...) = (true, index_ndims(I...)...)
|
||
@inline function index_ndims(i1::CartesianIndex, I...)
|
||
(map(x->true, i1.I)..., index_ndims(I...)...)
|
||
end
|
||
@inline function index_ndims(i1::AbstractArray{CartesianIndex{N}}, I...) where N
|
||
(ntuple(x->true, Val{N})..., index_ndims(I...)...)
|
||
end
|
||
index_ndims() = ()
|
||
|
||
# combined dimensionality of all indices
|
||
# rather than returning N, it returns an NTuple{N,Bool} so the result is inferrable
|
||
@inline index_dimsum(i1, I...) = (index_dimsum(I...)...)
|
||
@inline index_dimsum(::Colon, I...) = (true, index_dimsum(I...)...)
|
||
@inline index_dimsum(::AbstractArray{Bool}, I...) = (true, index_dimsum(I...)...)
|
||
@inline function index_dimsum(::AbstractArray{<:Any,N}, I...) where N
|
||
(ntuple(x->true, Val{N})..., index_dimsum(I...)...)
|
||
end
|
||
index_dimsum() = ()
|
||
|
||
# Recursively compute the lengths of a list of indices, without dropping scalars
|
||
index_lengths() = ()
|
||
@inline index_lengths(::Real, rest...) = (1, index_lengths(rest...)...)
|
||
@inline index_lengths(A::AbstractArray, rest...) = (length(A), index_lengths(rest...)...)
|
||
@inline index_lengths(A::Slice, rest...) = (length(indices1(A)), index_lengths(rest...)...)
|
||
|
||
# shape of array to create for getindex() with indexes I, dropping scalars
|
||
# returns a Tuple{Vararg{AbstractUnitRange}} of indices
|
||
index_shape() = ()
|
||
@inline index_shape(::Real, rest...) = index_shape(rest...)
|
||
@inline index_shape(A::AbstractArray, rest...) = (indices(A)..., index_shape(rest...)...)
|
||
|
||
"""
|
||
LogicalIndex(mask)
|
||
|
||
The `LogicalIndex` type is a special vector that simply contains all indices I
|
||
where `mask[I]` is true. This specialized type does not support indexing
|
||
directly as doing so would require O(n) lookup time. `AbstractArray{Bool}` are
|
||
wrapped with `LogicalIndex` upon calling `to_indices`.
|
||
"""
|
||
struct LogicalIndex{T, A<:AbstractArray{Bool}} <: AbstractVector{T}
|
||
mask::A
|
||
sum::Int
|
||
LogicalIndex{T,A}(mask::A) where {T,A<:AbstractArray{Bool}} = new(mask, countnz(mask))
|
||
end
|
||
LogicalIndex(mask::AbstractVector{Bool}) = LogicalIndex{Int, typeof(mask)}(mask)
|
||
LogicalIndex(mask::AbstractArray{Bool, N}) where {N} = LogicalIndex{CartesianIndex{N}, typeof(mask)}(mask)
|
||
(::Type{LogicalIndex{Int}})(mask::AbstractArray) = LogicalIndex{Int, typeof(mask)}(mask)
|
||
size(L::LogicalIndex) = (L.sum,)
|
||
length(L::LogicalIndex) = L.sum
|
||
collect(L::LogicalIndex) = [i for i in L]
|
||
show(io::IO, r::LogicalIndex) = print(io, "Base.LogicalIndex(", r.mask, ")")
|
||
# Iteration over LogicalIndex is very performance-critical, but it also must
|
||
# support arbitrary AbstractArray{Bool}s with both Int and CartesianIndex.
|
||
# Thus the iteration state contains an index iterator and its state. We also
|
||
# keep track of the count of elements since we already know how many there
|
||
# should be -- this way we don't need to look at future indices to check done.
|
||
@inline function start(L::LogicalIndex{Int})
|
||
r = linearindices(L.mask)
|
||
return (r, start(r), 1)
|
||
end
|
||
@inline function start(L::LogicalIndex{<:CartesianIndex})
|
||
r = CartesianRange(indices(L.mask))
|
||
return (r, start(r), 1)
|
||
end
|
||
@inline function next(L::LogicalIndex, s)
|
||
# We're looking for the n-th true element, using iterator r at state i
|
||
r, i, n = s
|
||
while true
|
||
done(r, i) # Call done(r, i) for the iteration protocol, but trust done(L, s) was called
|
||
idx, i = next(r, i)
|
||
L.mask[idx] && return (idx, (r, i, n+1))
|
||
end
|
||
end
|
||
done(L::LogicalIndex, s) = s[3] > length(L)
|
||
# When wrapping a BitArray, lean heavily upon its internals -- this is a common
|
||
# case. Just use the Int index and count as its state.
|
||
@inline start(L::LogicalIndex{Int,<:BitArray}) = (0, 1)
|
||
@inline function next(L::LogicalIndex{Int,<:BitArray}, s)
|
||
i, n = s
|
||
Bc = L.mask.chunks
|
||
while true
|
||
if Bc[_div64(i)+1] & (UInt64(1)<<_mod64(i)) != 0
|
||
i += 1
|
||
return (i, (i, n+1))
|
||
end
|
||
i += 1
|
||
end
|
||
end
|
||
@inline done(L::LogicalIndex{Int,<:BitArray}, s) = s[2] > length(L)
|
||
|
||
# Checking bounds with LogicalIndex{Int} is tricky since we allow linear indexing over trailing dimensions
|
||
@inline checkbounds_indices(::Type{Bool},IA::Tuple{},I::Tuple{LogicalIndex{Int,AbstractArray{Bool,N}}}) where {N} =
|
||
checkindex(Bool, IA, I[1])
|
||
@inline checkbounds_indices(::Type{Bool},IA::Tuple{Any},I::Tuple{LogicalIndex{Int,AbstractArray{Bool,N}}}) where {N} =
|
||
checkindex(Bool, IA[1], I[1])
|
||
@inline function checkbounds_indices(::Type{Bool}, IA::Tuple, I::Tuple{LogicalIndex{Int,AbstractArray{Bool,N}}}) where N
|
||
IA1, IArest = IteratorsMD.split(IA, Val{N})
|
||
checkindex(Bool, IA1, I[1])
|
||
end
|
||
@inline checkbounds(::Type{Bool}, A::AbstractArray, I::LogicalIndex{<:Any,<:AbstractArray{Bool,1}}) =
|
||
linearindices(A) == linearindices(I.mask)
|
||
@inline checkbounds(::Type{Bool}, A::AbstractArray, I::LogicalIndex) = indices(A) == indices(I.mask)
|
||
@inline checkindex(::Type{Bool}, indx::AbstractUnitRange, I::LogicalIndex) = (indx,) == indices(I.mask)
|
||
checkindex(::Type{Bool}, inds::Tuple, I::LogicalIndex) = false
|
||
|
||
ensure_indexable(I::Tuple{}) = ()
|
||
@inline ensure_indexable(I::Tuple{Any, Vararg{Any}}) = (I[1], ensure_indexable(tail(I))...)
|
||
@inline ensure_indexable(I::Tuple{LogicalIndex, Vararg{Any}}) = (collect(I[1]), ensure_indexable(tail(I))...)
|
||
|
||
# In simple cases, we know that we don't need to use indices(A). Optimize those
|
||
# until Julia gets smart enough to elide the call on its own:
|
||
to_indices(A, I::Tuple{}) = ()
|
||
@inline to_indices(A, I::Tuple{Vararg{Union{Integer, CartesianIndex}}}) = to_indices(A, (), I)
|
||
# But some index types require more context spanning multiple indices
|
||
# CartesianIndexes are simple; they just splat out
|
||
@inline to_indices(A, inds, I::Tuple{CartesianIndex, Vararg{Any}}) =
|
||
to_indices(A, inds, (I[1].I..., tail(I)...))
|
||
# But for arrays of CartesianIndex, we just skip the appropriate number of inds
|
||
@inline function to_indices(A, inds, I::Tuple{AbstractArray{CartesianIndex{N}}, Vararg{Any}}) where N
|
||
_, indstail = IteratorsMD.split(inds, Val{N})
|
||
(to_index(A, I[1]), to_indices(A, indstail, tail(I))...)
|
||
end
|
||
# And boolean arrays behave similarly; they also skip their number of dimensions
|
||
@inline function to_indices(A, inds, I::Tuple{AbstractArray{Bool, N}, Vararg{Any}}) where N
|
||
_, indstail = IteratorsMD.split(inds, Val{N})
|
||
(to_index(A, I[1]), to_indices(A, indstail, tail(I))...)
|
||
end
|
||
# As an optimization, we allow trailing Array{Bool} and BitArray to be linear over trailing dimensions
|
||
@inline to_indices(A, inds, I::Tuple{Union{Array{Bool,N}, BitArray{N}}}) where {N} =
|
||
(_maybe_linear_logical_index(IndexStyle(A), A, I[1]),)
|
||
_maybe_linear_logical_index(::IndexStyle, A, i) = to_index(A, i)
|
||
_maybe_linear_logical_index(::IndexLinear, A, i) = LogicalIndex{Int}(i)
|
||
|
||
# Colons get converted to slices by `uncolon`
|
||
@inline to_indices(A, inds, I::Tuple{Colon, Vararg{Any}}) =
|
||
(uncolon(inds, I), to_indices(A, _maybetail(inds), tail(I))...)
|
||
|
||
const CI0 = Union{CartesianIndex{0}, AbstractArray{CartesianIndex{0}}}
|
||
uncolon(inds::Tuple{}, I::Tuple{Colon}) = Slice(OneTo(1))
|
||
uncolon(inds::Tuple{}, I::Tuple{Colon, Vararg{Any}}) = Slice(OneTo(1))
|
||
uncolon(inds::Tuple{}, I::Tuple{Colon, Vararg{CI0}}) = Slice(OneTo(1))
|
||
uncolon(inds::Tuple{Any}, I::Tuple{Colon}) = Slice(inds[1])
|
||
uncolon(inds::Tuple{Any}, I::Tuple{Colon, Vararg{Any}}) = Slice(inds[1])
|
||
uncolon(inds::Tuple{Any}, I::Tuple{Colon, Vararg{CI0}}) = Slice(inds[1])
|
||
uncolon(inds::Tuple, I::Tuple{Colon, Vararg{Any}}) = Slice(inds[1])
|
||
uncolon(inds::Tuple, I::Tuple{Colon}) = Slice(OneTo(trailingsize(inds)))
|
||
uncolon(inds::Tuple, I::Tuple{Colon, Vararg{CI0}}) = Slice(OneTo(trailingsize(inds)))
|
||
|
||
### From abstractarray.jl: Internal multidimensional indexing definitions ###
|
||
getindex(x::Number, i::CartesianIndex{0}) = x
|
||
getindex(t::Tuple, i::CartesianIndex{1}) = getindex(t, i.I[1])
|
||
|
||
# These are not defined on directly on getindex to avoid
|
||
# ambiguities for AbstractArray subtypes. See the note in abstractarray.jl
|
||
|
||
@generated function _getindex(l::IndexStyle, A::AbstractArray, I::Union{Real, AbstractArray}...)
|
||
N = length(I)
|
||
quote
|
||
@_inline_meta
|
||
@boundscheck checkbounds(A, I...)
|
||
_unsafe_getindex(l, _maybe_reshape(l, A, I...), I...)
|
||
end
|
||
end
|
||
# But we can speed up IndexCartesian arrays by reshaping them to the appropriate dimensionality:
|
||
_maybe_reshape(::IndexLinear, A::AbstractArray, I...) = A
|
||
_maybe_reshape(::IndexCartesian, A::AbstractVector, I...) = A
|
||
@inline _maybe_reshape(::IndexCartesian, A::AbstractArray, I...) = __maybe_reshape(A, index_ndims(I...))
|
||
@inline __maybe_reshape(A::AbstractArray{T,N}, ::NTuple{N,Any}) where {T,N} = A
|
||
@inline __maybe_reshape(A::AbstractArray, ::NTuple{N,Any}) where {N} = reshape(A, Val{N})
|
||
|
||
@generated function _unsafe_getindex(::IndexStyle, A::AbstractArray, I::Union{Real, AbstractArray}...)
|
||
N = length(I)
|
||
quote
|
||
# This is specifically not inlined to prevent exessive allocations in type unstable code
|
||
@nexprs $N d->(I_d = I[d])
|
||
shape = @ncall $N index_shape I
|
||
dest = similar(A, shape)
|
||
map(unsafe_length, indices(dest)) == map(unsafe_length, shape) || throw_checksize_error(dest, shape)
|
||
@ncall $N _unsafe_getindex! dest A I
|
||
end
|
||
end
|
||
|
||
# Always index with the exactly indices provided.
|
||
@generated function _unsafe_getindex!(dest::AbstractArray, src::AbstractArray, I::Union{Real, AbstractArray}...)
|
||
N = length(I)
|
||
quote
|
||
$(Expr(:meta, :inline))
|
||
@nexprs $N d->(J_d = I[d])
|
||
D = eachindex(dest)
|
||
Ds = start(D)
|
||
@inbounds @nloops $N j d->J_d begin
|
||
d, Ds = next(D, Ds)
|
||
dest[d] = @ncall $N getindex src j
|
||
end
|
||
dest
|
||
end
|
||
end
|
||
|
||
@noinline throw_checksize_error(A, sz) = throw(DimensionMismatch("output array is the wrong size; expected $sz, got $(size(A))"))
|
||
|
||
## setindex! ##
|
||
@generated function _setindex!(l::IndexStyle, A::AbstractArray, x, I::Union{Real, AbstractArray}...)
|
||
N = length(I)
|
||
quote
|
||
@_inline_meta
|
||
@boundscheck checkbounds(A, I...)
|
||
_unsafe_setindex!(l, _maybe_reshape(l, A, I...), x, I...)
|
||
A
|
||
end
|
||
end
|
||
|
||
_iterable(v::AbstractArray) = v
|
||
_iterable(v) = Iterators.repeated(v)
|
||
@generated function _unsafe_setindex!(::IndexStyle, A::AbstractArray, x, I::Union{Real,AbstractArray}...)
|
||
N = length(I)
|
||
quote
|
||
X = _iterable(x)
|
||
@nexprs $N d->(I_d = I[d])
|
||
idxlens = @ncall $N index_lengths I
|
||
@ncall $N setindex_shape_check X (d->idxlens[d])
|
||
Xs = start(X)
|
||
@inbounds @nloops $N i d->I_d begin
|
||
v, Xs = next(X, Xs)
|
||
@ncall $N setindex! A v i
|
||
end
|
||
A
|
||
end
|
||
end
|
||
|
||
##
|
||
|
||
@generated function findn(A::AbstractArray{T,N}) where {T,N}
|
||
quote
|
||
nnzA = countnz(A)
|
||
@nexprs $N d->(I_d = Vector{Int}(nnzA))
|
||
k = 1
|
||
@nloops $N i A begin
|
||
@inbounds if (@nref $N A i) != zero(T)
|
||
@nexprs $N d->(I_d[k] = i_d)
|
||
k += 1
|
||
end
|
||
end
|
||
@ntuple $N I
|
||
end
|
||
end
|
||
|
||
|
||
# see discussion in #18364 ... we try not to widen type of the resulting array
|
||
# from cumsum or cumprod, but in some cases (+, Bool) we may not have a choice.
|
||
rcum_promote_type(op, ::Type{T}, ::Type{S}) where {T,S<:Number} = promote_op(op, T, S)
|
||
rcum_promote_type(op, ::Type{T}) where {T<:Number} = rcum_promote_type(op, T,T)
|
||
rcum_promote_type(op, ::Type{T}) where {T} = T
|
||
|
||
# handle sums of Vector{Bool} and similar. it would be nice to handle
|
||
# any AbstractArray here, but it's not clear how that would be possible
|
||
rcum_promote_type(op, ::Type{Array{T,N}}) where {T,N} = Array{rcum_promote_type(op,T), N}
|
||
|
||
# accumulate_pairwise slightly slower then accumulate, but more numerically
|
||
# stable in certain situations (e.g. sums).
|
||
# it does double the number of operations compared to accumulate,
|
||
# though for cheap operations like + this does not have much impact (20%)
|
||
function _accumulate_pairwise!(op::Op, c::AbstractVector{T}, v::AbstractVector, s, i1, n)::T where {T,Op}
|
||
@inbounds if n < 128
|
||
s_ = v[i1]
|
||
c[i1] = op(s, s_)
|
||
for i = i1+1:i1+n-1
|
||
s_ = op(s_, v[i])
|
||
c[i] = op(s, s_)
|
||
end
|
||
else
|
||
n2 = n >> 1
|
||
s_ = _accumulate_pairwise!(op, c, v, s, i1, n2)
|
||
s_ = op(s_, _accumulate_pairwise!(op, c, v, op(s, s_), i1+n2, n-n2))
|
||
end
|
||
return s_
|
||
end
|
||
|
||
function accumulate_pairwise!(op::Op, result::AbstractVector, v::AbstractVector) where Op
|
||
li = linearindices(v)
|
||
li != linearindices(result) && throw(DimensionMismatch("input and output array sizes and indices must match"))
|
||
n = length(li)
|
||
n == 0 && return result
|
||
i1 = first(li)
|
||
@inbounds result[i1] = v1 = v[i1]
|
||
n == 1 && return result
|
||
_accumulate_pairwise!(op, result, v, v1, i1+1, n-1)
|
||
return result
|
||
end
|
||
|
||
function accumulate_pairwise(op, v::AbstractVector{T}) where T
|
||
out = similar(v, rcum_promote_type(op, T))
|
||
return accumulate_pairwise!(op, out, v)
|
||
end
|
||
|
||
function cumsum!(out, v::AbstractVector, axis::Integer=1)
|
||
# for types prone to numerical stability issues, we want
|
||
# accumulate_pairwise.
|
||
axis == 1 ? accumulate_pairwise!(+, out, v) : copy!(out,v)
|
||
end
|
||
|
||
function cumsum!(out, v::AbstractVector{<:Integer}, axis::Integer=1)
|
||
axis == 1 ? accumulate!(+, out, v) : copy!(out,v)
|
||
end
|
||
|
||
"""
|
||
cumsum(A, dim=1)
|
||
|
||
Cumulative sum along a dimension `dim` (defaults to 1). See also [`cumsum!`](@ref)
|
||
to use a preallocated output array, both for performance and to control the precision of the
|
||
output (e.g. to avoid overflow).
|
||
|
||
```jldoctest
|
||
julia> a = [1 2 3; 4 5 6]
|
||
2×3 Array{Int64,2}:
|
||
1 2 3
|
||
4 5 6
|
||
|
||
julia> cumsum(a,1)
|
||
2×3 Array{Int64,2}:
|
||
1 2 3
|
||
5 7 9
|
||
|
||
julia> cumsum(a,2)
|
||
2×3 Array{Int64,2}:
|
||
1 3 6
|
||
4 9 15
|
||
```
|
||
"""
|
||
function cumsum(A::AbstractArray{T}, axis::Integer=1) where T
|
||
out = similar(A, rcum_promote_type(+, T))
|
||
cumsum!(out, A, axis)
|
||
end
|
||
|
||
"""
|
||
cumsum!(B, A, dim::Integer=1)
|
||
|
||
Cumulative sum of `A` along a dimension, storing the result in `B`. The dimension defaults
|
||
to 1. See also [`cumsum`](@ref).
|
||
"""
|
||
cumsum!(B, A, axis::Integer=1) = accumulate!(+, B, A, axis)
|
||
|
||
"""
|
||
cumprod(A, dim=1)
|
||
|
||
Cumulative product along a dimension `dim` (defaults to 1). See also
|
||
[`cumprod!`](@ref) to use a preallocated output array, both for performance and
|
||
to control the precision of the output (e.g. to avoid overflow).
|
||
|
||
```jldoctest
|
||
julia> a = [1 2 3; 4 5 6]
|
||
2×3 Array{Int64,2}:
|
||
1 2 3
|
||
4 5 6
|
||
|
||
julia> cumprod(a,1)
|
||
2×3 Array{Int64,2}:
|
||
1 2 3
|
||
4 10 18
|
||
|
||
julia> cumprod(a,2)
|
||
2×3 Array{Int64,2}:
|
||
1 2 6
|
||
4 20 120
|
||
```
|
||
"""
|
||
cumprod(A::AbstractArray, axis::Integer=1) = accumulate(*, A, axis)
|
||
|
||
"""
|
||
cumprod!(B, A, dim::Integer=1)
|
||
|
||
Cumulative product of `A` along a dimension, storing the result in `B`. The dimension defaults to 1.
|
||
See also [`cumprod`](@ref).
|
||
"""
|
||
cumprod!(B, A, axis::Integer=1) = accumulate!(*, B, A, axis)
|
||
|
||
"""
|
||
accumulate(op, A, dim=1)
|
||
|
||
Cumulative operation `op` along a dimension `dim` (defaults to 1). See also
|
||
[`accumulate!`](@ref) to use a preallocated output array, both for performance and
|
||
to control the precision of the output (e.g. to avoid overflow). For common operations
|
||
there are specialized variants of `accumulate`, see:
|
||
[`cumsum`](@ref), [`cumprod`](@ref)
|
||
|
||
```jldoctest
|
||
julia> accumulate(+, [1,2,3])
|
||
3-element Array{Int64,1}:
|
||
1
|
||
3
|
||
6
|
||
|
||
julia> accumulate(*, [1,2,3])
|
||
3-element Array{Int64,1}:
|
||
1
|
||
2
|
||
6
|
||
```
|
||
"""
|
||
function accumulate(op, A, axis::Integer=1)
|
||
out = similar(A, rcum_promote_type(op, eltype(A)))
|
||
accumulate!(op, out, A, axis)
|
||
end
|
||
|
||
|
||
"""
|
||
accumulate(op, v0, A)
|
||
|
||
Like `accumulate`, but using a starting element `v0`. The first entry of the result will be
|
||
`op(v0, first(A))`. For example:
|
||
|
||
```jldoctest
|
||
julia> accumulate(+, 100, [1,2,3])
|
||
3-element Array{Int64,1}:
|
||
101
|
||
103
|
||
106
|
||
|
||
julia> accumulate(min, 0, [1,2,-1])
|
||
3-element Array{Int64,1}:
|
||
0
|
||
0
|
||
-1
|
||
```
|
||
"""
|
||
function accumulate(op, v0, A, axis::Integer=1)
|
||
T = rcum_promote_type(op, typeof(v0), eltype(A))
|
||
out = similar(A, T)
|
||
accumulate!(op, out, v0, A, 1)
|
||
end
|
||
|
||
function accumulate!(op::Op, B, A::AbstractVector, axis::Integer=1) where Op
|
||
isempty(A) && return B
|
||
v1 = first(A)
|
||
_accumulate1!(op, B, v1, A, axis)
|
||
end
|
||
|
||
function accumulate!(op, B, v0, A::AbstractVector, axis::Integer=1)
|
||
isempty(A) && return B
|
||
v1 = op(v0, first(A))
|
||
_accumulate1!(op, B, v1, A, axis)
|
||
end
|
||
|
||
|
||
function _accumulate1!(op, B, v1, A::AbstractVector, axis::Integer=1)
|
||
axis > 0 || throw(ArgumentError("axis must be a positive integer"))
|
||
inds = linearindices(A)
|
||
inds == linearindices(B) || throw(DimensionMismatch("linearindices of A and B don't match"))
|
||
axis > 1 && return copy!(B, A)
|
||
i1 = inds[1]
|
||
cur_val = v1
|
||
B[i1] = cur_val
|
||
@inbounds for i in inds[2:end]
|
||
cur_val = op(cur_val, A[i])
|
||
B[i] = cur_val
|
||
end
|
||
return B
|
||
end
|
||
|
||
"""
|
||
accumulate!(op, B, A, dim=1)
|
||
|
||
Cumulative operation `op` on `A` along a dimension, storing the result in `B`.
|
||
The dimension defaults to 1. See also [`accumulate`](@ref).
|
||
"""
|
||
function accumulate!(op, B, A, axis::Integer=1)
|
||
axis > 0 || throw(ArgumentError("axis must be a positive integer"))
|
||
inds_t = indices(A)
|
||
indices(B) == inds_t || throw(DimensionMismatch("shape of B must match A"))
|
||
axis > ndims(A) && return copy!(B, A)
|
||
isempty(inds_t[axis]) && return B
|
||
if axis == 1
|
||
# We can accumulate to a temporary variable, which allows
|
||
# register usage and will be slightly faster
|
||
ind1 = inds_t[1]
|
||
@inbounds for I in CartesianRange(tail(inds_t))
|
||
tmp = convert(eltype(B), A[first(ind1), I])
|
||
B[first(ind1), I] = tmp
|
||
for i_1 = first(ind1)+1:last(ind1)
|
||
tmp = op(tmp, A[i_1, I])
|
||
B[i_1, I] = tmp
|
||
end
|
||
end
|
||
else
|
||
R1 = CartesianRange(indices(A)[1:axis-1]) # not type-stable
|
||
R2 = CartesianRange(indices(A)[axis+1:end])
|
||
_accumulate!(op, B, A, R1, inds_t[axis], R2) # use function barrier
|
||
end
|
||
return B
|
||
end
|
||
|
||
@noinline function _accumulate!(op, B, A, R1, ind, R2)
|
||
# Copy the initial element in each 1d vector along dimension `axis`
|
||
i = first(ind)
|
||
@inbounds for J in R2, I in R1
|
||
B[I, i, J] = A[I, i, J]
|
||
end
|
||
# Accumulate
|
||
@inbounds for J in R2, i in first(ind)+1:last(ind), I in R1
|
||
B[I, i, J] = op(B[I, i-1, J], A[I, i, J])
|
||
end
|
||
B
|
||
end
|
||
|
||
### from abstractarray.jl
|
||
|
||
function fill!(A::AbstractArray{T}, x) where T
|
||
xT = convert(T, x)
|
||
for I in eachindex(A)
|
||
@inbounds A[I] = xT
|
||
end
|
||
A
|
||
end
|
||
|
||
"""
|
||
copy!(dest, src) -> dest
|
||
|
||
Copy all elements from collection `src` to array `dest`.
|
||
"""
|
||
copy!(dest, src)
|
||
|
||
function copy!(dest::AbstractArray{T,N}, src::AbstractArray{T,N}) where {T,N}
|
||
@boundscheck checkbounds(dest, indices(src)...)
|
||
for I in eachindex(IndexStyle(src,dest), src)
|
||
@inbounds dest[I] = src[I]
|
||
end
|
||
dest
|
||
end
|
||
|
||
@generated function copy!{T1,T2,N}(dest::AbstractArray{T1,N},
|
||
Rdest::CartesianRange{CartesianIndex{N}},
|
||
src::AbstractArray{T2,N},
|
||
Rsrc::CartesianRange{CartesianIndex{N}})
|
||
quote
|
||
isempty(Rdest) && return dest
|
||
if size(Rdest) != size(Rsrc)
|
||
throw(ArgumentError("source and destination must have same size (got $(size(Rsrc)) and $(size(Rdest)))"))
|
||
end
|
||
@boundscheck checkbounds(dest, Rdest.start)
|
||
@boundscheck checkbounds(dest, Rdest.stop)
|
||
@boundscheck checkbounds(src, Rsrc.start)
|
||
@boundscheck checkbounds(src, Rsrc.stop)
|
||
ΔI = Rdest.start - Rsrc.start
|
||
# TODO: restore when #9080 is fixed
|
||
# for I in Rsrc
|
||
# @inbounds dest[I+ΔI] = src[I]
|
||
@nloops $N i (n->Rsrc.start[n]:Rsrc.stop[n]) begin
|
||
@inbounds @nref($N,dest,n->i_n+ΔI[n]) = @nref($N,src,i)
|
||
end
|
||
dest
|
||
end
|
||
end
|
||
|
||
"""
|
||
copy!(dest, Rdest::CartesianRange, src, Rsrc::CartesianRange) -> dest
|
||
|
||
Copy the block of `src` in the range of `Rsrc` to the block of `dest`
|
||
in the range of `Rdest`. The sizes of the two regions must match.
|
||
"""
|
||
copy!(::AbstractArray, ::CartesianRange, ::AbstractArray, ::CartesianRange)
|
||
|
||
# circshift!
|
||
circshift!(dest::AbstractArray, src, ::Tuple{}) = copy!(dest, src)
|
||
"""
|
||
circshift!(dest, src, shifts)
|
||
|
||
Circularly shift the data in `src`, storing the result in
|
||
`dest`. `shifts` specifies the amount to shift in each dimension.
|
||
|
||
The `dest` array must be distinct from the `src` array (they cannot
|
||
alias each other).
|
||
|
||
See also [`circshift`](@ref).
|
||
"""
|
||
@noinline function circshift!(dest::AbstractArray{T,N}, src, shiftamt::DimsInteger) where {T,N}
|
||
dest === src && throw(ArgumentError("dest and src must be separate arrays"))
|
||
inds = indices(src)
|
||
indices(dest) == inds || throw(ArgumentError("indices of src and dest must match (got $inds and $(indices(dest)))"))
|
||
_circshift!(dest, (), src, (), inds, fill_to_length(shiftamt, 0, Val{N}))
|
||
end
|
||
circshift!(dest::AbstractArray, src, shiftamt) = circshift!(dest, src, (shiftamt...,))
|
||
|
||
# For each dimension, we copy the first half of src to the second half
|
||
# of dest, and the second half of src to the first half of dest. This
|
||
# uses a recursive bifurcation strategy so that these splits can be
|
||
# encoded by ranges, which means that we need only one call to `mod`
|
||
# per dimension rather than one call per index.
|
||
# `rdest` and `rsrc` are tuples-of-ranges that grow one dimension at a
|
||
# time; when all the dimensions have been filled in, you call `copy!`
|
||
# for that block. In other words, in two dimensions schematically we
|
||
# have the following call sequence (--> means a call):
|
||
# circshift!(dest, src, shiftamt) -->
|
||
# _circshift!(dest, src, ("first half of dim1",)) -->
|
||
# _circshift!(dest, src, ("first half of dim1", "first half of dim2")) --> copy!
|
||
# _circshift!(dest, src, ("first half of dim1", "second half of dim2")) --> copy!
|
||
# _circshift!(dest, src, ("second half of dim1",)) -->
|
||
# _circshift!(dest, src, ("second half of dim1", "first half of dim2")) --> copy!
|
||
# _circshift!(dest, src, ("second half of dim1", "second half of dim2")) --> copy!
|
||
@inline function _circshift!(dest, rdest, src, rsrc,
|
||
inds::Tuple{AbstractUnitRange,Vararg{Any}},
|
||
shiftamt::Tuple{Integer,Vararg{Any}})
|
||
ind1, d = inds[1], shiftamt[1]
|
||
s = mod(d, length(ind1))
|
||
sf, sl = first(ind1)+s, last(ind1)-s
|
||
r1, r2 = first(ind1):sf-1, sf:last(ind1)
|
||
r3, r4 = first(ind1):sl, sl+1:last(ind1)
|
||
tinds, tshiftamt = tail(inds), tail(shiftamt)
|
||
_circshift!(dest, (rdest..., r1), src, (rsrc..., r4), tinds, tshiftamt)
|
||
_circshift!(dest, (rdest..., r2), src, (rsrc..., r3), tinds, tshiftamt)
|
||
end
|
||
# At least one of inds, shiftamt is empty
|
||
function _circshift!(dest, rdest, src, rsrc, inds, shiftamt)
|
||
copy!(dest, CartesianRange(rdest), src, CartesianRange(rsrc))
|
||
end
|
||
|
||
# circcopy!
|
||
"""
|
||
circcopy!(dest, src)
|
||
|
||
Copy `src` to `dest`, indexing each dimension modulo its length.
|
||
`src` and `dest` must have the same size, but can be offset in
|
||
their indices; any offset results in a (circular) wraparound. If the
|
||
arrays have overlapping indices, then on the domain of the overlap
|
||
`dest` agrees with `src`.
|
||
|
||
# Example
|
||
```julia-repl
|
||
julia> src = reshape(collect(1:16), (4,4))
|
||
4×4 Array{Int64,2}:
|
||
1 5 9 13
|
||
2 6 10 14
|
||
3 7 11 15
|
||
4 8 12 16
|
||
|
||
julia> dest = OffsetArray{Int}((0:3,2:5))
|
||
|
||
julia> circcopy!(dest, src)
|
||
OffsetArrays.OffsetArray{Int64,2,Array{Int64,2}} with indices 0:3×2:5:
|
||
8 12 16 4
|
||
5 9 13 1
|
||
6 10 14 2
|
||
7 11 15 3
|
||
|
||
julia> dest[1:3,2:4] == src[1:3,2:4]
|
||
true
|
||
```
|
||
"""
|
||
function circcopy!(dest, src)
|
||
dest === src && throw(ArgumentError("dest and src must be separate arrays"))
|
||
indssrc, indsdest = indices(src), indices(dest)
|
||
if (szsrc = map(length, indssrc)) != (szdest = map(length, indsdest))
|
||
throw(DimensionMismatch("src and dest must have the same sizes (got $szsrc and $szdest)"))
|
||
end
|
||
shift = map((isrc, idest)->first(isrc)-first(idest), indssrc, indsdest)
|
||
all(x->x==0, shift) && return copy!(dest, src)
|
||
_circcopy!(dest, (), indsdest, src, (), indssrc)
|
||
end
|
||
|
||
# This uses the same strategy described above for _circshift!
|
||
@inline function _circcopy!(dest, rdest, indsdest::Tuple{AbstractUnitRange,Vararg{Any}},
|
||
src, rsrc, indssrc::Tuple{AbstractUnitRange,Vararg{Any}})
|
||
indd1, inds1 = indsdest[1], indssrc[1]
|
||
l = length(indd1)
|
||
s = mod(first(inds1)-first(indd1), l)
|
||
sdf = first(indd1)+s
|
||
rd1, rd2 = first(indd1):sdf-1, sdf:last(indd1)
|
||
ssf = last(inds1)-s
|
||
rs1, rs2 = first(inds1):ssf, ssf+1:last(inds1)
|
||
tindsd, tindss = tail(indsdest), tail(indssrc)
|
||
_circcopy!(dest, (rdest..., rd1), tindsd, src, (rsrc..., rs2), tindss)
|
||
_circcopy!(dest, (rdest..., rd2), tindsd, src, (rsrc..., rs1), tindss)
|
||
end
|
||
|
||
# At least one of indsdest, indssrc are empty (and both should be, since we've checked)
|
||
function _circcopy!(dest, rdest, indsdest, src, rsrc, indssrc)
|
||
copy!(dest, CartesianRange(rdest), src, CartesianRange(rsrc))
|
||
end
|
||
|
||
### BitArrays
|
||
|
||
## getindex
|
||
|
||
# contiguous multidimensional indexing: if the first dimension is a range,
|
||
# we can get some performance from using copy_chunks!
|
||
@inline function _unsafe_getindex!(X::BitArray, B::BitArray, I0::Union{UnitRange{Int},Slice})
|
||
copy_chunks!(X.chunks, 1, B.chunks, indexoffset(I0)+1, length(I0))
|
||
return X
|
||
end
|
||
|
||
# Optimization where the inner dimension is contiguous improves perf dramatically
|
||
@generated function _unsafe_getindex!(X::BitArray, B::BitArray,
|
||
I0::Union{Slice,UnitRange{Int}}, I::Union{Int,UnitRange{Int},Slice}...)
|
||
N = length(I)
|
||
quote
|
||
$(Expr(:meta, :inline))
|
||
@nexprs $N d->(I_d = I[d])
|
||
|
||
idxlens = @ncall $N index_lengths I0 I
|
||
|
||
f0 = indexoffset(I0)+1
|
||
l0 = idxlens[1]
|
||
|
||
gap_lst_1 = 0
|
||
@nexprs $N d->(gap_lst_{d+1} = idxlens[d+1])
|
||
stride = 1
|
||
ind = f0
|
||
@nexprs $N d->begin
|
||
stride *= size(B, d)
|
||
stride_lst_d = stride
|
||
ind += stride * indexoffset(I_d)
|
||
gap_lst_{d+1} *= stride
|
||
end
|
||
|
||
storeind = 1
|
||
Xc, Bc = X.chunks, B.chunks
|
||
@nloops($N, i, d->(1:idxlens[d+1]),
|
||
d->nothing, # PRE
|
||
d->(ind += stride_lst_d - gap_lst_d), # POST
|
||
begin # BODY
|
||
copy_chunks!(Xc, storeind, Bc, ind, l0)
|
||
storeind += l0
|
||
end)
|
||
return X
|
||
end
|
||
end
|
||
|
||
# in the general multidimensional non-scalar case, can we do about 10% better
|
||
# in most cases by manually hoisting the bitarray chunks access out of the loop
|
||
# (This should really be handled by the compiler or with an immutable BitArray)
|
||
@generated function _unsafe_getindex!(X::BitArray, B::BitArray, I::Union{Int,AbstractArray{Int}}...)
|
||
N = length(I)
|
||
quote
|
||
$(Expr(:meta, :inline))
|
||
stride_1 = 1
|
||
@nexprs $N d->(stride_{d+1} = stride_d*size(B, d))
|
||
$(Symbol(:offset_, N)) = 1
|
||
ind = 0
|
||
Xc, Bc = X.chunks, B.chunks
|
||
@nloops $N i d->I[d] d->(@inbounds offset_{d-1} = offset_d + (i_d-1)*stride_d) begin
|
||
ind += 1
|
||
unsafe_bitsetindex!(Xc, unsafe_bitgetindex(Bc, offset_0), ind)
|
||
end
|
||
return X
|
||
end
|
||
end
|
||
|
||
## setindex!
|
||
|
||
function copy_to_bitarray_chunks!(Bc::Vector{UInt64}, pos_d::Int, C::StridedArray, pos_s::Int, numbits::Int)
|
||
bind = pos_d
|
||
cind = pos_s
|
||
lastind = pos_d + numbits - 1
|
||
@inbounds while bind ≤ lastind
|
||
unsafe_bitsetindex!(Bc, Bool(C[cind]), bind)
|
||
bind += 1
|
||
cind += 1
|
||
end
|
||
end
|
||
|
||
# Note: the next two functions rely on the following definition of the conversion to Bool:
|
||
# convert(::Type{Bool}, x::Real) = x==0 ? false : x==1 ? true : throw(InexactError())
|
||
# they're used to pre-emptively check in bulk when possible, which is much faster.
|
||
# Also, the functions can be overloaded for custom types T<:Real :
|
||
# a) in the unlikely eventuality that they use a different logic for Bool conversion
|
||
# b) to skip the check if not necessary
|
||
@inline try_bool_conversion(x::Real) = x == 0 || x == 1 || throw(InexactError())
|
||
@inline unchecked_bool_convert(x::Real) = x == 1
|
||
|
||
function copy_to_bitarray_chunks!(Bc::Vector{UInt64}, pos_d::Int, C::StridedArray{<:Real}, pos_s::Int, numbits::Int)
|
||
@inbounds for i = (1:numbits) + pos_s - 1
|
||
try_bool_conversion(C[i])
|
||
end
|
||
|
||
kd0, ld0 = get_chunks_id(pos_d)
|
||
kd1, ld1 = get_chunks_id(pos_d + numbits - 1)
|
||
|
||
delta_kd = kd1 - kd0
|
||
|
||
u = _msk64
|
||
if delta_kd == 0
|
||
msk_d0 = msk_d1 = ~(u << ld0) | (u << (ld1+1))
|
||
lt0 = ld1
|
||
else
|
||
msk_d0 = ~(u << ld0)
|
||
msk_d1 = (u << (ld1+1))
|
||
lt0 = 63
|
||
end
|
||
|
||
bind = kd0
|
||
ind = pos_s
|
||
@inbounds if ld0 > 0
|
||
c = UInt64(0)
|
||
for j = ld0:lt0
|
||
c |= (UInt64(unchecked_bool_convert(C[ind])) << j)
|
||
ind += 1
|
||
end
|
||
Bc[kd0] = (Bc[kd0] & msk_d0) | (c & ~msk_d0)
|
||
bind += 1
|
||
end
|
||
|
||
nc = _div64(numbits - ind + pos_s)
|
||
@inbounds for i = 1:nc
|
||
c = UInt64(0)
|
||
for j = 0:63
|
||
c |= (UInt64(unchecked_bool_convert(C[ind])) << j)
|
||
ind += 1
|
||
end
|
||
Bc[bind] = c
|
||
bind += 1
|
||
end
|
||
|
||
@inbounds if bind ≤ kd1
|
||
@assert bind == kd1
|
||
c = UInt64(0)
|
||
for j = 0:ld1
|
||
c |= (UInt64(unchecked_bool_convert(C[ind])) << j)
|
||
ind += 1
|
||
end
|
||
Bc[kd1] = (Bc[kd1] & msk_d1) | (c & ~msk_d1)
|
||
end
|
||
end
|
||
|
||
# contiguous multidimensional indexing: if the first dimension is a range,
|
||
# we can get some performance from using copy_chunks!
|
||
|
||
@inline function setindex!(B::BitArray, X::StridedArray, J0::Union{Colon,UnitRange{Int}})
|
||
I0 = to_indices(B, (J0,))[1]
|
||
@boundscheck checkbounds(B, I0)
|
||
l0 = length(I0)
|
||
setindex_shape_check(X, l0)
|
||
l0 == 0 && return B
|
||
f0 = indexoffset(I0)+1
|
||
copy_to_bitarray_chunks!(B.chunks, f0, X, 1, l0)
|
||
return B
|
||
end
|
||
|
||
@inline function setindex!(B::BitArray, X::StridedArray,
|
||
I0::Union{Colon,UnitRange{Int}}, I::Union{Int,UnitRange{Int},Colon}...)
|
||
J = to_indices(B, (I0, I...))
|
||
@boundscheck checkbounds(B, J...)
|
||
_unsafe_setindex!(B, X, J...)
|
||
end
|
||
@generated function _unsafe_setindex!(B::BitArray, X::StridedArray,
|
||
I0::Union{Slice,UnitRange{Int}}, I::Union{Int,UnitRange{Int},Slice}...)
|
||
N = length(I)
|
||
quote
|
||
idxlens = @ncall $N index_lengths I0 d->I[d]
|
||
@ncall $N setindex_shape_check X idxlens[1] d->idxlens[d+1]
|
||
isempty(X) && return B
|
||
f0 = indexoffset(I0)+1
|
||
l0 = idxlens[1]
|
||
|
||
gap_lst_1 = 0
|
||
@nexprs $N d->(gap_lst_{d+1} = idxlens[d+1])
|
||
stride = 1
|
||
ind = f0
|
||
@nexprs $N d->begin
|
||
stride *= size(B, d)
|
||
stride_lst_d = stride
|
||
ind += stride * indexoffset(I[d])
|
||
gap_lst_{d+1} *= stride
|
||
end
|
||
|
||
refind = 1
|
||
Bc = B.chunks
|
||
@nloops($N, i, d->I[d],
|
||
d->nothing, # PRE
|
||
d->(ind += stride_lst_d - gap_lst_d), # POST
|
||
begin # BODY
|
||
copy_to_bitarray_chunks!(Bc, ind, X, refind, l0)
|
||
refind += l0
|
||
end)
|
||
|
||
return B
|
||
end
|
||
end
|
||
|
||
@inline function setindex!(B::BitArray, x,
|
||
I0::Union{Colon,UnitRange{Int}}, I::Union{Int,UnitRange{Int},Colon}...)
|
||
J = to_indices(B, (I0, I...))
|
||
@boundscheck checkbounds(B, J...)
|
||
_unsafe_setindex!(B, x, J...)
|
||
end
|
||
@generated function _unsafe_setindex!(B::BitArray, x,
|
||
I0::Union{Slice,UnitRange{Int}}, I::Union{Int,UnitRange{Int},Slice}...)
|
||
N = length(I)
|
||
quote
|
||
y = Bool(x)
|
||
idxlens = @ncall $N index_lengths I0 d->I[d]
|
||
|
||
f0 = indexoffset(I0)+1
|
||
l0 = idxlens[1]
|
||
l0 == 0 && return B
|
||
@nexprs $N d->(isempty(I[d]) && return B)
|
||
|
||
gap_lst_1 = 0
|
||
@nexprs $N d->(gap_lst_{d+1} = idxlens[d+1])
|
||
stride = 1
|
||
ind = f0
|
||
@nexprs $N d->begin
|
||
stride *= size(B, d)
|
||
stride_lst_d = stride
|
||
ind += stride * indexoffset(I[d])
|
||
gap_lst_{d+1} *= stride
|
||
end
|
||
|
||
@nloops($N, i, d->I[d],
|
||
d->nothing, # PRE
|
||
d->(ind += stride_lst_d - gap_lst_d), # POST
|
||
fill_chunks!(B.chunks, y, ind, l0) # BODY
|
||
)
|
||
|
||
return B
|
||
end
|
||
end
|
||
|
||
## findn
|
||
|
||
@generated function findn(B::BitArray{N}) where N
|
||
quote
|
||
nnzB = countnz(B)
|
||
I = ntuple(x->Vector{Int}(nnzB), Val{$N})
|
||
if nnzB > 0
|
||
count = 1
|
||
@nloops $N i B begin
|
||
if (@nref $N B i) # TODO: should avoid bounds checking
|
||
@nexprs $N d->(I[d][count] = i_d)
|
||
count += 1
|
||
end
|
||
end
|
||
end
|
||
return I
|
||
end
|
||
end
|
||
|
||
## isassigned
|
||
|
||
@generated function isassigned(B::BitArray, I_0::Int, I::Int...)
|
||
N = length(I)
|
||
quote
|
||
@nexprs $N d->(I_d = I[d])
|
||
stride = 1
|
||
index = I_0
|
||
@nexprs $N d->begin
|
||
l = size(B,d)
|
||
stride *= l
|
||
1 <= I_{d-1} <= l || return false
|
||
index += (I_d - 1) * stride
|
||
end
|
||
return isassigned(B, index)
|
||
end
|
||
end
|
||
|
||
## permutedims
|
||
|
||
## Permute array dims ##
|
||
|
||
function permutedims(B::StridedArray, perm)
|
||
dimsB = size(B)
|
||
ndimsB = length(dimsB)
|
||
(ndimsB == length(perm) && isperm(perm)) || throw(ArgumentError("no valid permutation of dimensions"))
|
||
dimsP = ntuple(i->dimsB[perm[i]], ndimsB)::typeof(dimsB)
|
||
P = similar(B, dimsP)
|
||
permutedims!(P, B, perm)
|
||
end
|
||
|
||
function checkdims_perm{TP,TB,N}(P::AbstractArray{TP,N}, B::AbstractArray{TB,N}, perm)
|
||
indsB = indices(B)
|
||
length(perm) == N || throw(ArgumentError("expected permutation of size $N, but length(perm)=$(length(perm))"))
|
||
isperm(perm) || throw(ArgumentError("input is not a permutation"))
|
||
indsP = indices(P)
|
||
for i = 1:length(perm)
|
||
indsP[i] == indsB[perm[i]] || throw(DimensionMismatch("destination tensor of incorrect size"))
|
||
end
|
||
nothing
|
||
end
|
||
|
||
for (V, PT, BT) in [((:N,), BitArray, BitArray), ((:T,:N), Array, StridedArray)]
|
||
@eval @generated function permutedims!(P::$PT{$(V...)}, B::$BT{$(V...)}, perm) where $(V...)
|
||
quote
|
||
checkdims_perm(P, B, perm)
|
||
|
||
#calculates all the strides
|
||
strides_1 = 0
|
||
@nexprs $N d->(strides_{d+1} = stride(B, perm[d]))
|
||
|
||
#Creates offset, because indexing starts at 1
|
||
offset = 1 - sum(@ntuple $N d->strides_{d+1})
|
||
|
||
if isa(B, SubArray)
|
||
offset += first_index(B::SubArray) - 1
|
||
B = B.parent
|
||
end
|
||
|
||
ind = 1
|
||
@nexprs 1 d->(counts_{$N+1} = strides_{$N+1}) # a trick to set counts_($N+1)
|
||
@nloops($N, i, P,
|
||
d->(counts_d = strides_d), # PRE
|
||
d->(counts_{d+1} += strides_{d+1}), # POST
|
||
begin # BODY
|
||
sumc = sum(@ntuple $N d->counts_{d+1})
|
||
@inbounds P[ind] = B[sumc+offset]
|
||
ind += 1
|
||
end)
|
||
|
||
return P
|
||
end
|
||
end
|
||
end
|
||
|
||
## unique across dim
|
||
|
||
# TODO: this doesn't fit into the new hashing scheme in any obvious way
|
||
|
||
struct Prehashed
|
||
hash::UInt
|
||
end
|
||
hash(x::Prehashed) = x.hash
|
||
|
||
"""
|
||
unique(itr[, dim])
|
||
|
||
Returns an array containing only the unique elements of the iterable `itr`, in
|
||
the order that the first of each set of equivalent elements originally appears.
|
||
If `dim` is specified, returns unique regions of the array `itr` along `dim`.
|
||
|
||
```jldoctest
|
||
julia> A = map(isodd, reshape(collect(1:8), (2,2,2)))
|
||
2×2×2 Array{Bool,3}:
|
||
[:, :, 1] =
|
||
true true
|
||
false false
|
||
|
||
[:, :, 2] =
|
||
true true
|
||
false false
|
||
|
||
julia> unique(A)
|
||
2-element Array{Bool,1}:
|
||
true
|
||
false
|
||
|
||
julia> unique(A, 2)
|
||
2×1×2 Array{Bool,3}:
|
||
[:, :, 1] =
|
||
true
|
||
false
|
||
|
||
[:, :, 2] =
|
||
true
|
||
false
|
||
|
||
julia> unique(A, 3)
|
||
2×2×1 Array{Bool,3}:
|
||
[:, :, 1] =
|
||
true true
|
||
false false
|
||
```
|
||
"""
|
||
@generated function unique(A::AbstractArray{T,N}, dim::Int) where {T,N}
|
||
inds = inds -> zeros(UInt, inds)
|
||
quote
|
||
1 <= dim <= $N || return copy(A)
|
||
hashes = similar($inds, indices(A, dim))
|
||
|
||
# Compute hash for each row
|
||
k = 0
|
||
@nloops $N i A d->(if d == dim; k = i_d; end) begin
|
||
@inbounds hashes[k] = hash(hashes[k], hash((@nref $N A i)))
|
||
end
|
||
|
||
# Collect index of first row for each hash
|
||
uniquerow = similar(Array{Int}, indices(A, dim))
|
||
firstrow = Dict{Prehashed,Int}()
|
||
for k = indices(A, dim)
|
||
uniquerow[k] = get!(firstrow, Prehashed(hashes[k]), k)
|
||
end
|
||
uniquerows = collect(values(firstrow))
|
||
|
||
# Check for collisions
|
||
collided = similar(falses, indices(A, dim))
|
||
@inbounds begin
|
||
@nloops $N i A d->(if d == dim
|
||
k = i_d
|
||
j_d = uniquerow[k]
|
||
else
|
||
j_d = i_d
|
||
end) begin
|
||
if (@nref $N A j) != (@nref $N A i)
|
||
collided[k] = true
|
||
end
|
||
end
|
||
end
|
||
|
||
if any(collided)
|
||
nowcollided = similar(BitArray, indices(A, dim))
|
||
while any(collided)
|
||
# Collect index of first row for each collided hash
|
||
empty!(firstrow)
|
||
for j = indices(A, dim)
|
||
collided[j] || continue
|
||
uniquerow[j] = get!(firstrow, Prehashed(hashes[j]), j)
|
||
end
|
||
for v in values(firstrow)
|
||
push!(uniquerows, v)
|
||
end
|
||
|
||
# Check for collisions
|
||
fill!(nowcollided, false)
|
||
@nloops $N i A d->begin
|
||
if d == dim
|
||
k = i_d
|
||
j_d = uniquerow[k]
|
||
(!collided[k] || j_d == k) && continue
|
||
else
|
||
j_d = i_d
|
||
end
|
||
end begin
|
||
if (@nref $N A j) != (@nref $N A i)
|
||
nowcollided[k] = true
|
||
end
|
||
end
|
||
(collided, nowcollided) = (nowcollided, collided)
|
||
end
|
||
end
|
||
|
||
@nref $N A d->d == dim ? sort!(uniquerows) : (indices(A, d))
|
||
end
|
||
end
|
||
|
||
"""
|
||
extrema(A, dims) -> Array{Tuple}
|
||
|
||
Compute the minimum and maximum elements of an array over the given dimensions.
|
||
|
||
# Example
|
||
```jldoctest
|
||
julia> A = reshape(collect(1:2:16), (2,2,2))
|
||
2×2×2 Array{Int64,3}:
|
||
[:, :, 1] =
|
||
1 5
|
||
3 7
|
||
|
||
[:, :, 2] =
|
||
9 13
|
||
11 15
|
||
|
||
julia> extrema(A, (1,2))
|
||
1×1×2 Array{Tuple{Int64,Int64},3}:
|
||
[:, :, 1] =
|
||
(1, 7)
|
||
|
||
[:, :, 2] =
|
||
(9, 15)
|
||
```
|
||
"""
|
||
function extrema(A::AbstractArray, dims)
|
||
sz = [size(A)...]
|
||
sz[[dims...]] = 1
|
||
B = Array{Tuple{eltype(A),eltype(A)}}(sz...)
|
||
return extrema!(B, A)
|
||
end
|
||
|
||
@noinline function extrema!(B, A)
|
||
sA = size(A)
|
||
sB = size(B)
|
||
for I in CartesianRange(sB)
|
||
AI = A[I]
|
||
B[I] = (AI, AI)
|
||
end
|
||
Bmax = CartesianIndex(sB)
|
||
@inbounds @simd for I in CartesianRange(sA)
|
||
J = min(Bmax,I)
|
||
BJ = B[J]
|
||
AI = A[I]
|
||
if AI < BJ[1]
|
||
B[J] = (AI, BJ[2])
|
||
elseif AI > BJ[2]
|
||
B[J] = (BJ[1], AI)
|
||
end
|
||
end
|
||
return B
|
||
end
|