diag with a Val band index

[jishnub]

This provides a public API to access the underlying bands of a structured matrix without having to access its fields. These bands are often used by packages, so this method seems useful.

julia> D = Diagonal(1:3)
3×3 Diagonal{Int64, UnitRange{Int64}}:
 1  ⋅  ⋅
 ⋅  2  ⋅
 ⋅  ⋅  3

julia> diag(D, Val(0))
1:3

julia> diag(D) # materializes the range for type-stability
3-element Vector{Int64}:
 1
 2
 3

This generalizes better to other structured matrix types for which we can't specialize the 1-argument diag.

julia> T = Tridiagonal(1:2, 1:3, 1:2)
3×3 Tridiagonal{Int64, UnitRange{Int64}}:
 1  1  ⋅
 1  2  2
 ⋅  2  3

julia> diag(T,1)
2-element Vector{Int64}:
 1
 2

julia> diag(T,Val(1))
1:2

Currently, diag(M, ::Val{k}) where {k} falls back to diag(M, k::Integer) when a specialized implementation isn't available.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 93.89%. Comparing base (ed53855) to head (b11d345).
⚠️ Report is 7 commits behind head on master.

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #1424   +/-   ##
=======================================
  Coverage   93.88%   93.89%           
=======================================
  Files          34       34           
  Lines       15891    15913   +22     
=======================================
+ Hits        14920    14942   +22     
  Misses        971      971           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[dkarrasch]

Nice stuff! Should we use it in multiplication code? Looking at your example, we should perhaps do it only in cases where we only read from the bands, but don't write to them?

[dkarrasch]

This may increase latency a bit, but should decrease allocations and potentially speed-up reading from the bands?

[jishnub]

Yes, I see this primarily as an optimization that downstream packages may find useful. We may be able to write to the bands as well, but this isn't guaranteed in general. Reading should always be possible.

An example use case may be that diag(s::StaticMatrix, ::Val{k}) where {k} may return a StaticVector, unlike the case where we don't specialize on the value of k.

We are already using something similar in our multiplication code that unwraps the fields (_diag), so this isn't necessary.

[nsajko]

IMO it's ugly to introduce more ::Val methods, instead of relying on the "static integer" types present in some packages, possibly through package extensions.

If you really do want Val, I'd suggest doing ::Int typeasserts on the method static parameter in the beginning of the body of each method with a Val{k} method static parameter.

Why: The value of the static parameter is intended to be Int. This is both usual with Val usage in general and implied by the fact that you define methods for Val{0}, but not for, e.g., Val{0x0}. Adding typeasserts provides some measure of type safety, protecting users who might choose the wrong type for the static parameter isbits value from performance gotchas/unexpected dispatch.

The typeasserts might also improve worst-case inference results.

[nsajko]

Suggested change

	diag(A::Transpose, ::Val{k}) where {k} = _vectranspose(diag(parent(A), Val(-k)))
	diag(A::Transpose, ::Val{k}) where {k} = _vectranspose(diag(parent(A), Val(-k::Int)))

[nsajko]

Suggested change

	diag(A::Adjoint, ::Val{k}) where {k} = _vecadjoint(diag(parent(A), Val(-k)))
	diag(A::Adjoint, ::Val{k}) where {k} = _vecadjoint(diag(parent(A), Val(-k::Int)))

[nsajko]

Suggested change

	function diag(A::AbstractMatrix, ::Val{k}) where {k}
	function diag(A::AbstractMatrix, ::Val{K}) where {K}
	k = K::Int

[nsajko]

Suggested change

	diag(A::UpperHessenberg, ::Val{k}) where {k} = k >= -1 ? diag(A.data, Val(k)) : diag(A, k)
	function diag(A::UpperHessenberg, ::Val{K}) where {K}
	k = K::Int
	k >= -1 ? diag(A.data, Val(k)) : diag(A, k)
	end

[nsajko]

Suggested change

	diag(A::UpperTriangular, ::Val{k}) where {k} = k >= 0 ? diag(A.data, Val(k)) : diag(A, k)
	function diag(A::UpperTriangular, ::Val{K}) where {K}
	k = K::Int
	k >= 0 ? diag(A.data, Val(k)) : diag(A, k)
	end

[nsajko]

Suggested change

	diag(A::LowerTriangular, ::Val{k}) where {k} = k <= 0 ? diag(A.data, Val(k)) : diag(A, k)
	function diag(A::LowerTriangular, ::Val{K}) where {K}
	k = K::Int
	k <= 0 ? diag(A.data, Val(k)) : diag(A, k)
	end

[nsajko]

Suggested change

	diag(A::UnitUpperTriangular, ::Val{k}) where {k} = k > 0 ? diag(A.data, Val(k)) : diag(A, k)
	function diag(A::UnitUpperTriangular, ::Val{K}) where {K}
	k = K::Int
	k > 0 ? diag(A.data, Val(k)) : diag(A, k)
	end

[nsajko]

Suggested change

	diag(A::UnitLowerTriangular, ::Val{k}) where {k} = k < 0 ? diag(A.data, Val(k)) : diag(A, k)
	function diag(A::UnitLowerTriangular, ::Val{K}) where {K}
	k = K::Int
	k < 0 ? diag(A.data, Val(k)) : diag(A, k)
	end

[putianyi889]

Would it make sense to let diag(A) = diag(A, Val(0))?

[jishnub]

That might be too breaking. It's usually reasonable to assume that diag(A), diag(A, 0) and diag(A, 1) have the same types. Also, it's common practice to assume that diag(A) doesn't share memory with A.

We may try something like that to see how much it impacts the ecosystem, but let's do it in another PR.

[fredrikekre]

I think this should be based on view with a custom Diag index type, similarly to how you extract blockarray views (@view A[Block(1)]). Using Val-indices to essentially mean aliased getindex is a really strange interface IMO.

[jishnub]

Are you suggesting a Band type to be introduced here that view might dispatch on? That certainly makes sense to me as the fallback may be diagview instead of diag, which is consistent with the non-allocating semantics.

The only concern there is type-stability. The current PR was embedding the band index in the Val to avoid this issue, and perhaps we need something similar for Band.

[jishnub]

Looking at what BandedMatrices does, it seems analogous to what diagview does in LinearAlgebra.

julia> B = brand(Int8, 5,5, 1,1)
5×5 BandedMatrix{Int8} with bandwidths (1, 1):
 -81   -73     ⋅     ⋅    ⋅
  92  -122   -80     ⋅    ⋅
   ⋅  -120    37  -100    ⋅
   ⋅     ⋅  -111   -80  -32
   ⋅     ⋅     ⋅    69   28

julia> @view B[Band(1)]
4-element view(reshape(::BandedMatrix{Int8, Matrix{Int8}, Base.OneTo{Int64}}, 25), BandedMatrices.BandSlice(Band(1), 6:6:24)) with eltype Int8:
  -73
  -80
 -100
  -32

This returns a view of the BandedMatrix instead of a view into the underlying storage.