Exponential

ext/MatrixAlgebraKitGenericSchurExt.jl

@@ -22,4 +22,8 @@ function MatrixAlgebraKit.eig_vals!(A::AbstractMatrix, D, ::GS_QRIteration)
     return GenericSchur.eigvals!(A)
 end
 
+function MatrixAlgebraKit.default_exponential_algorithm(E::Type{T}; kwargs...) where {T <: StridedMatrix{<:Union{BigFloat, Complex{BigFloat}}}}
+    return ExponentialViaEig(GS_QRIteration(; kwargs...))
+end
+
 end

src/MatrixAlgebraKit.jl

@@ -29,13 +29,15 @@ export left_polar, right_polar
 export left_polar!, right_polar!
 export left_orth, right_orth, left_null, right_null
 export left_orth!, right_orth!, left_null!, right_null!
+export exponential, exponential!
 
 export LAPACK_HouseholderQR, LAPACK_HouseholderLQ, LAPACK_Simple, LAPACK_Expert,
     LAPACK_QRIteration, LAPACK_Bisection, LAPACK_MultipleRelativelyRobustRepresentations,
     LAPACK_DivideAndConquer, LAPACK_Jacobi
 export GLA_HouseholderQR, GLA_QRIteration, GS_QRIteration
 export LQViaTransposedQR
 export PolarViaSVD, PolarNewton
+export ExponentialViaLA, ExponentialViaEig, ExponentialViaEigh
 export DiagonalAlgorithm
 export NativeBlocked
 export CUSOLVER_Simple, CUSOLVER_HouseholderQR, CUSOLVER_QRIteration, CUSOLVER_SVDPolar,
@@ -92,6 +94,7 @@ include("interface/gen_eig.jl")
 include("interface/schur.jl")
 include("interface/polar.jl")
 include("interface/orthnull.jl")
+include("interface/exponential.jl")
 
 include("implementations/projections.jl")
 include("implementations/truncation.jl")
@@ -104,6 +107,7 @@ include("implementations/gen_eig.jl")
 include("implementations/schur.jl")
 include("implementations/polar.jl")
 include("implementations/orthnull.jl")
+include("implementations/exponential.jl")
 
 include("pullbacks/qr.jl")
 include("pullbacks/lq.jl")

src/implementations/exponential.jl

@@ -0,0 +1,66 @@
+# Inputs
+# ------
+function copy_input(::typeof(exponential), A::AbstractMatrix)
+    return copy!(similar(A, float(eltype(A))), A)
+end
+
+copy_input(::typeof(exponential), A::Diagonal) = copy(A)
+
+function check_input(::typeof(exponential!), A::AbstractMatrix, expA::AbstractMatrix, alg::AbstractAlgorithm)
+    m, n = size(A)
+    m == n || throw(DimensionMismatch("square input matrix expected. Got ($m,$n)"))
+    @check_size(expA, (m, m))
+    return @check_scalar(expA, A)
+end
+
+function check_input(::typeof(exponential!), A::AbstractMatrix, expA::AbstractMatrix, ::DiagonalAlgorithm)
+    m, n = size(A)
+    @assert m == n && isdiag(A)
+    @assert expA isa Diagonal
+    @check_size(expA, (m, m))
+    @check_scalar(expA, A)
+    return nothing
+end
+
+# Outputs
+# -------
+function initialize_output(::typeof(exponential!), A::AbstractMatrix, ::AbstractAlgorithm)
+    n = size(A, 1) # square check will happen later
+    expA = similar(A, (n, n))
+    return expA
+end
+
+function initialize_output(::typeof(exponential!), A::Diagonal, ::DiagonalAlgorithm)
+    return similar(A)
+end
+
+# Implementation
+# --------------
+function exponential!(A::AbstractMatrix, expA::AbstractMatrix, alg::ExponentialViaLA)
+    copyto!(expA, LinearAlgebra.exp(A))
+    return expA
+end
+
+function exponential!(A::AbstractMatrix, expA::AbstractMatrix, alg::ExponentialViaEigh)
+    D, V = eigh_full(A, alg.eigh_alg)
+    iV = inv(V)
+    map!(exp, diagview(D), diagview(D))
+    mul!(expA, rmul!(V, D), iV)
+    return expA
+end
+
+function exponential!(A::AbstractMatrix, expA::AbstractMatrix, alg::ExponentialViaEig)
+    D, V = eig_full(A, alg.eig_alg)
+    iV = inv(V)
+    map!(exp, diagview(D), diagview(D))
+    mul!(expA, rmul!(V, D), iV)
+    return expA
+end
+
+# Diagonal logic
+# --------------
+function exponential!(A::Diagonal, expA, alg::DiagonalAlgorithm)
+    check_input(exponential!, A, expA, alg)
+    map!(exp, diagview(expA), diagview(A))
+    return expA
+end

src/interface/decompositions.jl

@@ -319,6 +319,46 @@ Divide and Conquer algorithm.
 
 const ROCSOLVER_SVDAlgorithm = Union{ROCSOLVER_QRIteration, ROCSOLVER_Jacobi}
 
+# ================================
+# EXPONENTIAL ALGORITHMS
+# ================================
+"""
+    ExponentialViaLA()
+
+Algorithm type to denote finding the exponential of `A` via the implementation of `LinearAlgebra`.
+"""
+@algdef ExponentialViaLA
+
+"""
+    ExponentialViaEigh()
+
+Algorithm type to denote finding the exponential `A` by computing the hermitian eigendecomposition of `A`.
+The `eigh_alg` specifies which hermitian eigendecomposition implementation to use.
+"""
+struct ExponentialViaEigh{A <: AbstractAlgorithm} <: AbstractAlgorithm
+    eigh_alg::A
+end
+function Base.show(io::IO, alg::ExponentialViaEigh)
+    print(io, "ExponentialViaEigh(")
+    _show_alg(io, alg.eigh_alg)
+    return print(io, ")")
+end
+
+"""
+    ExponentialViaEig()
+
+Algorithm type to denote finding the exponential `A` by computing the eigendecomposition of `A`.
+The `eig_alg` specifies which eigendecomposition implementation to use.
+"""
+struct ExponentialViaEig{A <: AbstractAlgorithm} <: AbstractAlgorithm
+    eig_alg::A
+end
+function Base.show(io::IO, alg::ExponentialViaEig)
+    print(io, "ExponentialViaEigh(")
+    _show_alg(io, alg.eig_alg)
+    return print(io, ")")
+end
+
 # Various consts and unions
 # -------------------------
 

src/interface/exponential.jl

@@ -0,0 +1,16 @@
+# Exponential functions
+# --------------
+@functiondef exponential
+
+# Algorithm selection
+# -------------------
+default_exponential_algorithm(A; kwargs...) = default_exponential_algorithm(typeof(A); kwargs...)
+function default_exponential_algorithm(T::Type; kwargs...)
+    return ExponentialViaLA(; kwargs...)
+end
+
+for f in (:exponential!,)
+    @eval function default_algorithm(::typeof($f), ::Type{A}; kwargs...) where {A}
+        return default_exponential_algorithm(A; kwargs...)
+    end
+end

test/exponential.jl

@@ -0,0 +1,48 @@
+using MatrixAlgebraKit
+using Test
+using TestExtras
+using StableRNGs
+using MatrixAlgebraKit: diagview
+using LinearAlgebra
+
+BLASFloats = (Float32, Float64, ComplexF32, ComplexF64)
+GenericFloats = (Float16, ComplexF16, BigFloat, Complex{BigFloat})
+
+@testset "exponential! for T = $T" for T in BLASFloats
+    rng = StableRNG(123)
+    m = 2
+
+    A = randn(rng, T, m, m)
+    A = (A + A') / 2
+    D, V = @constinferred eigh_full(A)
+    algs = (ExponentialViaLA(), ExponentialViaEig(LAPACK_Simple()), ExponentialViaEigh(LAPACK_QRIteration()))
+    expA_LA = @constinferred exp(A)
+    @testset "algorithm $alg" for alg in algs
+        expA = similar(A)
+
+        @constinferred exponential!(copy(A), expA)
+        expA2 = @constinferred exponential(A; alg = alg)
+        @test expA ≈ expA_LA
+        @test expA2 ≈ expA
+
+        Dexp, Vexp = @constinferred eigh_full(expA)
+        @test diagview(Dexp) ≈ LinearAlgebra.exp.(diagview(D))
+    end
+end
+
+@testset "exponential! for Diagonal{$T}" for T in (BLASFloats..., GenericFloats...)
+    rng = StableRNG(123)
+    atol = sqrt(eps(real(T)))
+    m = 54
+    Ad = randn(T, m)
+    A = Diagonal(Ad)
+
+    expA = similar(A)
+    @constinferred exponential!(copy(A), expA)
+    expA2 = @constinferred exponential(A; alg = DiagonalAlgorithm())
+    @test expA2 ≈ expA
+
+    D, V = @constinferred eig_full(A)
+    Dexp, Vexp = @constinferred eig_full(expA)
+    @test diagview(Dexp) ≈ LinearAlgebra.exp.(diagview(D))
+end

test/genericlinearalgebra/exponential.jl

@@ -0,0 +1,28 @@
+using MatrixAlgebraKit
+using Test
+using TestExtras
+using StableRNGs
+using MatrixAlgebraKit: diagview
+using LinearAlgebra
+
+GenericFloats = (BigFloat, Complex{BigFloat})
+
+@testset "exp! for T = $T" for T in GenericFloats
+    rng = StableRNG(123)
+    m = 2
+
+    A = randn(rng, T, m, m)
+    A = (A + A') / 2
+    D, V = @constinferred eigh_full(A)
+    algs = (ExponentialViaEigh(GLA_QRIteration()),)
+    @testset "algorithm $alg" for alg in algs
+        expA = similar(A)
+
+        @constinferred exponential!(copy(A), expA; alg)
+        expA2 = @constinferred exponential(A; alg)
+        @test expA2 ≈ expA
+
+        Dexp, Vexp = @constinferred eigh_full(expA)
+        @test diagview(Dexp) ≈ LinearAlgebra.exp.(diagview(D))
+    end
+end

test/genericschur/exponential.jl

@@ -0,0 +1,29 @@
+using MatrixAlgebraKit
+using Test
+using TestExtras
+using StableRNGs
+using MatrixAlgebraKit: diagview
+using LinearAlgebra
+
+GenericFloats = (BigFloat, Complex{BigFloat})
+
+@testset "exp! for T = $T" for T in GenericFloats
+    rng = StableRNG(123)
+    m = 2
+
+    A = randn(rng, T, m, m)
+    D, V = @constinferred eig_full(A)
+    algs = (ExponentialViaEig(GS_QRIteration()),)
+    expA_LA = @constinferred exponential(A)
+    @testset "algorithm $alg" for alg in algs
+        expA = similar(A)
+
+        @constinferred exponential!(copy(A), expA)
+        expA2 = @constinferred exponential(A; alg = alg)
+        @test expA ≈ expA_LA
+        @test expA2 ≈ expA
+
+        Dexp, Vexp = @constinferred eig_full(expA)
+        @test diagview(Dexp) ≈ LinearAlgebra.exp.(diagview(D))
+    end
+end

test/runtests.jl

@@ -38,6 +38,9 @@ if !is_buildkite
     @safetestset "Image and Null Space" begin
         include("orthnull.jl")
     end
+    @safetestset "Exponential" begin
+        include("exponential.jl")
+    end
     @safetestset "ChainRules" begin
         include("chainrules.jl")
     end
@@ -119,8 +122,14 @@ end
 @safetestset "Hermitian Eigenvalue Decomposition" begin
     include("genericlinearalgebra/eigh.jl")
 end
+@safetestset "Exponential" begin
+    include("genericlinearalgebra/exponential.jl")
+end
 
 using GenericSchur
 @safetestset "General Eigenvalue Decomposition" begin
     include("genericschur/eig.jl")
 end
+@safetestset "Exponential" begin
+    include("genericschur/exponential.jl")
+end