Add diagview to obtain a view along a diagonal

Author: jishnub

NEWS.md

@@ -151,7 +151,8 @@ Standard library changes
 * The matrix multiplication `A * B` calls `matprod_dest(A, B, T::Type)` to generate the destination.
   This function is now public ([#55537]).
 * The function `haszero(T::Type)` is used to check if a type `T` has a unique zero element defined as `zero(T)`.
-  This is now public.
+  This is now public ([#56223]).
+* A new function `diagview` is added that returns a view into a specific band of an `AbstractMatrix` ([#56175]).
 
 #### Logging
 

stdlib/LinearAlgebra/src/LinearAlgebra.jl

@@ -87,6 +87,7 @@ export
     diag,
     diagind,
     diagm,
+    diagview,
     dot,
     eigen!,
     eigen,

stdlib/LinearAlgebra/src/bidiag.jl

@@ -191,8 +191,8 @@ function Matrix{T}(A::Bidiagonal) where T
     B = Matrix{T}(undef, size(A))
     if haszero(T) # optimized path for types with zero(T) defined
         size(B,1) > 1 && fill!(B, zero(T))
-        copyto!(view(B, diagind(B)), A.dv)
-        copyto!(view(B, diagind(B, _offdiagind(A.uplo))), A.ev)
+        copyto!(diagview(B), A.dv)
+        copyto!(diagview(B, _offdiagind(A.uplo)), A.ev)
     else
         copyto!(B, A)
     end
@@ -570,7 +570,7 @@ end
 # to avoid allocations in _mul! below (#24324, #24578)
 _diag(A::Tridiagonal, k) = k == -1 ? A.dl : k == 0 ? A.d : A.du
 _diag(A::SymTridiagonal{<:Number}, k) = k == 0 ? A.dv : A.ev
-_diag(A::SymTridiagonal, k) = k == 0 ? view(A, diagind(A, IndexStyle(A))) : view(A, diagind(A, 1, IndexStyle(A)))
+_diag(A::SymTridiagonal, k) = diagview(A,k)
 function _diag(A::Bidiagonal, k)
     if k == 0
         return A.dv

stdlib/LinearAlgebra/src/dense.jl

@@ -290,6 +290,35 @@ julia> diag(A,1)
 """
 diag(A::AbstractMatrix, k::Integer=0) = A[diagind(A, k, IndexStyle(A))]
 
+"""
+    diagview(M, k::Integer=0)
+
+Return a view into the `k`th diagonal of the matrix `M`.
+
+See also [`diag`](@ref), [`diagind`](@ref).
+
+# Examples
+```jldoctest
+julia> A = [1 2 3; 4 5 6; 7 8 9]
+3×3 Matrix{Int64}:
+ 1  2  3
+ 4  5  6
+ 7  8  9
+
+julia> diagview(A)
+3-element view(::Vector{Int64}, 1:4:9) with eltype Int64:
+ 1
+ 5
+ 9
+
+julia> diagview(A, 1)
+2-element view(::Vector{Int64}, 4:4:8) with eltype Int64:
+ 2
+ 6
+```
+"""
+diagview(A::AbstractMatrix, k::Integer=0) = @view A[diagind(A, k, IndexStyle(A))]
+
 """
     diagm(kv::Pair{<:Integer,<:AbstractVector}...)
     diagm(m::Integer, n::Integer, kv::Pair{<:Integer,<:AbstractVector}...)

stdlib/LinearAlgebra/src/diagonal.jl

@@ -120,7 +120,7 @@ function Matrix{T}(D::Diagonal) where {T}
     B = Matrix{T}(undef, size(D))
     if haszero(T) # optimized path for types with zero(T) defined
         size(B,1) > 1 && fill!(B, zero(T))
-        copyto!(view(B, diagind(B)), D.diag)
+        copyto!(diagview(B), D.diag)
     else
         copyto!(B, D)
     end
@@ -1041,7 +1041,7 @@ dot(x::AbstractVector, D::Diagonal, y::AbstractVector) = _mapreduce_prod(dot, x,
 dot(A::Diagonal, B::Diagonal) = dot(A.diag, B.diag)
 function dot(D::Diagonal, B::AbstractMatrix)
     size(D) == size(B) || throw(DimensionMismatch(lazy"Matrix sizes $(size(D)) and $(size(B)) differ"))
-    return dot(D.diag, view(B, diagind(B, IndexStyle(B))))
+    return dot(D.diag, diagview(B))
 end
 
 dot(A::AbstractMatrix, B::Diagonal) = conj(dot(B, A))

stdlib/LinearAlgebra/src/special.jl

@@ -22,7 +22,7 @@ function Tridiagonal(A::Bidiagonal)
 end
 
 _diagview(S::SymTridiagonal{<:Number}) = S.dv
-_diagview(S::SymTridiagonal) = view(S, diagind(S, IndexStyle(S)))
+_diagview(S::SymTridiagonal) = diagview(S)
 
 # conversions from SymTridiagonal to other special matrix types
 Diagonal(A::SymTridiagonal) = Diagonal(_diagview(A))
@@ -370,20 +370,20 @@ function copyto!(dest::BandedMatrix, src::BandedMatrix)
 end
 function _copyto_banded!(T::Tridiagonal, D::Diagonal)
     T.d .= D.diag
-    T.dl .= view(D, diagind(D, -1, IndexStyle(D)))
-    T.du .= view(D, diagind(D,  1, IndexStyle(D)))
+    T.dl .= diagview(D, -1)
+    T.du .= diagview(D,  1)
     return T
 end
 function _copyto_banded!(SymT::SymTridiagonal, D::Diagonal)
     issymmetric(D) || throw(ArgumentError("cannot copy a non-symmetric Diagonal matrix to a SymTridiagonal"))
     SymT.dv .= D.diag
     _ev = _evview(SymT)
-    _ev .= view(D, diagind(D,  1, IndexStyle(D)))
+    _ev .= diagview(D,  1)
     return SymT
 end
 function _copyto_banded!(B::Bidiagonal, D::Diagonal)
     B.dv .= D.diag
-    B.ev .= view(D, diagind(D,  B.uplo == 'U' ? 1 : -1, IndexStyle(D)))
+    B.ev .= diagview(D, _offdiagind(B.uplo))
     return B
 end
 function _copyto_banded!(D::Diagonal, B::Bidiagonal)
@@ -411,10 +411,10 @@ function _copyto_banded!(T::Tridiagonal, B::Bidiagonal)
     T.d .= B.dv
     if B.uplo == 'U'
         T.du .= B.ev
-        T.dl .= view(B, diagind(B, -1, IndexStyle(B)))
+        T.dl .= diagview(B,-1)
     else
         T.dl .= B.ev
-        T.du .= view(B, diagind(B,  1, IndexStyle(B)))
+        T.du .= diagview(B, 1)
     end
     return T
 end

stdlib/LinearAlgebra/src/triangular.jl

@@ -2041,7 +2041,7 @@ function _find_params_log_quasitriu!(A)
 
     # Find s0, the smallest s such that the ρ(triu(A)^(1/2^s) - I) ≤ theta[tmax], where ρ(X)
     # is the spectral radius of X
-    d = complex.(@view(A[diagind(A)]))
+    d = complex.(diagview(A))
     dm1 = d .- 1
     s = 0
     while norm(dm1, Inf) > theta[tmax] && s < maxsqrt

stdlib/LinearAlgebra/src/tridiag.jl

@@ -612,9 +612,9 @@ function Matrix{T}(M::Tridiagonal) where {T}
     A = Matrix{T}(undef, size(M))
     if haszero(T) # optimized path for types with zero(T) defined
         size(A,1) > 2 && fill!(A, zero(T))
-        copyto!(view(A, diagind(A)), M.d)
-        copyto!(view(A, diagind(A,1)), M.du)
-        copyto!(view(A, diagind(A,-1)), M.dl)
+        copyto!(diagview(A), M.d)
+        copyto!(diagview(A,1), M.du)
+        copyto!(diagview(A,-1), M.dl)
     else
         copyto!(A, M)
     end
@@ -1092,7 +1092,7 @@ function show(io::IO, T::Tridiagonal)
 end
 function show(io::IO, S::SymTridiagonal)
     print(io, "SymTridiagonal(")
-    show(io, eltype(S) <: Number ? S.dv : view(S, diagind(S, IndexStyle(S))))
+    show(io, _diagview(S))
     print(io, ", ")
     show(io, S.ev)
     print(io, ")")

stdlib/LinearAlgebra/test/dense.jl

@@ -1024,6 +1024,15 @@ end
     @test diag(zeros(0,1),2) == []
 end
 
+@testset "diagview" begin
+    for sz in ((3,3), (3,5), (5,3))
+        A = rand(sz...)
+        for k in -5:5
+            @test diagview(A,k) == diag(A,k)
+        end
+    end
+end
+
 @testset "issue #39857" begin
     @test lyap(1.0+2.0im, 3.0+4.0im) == -1.5 - 2.0im
 end

stdlib/LinearAlgebra/test/tridiag.jl

@@ -1065,4 +1065,14 @@ end
     end
 end
 
+@testset "diagview" begin
+    A = Tridiagonal(rand(3), rand(4), rand(3))
+    for k in -5:5
+        @test diagview(A,k) == diag(A,k)
+    end
+    v = diagview(A,1)
+    v .= 0
+    @test all(iszero, diag(A,1))
+end
+
 end # module TestTridiagonal