Skip to content

Commit

Permalink
Merge pull request #17861 from JuliaLang/teh/circshift
Browse files Browse the repository at this point in the history
Faster, indices-aware circshift (and non-allocating circshift!)
  • Loading branch information
timholy authored Aug 9, 2016
2 parents 59f158c + 60660b5 commit b378ece
Show file tree
Hide file tree
Showing 6 changed files with 86 additions and 14 deletions.
22 changes: 10 additions & 12 deletions base/abstractarraymath.jl
Original file line number Diff line number Diff line change
Expand Up @@ -158,8 +158,10 @@ function flipdim(A::AbstractArray, d::Integer)
return B
end

circshift(a::AbstractArray, shiftamt::Real) = circshift(a, [Integer(shiftamt)])

function circshift(a::AbstractArray, shiftamt::Real)
circshift!(similar(a), a, (Integer(shiftamt),))
end
circshift(a::AbstractArray, shiftamt::DimsInteger) = circshift!(similar(a), a, shiftamt)
"""
circshift(A, shifts)
Expand All @@ -174,29 +176,25 @@ julia> b = reshape(collect(1:16), (4,4))
3 7 11 15
4 8 12 16
julia> circshift(b, [0,2])
julia> circshift(b, (0,2))
4×4 Array{Int64,2}:
9 13 1 5
10 14 2 6
11 15 3 7
12 16 4 8
julia> circshift(b, [-1,0])
julia> circshift(b, (-1,0))
4×4 Array{Int64,2}:
2 6 10 14
3 7 11 15
4 8 12 16
1 5 9 13
```
See also `circshift!`.
"""
function circshift{T,N}(a::AbstractArray{T,N}, shiftamts)
I = ()
for i=1:N
s = size(a,i)
d = i<=length(shiftamts) ? shiftamts[i] : 0
I = tuple(I..., d==0 ? [1:s;] : mod([-d:s-1-d;], s).+1)
end
a[(I::NTuple{N,Vector{Int}})...]
function circshift(a::AbstractArray, shiftamt)
circshift!(similar(a), a, map(Integer, (shiftamt...,)))
end

# Uses K-B-N summation
Expand Down
1 change: 1 addition & 0 deletions base/exports.jl
Original file line number Diff line number Diff line change
Expand Up @@ -490,6 +490,7 @@ export
checkbounds,
checkindex,
circshift,
circshift!,
clamp!,
colon,
conj!,
Expand Down
51 changes: 51 additions & 0 deletions base/multidimensional.jl
Original file line number Diff line number Diff line change
Expand Up @@ -610,6 +610,57 @@ function copy!{T,N}(dest::AbstractArray{T,N}, src::AbstractArray{T,N})
dest
end

function copy!(dest::AbstractArray, Rdest::CartesianRange, src::AbstractArray, Rsrc::CartesianRange)
isempty(Rdest) && return dest
size(Rdest) == size(Rsrc) || throw(ArgumentError("source and destination must have same size (got $(size(Rsrc)) and $(size(Rdest)))"))
@boundscheck checkbounds(dest, Rdest.start)
@boundscheck checkbounds(dest, Rdest.stop)
@boundscheck checkbounds(src, Rsrc.start)
@boundscheck checkbounds(src, Rsrc.stop)
deltaI = Rdest.start - Rsrc.start
for I in Rsrc
@inbounds dest[I+deltaI] = src[I]
end
dest
end

# 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`.
"""
@noinline function circshift!{T,N}(dest::AbstractArray{T,N}, src, shiftamt::DimsInteger)
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...,))

@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

### BitArrays

Expand Down
16 changes: 14 additions & 2 deletions doc/stdlib/arrays.rst
Original file line number Diff line number Diff line change
Expand Up @@ -590,20 +590,32 @@ Indexing, Assignment, and Concatenation
3 7 11 15
4 8 12 16

julia> circshift(b, [0,2])
julia> circshift(b, (0,2))
4×4 Array{Int64,2}:
9 13 1 5
10 14 2 6
11 15 3 7
12 16 4 8

julia> circshift(b, [-1,0])
julia> circshift(b, (-1,0))
4×4 Array{Int64,2}:
2 6 10 14
3 7 11 15
4 8 12 16
1 5 9 13

See also ``circshift!``\ .

.. function:: circshift!(dest, src, shifts)

.. Docstring generated from Julia source
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``\ .

.. function:: find(A)

.. Docstring generated from Julia source
Expand Down
9 changes: 9 additions & 0 deletions test/arrayops.jl
Original file line number Diff line number Diff line change
Expand Up @@ -462,6 +462,15 @@ for i = tensors
@test isequal(i,permutedims(ipermutedims(i,perm),perm))
end

## circshift

@test circshift(1:5, -1) == circshift(1:5, 4) == circshift(1:5, -6) == [2,3,4,5,1]
@test circshift(1:5, 1) == circshift(1:5, -4) == circshift(1:5, 6) == [5,1,2,3,4]
a = [1:5;]
@test_throws ArgumentError Base.circshift!(a, a, 1)
b = copy(a)
@test Base.circshift!(b, a, 1) == [5,1,2,3,4]

## unique across dim ##

# All rows and columns unique
Expand Down
1 change: 1 addition & 0 deletions test/offsetarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -411,6 +411,7 @@ v = OffsetArray(rand(8), (-2,))
@test rotr90(A) == OffsetArray(rotr90(parent(A)), A.offsets[[2,1]])
@test flipdim(A, 1) == OffsetArray(flipdim(parent(A), 1), A.offsets)
@test flipdim(A, 2) == OffsetArray(flipdim(parent(A), 2), A.offsets)
@test circshift(A, (-1,2)) == OffsetArray(circshift(parent(A), (-1,2)), A.offsets)

@test A+1 == OffsetArray(parent(A)+1, A.offsets)
@test 2*A == OffsetArray(2*parent(A), A.offsets)
Expand Down

0 comments on commit b378ece

Please sign in to comment.