Allow indexing of SparseMatrixCSC by array of CartesianIndex (#33225

)
JuliaLang · Sep 14, 2019 · 77fa781 · 77fa781
1 parent 04234fb
commit 77fa781
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 1 deletion.
diff --git a/base/multidimensional.jl b/base/multidimensional.jl
@@ -156,6 +156,8 @@ module IteratorsMD
         return I
     end
 
+    Base._ind2sub(t::Tuple, ind::CartesianIndex) = Tuple(ind)
+
     # Iteration over the elements of CartesianIndex cannot be supported until its length can be inferred,
     # see #23719
     Base.iterate(::CartesianIndex) =

diff --git a/stdlib/SparseArrays/src/sparsevector.jl b/stdlib/SparseArrays/src/sparsevector.jl
@@ -663,6 +663,7 @@ end
 
 function getindex(A::AbstractSparseMatrixCSC{Tv,Ti}, I::AbstractVector) where {Tv,Ti}
     require_one_based_indexing(A, I)
+    @boundscheck checkbounds(A, I)
     szA = size(A)
     nA = szA[1]*szA[2]
     colptrA = getcolptr(A)
@@ -676,7 +677,6 @@ function getindex(A::AbstractSparseMatrixCSC{Tv,Ti}, I::AbstractVector) where {T
 
     idxB = 1
     for i in 1:n
-        ((I[i] < 1) | (I[i] > nA)) && throw(BoundsError(A, I))
         row,col = Base._ind2sub(szA, I[i])
         for r in colptrA[col]:(colptrA[col+1]-1)
             @inbounds if rowvalA[r] == row

diff --git a/stdlib/SparseArrays/test/sparse.jl b/stdlib/SparseArrays/test/sparse.jl
@@ -870,6 +870,21 @@ end
         end
     end
 
+    # indexing by array of CartesianIndex (issue #30981)
+    S = sprand(10, 10, 0.4)
+    inds_sparse = S[findall(S .> 0.2)]
+    M = Matrix(S)
+    inds_dense = M[findall(M .> 0.2)]
+    @test Array(inds_sparse) == inds_dense
+    inds_out = Array([CartesianIndex(1, 1), CartesianIndex(0, 1)])
+    @test_throws BoundsError S[inds_out]
+    pop!(inds_out); push!(inds_out, CartesianIndex(1, 0))
+    @test_throws BoundsError S[inds_out]
+    pop!(inds_out); push!(inds_out, CartesianIndex(11, 1))
+    @test_throws BoundsError S[inds_out]
+    pop!(inds_out); push!(inds_out, CartesianIndex(1, 11))
+    @test_throws BoundsError S[inds_out]
+
     # workaround issue #7197: comment out let-block
     #let S = SparseMatrixCSC(3, 3, UInt8[1,1,1,1], UInt8[], Int64[])
     S1290 = SparseMatrixCSC(3, 3, UInt8[1,1,1,1], UInt8[], Int64[])