Compute real matrix logarithm and matrix square root using real arithmetic (JuliaLang#39973)

* Add failing test

* Add sylvester methods for small matrices

* Add 2x2 real matrix square root

* Add real square root of quasitriangular matrix

* Simplify 2x2 real square root

* Rename functions to use quasitriu

* Avoid NaNs when eigenvalues are all zero

* Reuse ranges

* Add clarifying comments

* Unify real and complex matrix square root

* Add reference for real sqrt

* Move quasitriu auxiliary functions to triangular.jl

* Ensure loops are type-stable and use simd

* Remove duplicate computation

* Correctly promote for dimensionful A

* Use simd directive

* Test that UpperTriangular is returned by sqrt

* Test sqrt for UnitUpperTriangular

* Test that return type is complex when input type is

* Test that output is complex when input is

* Add failing test

* Separate type-stable from type-unstable part

* Use generic sqrt_quasitriu for sqrt triu

* Avoid redundant matmul

* Clarify comment

* Return complex output for complex input

* Call log_quasitriu

* Add failing test for log type-inferrability

* Realify or complexify as necessary

* Call sqrt_quasitriu directly

* Refactor sqrt_diag!

* Simplify utility function

* Add comment

* Compute accurate block-diagonal

* Compute superdiagonal for quasi triu A0

* Compute accurate block superdiagonal

* Avoid full LU decomposition in inner loop

* Avoid promotion to improve type-stability

* Modify return type if necessary

* Clarify comment

* Add comments

* Call log_quasitriu on quasitriu matrices

* Document quasi-triangular algorithm

* Remove test

This matrix has eigenvalues to close to zero that its eltype is not stable

* Rearrange definition

* Add compatibility for unit triangular matrices

* Release constraints on tests

* Separate copying of A from log computation

* Revert "Separate copying of A from log computation"

This reverts commit 23becc5.

* Use Givens rotations

* Compute Schur in-place when possible

* Always allocate a copy

* Fix block indexing

* Compute sqrt in-place

* Overwrite AmI

* Reduce allocations in Pade approximation

* Use T

* Don't unnecessarily unwrap

* Test remaining log branches

* Add additional matrix square root tests

* Separate type-unstable from type-stable part

This substantially reduces allocations for some reason

* Use Ref instead of a Vector

* Eliminate allocation in checksquare

* Refactor param choosing code to own function

* Comment section

* Use more descriptive variable name

* Reuse temporaries

* Add reference

* More accurately describe condition

Author: sethaxen

stdlib/LinearAlgebra/src/dense.jl

@@ -679,7 +679,7 @@ function rcswap!(i::Integer, j::Integer, X::StridedMatrix{<:Number})
 end
 
 """
-    log(A{T}::StridedMatrix{T})
+    log(A::StridedMatrix)
 
 If `A` has no negative real eigenvalue, compute the principal matrix logarithm of `A`, i.e.
 the unique matrix ``X`` such that ``e^X = A`` and ``-\\pi < Im(\\lambda) < \\pi`` for all
@@ -688,9 +688,10 @@ matrix function is returned whenever possible.
 
 If `A` is symmetric or Hermitian, its eigendecomposition ([`eigen`](@ref)) is
 used, if `A` is triangular an improved version of the inverse scaling and squaring method is
-employed (see [^AH12] and [^AHR13]). For general matrices, the complex Schur form
-([`schur`](@ref)) is computed and the triangular algorithm is used on the
-triangular factor.
+employed (see [^AH12] and [^AHR13]). If `A` is real with no negative eigenvalues, then
+the real Schur form is computed. Otherwise, the complex Schur form is computed. Then
+the upper (quasi-)triangular algorithm in [^AHR13] is used on the upper (quasi-)triangular
+factor.
 
 [^AH12]: Awad H. Al-Mohy and Nicholas J. Higham, "Improved inverse  scaling and squaring algorithms for the matrix logarithm", SIAM Journal on Scientific Computing, 34(4), 2012, C153-C169. [doi:10.1137/110852553](https://doi.org/10.1137/110852553)
 
@@ -713,27 +714,28 @@ function log(A::StridedMatrix)
     # If possible, use diagonalization
     if ishermitian(A)
         logHermA = log(Hermitian(A))
-        return isa(logHermA, Hermitian) ? copytri!(parent(logHermA), 'U', true) : parent(logHermA)
-    end
-
-    # Use Schur decomposition
-    n = checksquare(A)
-    if istriu(A)
-        return triu!(parent(log(UpperTriangular(complex(A)))))
-    else
-        if isreal(A)
-            SchurF = schur(real(A))
+        return ishermitian(logHermA) ? copytri!(parent(logHermA), 'U', true) : parent(logHermA)
+    elseif istriu(A)
+        return triu!(parent(log(UpperTriangular(A))))
+    elseif isreal(A)
+        SchurF = schur(real(A))
+        if istriu(SchurF.T)
+            logA = SchurF.Z * log(UpperTriangular(SchurF.T)) * SchurF.Z'
         else
-            SchurF = schur(A)
-        end
-        if !istriu(SchurF.T)
-            SchurS = schur(complex(SchurF.T))
-            logT = SchurS.Z * log(UpperTriangular(SchurS.T)) * SchurS.Z'
-            return SchurF.Z * logT * SchurF.Z'
-        else
-            R = log(UpperTriangular(complex(SchurF.T)))
-            return SchurF.Z * R * SchurF.Z'
+            # real log exists whenever all eigenvalues are positive
+            is_log_real = !any(x -> isreal(x) && real(x) ≤ 0, SchurF.values)
+            if is_log_real
+                logA = SchurF.Z * log_quasitriu(SchurF.T) * SchurF.Z'
+            else
+                SchurS = schur!(complex(SchurF.T))
+                Z = SchurF.Z * SchurS.Z
+                logA = Z * log(UpperTriangular(SchurS.T)) * Z'
+            end
         end
+        return eltype(A) <: Complex ? complex(logA) : logA
+    else
+        SchurF = schur(A)
+        return SchurF.vectors * log(UpperTriangular(SchurF.T)) * SchurF.vectors'
     end
 end
 
@@ -755,13 +757,21 @@ defaults to machine precision scaled by `size(A,1)`.
 Otherwise, the square root is determined by means of the
 Björck-Hammarling method [^BH83], which computes the complex Schur form ([`schur`](@ref))
 and then the complex square root of the triangular factor.
+If a real square root exists, then an extension of this method [^H87] that computes the real
+Schur form and then the real square root of the quasi-triangular factor is instead used.
 
 [^BH83]:
 
     Åke Björck and Sven Hammarling, "A Schur method for the square root of a matrix",
     Linear Algebra and its Applications, 52-53, 1983, 127-140.
     [doi:10.1016/0024-3795(83)80010-X](https://doi.org/10.1016/0024-3795(83)80010-X)
 
+[^H87]:
+
+    Nicholas J. Higham, "Computing real square roots of a real matrix",
+    Linear Algebra and its Applications, 88-89, 1987, 405-430.
+    [doi:10.1016/0024-3795(87)90118-2](https://doi.org/10.1016/0024-3795(87)90118-2)
+
 # Examples
 ```jldoctest
 julia> A = [4 0; 0 4]
@@ -775,31 +785,32 @@ julia> sqrt(A)
  0.0  2.0
 ```
 """
-function sqrt(A::StridedMatrix{<:Real})
-    if issymmetric(A)
-        return copytri!(parent(sqrt(Symmetric(A))), 'U')
-    end
-    n = checksquare(A)
-    if istriu(A)
-        return triu!(parent(sqrt(UpperTriangular(A))))
-    else
-        SchurF = schur(complex(A))
-        R = triu!(parent(sqrt(UpperTriangular(SchurF.T)))) # unwrapping unnecessary?
-        return SchurF.vectors * R * SchurF.vectors'
-    end
-end
-function sqrt(A::StridedMatrix{<:Complex})
+function sqrt(A::StridedMatrix{T}) where {T<:Union{Real,Complex}}
     if ishermitian(A)
         sqrtHermA = sqrt(Hermitian(A))
-        return isa(sqrtHermA, Hermitian) ? copytri!(parent(sqrtHermA), 'U', true) : parent(sqrtHermA)
-    end
-    n = checksquare(A)
-    if istriu(A)
+        return ishermitian(sqrtHermA) ? copytri!(parent(sqrtHermA), 'U', true) : parent(sqrtHermA)
+    elseif istriu(A)
         return triu!(parent(sqrt(UpperTriangular(A))))
+    elseif isreal(A)
+        SchurF = schur(real(A))
+        if istriu(SchurF.T)
+            sqrtA = SchurF.Z * sqrt(UpperTriangular(SchurF.T)) * SchurF.Z'
+        else
+            # real sqrt exists whenever no eigenvalues are negative
+            is_sqrt_real = !any(x -> isreal(x) && real(x) < 0, SchurF.values)
+            # sqrt_quasitriu uses LAPACK functions for non-triu inputs
+            if typeof(sqrt(zero(T))) <: BlasFloat && is_sqrt_real
+                sqrtA = SchurF.Z * sqrt_quasitriu(SchurF.T) * SchurF.Z'
+            else
+                SchurS = schur!(complex(SchurF.T))
+                Z = SchurF.Z * SchurS.Z
+                sqrtA = Z * sqrt(UpperTriangular(SchurS.T)) * Z'
+            end
+        end
+        return eltype(A) <: Complex ? complex(sqrtA) : sqrtA
     else
         SchurF = schur(A)
-        R = triu!(parent(sqrt(UpperTriangular(SchurF.T)))) # unwrapping unnecessary?
-        return SchurF.vectors * R * SchurF.vectors'
+        return SchurF.vectors * sqrt(UpperTriangular(SchurF.T)) * SchurF.vectors'
     end
 end
 
@@ -1526,6 +1537,34 @@ function sylvester(A::StridedMatrix{T},B::StridedMatrix{T},C::StridedMatrix{T})
 end
 sylvester(A::StridedMatrix{T}, B::StridedMatrix{T}, C::StridedMatrix{T}) where {T<:Integer} = sylvester(float(A), float(B), float(C))
 
+Base.@propagate_inbounds function _sylvester_2x1!(A, B, C)
+    b = B[1]
+    a21, a12 = A[2, 1], A[1, 2]
+    m11 = b + A[1, 1]
+    m22 = b + A[2, 2]
+    d = m11 * m22 - a12 * a21
+    c1, c2 = C
+    C[1] = (a12 * c2 - m22 * c1) / d
+    C[2] = (a21 * c1 - m11 * c2) / d
+    return C
+end
+Base.@propagate_inbounds function _sylvester_1x2!(A, B, C)
+    a = A[1]
+    b21, b12 = B[2, 1], B[1, 2]
+    m11 = a + B[1, 1]
+    m22 = a + B[2, 2]
+    d = m11 * m22 - b21 * b12
+    c1, c2 = C
+    C[1] = (b21 * c2 - m22 * c1) / d
+    C[2] = (b12 * c1 - m11 * c2) / d
+    return C
+end
+function _sylvester_2x2!(A, B, C)
+    _, scale = LAPACK.trsyl!('N', 'N', A, B, C)
+    rmul!(C, -inv(scale))
+    return C
+end
+
 sylvester(a::Union{Real,Complex}, b::Union{Real,Complex}, c::Union{Real,Complex}) = -c / (a + b)
 
 # AX + XA' + C = 0

stdlib/LinearAlgebra/src/lapack.jl

@@ -6449,15 +6449,15 @@ for (fn, elty, relty) in ((:dtrsyl_, :Float64, :Float64),
                         B::AbstractMatrix{$elty}, C::AbstractMatrix{$elty}, isgn::Int=1)
             require_one_based_indexing(A, B, C)
             chkstride1(A, B, C)
-            m, n = checksquare(A, B)
+            m, n = checksquare(A), checksquare(B)
             lda = max(1, stride(A, 2))
             ldb = max(1, stride(B, 2))
             m1, n1 = size(C)
             if m != m1 || n != n1
                 throw(DimensionMismatch("dimensions of A, ($m,$n), and C, ($m1,$n1), must match"))
             end
             ldc = max(1, stride(C, 2))
-            scale = Vector{$relty}(undef, 1)
+            scale = Ref{$relty}()
             info  = Ref{BlasInt}()
             ccall((@blasfunc($fn), libblastrampoline), Cvoid,
                 (Ref{UInt8}, Ref{UInt8}, Ref{BlasInt}, Ref{BlasInt}, Ref{BlasInt},
@@ -6467,7 +6467,7 @@ for (fn, elty, relty) in ((:dtrsyl_, :Float64, :Float64),
                 A, lda, B, ldb, C, ldc,
                 scale, info, 1, 1)
             chklapackerror(info[])
-            C, scale[1]
+            C, scale[]
         end
     end
 end

stdlib/LinearAlgebra/src/triangular.jl

@@ -148,6 +148,14 @@ end
             asym = a + a' # symmetric indefinite
             asymsq = sqrt(asym)
             @test asymsq*asymsq ≈ asym
+            if eltype(a) <: Real  # real square root
+                apos = a * a
+                @test sqrt(apos)^2 ≈ apos
+                @test eltype(sqrt(apos)) <: Real
+                # test that real but Complex input produces Complex output
+                @test sqrt(complex(apos)) ≈ sqrt(apos)
+                @test eltype(sqrt(complex(apos))) <: Complex
+            end
         end
 
         @testset "Powers" begin
@@ -708,9 +716,6 @@ end
     A11 = convert(Matrix{elty}, [3 2; -5 -3])
     @test exp(log(A11)) ≈ A11
 
-    A12 = convert(Matrix{elty}, [1 -1; 1 -1])
-    @test typeof(log(A12)) == Array{ComplexF64, 2}
-
     A13 = convert(Matrix{elty}, [2 0; 0 2])
     @test typeof(log(A13)) == Array{elty, 2}
 
@@ -723,6 +728,7 @@ end
                                     0.2310490602 0.1969543025 1.363756107])
     @test log(A1) ≈ logA1
     @test exp(log(A1)) ≈ A1
+    @test typeof(log(A1)) == Matrix{elty}
 
     A4  = convert(Matrix{elty}, [1/2 1/3 1/4 1/5+eps();
                                  1/3 1/4 1/5 1/6;
@@ -734,6 +740,134 @@ end
                                     0.2414170219 0.5865285289 3.318413247 -5.444632124])
     @test log(A4) ≈ logA4
     @test exp(log(A4)) ≈ A4
+    @test typeof(log(A4)) == Matrix{elty}
+
+    # real triu matrix
+    A5  = convert(Matrix{elty}, [1 2 3; 0 4 5; 0 0 6])  # triu
+    logA5 = convert(Matrix{elty}, [0.0 0.9241962407465937 0.5563245488984037;
+                                   0.0 1.3862943611198906 1.0136627702704109;
+                                   0.0 0.0 1.791759469228055])
+    @test log(A5) ≈ logA5
+    @test exp(log(A5)) ≈ A5
+    @test typeof(log(A5)) == Matrix{elty}
+
+    # real quasitriangular schur form with 2 2x2 blocks, 2 1x1 blocks, and all positive eigenvalues
+    A6 = convert(Matrix{elty}, [2 3 2 2 3 1;
+                                1 3 3 2 3 1;
+                                3 3 3 1 1 2;
+                                2 1 2 2 2 2;
+                                1 1 2 2 3 1;
+                                2 2 2 2 1 3])
+    @test exp(log(A6)) ≈ A6
+    @test typeof(log(A6)) == Matrix{elty}
+
+    # real quasitriangular schur form with a negative eigenvalue
+    A7 = convert(Matrix{elty}, [1 3 3 2 2 2;
+                                1 2 1 3 1 2;
+                                3 1 2 3 2 1;
+                                3 1 2 2 2 1;
+                                3 1 3 1 2 1;
+                                1 1 3 1 1 3])
+    @test exp(log(A7)) ≈ A7
+    @test typeof(log(A7)) == Matrix{complex(elty)}
+
+    if elty <: Complex
+        A8 = convert(Matrix{elty}, [1 + 1im 1 + 1im 1 - 1im;
+                                    1 + 1im -1 + 1im 1 + 1im;
+                                    1 - 1im 1 + 1im -1 - 1im])
+        logA8 = convert(
+            Matrix{elty},
+            [0.9478628953131517 + 1.3725201223387407im -0.2547157147532057 + 0.06352318334299434im 0.8560050197863862 - 1.0471975511965979im;
+             -0.2547157147532066 + 0.06352318334299467im -0.16285783922644065 + 0.2617993877991496im 0.2547157147532063 + 2.1579182857361894im;
+             0.8560050197863851 - 1.0471975511965974im 0.25471571475320665 + 2.1579182857361903im 0.9478628953131519 - 0.8489213467404436im],
+        )
+        @test log(A8) ≈ logA8
+        @test exp(log(A8)) ≈ A8
+        @test typeof(log(A8)) == Matrix{elty}
+    end
+end
+
+@testset "matrix logarithm is type-inferrable" for elty in (Float32,Float64,ComplexF32,ComplexF64)
+    A1 = randn(elty, 4, 4)
+    @inferred Union{Matrix{elty},Matrix{complex(elty)}} log(A1)
+end
+
+@testset "Additional matrix square root tests" for elty in (Float64, ComplexF64)
+    A11 = convert(Matrix{elty}, [3 2; -5 -3])
+    @test sqrt(A11)^2 ≈ A11
+
+    A13 = convert(Matrix{elty}, [2 0; 0 2])
+    @test typeof(sqrt(A13)) == Array{elty, 2}
+
+    T = elty == Float64 ? Symmetric : Hermitian
+    @test typeof(sqrt(T(A13))) == T{elty, Array{elty, 2}}
+
+    A1  = convert(Matrix{elty}, [4 2 0; 1 4 1; 1 1 4])
+    sqrtA1 = convert(Matrix{elty}, [1.971197119306979 0.5113118387140085 -0.03301921523780871;
+                                   0.23914631173809942 1.9546875116880718 0.2556559193570036;
+                                   0.23914631173810008 0.22263670411919556 1.9877067269258815])
+    @test sqrt(A1) ≈ sqrtA1
+    @test sqrt(A1)^2 ≈ A1
+    @test typeof(sqrt(A1)) == Matrix{elty}
+
+    A4  = convert(Matrix{elty}, [1/2 1/3 1/4 1/5+eps();
+                                 1/3 1/4 1/5 1/6;
+                                 1/4 1/5 1/6 1/7;
+                                 1/5 1/6 1/7 1/8])
+                                 sqrtA4 = convert(
+        Matrix{elty},
+        [0.590697761556362 0.3055006800405779 0.19525404749300546 0.14007621469988107;
+         0.30550068004057784 0.2825388389385975 0.21857572599211642 0.17048692323164674;
+         0.19525404749300565 0.21857572599211622 0.21155429252242863 0.18976816626246887;
+         0.14007621469988046 0.17048692323164724 0.1897681662624689 0.20075085592778794],
+    )
+    @test sqrt(A4) ≈ sqrtA4
+    @test sqrt(A4)^2 ≈ A4
+    @test typeof(sqrt(A4)) == Matrix{elty}
+
+    # real triu matrix
+    A5  = convert(Matrix{elty}, [1 2 3; 0 4 5; 0 0 6])  # triu
+    sqrtA5 = convert(Matrix{elty}, [1.0 0.6666666666666666 0.6525169217864183;
+                                   0.0 2.0 1.1237243569579454;
+                                   0.0 0.0 2.449489742783178])
+    @test sqrt(A5) ≈ sqrtA5
+    @test sqrt(A5)^2 ≈ A5
+    @test typeof(sqrt(A5)) == Matrix{elty}
+
+    # real quasitriangular schur form with 2 2x2 blocks, 2 1x1 blocks, and all positive eigenvalues
+    A6 = convert(Matrix{elty}, [2 3 2 2 3 1;
+                                1 3 3 2 3 1;
+                                3 3 3 1 1 2;
+                                2 1 2 2 2 2;
+                                1 1 2 2 3 1;
+                                2 2 2 2 1 3])
+    @test sqrt(A6)^2 ≈ A6
+    @test typeof(sqrt(A6)) == Matrix{elty}
+
+    # real quasitriangular schur form with a negative eigenvalue
+    A7 = convert(Matrix{elty}, [1 3 3 2 2 2;
+                                1 2 1 3 1 2;
+                                3 1 2 3 2 1;
+                                3 1 2 2 2 1;
+                                3 1 3 1 2 1;
+                                1 1 3 1 1 3])
+    @test sqrt(A7)^2 ≈ A7
+    @test typeof(sqrt(A7)) == Matrix{complex(elty)}
+
+    if elty <: Complex
+        A8 = convert(Matrix{elty}, [1 + 1im 1 + 1im 1 - 1im;
+                                    1 + 1im -1 + 1im 1 + 1im;
+                                    1 - 1im 1 + 1im -1 - 1im])
+        sqrtA8 = convert(
+            Matrix{elty},
+            [1.2559748527474284 + 0.6741878819930323im 0.20910077991005582 + 0.24969165051825476im 0.591784212275146 - 0.6741878819930327im;
+             0.2091007799100553 + 0.24969165051825515im 0.3320953202361413 + 0.2915044496279425im 0.33209532023614136 + 1.0568713143581219im;
+             0.5917842122751455 - 0.674187881993032im 0.33209532023614147 + 1.0568713143581223im 0.7147787526012315 - 0.6323750828833452im],
+        )
+        @test sqrt(A8) ≈ sqrtA8
+        @test sqrt(A8)^2 ≈ A8
+        @test typeof(sqrt(A8)) == Matrix{elty}
+    end
 end
 
 @testset "issue #7181" begin

stdlib/LinearAlgebra/test/dense.jl

@@ -198,7 +198,7 @@ for elty1 in (Float32, Float64, BigFloat, ComplexF32, ComplexF64, Complex{BigFlo
         end
 
         #exp/log
-        if (elty1 == Float64 || elty1 == ComplexF64) && (t1 == UpperTriangular || t1 == LowerTriangular)
+        if elty1 ∈ (Float32,Float64,ComplexF32,ComplexF64)
             @test exp(Matrix(log(A1))) ≈ A1
         end
 
@@ -496,10 +496,40 @@ end
 # Matrix square root
 Atn = UpperTriangular([-1 1 2; 0 -2 2; 0 0 -3])
 Atp = UpperTriangular([1 1 2; 0 2 2; 0 0 3])
+Atu = UnitUpperTriangular([1 1 2; 0 1 2; 0 0 1])
 @test sqrt(Atn) |> t->t*t ≈ Atn
+@test sqrt(Atn) isa UpperTriangular
 @test typeof(sqrt(Atn)[1,1]) <: Complex
 @test sqrt(Atp) |> t->t*t ≈ Atp
+@test sqrt(Atp) isa UpperTriangular
 @test typeof(sqrt(Atp)[1,1]) <: Real
+@test typeof(sqrt(complex(Atp))[1,1]) <: Complex
+@test sqrt(Atu) |> t->t*t ≈ Atu
+@test sqrt(Atu) isa UnitUpperTriangular
+@test typeof(sqrt(Atu)[1,1]) <: Real
+@test typeof(sqrt(complex(Atu))[1,1]) <: Complex
+
+@testset "check matrix logarithm type-inferrable" for elty in (Float32,Float64,ComplexF32,ComplexF64)
+    A = UpperTriangular(exp(triu(randn(elty, n, n))))
+    @inferred Union{typeof(A),typeof(complex(A))} log(A)
+    @test exp(Matrix(log(A))) ≈ A
+    if elty <: Real
+        @test typeof(log(A)) <: UpperTriangular{elty}
+        @test typeof(log(complex(A))) <: UpperTriangular{complex(elty)}
+        @test isreal(log(complex(A)))
+        @test log(complex(A)) ≈ log(A)
+    end
+
+    Au = UnitUpperTriangular(exp(triu(randn(elty, n, n), 1)))
+    @inferred Union{typeof(A),typeof(complex(A))} log(Au)
+    @test exp(Matrix(log(Au))) ≈ Au
+    if elty <: Real
+        @test typeof(log(Au)) <: UpperTriangular{elty}
+        @test typeof(log(complex(Au))) <: UpperTriangular{complex(elty)}
+        @test isreal(log(complex(Au)))
+        @test log(complex(Au)) ≈ log(Au)
+    end
+end
 
 Areal   = randn(n, n)/2
 Aimg    = randn(n, n)/2