Merge pull request #15423 from martinholters/sparsecopy

Specialized copy! for sparse matrices and vectors

Author: timholy

base/sparse/sparsematrix.jl

@@ -193,20 +193,51 @@ end
 copy(S::SparseMatrixCSC) =
     SparseMatrixCSC(S.m, S.n, copy(S.colptr), copy(S.rowval), copy(S.nzval))
 
-function copy!{TvA, TiA, TvB, TiB}(A::SparseMatrixCSC{TvA,TiA},
-                                   B::SparseMatrixCSC{TvB,TiB})
-    # If the two matrices have the same size then all the
-    # elements in A will be overwritten and we can simply copy the
-    # internal fields of B to A.
-    if size(A) == size(B)
-        copy!(A.colptr, B.colptr)
+function copy!(A::SparseMatrixCSC, B::SparseMatrixCSC)
+    # If the two matrices have the same length then all the
+    # elements in A will be overwritten.
+    if length(A) == length(B)
         resize!(A.nzval, length(B.nzval))
         resize!(A.rowval, length(B.rowval))
-        copy!(A.rowval, B.rowval)
-        copy!(A.nzval, B.nzval)
+        if size(A) == size(B)
+            # Simple case: we can simply copy the internal fields of B to A.
+            copy!(A.colptr, B.colptr)
+            copy!(A.rowval, B.rowval)
+        else
+            # This is like a "reshape B into A".
+            sparse_compute_reshaped_colptr_and_rowval(A.colptr, A.rowval, A.m, A.n, B.colptr, B.rowval, B.m, B.n)
+        end
     else
-        invoke(Base.copy!, Tuple{AbstractMatrix{TvA}, AbstractMatrix{TvB}}, A, B)
+        length(A) >= length(B) || throw(BoundsError())
+        lB = length(B)
+        nnzA = nnz(A)
+        nnzB = nnz(B)
+        # Up to which col, row, and ptr in rowval/nzval will A be overwritten?
+        lastmodcolA = div(lB - 1, A.m) + 1
+        lastmodrowA = mod(lB - 1, A.m) + 1
+        lastmodptrA = A.colptr[lastmodcolA]
+        while lastmodptrA < A.colptr[lastmodcolA+1] && A.rowval[lastmodptrA] <= lastmodrowA
+            lastmodptrA += 1
+        end
+        lastmodptrA -= 1
+        if lastmodptrA >= nnzB
+            # A will have fewer non-zero elements; unmodified elements are kept at the end.
+            deleteat!(A.rowval, nnzB+1:lastmodptrA)
+            deleteat!(A.nzval, nnzB+1:lastmodptrA)
+        else
+            # A will have more non-zero elements; unmodified elements are kept at the end.
+            resize!(A.rowval, nnzB + nnzA - lastmodptrA)
+            resize!(A.nzval, nnzB + nnzA - lastmodptrA)
+            copy!(A.rowval, nnzB+1, A.rowval, lastmodptrA+1, nnzA-lastmodptrA)
+            copy!(A.nzval, nnzB+1, A.nzval, lastmodptrA+1, nnzA-lastmodptrA)
+        end
+        # Adjust colptr accordingly.
+        @inbounds for i in 2:length(A.colptr)
+            A.colptr[i] += nnzB - lastmodptrA
+        end
+        sparse_compute_reshaped_colptr_and_rowval(A.colptr, A.rowval, A.m, lastmodcolA-1, B.colptr, B.rowval, B.m, B.n)
     end
+    copy!(A.nzval, B.nzval)
     return A
 end
 

base/sparse/sparsevector.jl

@@ -320,6 +320,60 @@ convert{Tv,TvS,TiS}(::Type{SparseVector{Tv}}, s::SparseVector{TvS,TiS}) =
     SparseVector{Tv,TiS}(s.n, s.nzind, convert(Vector{Tv}, s.nzval))
 
 
+### copying
+function prep_sparsevec_copy_dest!(A::SparseVector, lB, nnzB)
+    lA = length(A)
+    lA >= lB || throw(BoundsError())
+    # If the two vectors have the same length then all the elements in A will be overwritten.
+    if length(A) == lB
+        resize!(A.nzval, nnzB)
+        resize!(A.nzind, nnzB)
+    else
+        nnzA = nnz(A)
+
+        lastmodindA = searchsortedlast(A.nzind, lB)
+        if lastmodindA >= nnzB
+            # A will have fewer non-zero elements; unmodified elements are kept at the end.
+            deleteat!(A.nzind, nnzB+1:lastmodindA)
+            deleteat!(A.nzval, nnzB+1:lastmodindA)
+        else
+            # A will have more non-zero elements; unmodified elements are kept at the end.
+            resize!(A.nzind, nnzB + nnzA - lastmodindA)
+            resize!(A.nzval, nnzB + nnzA - lastmodindA)
+            copy!(A.nzind, nnzB+1, A.nzind, lastmodindA+1, nnzA-lastmodindA)
+            copy!(A.nzval, nnzB+1, A.nzval, lastmodindA+1, nnzA-lastmodindA)
+        end
+    end
+end
+
+function copy!(A::SparseVector, B::SparseVector)
+    prep_sparsevec_copy_dest!(A, length(B), nnz(B))
+    copy!(A.nzind, B.nzind)
+    copy!(A.nzval, B.nzval)
+    return A
+end
+
+function copy!(A::SparseVector, B::SparseMatrixCSC)
+    prep_sparsevec_copy_dest!(A, length(B), nnz(B))
+
+    ptr = 1
+    @assert length(A.nzind) >= length(B.rowval)
+    maximum(B.colptr)-1 <= length(B.rowval) || throw(BoundsError())
+    @inbounds for col=1:length(B.colptr)-1
+        offsetA = (col - 1) * B.m
+        while ptr <= B.colptr[col+1]-1
+            A.nzind[ptr] = B.rowval[ptr] + offsetA
+            ptr += 1
+        end
+    end
+    copy!(A.nzval, B.nzval)
+    return A
+end
+
+copy!{TvB, TiB}(A::SparseMatrixCSC, B::SparseVector{TvB,TiB}) =
+    copy!(A, SparseMatrixCSC{TvB,TiB}(B.n, 1, TiB[1, length(B.nzind)+1], B.nzind, B.nzval))
+
+
 ### Rand Construction
 sprand{T}(n::Integer, p::AbstractFloat, rfn::Function, ::Type{T}) = sprand(GLOBAL_RNG, n, p, rfn, T)
 function sprand{T}(r::AbstractRNG, n::Integer, p::AbstractFloat, rfn::Function, ::Type{T})

test/sparsedir/sparse.jl

@@ -253,14 +253,30 @@ let
     @test pointer(A.nzval) != pointer(B.nzval)
     @test pointer(A.rowval) != pointer(B.rowval)
     @test pointer(A.colptr) != pointer(B.colptr)
-    # Test size(A) != size(B)
-    B = sprand(3, 3, 0.2)
+    # Test size(A) != size(B), but length(A) == length(B)
+    B = sprand(25, 1, 0.2)
     copy!(A, B)
-    @test A[1:9] == B[:]
+    @test A[:] == B[:]
+    # Test various size(A) / size(B) combinations
+    for mA in [5, 10, 20], nA in [5, 10, 20], mB in [5, 10, 20], nB in [5, 10, 20]
+        A = sprand(mA,nA,0.4)
+        Aorig = copy(A)
+        B = sprand(mB,nB,0.4)
+        if mA*nA >= mB*nB
+            copy!(A,B)
+            @assert(A[1:length(B)] == B[:])
+            @assert(A[length(B)+1:end] == Aorig[length(B)+1:end])
+        else
+            @test_throws BoundsError copy!(A,B)
+        end
+    end
     # Test eltype(A) != eltype(B), size(A) != size(B)
+    A = sprand(5, 5, 0.2)
+    Aorig = copy(A)
     B = sparse(rand(Float32, 3, 3))
     copy!(A, B)
     @test A[1:9] == B[:]
+    @test A[10:end] == Aorig[10:end]
     # Test eltype(A) != eltype(B), size(A) == size(B)
     A = sparse(rand(Float64, 3, 3))
     B = sparse(rand(Float32, 3, 3))

test/sparsedir/sparsevector.jl

@@ -243,6 +243,90 @@ let a = SparseVector(8, [2, 5, 6], Int32[12, 35, 72])
     @test sparsevec(ctranspose(ctranspose(acp))) == acp
 end
 
+let x1 = SparseVector(8, [2, 5, 6], [12.2, 1.4, 5.0])
+    x2 = SparseVector(8, [3, 4], [1.2, 3.4])
+    copy!(x2, x1)
+    @test x2 == x1
+    x2 = SparseVector(8, [2, 4, 8], [10.3, 7.4, 3.1])
+    copy!(x2, x1)
+    @test x2 == x1
+    x2 = SparseVector(8, [1, 3, 4, 7], [0.3, 1.2, 3.4, 0.1])
+    copy!(x2, x1)
+    @test x2 == x1
+    x2 = SparseVector(10, [3, 4], [1.2, 3.4])
+    copy!(x2, x1)
+    @test x2[1:8] == x1
+    @test x2[9:10] == spzeros(2)
+    x2 = SparseVector(10, [3, 4, 9], [1.2, 3.4, 17.8])
+    copy!(x2, x1)
+    @test x2[1:8] == x1
+    @test x2[9] == 17.8
+    @test x2[10] == 0
+    x2 = SparseVector(10, [3, 4, 5, 6, 9], [8.3, 7.2, 1.2, 3.4, 17.8])
+    copy!(x2, x1)
+    @test x2[1:8] == x1
+    @test x2[9] == 17.8
+    @test x2[10] == 0
+    x2 = SparseVector(6, [3, 4], [1.2, 3.4])
+    @test_throws BoundsError copy!(x2, x1)
+end
+
+let x1 = sparse([2, 1, 2], [1, 3, 3], [12.2, 1.4, 5.0], 2, 4)
+    x2 = SparseVector(8, [3, 4], [1.2, 3.4])
+    copy!(x2, x1)
+    @test x2[:] == x1[:]
+    x2 = SparseVector(8, [2, 4, 8], [10.3, 7.4, 3.1])
+    copy!(x2, x1)
+    @test x2[:] == x1[:]
+    x2 = SparseVector(8, [1, 3, 4, 7], [0.3, 1.2, 3.4, 0.1])
+    copy!(x2, x1)
+    @test x2[:] == x1[:]
+    x2 = SparseVector(10, [3, 4], [1.2, 3.4])
+    copy!(x2, x1)
+    @test x2[1:8] == x1[:]
+    @test x2[9:10] == spzeros(2)
+    x2 = SparseVector(10, [3, 4, 9], [1.2, 3.4, 17.8])
+    copy!(x2, x1)
+    @test x2[1:8] == x1[:]
+    @test x2[9] == 17.8
+    @test x2[10] == 0
+    x2 = SparseVector(10, [3, 4, 5, 6, 9], [8.3, 7.2, 1.2, 3.4, 17.8])
+    copy!(x2, x1)
+    @test x2[1:8] == x1[:]
+    @test x2[9] == 17.8
+    @test x2[10] == 0
+    x2 = SparseVector(6, [3, 4], [1.2, 3.4])
+    @test_throws BoundsError copy!(x2, x1)
+end
+
+let x1 = SparseVector(8, [2, 5, 6], [12.2, 1.4, 5.0])
+    x2 = sparse([1, 2], [2, 2], [1.2, 3.4], 2, 4)
+    copy!(x2, x1)
+    @test x2[:] == x1[:]
+    x2 = sparse([2, 2, 2], [1, 3, 4], [10.3, 7.4, 3.1], 2, 4)
+    copy!(x2, x1)
+    @test x2[:] == x1[:]
+    x2 = sparse([1, 1, 2, 1], [1, 2, 2, 4], [0.3, 1.2, 3.4, 0.1], 2, 4)
+    copy!(x2, x1)
+    @test x2[:] == x1[:]
+    x2 = sparse([1, 2], [2, 2], [1.2, 3.4], 2, 5)
+    copy!(x2, x1)
+    @test x2[1:8] == x1
+    @test x2[9:10] == spzeros(2)
+    x2 = sparse([1, 2, 1], [2, 2, 5], [1.2, 3.4, 17.8], 2, 5)
+    copy!(x2, x1)
+    @test x2[1:8] == x1
+    @test x2[9] == 17.8
+    @test x2[10] == 0
+    x2 = sparse([1, 2, 1, 2, 1], [2, 2, 3, 3, 5], [8.3, 7.2, 1.2, 3.4, 17.8], 2, 5)
+    copy!(x2, x1)
+    @test x2[1:8] == x1
+    @test x2[9] == 17.8
+    @test x2[10] == 0
+    x2 = sparse([1, 2], [2, 2], [1.2, 3.4], 2, 3)
+    @test_throws BoundsError copy!(x2, x1)
+end
+
 ### Type conversion
 
 let x = convert(SparseVector, sparse([2, 5, 6], [1, 1, 1], [1.25, -0.75, 3.5], 8, 1))