Skip to content

Commit

Permalink
Improve similar for SparseMatrixCSC.
Browse files Browse the repository at this point in the history 
Makes similar methods for SparseMatrixCSC more consistent both among themselves
and with similar methods for other types. Also fleshes those methods out in full.
  • Loading branch information
@Sacha0
Sacha0 committed Oct 17, 2017
1 parent 26c7ecb commit 85fb550
Show file tree
Hide file tree
Showing 2 changed files with 120 additions and 29 deletions.
44 changes: 34 additions & 10 deletions base/sparse/sparsematrix.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -313,17 +313,41 @@ function copy!(A::SparseMatrixCSC, B::SparseMatrixCSC)
return A
end

function similar(S::SparseMatrixCSC, ::Type{Tv} = eltype(S)) where Tv
SparseMatrixCSC(S.m, S.n, copy(S.colptr), copy(S.rowval), Vector{Tv}(length(S.nzval)))
end
## similar
#
# parent method for similar that preserves stored-entry structure (for when new and old dims match)
function _sparsesimilar(S::SparseMatrixCSC, ::Type{TvNew}, ::Type{TiNew}) where {TvNew,TiNew}
newcolptr = copy!(similar(S.colptr, TiNew), S.colptr)
newrowval = copy!(similar(S.rowval, TiNew), S.rowval)
return SparseMatrixCSC(S.m, S.n, newcolptr, newrowval, similar(S.nzval, TvNew))
end
# parent methods for similar that preserves only storage space (for when new and old dims differ)
_sparsesimilar(S::SparseMatrixCSC, ::Type{TvNew}, ::Type{TiNew}, dims::Dims{2}) where {TvNew,TiNew} =
SparseMatrixCSC(dims..., ones(TiNew, last(dims)+1), similar(S.rowval, TiNew), similar(S.nzval, TvNew))
# parent method for similar that allocates an empty sparse vector (when new dims are single)
_sparsesimilar(S::SparseMatrixCSC, ::Type{TvNew}, ::Type{TiNew}, dims::Dims{1}) where {TvNew,TiNew} =
SparseVector(dims..., similar(S.rowval, TiNew, 0), similar(S.nzval, TvNew, 0))
#
# The following methods hook into the AbstractArray similar hierarchy. The first method
# covers similar(A[, Tv]) calls, which preserve stored-entry structure, and the latter
# methods cover similar(A[, Tv], shape...) calls, which preserve storage space when the shape
# calls for a two-dimensional result.
similar(S::SparseMatrixCSC{<:Any,Ti}, ::Type{TvNew}) where {Ti,TvNew} = _sparsesimilar(S, TvNew, Ti)
similar(S::SparseMatrixCSC{<:Any,Ti}, ::Type{TvNew}, dims::Union{Dims{1},Dims{2}}) where {Ti,TvNew} =
_sparsesimilar(S, TvNew, Ti, dims)
# The following methods cover similar(A, Tv, Ti[, shape...]) calls, which specify the
# result's index type in addition to its entry type, and aren't covered by the hooks above.
# The calls without shape again preserve stored-entry structure, whereas those with shape
# preserve storage space when the shape calls for a two-dimensional result.
similar(S::SparseMatrixCSC, ::Type{TvNew}, ::Type{TiNew}) where{TvNew,TiNew} =
_sparsesimilar(S, TvNew, TiNew)
similar(S::SparseMatrixCSC, ::Type{TvNew}, ::Type{TiNew}, dims::Union{Dims{1},Dims{2}}) where {TvNew,TiNew} =
_sparsesimilar(S, TvNew, TiNew, dims)
similar(S::SparseMatrixCSC, ::Type{TvNew}, ::Type{TiNew}, m::Integer) where {TvNew,TiNew} =
_sparsesimilar(S, TvNew, TiNew, (m,))
similar(S::SparseMatrixCSC, ::Type{TvNew}, ::Type{TiNew}, m::Integer, n::Integer) where {TvNew,TiNew} =
_sparsesimilar(S, TvNew, TiNew, (m, n))

function similar(S::SparseMatrixCSC, ::Type{Tv}, ::Type{Ti}) where {Tv,Ti}
new_colptr = copy!(similar(S.colptr, Ti), S.colptr)
new_rowval = copy!(similar(S.rowval, Ti), S.rowval)
new_nzval = copy!(similar(S.nzval, Tv), S.nzval)
SparseMatrixCSC(S.m, S.n, new_colptr, new_rowval, new_nzval)
end
@inline similar(S::SparseMatrixCSC, ::Type{Tv}, d::Dims) where {Tv} = spzeros(Tv, d...)

# convert'ing between SparseMatrixCSC types
convert(::Type{AbstractMatrix{Tv}}, A::SparseMatrixCSC{Tv}) where {Tv} = A
Expand Down
105 changes: 86 additions & 19 deletions test/sparse/sparse.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1302,16 +1302,6 @@ end
@test_throws ArgumentError squeeze(A,(1, 1))
end

@testset "similar with type conversion" begin
local A = speye(5)
@test size(similar(A, Complex128, Int)) == (5, 5)
@test typeof(similar(A, Complex128, Int)) == SparseMatrixCSC{Complex128, Int}
@test size(similar(A, Complex128, Int8)) == (5, 5)
@test typeof(similar(A, Complex128, Int8)) == SparseMatrixCSC{Complex128, Int8}
@test similar(A, Complex128,(6, 6)) == spzeros(Complex128, 6, 6)
@test convert(Matrix, A) == Array(A)
end

@testset "float" begin
local A
A = sprand(Bool, 5, 5, 0.0)
Expand Down Expand Up @@ -1945,15 +1935,6 @@ end
" [2, 1] = 2.0\n [3, 2] = 3.0\n [4, 2] = 4.0\n [5, 3] = 5.0\n [6, 3] = 6.0")
end

@testset "similar aliasing" begin
a = sparse(rand(3,3) .+ 0.1)
b = similar(a, Float32, Int32)
c = similar(b, Float32, Int32)
Base.SparseArrays.dropstored!(b, 1, 1)
@test length(c.rowval) == 9
@test length(c.nzval) == 9
end

@testset "check buffers" for n in 1:3
local A
colptr = [1,2,3,4]
Expand Down Expand Up @@ -2035,3 +2016,89 @@ end
@test isfinite.(cov_sparse) == isfinite.(cov_dense)
end
end

@testset "similar should not alias the input sparse array" begin
a = sparse(rand(3,3) .+ 0.1)
b = similar(a, Float32, Int32)
c = similar(b, Float32, Int32)
Base.SparseArrays.dropstored!(b, 1, 1)
@test length(c.rowval) == 9
@test length(c.nzval) == 9
end

@testset "similar with type conversion" begin
local A = speye(5)
@test size(similar(A, Complex128, Int)) == (5, 5)
@test typeof(similar(A, Complex128, Int)) == SparseMatrixCSC{Complex128, Int}
@test size(similar(A, Complex128, Int8)) == (5, 5)
@test typeof(similar(A, Complex128, Int8)) == SparseMatrixCSC{Complex128, Int8}
@test similar(A, Complex128,(6, 6)) == spzeros(Complex128, 6, 6)
@test convert(Matrix, A) == Array(A) # lolwut, are you lost, test?
end

@testset "similar for SparseMatrixCSC" begin
A = speye(5)
# test similar without specifications (preserves stored-entry structure)
simA = similar(A)
@test typeof(simA) == typeof(A)
@test size(simA) == size(A)
@test simA.colptr == A.colptr
@test simA.rowval == A.rowval
@test length(simA.nzval) == length(A.nzval)
# test similar with entry type specification (preserves stored-entry structure)
simA = similar(A, Float32)
@test typeof(simA) == SparseMatrixCSC{Float32,eltype(A.colptr)}
@test size(simA) == size(A)
@test simA.colptr == A.colptr
@test simA.rowval == A.rowval
@test length(simA.nzval) == length(A.nzval)
# test similar with entry and index type specification (preserves stored-entry structure)
simA = similar(A, Float32, Int8)
@test typeof(simA) == SparseMatrixCSC{Float32,Int8}
@test size(simA) == size(A)
@test simA.colptr == A.colptr
@test simA.rowval == A.rowval
@test length(simA.nzval) == length(A.nzval)
# test similar with Dims{2} specification (preserves storage space only, not stored-entry structure)
simA = similar(A, (6,6))
@test typeof(simA) == typeof(A)
@test size(simA) == (6,6)
@test simA.colptr == ones(eltype(A.colptr), 6+1)
@test length(simA.rowval) == length(A.rowval)
@test length(simA.nzval) == length(A.nzval)
# test similar with entry type and Dims{2} specification (preserves storage space only)
simA = similar(A, Float32, (6,6))
@test typeof(simA) == SparseMatrixCSC{Float32,eltype(A.colptr)}
@test size(simA) == (6,6)
@test simA.colptr == ones(eltype(A.colptr), 6+1)
@test length(simA.rowval) == length(A.rowval)
@test length(simA.nzval) == length(A.nzval)
# test similar with entry type, index type, and Dims{2} specification (preserves storage space only)
simA = similar(A, Float32, Int8, (6,6))
@test typeof(simA) == SparseMatrixCSC{Float32, Int8}
@test size(simA) == (6,6)
@test simA.colptr == ones(eltype(A.colptr), 6+1)
@test length(simA.rowval) == length(A.rowval)
@test length(simA.nzval) == length(A.nzval)
# test similar with Dims{1} specification (preserves nothing)
simA = similar(A, (6,))
@test typeof(simA) == SparseVector{eltype(A.nzval),eltype(A.colptr)}
@test size(simA) == (6,)
@test length(simA.nzind) == 0
@test length(simA.nzval) == 0
# test similar with entry type and Dims{1} specification (preserves nothing)
simA = similar(A, Float32, (6,))
@test typeof(simA) == SparseVector{Float32,eltype(A.colptr)}
@test size(simA) == (6,)
@test length(simA.nzind) == 0
@test length(simA.nzval) == 0
# test similar with entry type, index type, and Dims{1} specification (preserves nothing)
simA = similar(A, Float32, Int8, (6,))
@test typeof(simA) == SparseVector{Float32,Int8}
@test size(simA) == (6,)
@test length(simA.nzind) == 0
@test length(simA.nzval) == 0
# test entry points to similar with entry type, index type, and non-Dims shape specification
@test similar(A, Float32, Int8, 6, 6) == similar(A, Float32, Int8, (6, 6))
@test similar(A, Float32, Int8, 6) == similar(A, Float32, Int8, (6,))
end

0 comments on commit 85fb550

Please sign in to comment.