symmetrize/transpose in getindex of SymTridiagonal (#31327)

* symmetrize/transpose in SymTridiagonal

added block matrix docstring

more consistency in SymTridiagonal(A) and diag

fix x-ref

fix and test diag

remove x-ref for Transpose

fix and test for type inference

simplify constructor, fix use of symmetric

update docstrings, add comment

* don't symmetrize diagonal elements at construction, but at getindex

* remove symmetry check at construction

* don't symmetrize in SymTridiagonal(::Matrix)

* fix doctests

Author: dkarrasch

stdlib/LinearAlgebra/src/tridiag.jl

@@ -3,15 +3,15 @@
 #### Specialized matrix types ####
 
 ## (complex) symmetric tridiagonal matrices
-struct SymTridiagonal{T,V<:AbstractVector{T}} <: AbstractMatrix{T}
+struct SymTridiagonal{T, V<:AbstractVector{T}} <: AbstractMatrix{T}
     dv::V                        # diagonal
-    ev::V                        # subdiagonal
-    function SymTridiagonal{T,V}(dv, ev) where {T,V<:AbstractVector{T}}
+    ev::V                        # superdiagonal
+    function SymTridiagonal{T, V}(dv, ev) where {T, V<:AbstractVector{T}}
         require_one_based_indexing(dv, ev)
         if !(length(dv) - 1 <= length(ev) <= length(dv))
             throw(DimensionMismatch("subdiagonal has wrong length. Has length $(length(ev)), but should be either $(length(dv) - 1) or $(length(dv))."))
         end
-        new{T,V}(dv,ev)
+        new{T, V}(dv, ev)
     end
 end
 
@@ -23,6 +23,10 @@ sub/super-diagonal (`ev`), respectively. The result is of type `SymTridiagonal`
 and provides efficient specialized eigensolvers, but may be converted into a
 regular matrix with [`convert(Array, _)`](@ref) (or `Array(_)` for short).
 
+For `SymTridiagonal` block matrices, the elements of `dv` are symmetrized.
+The argument `ev` is interpreted as the superdiagonal. Blocks from the
+subdiagonal are (materialized) transpose of the corresponding superdiagonal blocks.
+
 # Examples
 ```jldoctest
 julia> dv = [1, 2, 3, 4]
@@ -44,6 +48,23 @@ julia> SymTridiagonal(dv, ev)
  7  2  8  ⋅
  ⋅  8  3  9
  ⋅  ⋅  9  4
+
+julia> A = SymTridiagonal(fill([1 2; 3 4], 3), fill([1 2; 3 4], 2));
+
+julia> A[1,1]
+2×2 Symmetric{Int64,Array{Int64,2}}:
+ 1  2
+ 2  4
+
+julia> A[1,2]
+2×2 Array{Int64,2}:
+ 1  2
+ 3  4
+
+julia> A[2,1]
+2×2 Array{Int64,2}:
+ 1  3
+ 2  4
 ```
 """
 SymTridiagonal(dv::V, ev::V) where {T,V<:AbstractVector{T}} = SymTridiagonal{T}(dv, ev)
@@ -56,8 +77,8 @@ end
 """
     SymTridiagonal(A::AbstractMatrix)
 
-Construct a symmetric tridiagonal matrix from the diagonal and
-first sub/super-diagonal, of the symmetric matrix `A`.
+Construct a symmetric tridiagonal matrix from the diagonal and first superdiagonal
+of the symmetric matrix `A`.
 
 # Examples
 ```jldoctest
@@ -72,11 +93,18 @@ julia> SymTridiagonal(A)
  1  2  ⋅
  2  4  5
  ⋅  5  6
+
+julia> B = reshape([[1 2; 2 3], [1 2; 3 4], [1 3; 2 4], [1 2; 2 3]], 2, 2);
+
+julia> SymTridiagonal(B)
+2×2 SymTridiagonal{Array{Int64,2},Array{Array{Int64,2},1}}:
+ [1 2; 2 3]  [1 3; 2 4]
+ [1 2; 3 4]  [1 2; 2 3]
 ```
 """
 function SymTridiagonal(A::AbstractMatrix)
-    if diag(A,1) == diag(A,-1)
-        SymTridiagonal(diag(A,0), diag(A,1))
+    if (diag(A, 1) == transpose.(diag(A, -1))) && all(issymmetric.(diag(A, 0)))
+        SymTridiagonal(diag(A, 0), diag(A, 1))
     else
         throw(ArgumentError("matrix is not symmetric; cannot convert to SymTridiagonal"))
     end
@@ -139,13 +167,13 @@ adjoint(S::SymTridiagonal) = Adjoint(S)
 Base.copy(S::Adjoint{<:Any,<:SymTridiagonal}) = SymTridiagonal(map(x -> copy.(adjoint.(x)), (S.parent.dv, S.parent.ev))...)
 Base.copy(S::Transpose{<:Any,<:SymTridiagonal}) = SymTridiagonal(map(x -> copy.(transpose.(x)), (S.parent.dv, S.parent.ev))...)
 
-function diag(M::SymTridiagonal, n::Integer=0)
+function diag(M::SymTridiagonal{<:Number}, n::Integer=0)
     # every branch call similar(..., ::Int) to make sure the
     # same vector type is returned independent of n
     absn = abs(n)
     if absn == 0
         return copyto!(similar(M.dv, length(M.dv)), M.dv)
-    elseif absn==1
+    elseif absn == 1
         return copyto!(similar(M.ev, length(M.ev)), M.ev)
     elseif absn <= size(M,1)
         return fill!(similar(M.dv, size(M,1)-absn), 0)
@@ -154,6 +182,22 @@ function diag(M::SymTridiagonal, n::Integer=0)
             "and at most $(size(M, 2)) for an $(size(M, 1))-by-$(size(M, 2)) matrix")))
     end
 end
+function diag(M::SymTridiagonal, n::Integer=0)
+    # every branch call similar(..., ::Int) to make sure the
+    # same vector type is returned independent of n
+    if n == 0
+        return copyto!(similar(M.dv, length(M.dv)), symmetric.(M.dv, :U))
+    elseif n == 1
+        return copyto!(similar(M.ev, length(M.ev)), M.ev)
+    elseif n == -1
+        return copyto!(similar(M.ev, length(M.ev)), transpose.(M.ev))
+    elseif n <= size(M,1)
+        throw(ArgumentError("requested diagonal contains undefined zeros of an array type"))
+    else
+        throw(ArgumentError(string("requested diagonal, $n, must be at least $(-size(M, 1)) ",
+            "and at most $(size(M, 2)) for an $(size(M, 1))-by-$(size(M, 2)) matrix")))
+    end
+end
 
 +(A::SymTridiagonal, B::SymTridiagonal) = SymTridiagonal(A.dv+B.dv, A.ev+B.ev)
 -(A::SymTridiagonal, B::SymTridiagonal) = SymTridiagonal(A.dv-B.dv, A.ev-B.ev)
@@ -390,9 +434,9 @@ function getindex(A::SymTridiagonal{T}, i::Integer, j::Integer) where T
         throw(BoundsError(A, (i,j)))
     end
     if i == j
-        return A.dv[i]
+        return symmetric(A.dv[i], :U)::symmetric_type(eltype(A.dv))
     elseif i == j + 1
-        return A.ev[j]
+        return copy(transpose(A.ev[j])) # materialized for type stability
     elseif i + 1 == j
         return A.ev[i]
     else

stdlib/LinearAlgebra/test/tridiag.jl

@@ -398,6 +398,16 @@ end
     end
 end
 
+@testset "SymTridiagonal block matrix" begin
+    M = [1 2; 2 4]
+    A = SymTridiagonal(fill(M, 3), fill(M, 2))
+    @test @inferred A[1,1] == Symmetric(M)
+    @test @inferred A[1,2] == M
+    @test @inferred A[2,1] == transpose(M)
+    @test @inferred diag(A, 1) == fill(M, 2)
+    @test @inferred diag(A, 0) == fill(Symmetric(M), 3)
+    @test @inferred diag(A, -1) == fill(transpose(M), 2)
+end
 
 @testset "Issue 12068" begin
     @test SymTridiagonal([1, 2], [0])^3 == [1 0; 0 8]