Make (Sym)Tridiag-Diagonal solves return TriDiagonal (JuliaLang#42744)

Author: dkarrasch

stdlib/LinearAlgebra/src/diagonal.jl

@@ -363,7 +363,7 @@ function _rdiv!(B::AbstractVecOrMat, A::AbstractVecOrMat, D::Diagonal)
     require_one_based_indexing(A)
     dd = D.diag
     m, n = size(A, 1), size(A, 2)
-    if (k = length(dd)) ≠ n
+    if (k = length(dd)) != n
         throw(DimensionMismatch("left hand side has $n columns but D is $k by $k"))
     end
     @inbounds for j in 1:n
@@ -395,7 +395,7 @@ end
 \(D::Diagonal, B::Diagonal) = ldiv!(similar(B, promote_op(\, eltype(D), eltype(B))), D, B)
 /(A::Diagonal, D::Diagonal) = _rdiv!(similar(A, promote_op(/, eltype(A), eltype(D))), A, D)
 function _rdiv!(Dc::Diagonal, Db::Diagonal, Da::Diagonal)
-    n, k = length(Db.diag), length(Db.diag)
+    n, k = length(Db.diag), length(Da.diag)
     n == k || throw(DimensionMismatch("left hand side has $n columns but D is $k by $k"))
     j = findfirst(iszero, Da.diag)
     isnothing(j) || throw(SingularException(j))
@@ -404,6 +404,88 @@ function _rdiv!(Dc::Diagonal, Db::Diagonal, Da::Diagonal)
 end
 ldiv!(Dc::Diagonal, Da::Diagonal, Db::Diagonal) = Diagonal(ldiv!(Dc.diag, Da, Db.diag))
 
+# optimizations for (Sym)Tridiagonal and Diagonal
+@propagate_inbounds _getudiag(T::Tridiagonal, i) = T.du[i]
+@propagate_inbounds _getudiag(S::SymTridiagonal, i) = S.ev[i]
+@propagate_inbounds _getdiag(T::Tridiagonal, i) = T.d[i]
+@propagate_inbounds _getdiag(S::SymTridiagonal, i) = symmetric(S.dv[i], :U)::symmetric_type(eltype(S.dv))
+@propagate_inbounds _getldiag(T::Tridiagonal, i) = T.dl[i]
+@propagate_inbounds _getldiag(S::SymTridiagonal, i) = transpose(S.ev[i])
+
+function (\)(D::Diagonal, S::SymTridiagonal)
+    T = promote_op(\, eltype(D), eltype(S))
+    du = similar(S.ev, T, max(length(S.dv)-1, 0))
+    d  = similar(S.dv, T, length(S.dv))
+    dl = similar(S.ev, T, max(length(S.dv)-1, 0))
+    ldiv!(Tridiagonal(dl, d, du), D, S)
+end
+(\)(D::Diagonal, T::Tridiagonal) = ldiv!(similar(T, promote_op(\, eltype(D), eltype(T))), D, T)
+function ldiv!(T::Tridiagonal, D::Diagonal, S::Union{SymTridiagonal,Tridiagonal})
+    m = size(S, 1)
+    dd = D.diag
+    if (k = length(dd)) != m
+        throw(DimensionMismatch("diagonal matrix is $k by $k but right hand side has $m rows"))
+    end
+    if length(T.d) != m
+        throw(DimensionMismatch("target matrix size $(size(T)) does not match input matrix size $(size(S))"))
+    end
+    m == 0 && return T
+    j = findfirst(iszero, dd)
+    isnothing(j) || throw(SingularException(j))
+    ddj = dd[1]
+    T.d[1] = ddj \ _getdiag(S, 1)
+    @inbounds if m > 1
+        T.du[1] = ddj \ _getudiag(S, 1)
+        for j in 2:m-1
+            ddj = dd[j]
+            T.dl[j-1] = ddj \ _getldiag(S, j-1)
+            T.d[j]  = ddj \ _getdiag(S, j)
+            T.du[j] = ddj \ _getudiag(S, j)
+        end
+        ddj = dd[m]
+        T.dl[m-1] = ddj \ _getldiag(S, m-1)
+        T.d[m] = ddj \ _getdiag(S, m)
+    end
+    return T
+end
+
+function (/)(S::SymTridiagonal, D::Diagonal)
+    T = promote_op(\, eltype(D), eltype(S))
+    du = similar(S.ev, T, max(length(S.dv)-1, 0))
+    d  = similar(S.dv, T, length(S.dv))
+    dl = similar(S.ev, T, max(length(S.dv)-1, 0))
+    _rdiv!(Tridiagonal(dl, d, du), S, D)
+end
+(/)(T::Tridiagonal, D::Diagonal) = _rdiv!(similar(T, promote_op(/, eltype(T), eltype(D))), T, D)
+function _rdiv!(T::Tridiagonal, S::Union{SymTridiagonal,Tridiagonal}, D::Diagonal)
+    n = size(S, 2)
+    dd = D.diag
+    if (k = length(dd)) != n
+        throw(DimensionMismatch("left hand side has $n columns but D is $k by $k"))
+    end
+    if length(T.d) != n
+        throw(DimensionMismatch("target matrix size $(size(T)) does not match input matrix size $(size(S))"))
+    end
+    n == 0 && return T
+    j = findfirst(iszero, dd)
+    isnothing(j) || throw(SingularException(j))
+    ddj = dd[1]
+    T.d[1] = _getdiag(S, 1) / ddj
+    @inbounds if n > 1
+        T.dl[1] = _getldiag(S, 1) / ddj
+        for j in 2:n-1
+            ddj = dd[j]
+            T.dl[j] = _getldiag(S, j) / ddj
+            T.d[j] = _getdiag(S, j) / ddj
+            T.du[j-1] = _getudiag(S, j-1) / ddj
+        end
+        ddj = dd[n]
+        T.d[n] = _getdiag(S, n) / ddj
+        T.du[n-1] = _getudiag(S, n-1) / ddj
+    end
+    return T
+end
+
 # Optimizations for [l/r]mul!, l/rdiv!, *, / and \ between Triangular and Diagonal.
 # These functions are generally more efficient if we calculate the whole data field.
 # The following code implements them in a unified pattern to avoid missing.
@@ -623,7 +705,7 @@ dot(A::AbstractMatrix, B::Diagonal) = conj(dot(B, A))
 
 function _mapreduce_prod(f, x, D::Diagonal, y)
     if isempty(x) && isempty(D) && isempty(y)
-        return zero(Base.promote_op(f, eltype(x), eltype(D), eltype(y)))
+        return zero(promote_op(f, eltype(x), eltype(D), eltype(y)))
     else
         return mapreduce(t -> f(t[1], t[2], t[3]), +, zip(x, D.diag, y))
     end

stdlib/LinearAlgebra/test/diagonal.jl

@@ -3,7 +3,7 @@
 module TestDiagonal
 
 using Test, LinearAlgebra, SparseArrays, Random
-using LinearAlgebra: mul!, mul!, rmul!, lmul!, ldiv!, rdiv!, BlasFloat, BlasComplex, SingularException
+using LinearAlgebra: BlasFloat, BlasComplex
 
 n=12 #Size of matrix problem to test
 Random.seed!(1)
@@ -776,6 +776,59 @@ end
     end
 end
 
+@testset "(Sym)Tridiagonal division by Diagonal" begin
+    for K in (5, 1), elty in (Float64, ComplexF32), overlength in (1, 0)
+        S = SymTridiagonal(randn(elty, K), randn(elty, K-overlength))
+        T = Tridiagonal(randn(elty, K-1), randn(elty, K), randn(elty, K-1))
+        D = Diagonal(randn(elty, K))
+        D0 = Diagonal(zeros(elty, K))
+        @test (D \ S)::Tridiagonal{elty} == Tridiagonal(Matrix(D) \ Matrix(S))
+        @test (D \ T)::Tridiagonal{elty} == Tridiagonal(Matrix(D) \ Matrix(T))
+        @test (S / D)::Tridiagonal{elty} == Tridiagonal(Matrix(S) / Matrix(D))
+        @test (T / D)::Tridiagonal{elty} == Tridiagonal(Matrix(T) / Matrix(D))
+        @test_throws SingularException D0 \ S
+        @test_throws SingularException D0 \ T
+        @test_throws SingularException S / D0
+        @test_throws SingularException T / D0
+    end
+    # 0-length case
+    S = SymTridiagonal(Float64[], Float64[])
+    T = Tridiagonal(Float64[], Float64[], Float64[])
+    D = Diagonal(Float64[])
+    @test (D \ S)::Tridiagonal{Float64} == T
+    @test (D \ T)::Tridiagonal{Float64} == T
+    @test (S / D)::Tridiagonal{Float64} == T
+    @test (T / D)::Tridiagonal{Float64} == T
+    # matrix eltype case
+    K = 5
+    for elty in (Float64, ComplexF32), overlength in (1, 0)
+        S = SymTridiagonal([rand(elty, 2, 2) for _ in 1:K], [rand(elty, 2, 2) for _ in 1:K-overlength])
+        T = Tridiagonal([rand(elty, 2, 2) for _ in 1:K-1], [rand(elty, 2, 2) for _ in 1:K], [rand(elty, 2, 2) for _ in 1:K-1])
+        D = Diagonal(randn(elty, K))
+        SM = fill(zeros(elty, 2, 2), K, K)
+        TM = copy(SM)
+        SM[1,1] = S[1,1]; TM[1,1] = T[1,1]
+        for j in 2:K
+            SM[j,j-1] = S[j,j-1]; SM[j,j] = S[j,j]; SM[j-1,j] = S[j-1,j]
+            TM[j,j-1] = T[j,j-1]; TM[j,j] = T[j,j]; TM[j-1,j] = T[j-1,j]
+        end
+        for (M, Mm) in ((S, SM), (T, TM))
+            DS = D \ M
+            @test DS isa Tridiagonal
+            DM = D \ Mm
+            for i in -1:1; @test diag(DS, i) ≈ diag(DM, i) end
+        end
+    end
+    # eltype promotion case
+    S = SymTridiagonal(rand(-20:20, K), rand(-20:20, K-1))
+    T = Tridiagonal(rand(-20:20, K-1), rand(-20:20, K), rand(-20:20, K-1))
+    D = Diagonal(rand(1:20, K))
+    @test (D \ S)::Tridiagonal{Float64} == Tridiagonal(Matrix(D) \ Matrix(S))
+    @test (D \ T)::Tridiagonal{Float64} == Tridiagonal(Matrix(D) \ Matrix(T))
+    @test (S / D)::Tridiagonal{Float64} == Tridiagonal(Matrix(S) / Matrix(D))
+    @test (T / D)::Tridiagonal{Float64} == Tridiagonal(Matrix(T) / Matrix(D))
+end
+
 @testset "eigenvalue sorting" begin
     D = Diagonal([0.4, 0.2, -1.3])
     @test eigvals(D) == eigen(D).values == [0.4, 0.2, -1.3] # not sorted by default