Add diagview to obtain a view along a diagonal (#56175)

jishnub · web-flow · commit 38e3d1433a24 · 2024-11-11T17:17:36.000+05:30
A function to obtain a view of a diagonal of a matrix is useful, and this is clearly being used widely within `LinearAlgebra`. The implementation here iterates according to the `IndexStyle` of the array: ```julia julia> using LinearAlgebra julia> A = reshape(1:9, 3, 3) 3×3 reshape(::UnitRange{Int64}, 3, 3) with eltype Int64: 1 4 7 2 5 8 3 6 9 julia> diagview(A,1) 2-element view(::UnitRange{Int64}, 4:4:8) with eltype Int64: 4 8 julia> T = Tridiagonal(1:3, 3:6, 4:6) 4×4 Tridiagonal{Int64, UnitRange{Int64}}: 3 4 ⋅ ⋅ 1 4 5 ⋅ ⋅ 2 5 6 ⋅ ⋅ 3 6 julia> diagview(T,1) 3-element view(::Tridiagonal{Int64, UnitRange{Int64}}, StepRangeLen(CartesianIndex(1, 2), CartesianIndex(1, 1), 3)) with eltype Int64: 4 5 6 ``` Closes #30250
diff --git a/NEWS.md b/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
 
diff --git a/stdlib/LinearAlgebra/src/LinearAlgebra.jl b/stdlib/LinearAlgebra/src/LinearAlgebra.jl
@@ -87,6 +87,7 @@ export
     diag,
     diagind,
     diagm,
+    diagview,
     dot,
     eigen!,
     eigen,
diff --git a/stdlib/LinearAlgebra/src/abstractq.jl b/stdlib/LinearAlgebra/src/abstractq.jl
@@ -456,11 +456,9 @@ end
 
 det(Q::QRPackedQ) = _det_tau(Q.τ)
 det(Q::QRCompactWYQ) =
-    prod(i -> _det_tau(_diagview(Q.T[:, i:min(i + size(Q.T, 1), size(Q.T, 2))])),
+    prod(i -> _det_tau(diagview(Q.T[:, i:min(i + size(Q.T, 1), size(Q.T, 2))])),
          1:size(Q.T, 1):size(Q.T, 2))
 
-_diagview(A) = @view A[diagind(A)]
-
 # Compute `det` from the number of Householder reflections.  Handle
 # the case `Q.τ` contains zeros.
 _det_tau(τs::AbstractVector{<:Real}) =
diff --git a/stdlib/LinearAlgebra/src/bidiag.jl b/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
diff --git a/stdlib/LinearAlgebra/src/dense.jl b/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}...)
@@ -1636,13 +1665,11 @@ function pinv(A::AbstractMatrix{T}; atol::Real = 0.0, rtol::Real = (eps(real(flo
         return similar(A, Tout, (n, m))
     end
     if isdiag(A)
-        indA = diagind(A)
-        dA = view(A, indA)
+        dA = diagview(A)
         maxabsA = maximum(abs, dA)
         tol = max(rtol * maxabsA, atol)
         B = fill!(similar(A, Tout, (n, m)), 0)
-        indB = diagind(B)
-        B[indB] .= (x -> abs(x) > tol ? pinv(x) : zero(x)).(dA)
+        diagview(B) .= (x -> abs(x) > tol ? pinv(x) : zero(x)).(dA)
         return B
     end
     SVD         = svd(A)
diff --git a/stdlib/LinearAlgebra/src/diagonal.jl b/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))
diff --git a/stdlib/LinearAlgebra/src/special.jl b/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
diff --git a/stdlib/LinearAlgebra/src/triangular.jl b/stdlib/LinearAlgebra/src/triangular.jl
@@ -188,7 +188,7 @@ end
 function full(A::UnitUpperOrUnitLowerTriangular)
     isupper = A isa UnitUpperTriangular
     Ap = _triangularize(A)(parent(A), isupper ? 1 : -1)
-    Ap[diagind(Ap, IndexStyle(Ap))] = @view A[diagind(A, IndexStyle(A))]
+    diagview(Ap) .= diagview(A)
     return Ap
 end
 
@@ -400,20 +400,20 @@ function tril!(A::UnitUpperTriangular{T}, k::Integer=0) where {T}
         return UpperTriangular(A.data)
     elseif k == 0
         fill!(A.data, zero(T))
-        for i in diagind(A)
+        for i in diagind(A.data, IndexStyle(A.data))
             A.data[i] = oneunit(T)
         end
         return UpperTriangular(A.data)
     else
-        for i in diagind(A)
+        for i in diagind(A.data, IndexStyle(A.data))
             A.data[i] = oneunit(T)
         end
         return UpperTriangular(tril!(A.data,k))
     end
 end
 
 function triu!(A::UnitUpperTriangular, k::Integer=0)
-    for i in diagind(A.data)
+    for i in diagind(A.data, IndexStyle(A.data))
         A.data[i] = oneunit(eltype(A))
     end
     return triu!(UpperTriangular(A.data), k)
@@ -448,20 +448,20 @@ function triu!(A::UnitLowerTriangular{T}, k::Integer=0) where T
         return LowerTriangular(A.data)
     elseif k == 0
         fill!(A.data, zero(T))
-        for i in diagind(A)
+        for i in diagind(A.data, IndexStyle(A.data))
             A.data[i] = oneunit(T)
         end
         return LowerTriangular(A.data)
     else
-        for i in diagind(A)
+        for i in diagind(A.data, IndexStyle(A.data))
             A.data[i] = oneunit(T)
         end
         return LowerTriangular(triu!(A.data, k))
     end
 end
 
 function tril!(A::UnitLowerTriangular, k::Integer=0)
-    for i in diagind(A.data)
+    for i in diagind(A.data, IndexStyle(A.data))
         A.data[i] = oneunit(eltype(A))
     end
     return tril!(LowerTriangular(A.data), k)
@@ -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
diff --git a/stdlib/LinearAlgebra/src/tridiag.jl b/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, ")")
diff --git a/stdlib/LinearAlgebra/src/uniformscaling.jl b/stdlib/LinearAlgebra/src/uniformscaling.jl
@@ -201,7 +201,7 @@ end
 function (+)(A::Hermitian, J::UniformScaling{<:Complex})
     TS = Base.promote_op(+, eltype(A), typeof(J))
     B = copytri!(copymutable_oftype(parent(A), TS), A.uplo, true)
-    for i in diagind(B)
+    for i in diagind(B, IndexStyle(B))
         B[i] = A[i] + J
     end
     return B
@@ -211,7 +211,7 @@ function (-)(J::UniformScaling{<:Complex}, A::Hermitian)
     TS = Base.promote_op(+, eltype(A), typeof(J))
     B = copytri!(copymutable_oftype(parent(A), TS), A.uplo, true)
     B .= .-B
-    for i in diagind(B)
+    for i in diagind(B, IndexStyle(B))
         B[i] = J - A[i]
     end
     return B
diff --git a/stdlib/LinearAlgebra/test/dense.jl b/stdlib/LinearAlgebra/test/dense.jl
@@ -1029,6 +1029,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
diff --git a/stdlib/LinearAlgebra/test/tridiag.jl b/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
