Exponential

ext/MatrixAlgebraKitGenericSchurExt.jl

@@ -14,6 +14,13 @@ function MatrixAlgebraKit.default_eig_algorithm(
     return QRIteration(; driver, kwargs...)
 end
 
+function MatrixAlgebraKit.default_exponential_algorithm(
+        type::Type{T}; kwargs...
+    ) where {T <: StridedMatrix{<:GSFloat}}
+    eig_alg = MatrixAlgebraKit.default_eig_algorithm(type; kwargs...)
+    return MatrixFunctionViaEig(eig_alg)
+end
+
 function geev!(::GS, A::AbstractMatrix, Dd::AbstractVector, V::AbstractMatrix; kwargs...)
     D, Vmat = GenericSchur.eigen!(A)
     copyto!(Dd, D)

src/MatrixAlgebraKit.jl

@@ -30,6 +30,7 @@ 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 Householder, Native_HouseholderQR, Native_HouseholderLQ
 export DivideAndConquer, SafeDivideAndConquer, QRIteration, Bisection, Jacobi, SVDViaPolar
@@ -40,6 +41,7 @@ export LAPACK_HouseholderQR, LAPACK_HouseholderLQ, LAPACK_Simple, LAPACK_Expert,
 export GLA_HouseholderQR, GLA_QRIteration, GS_QRIteration
 export LQViaTransposedQR
 export PolarViaSVD, PolarNewton
+export MatrixFunctionViaLA, MatrixFunctionViaEig, MatrixFunctionViaEigh
 export DefaultAlgorithm
 export DiagonalAlgorithm
 export NativeBlocked
@@ -96,9 +98,12 @@ include("common/utility.jl")
 
 include("yalapack.jl")
 include("algorithms.jl")
+
 include("interface/projections.jl")
 include("interface/decompositions.jl")
 include("interface/truncation.jl")
+include("interface/matrixfunctions.jl")
+
 include("interface/qr.jl")
 include("interface/lq.jl")
 include("interface/svd.jl")
@@ -108,6 +113,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")
@@ -120,6 +126,7 @@ include("implementations/gen_eig.jl")
 include("implementations/schur.jl")
 include("implementations/polar.jl")
 include("implementations/orthnull.jl")
+include("implementations/exponential.jl")
 
 include("common/gauge.jl") # needs to be defined after the functions are
 

src/common/view.jl

@@ -20,6 +20,26 @@ See also [`diagview`](@ref).
 
 diagonal(v::AbstractVector) = Diagonal(v)
 
+"""
+    map_diagonal!(f, dst, src...)
+
+Map the scalar function `f` over all elements of the diagonal of `src...`, returning
+a diagonal result.
+
+See also [`map_diagonal!`](@ref).
+"""
+map_diagonal(f, src, srcs...) = diagonal(f.(diagview(src), map(diagview, srcs)...))
+
+"""
+    map_diagonal!(f, dst, src...)
+
+Map the scalar function `f` over all elements of the diagonal of `src...`,
+into the diagonal elements of destination `dst`.
+
+See also [`map_diagonal`](@ref).
+"""
+map_diagonal!(f, dst, src, srcs...) = (diagview(dst) .= f.(diagview(src), map(diagview, srcs)...); dst)
+
 # triangularind
 function lowertriangularind(A::AbstractMatrix)
     Base.require_one_based_indexing(A)

src/implementations/exponential.jl

@@ -0,0 +1,109 @@
+# 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)
+copy_input(::typeof(exponential), (τ, A)::Tuple{Number, AbstractMatrix}) = (τ, copy!(similar(A, float(eltype(A))), A))
+copy_input(::typeof(exponential), (τ, A)::Tuple{Number, Diagonal}) = τ, copy(A)
+
+function check_input(::typeof(exponential!), A::AbstractMatrix, expA::AbstractMatrix, alg::AbstractAlgorithm)
+    m = LinearAlgebra.checksquare(A)
+    @check_size(expA, (m, m))
+    @check_scalar(expA, A)
+    return nothing
+end
+
+function check_input(::typeof(exponential!), A::AbstractMatrix, expA::AbstractMatrix, ::DiagonalAlgorithm)
+    m = LinearAlgebra.checksquare(A)
+    @assert isdiag(A)
+    @assert expA isa Diagonal
+    @check_size(expA, (m, m))
+    @check_scalar(expA, A)
+    return nothing
+end
+
+function check_input(::typeof(exponential!), (τ, A)::Tuple{Number, AbstractMatrix}, expA::AbstractMatrix, alg::AbstractAlgorithm)
+    m = LinearAlgebra.checksquare(A)
+    @check_size(expA, (m, m))
+    @check_scalar(expA, A, (τ isa Real) ? identity : complex)
+    return nothing
+end
+
+function check_input(::typeof(exponential!), (τ, A)::Tuple{Number, AbstractMatrix}, expA::AbstractMatrix, ::DiagonalAlgorithm)
+    m = LinearAlgebra.checksquare(A)
+    @assert isdiag(A)
+    @assert expA isa Diagonal
+    @check_size(expA, (m, m))
+    @check_scalar(expA, A, (τ isa Real) ? identity : complex)
+    return nothing
+end
+
+# Outputs
+# -------
+initialize_output(::typeof(exponential!), A::AbstractMatrix, ::AbstractAlgorithm) = A
+initialize_output(::typeof(exponential!), (τ, A)::Tuple{T, AbstractMatrix}, ::AbstractAlgorithm) where {T <: Real} = A
+initialize_output(::typeof(exponential!), (τ, A)::Tuple{Number, AbstractMatrix}, ::AbstractAlgorithm) = complex(A)
+
+# Implementation
+# --------------
+function exponential!(A::AbstractMatrix, expA::AbstractMatrix, alg::MatrixFunctionViaLA)
+    check_input(exponential!, A, expA, alg)
+    A = LinearAlgebra.exp!(A)
+    A === expA || copy!(expA, A)
+    return expA
+end
+
+exponential!(A, expA, alg::MatrixFunctionViaEigh) = exponential!((one(eltype(A)), A), expA, alg)
+exponential!(A::AbstractMatrix, expA::AbstractMatrix, alg::MatrixFunctionViaEig) = exponential!((one(eltype(A)), A), expA, alg)
+
+function exponential!((τ, A)::Tuple{Number, AbstractMatrix}, expA::AbstractMatrix, alg::AbstractAlgorithm)
+    expA .= A .* τ
+    return exponential!(expA, expA, alg)
+end
+
+function exponential!((τ, A)::Tuple{Number, AbstractMatrix}, expA::AbstractMatrix, alg::MatrixFunctionViaEigh)
+    check_input(exponential!, (τ, A), expA, alg)
+    D, V = eigh_full!(A, alg.eigh_alg)
+    if eltype(A) <: Real
+        if eltype(τ) <: Real
+            VexpD = rmul!(V, exponential!((τ / 2, D), D))
+        else
+            VexpD = V * exponential((τ / 2, D))
+        end
+        return mul!(expA, VexpD, transpose(VexpD))
+    else
+        if eltype(τ) <: Real
+            VexpD = V * exponential!((τ, D), D)
+        else
+            VexpD = V * exponential((τ, D))
+        end
+        return mul!(expA, VexpD, V')
+    end
+end
+
+function exponential!((τ, A)::Tuple{Number, AbstractMatrix}, expA::AbstractMatrix, alg::MatrixFunctionViaEig)
+    check_input(exponential!, (τ, A), expA, alg)
+    D, V = eig_full!(A, alg.eig_alg)
+    if eltype(A) <: Real && eltype(τ) <: Real
+        VexpD = V * exponential!((τ, D), D)
+        expAc = rdiv!(VexpD, LinearAlgebra.lu!(V))
+        return expA .= real.(expAc)
+    else
+        expA .= V .* transpose(diagview(exponential!((τ, D), D)))
+        return rdiv!(expA, LinearAlgebra.lu!(V))
+    end
+end
+
+# Diagonal logic
+# --------------
+function exponential!(A, expA, alg::DiagonalAlgorithm)
+    check_input(exponential!, A, expA, alg)
+    return map_diagonal!(exp, expA, A)
+end
+
+function exponential!((τ, A)::Tuple{Number, Any}, expA, alg::DiagonalAlgorithm)
+    check_input(exponential!, (τ, A), expA, alg)
+    return map_diagonal!(x -> exp(x * τ), expA, A)
+end

src/interface/exponential.jl

@@ -0,0 +1,39 @@
+# Exponential functions
+# --------------
+
+"""
+    exponential(A; kwargs...) -> expA
+    exponential(A, alg::AbstractAlgorithm) -> expA
+    exponential!(A, [expA]; kwargs...) -> expA
+    exponential!(A, [expA], alg::AbstractAlgorithm) -> expA
+    exponential((τ, A); kwargs...) -> expτA
+    exponential((τ, A), alg::AbstractAlgorithm) -> expτA
+    exponential!((τ, A), [expA]; kwargs...) -> expτA
+    exponential!((τ, A), [expA], alg::AbstractAlgorithm) -> expτA
+
+Compute the exponential of the square matrix `A` or `τ * A`,
+
+!!! note
+    The bang method `exponential!` optionally accepts the output structure and
+    possibly destroys the input matrix `A`. Always use the return value of the function
+    as it may not always be possible to use the provided `expA` as output.
+"""
+@functiondef exponential
+
+# Algorithm selection
+# -------------------
+default_exponential_algorithm(A; kwargs...) = default_exponential_algorithm(typeof(A); kwargs...)
+function default_exponential_algorithm(T::Type; kwargs...)
+    return MatrixFunctionViaLA(; kwargs...)
+end
+function default_exponential_algorithm(::Type{T}; kwargs...) where {T <: Diagonal}
+    return DiagonalAlgorithm(; kwargs...)
+end
+
+function default_algorithm(::typeof(exponential!), ::Type{A}; kwargs...) where {A}
+    return default_exponential_algorithm(A; kwargs...)
+end
+
+function default_algorithm(::typeof(exponential!), ::Tuple{A, B}; kwargs...) where {A, B}
+    return default_exponential_algorithm(B; kwargs...)
+end

src/interface/matrixfunctions.jl

@@ -0,0 +1,39 @@
+# ================================
+# EXPONENTIAL ALGORITHMS
+# ================================
+"""
+    MatrixFunctionViaLA()
+
+Algorithm type to denote finding the exponential of `A` via the implementation of `LinearAlgebra`.
+"""
+@algdef MatrixFunctionViaLA
+
+"""
+    MatrixFunctionViaEigh(eigh_alg)
+
+Algorithm type for computing a function of a matrix by computing its hermitian eigenvalue decomposition and applying the function to the eigenvalues.
+The `eigh_alg` specifies which hermitian eigendecomposition implementation to use.
+"""
+struct MatrixFunctionViaEigh{A <: AbstractAlgorithm} <: AbstractAlgorithm
+    eigh_alg::A
+end
+function Base.show(io::IO, alg::MatrixFunctionViaEigh)
+    print(io, "MatrixFunctionViaEigh(")
+    _show_alg(io, alg.eigh_alg)
+    return print(io, ")")
+end
+
+"""
+    MatrixFunctionViaEig(eig_alg)
+
+Algorithm type for computing a function of a matrix by computing its eigenvalue decomposition and applying the function to the eigenvalues.
+The `eig_alg` specifies which eigendecomposition implementation to use.
+"""
+struct MatrixFunctionViaEig{A <: AbstractAlgorithm} <: AbstractAlgorithm
+    eig_alg::A
+end
+function Base.show(io::IO, alg::MatrixFunctionViaEig)
+    print(io, "MatrixFunctionViaEig(")
+    _show_alg(io, alg.eig_alg)
+    return print(io, ")")
+end

test/exponential.jl

@@ -0,0 +1,88 @@
+using MatrixAlgebraKit
+using Test
+using TestExtras
+using StableRNGs
+using MatrixAlgebraKit: diagview
+using LinearAlgebra
+using LinearAlgebra: exp
+
+BLASFloats = (Float32, Float64, ComplexF32, ComplexF64)
+GenericFloats = (Float16, ComplexF16, BigFloat, Complex{BigFloat})
+
+@testset "exponential! for T = $T" for T in BLASFloats
+    rng = StableRNG(123)
+    m = 54
+
+    A = LinearAlgebra.normalize!(randn(rng, T, m, m))
+    Ac = copy(A)
+    expA = LinearAlgebra.exp(A)
+
+    expA2 = @constinferred exponential(A)
+    @test expA ≈ expA2
+    @test A == Ac
+
+    algs = (MatrixFunctionViaLA(), MatrixFunctionViaEig(LAPACK_Simple()))
+    @testset "algorithm $alg" for alg in algs
+        expA2 = @constinferred exponential(A, alg)
+        @test expA ≈ expA2
+        @test A == Ac
+    end
+
+    @test_throws DomainError exponential(A; alg = MatrixFunctionViaEigh(LAPACK_QRIteration()))
+end
+
+@testset "exponential! for T = $T" for T in BLASFloats
+    rng = StableRNG(123)
+    m = 54
+
+    A = randn(rng, T, m, m)
+    τ = randn(rng, T)
+    Ac = copy(A)
+
+    Aτ = A * τ
+    expAτ = LinearAlgebra.exp(Aτ)
+
+    expAτ2 = @constinferred exponential((τ, A))
+    @test expAτ ≈ expAτ2
+    @test A == Ac
+
+    algs = (MatrixFunctionViaLA(), MatrixFunctionViaEig(LAPACK_Simple()))
+    @testset "algorithm $alg" for alg in algs
+        expAτ2 = @constinferred exponential((τ, A), alg)
+        @test expAτ ≈ expAτ2
+        @test A == Ac
+    end
+
+    @test_throws DomainError exponential((τ, A); alg = MatrixFunctionViaEigh(LAPACK_QRIteration()))
+end
+
+@testset "exponential! for Diagonal{$T}" for T in (BLASFloats..., GenericFloats...)
+    rng = StableRNG(123)
+    m = 54
+
+    A = Diagonal(randn(rng, T, m))
+    τ = randn(rng, T)
+    Ac = copy(A)
+
+    expA = LinearAlgebra.exp(A)
+
+    expA2 = @constinferred exponential(A)
+    @test expA ≈ expA2
+    @test A == Ac
+end
+
+@testset "exponential! for Diagonal{$T}" for T in (BLASFloats..., GenericFloats...)
+    rng = StableRNG(123)
+    m = 1
+
+    A = Diagonal(randn(rng, T, m))
+    τ = randn(rng, T)
+    Ac = copy(A)
+
+    Aτ = A * τ
+    expAτ = LinearAlgebra.exp(Aτ)
+
+    expAτ2 = @constinferred exponential((τ, A))
+    @test expAτ ≈ expAτ2
+    @test A == Ac
+end