diag of SparseMatrixCSC should always return SparseVector

[fredrikekre]

diag(A::SparseMatrixCSC, k::Int=0) now dispatches here and then to getindex for sparse matrices. Will do some benchmarking if it is worth updating SpDiagIterator, or keep it as is.

fix #21064

[tkelman]

if it doesn't help performance to adapt SpDiagIterator, we should maybe deprecate that if it isn't used anywhere else

[fredrikekre]

It is used in trace too, but thats it. Will see what benchmarks say tomorrow.

[Sacha0]

superficially lgtm :).

[andreasnoack]

Is this to avoid changing the behavior? I guess it should still work with a sparse vector, right?

[fredrikekre]

Yea it works, but will result in a SparseMatrixCSC, so the result is not inferable.

[fredrikekre]

It seems that using SpDiagIterator is faster so we should rewrite that to work with arbitrary diagonal. I suppose this could be merged as is and we can make performance improvements later, but might as well leave it as WIP until I get around to fix it.

[fredrikekre]

I updated SpDiagIterator such that we can iterate over any diagonal.

[tkelman]

does SpDiagIterator skip structural zeros or produce them?

[fredrikekre]

Well, it is like getindex so it returns either the stored value at that position or zero(T). I am tempted to just iterate over stored values but that would be breaking I guess. Could also just remove SpDiagIterator and move it into the diag function.

[tkelman]

that would be breaking I guess

in what way?

[fredrikekre]

Something like

for i in SpDiagIterator(spzeros(3,3))
    # ...
end

Currently gives a loop of length 3, with i being 0 all three times. If the iterator only returned stored values, this loop would be done on the first iteration.

[tkelman]

Ah yes good point. I just grepped all registered packages and it's never used outside of Base, so I think as long as it still does what it's supposed to in the places that Base uses it, we can change its behavior if it would be an improvement.

[fredrikekre]

Yea I did some searching too, but found it nowhere else. In that case I think we can just remove it. I did some benchmarking for trace and it seems that

function trace2(A::SparseMatrixCSC{Tv,Ti}) where {Tv,Ti}
    n = checksquare(A)
    d = zero(Tv)
    for i in 1:n
        d += A[i,i]
    end
    return d
end

is faster than using the iterator. So I will remove it, and then include the iteration code inside the diag body instead.

[fredrikekre]

Take 3: Removed SpDiagIterator and imported the iteration code into the diag body. Now we collect all the stored values, so stored zeros will be preserved, which is kinda neat IMO

julia> A
3×3 SparseMatrixCSC{Float64,Int64} with 1 stored entry:
  [1, 1]  =  0.0

julia> diag(A)
3-element SparseVector{Float64,Int64} with 1 stored entry:
  [1]  =  0.0

[tkelman]

looks like most of the other methods allow Integer for the second argument

[Sacha0]

Perhaps either type r and c or use zero to on the right side to avoid potential type instability?

[fredrikekre]

Converted d directly to Int so should be ok now :)

[Sacha0]

Likewise here, either type r and c or use oneunit on the right side to avoid potential type instability?

[fredrikekre]

Also fixed by converting to Int directly.

[Sacha0]

Does this version perform similarly to the original? Sparse getindex is fairly complex / expensive.

[KristofferC]

Shouldn't it do pretty much exactly what the sparse iterator what doing? You still gotta search, right?

[Sacha0]

I expect so, at least roughly. But my conjectures often fail, so explicit verification has become my friend :).

[fredrikekre]

Some benchmarking suggested this was faster. The iterator was even allocating some stuff. I will get back with number :)

[fredrikekre]

using BenchmarkTools
function trace2(A::SparseMatrixCSC{Tv}) where Tv
    n = Base.LinAlg.checksquare(A)
    s = zero(Tv)
    for i in 1:n
        s += A[i,i]
    end
    return s
end

for s in (1000, 5000), p in (0.1, 0.01, 0.005)
    S1 = sprand(s, s, p)
    S2 = S1 + speye(s, s) # typical case with values on all diagonal positions
    println("trace")
    @btime trace($S1)
    @btime trace($S2)
    println("trace2")
    @btime trace2($S1)
    @btime trace2($S2)
end

with output:

trace
  24.820 μs (6 allocations: 176 bytes)
  26.273 μs (6 allocations: 176 bytes)
trace2
  25.804 μs (0 allocations: 0 bytes)
  25.010 μs (0 allocations: 0 bytes)
trace
  11.110 μs (6 allocations: 176 bytes)
  11.680 μs (6 allocations: 176 bytes)
trace2
  9.736 μs (0 allocations: 0 bytes)
  10.217 μs (0 allocations: 0 bytes)
trace
  9.370 μs (6 allocations: 176 bytes)
  9.862 μs (6 allocations: 176 bytes)
trace2
  7.443 μs (0 allocations: 0 bytes)
  7.938 μs (0 allocations: 0 bytes)
trace
  204.548 μs (6 allocations: 176 bytes)
  218.412 μs (6 allocations: 176 bytes)
trace2
  197.253 μs (0 allocations: 0 bytes)
  210.470 μs (0 allocations: 0 bytes)
trace
  117.837 μs (6 allocations: 176 bytes)
  120.147 μs (6 allocations: 176 bytes)
trace2
  112.441 μs (0 allocations: 0 bytes)
  113.419 μs (0 allocations: 0 bytes)
trace
  104.461 μs (6 allocations: 176 bytes)
  106.751 μs (6 allocations: 176 bytes)
trace2
  97.655 μs (0 allocations: 0 bytes)
  100.329 μs (0 allocations: 0 bytes)

[Sacha0]

IIRC searchsortedfirst always returns an Int?

[Sacha0]

Extra space before ==?

[Sacha0]

Extra space before ==?

[Sacha0]

Thanks @fredrikekre! :)

The Travis macOS failure appears unrelated.