implement circshift! for SparseMatrixCS

JuliaLang · Dec 13, 2018 · b48e876 · b48e876
1 parent 560e829
commit b48e876
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 1 deletion.
diff --git a/stdlib/SparseArrays/src/SparseArrays.jl b/stdlib/SparseArrays/src/SparseArrays.jl
@@ -27,7 +27,7 @@ import Base: @get!, acos, acosd, acot, acotd, acsch, asech, asin, asind, asinh,
     vcat, hcat, hvcat, cat, imag, argmax, kron, length, log, log1p, max, min,
     maximum, minimum, one, promote_eltype, real, reshape, rot180,
     rotl90, rotr90, round, setindex!, similar, size, transpose,
-    vec, permute!, map, map!, Array, diff
+    vec, permute!, map, map!, Array, diff, circshift!, circshift
 
 using Random: GLOBAL_RNG, AbstractRNG, randsubseq, randsubseq!
 

diff --git a/stdlib/SparseArrays/src/sparsematrix.jl b/stdlib/SparseArrays/src/sparsematrix.jl
@@ -3506,3 +3506,13 @@ end
 (+)(A::SparseMatrixCSC, J::UniformScaling) = A + sparse(J, size(A)...)
 (-)(A::SparseMatrixCSC, J::UniformScaling) = A - sparse(J, size(A)...)
 (-)(J::UniformScaling, A::SparseMatrixCSC) = sparse(J, size(A)...) - A
+
+## circular shift
+
+function circshift!(O::SparseMatrixCSC, X::SparseMatrixCSC, (r,c)::Base.DimsInteger{2})
+    I, J, V = findnz(X)
+    O .= sparse(mod1.(I .+ r, X.m), mod1.(J .+ c, X.n), V, X.m, X.n)
+end
+
+circshift!(O::SparseMatrixCSC, X::SparseMatrixCSC, (r,)::Base.DimsInteger{1}) = circshift!(O, X, (r,0))
+circshift!(O::SparseMatrixCSC, X::SparseMatrixCSC, r::Real) = circshift!(O, X, (Integer(r),0))
diff --git a/stdlib/SparseArrays/test/sparse.jl b/stdlib/SparseArrays/test/sparse.jl
@@ -2359,4 +2359,16 @@ end
     @test one(A) isa SparseMatrixCSC{Int}
 end
 
+@testset "circshift" begin
+    A = sprand(40, 30, 0.3)
+    @test nnz(A) == nnz(circshift(A,(1,2)))
+    @test circshift(A, (3,4)) == circshift(Matrix(A), (3,4))
+    @test circshift(A, (7,-4)) == circshift(Matrix(A), (7,-4))
+    @test circshift(A, 14) == circshift(Matrix(A), 14)
+    O = similar(A)
+    circshift!(O, A, (11,13))
+    @test nnz(O) == nnz(A)
+    @test O == circshift(A, (11,13))
+end
+
 end # module