Skip to content

Commit

Permalink
wip rewrite spdiagm with pairs
Browse files Browse the repository at this point in the history
instead of using one tuple with diagonals and one tuple with vectors
[ci skip]
  • Loading branch information
fredrikekre committed Sep 18, 2017
1 parent d668a95 commit 2be114e
Showing 1 changed file with 24 additions and 43 deletions.
67 changes: 24 additions & 43 deletions base/sparse/sparsematrix.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1117,7 +1117,7 @@ For expert drivers and additional information, see [`permute!`](@ref).
# Examples
```jldoctest
julia> A = spdiagm([1, 2, 3, 4], 0, 4, 4) + spdiagm([5, 6, 7], 1, 4, 4)
julia> A = spdiagm(0 => [1, 2, 3, 4], 1 => [5, 6, 7])
4×4 SparseMatrixCSC{Int64,Int64} with 7 stored entries:
[1, 1] = 1
[1, 2] = 5
Expand Down Expand Up @@ -3298,57 +3298,46 @@ function istril(A::SparseMatrixCSC)
return true
end

# Create a sparse diagonal matrix by specifying multiple diagonals
# packed into a tuple, alongside their diagonal offsets and matrix shape

function spdiagm_internal(B, d)
ndiags = length(d)
if length(B) != ndiags; throw(ArgumentError("first argument should be a tuple of length(d)=$ndiags arrays of diagonals")); end
function spdiagm_internal(kv::Pair{<:Integer, <:AbstractVector}...)
ncoeffs = 0
for vec in B
ncoeffs += length(vec)
for p in kv
ncoeffs += length(p.second)
end
I = Vector{Int}(ncoeffs)
J = Vector{Int}(ncoeffs)
V = Vector{promote_type(map(eltype, B)...)}(ncoeffs)
id = 0
V = Vector{promote_type(map(x -> eltype(x.second), kv)...)}(ncoeffs)
i = 0
for vec in B
id += 1
diag = d[id]
numel = length(vec)
if diag < 0
row = -diag
for p in kv
numel = length(p.second)
if p.first < 0
row = -p.first
col = 0
elseif diag > 0
elseif p.first > 0
row = 0
col = diag
col = p.first
else
row = 0
col = 0
end
range = 1+i:numel+i
I[range] = row+1:row+numel
J[range] = col+1:col+numel
copy!(view(V, range), vec)
r = 1+i:numel+i
I[r] = row+1:row+numel
J[r] = col+1:col+numel
copy!(view(V, r), p.second)
i += numel
end

return (I,J,V)
return I, J, V
end

"""
spdiagm(B, d[, m, n])
spdiagm(kv::Pair{<:Integer, <:AbstractVector}...)
Construct a sparse diagonal matrix. `B` is a tuple of vectors containing the diagonals and
`d` is a tuple containing the positions of the diagonals. In the case the input contains only
one diagonal, `B` can be a vector (instead of a tuple) and `d` can be the diagonal position
(instead of a tuple), defaulting to 0 (diagonal). Optionally, `m` and `n` specify the size
of the resulting sparse matrix.
Construct a sparse diagonal matrix from the diagonal-number/vector pair `kv`,
placing each vector `kv.second` on the `kv.first` diagonal.
# Examples
```jldoctest
julia> spdiagm(([1,2,3,4],[4,3,2,1]),(-1,1))
julia> spdiagm(-1 => [1,2,3,4], 1 => [4,3,2,1])
5×5 SparseMatrixCSC{Int64,Int64} with 8 stored entries:
[2, 1] = 1
[1, 2] = 4
Expand All @@ -3360,20 +3349,12 @@ julia> spdiagm(([1,2,3,4],[4,3,2,1]),(-1,1))
[4, 5] = 1
```
"""
function spdiagm(B, d, m::Integer, n::Integer)
(I,J,V) = spdiagm_internal(B, d)
return sparse(I,J,V,m,n)
end

function spdiagm(B, d)
(I,J,V) = spdiagm_internal(B, d)
return sparse(I,J,V)
function spdiagm(kv::Pair{<:Integer, <:AbstractVector}...)
I, J, V = spdiagm_internal(kv...)
n = max(dimlub(I), dimlub(J))
return sparse(I, J, V, n, n)
end

spdiagm(B::AbstractVector, d::Number, m::Integer, n::Integer) = spdiagm((B,), (d,), m, n)

spdiagm(B::AbstractVector, d::Number=0) = spdiagm((B,), (d,))

## expand a colptr or rowptr into a dense index vector
function expandptr(V::Vector{<:Integer})
if V[1] != 1 throw(ArgumentError("first index must be one")) end
Expand Down

0 comments on commit 2be114e

Please sign in to comment.