feat: exponential (#465)

* Exponential via MatrixAlgebraKit

* include case where t is real and tau is complex

* Add test

* Fix formatting

* Update correspondence with MAK and extend tests

* Update src/factorizations/matrixalgebrakit.jl

Co-authored-by: Lukas Devos <ldevos98@gmail.com>

* deprecate `exp!`

* Remove extra definitions that are now in MatrixAlgebraKit

* bump MAK version

* update implementations

* remove double test

* refactor exponential tests in the style of other tests

* fix `copy_input` shouldn't alias

* some more small fixes

---------

Co-authored-by: Lukas Devos <ldevos98@gmail.com>

Author: sanderdemeyer

Project.toml

@@ -54,7 +54,7 @@ FiniteDifferences = "0.12"
 GPUArrays = "11.4.1"
 LRUCache = "1.6"
 LinearAlgebra = "1"
-MatrixAlgebraKit = "0.6.8"
+MatrixAlgebraKit = "0.6.9"
 Mooncake = "0.5.27"
 OhMyThreads = "0.8.0"
 Printf = "1"

docs/src/man/linearalgebra.md

@@ -40,7 +40,7 @@ Note that, because the adjoint interchanges domain and codomain, we have `space(
 
 When tensor map instances are endomorphisms, i.e. they have the same domain and codomain, there is a multiplicative identity which can be obtained as `one(t)` or `one!(t)`, where the latter overwrites the contents of `t`.
 The multiplicative identity on a space `V` can also be obtained using `id(A, V)` as discussed [above](@ref ss_tensor_construction), such that for a general homomorphism `t′`, we have `t′ == id(codomain(t′)) * t′ == t′ * id(domain(t′))`.
-Returning to the case of endomorphisms `t`, we can compute the trace via `tr(t)` and exponentiate them using `exp(t)`, or if the contents of `t` can be destroyed in the process, `exp!(t)`.
+Returning to the case of endomorphisms `t`, we can compute the trace via `tr(t)` and exponentiate them using `exp(t)` (or the more general `exponential(t)`), or if the contents of `t` can be destroyed in the process, `exponential!(t)`.
 Furthermore, there are a number of tensor factorizations for both endomorphisms and general homomorphisms that we discuss on the [Tensor factorizations](@ref ss_tensor_factorization) page.
 
 Finally, there are a number of operations that also belong in this paragraph because of their analogy to common matrix operations.

src/TensorKit.jl

@@ -82,7 +82,7 @@ export left_orth, right_orth, left_null, right_null,
     qr_full, qr_compact, qr_null, lq_full, lq_compact, lq_null,
     qr_full!, qr_compact!, qr_null!, lq_full!, lq_compact!, lq_null!,
     svd_compact!, svd_full!, svd_trunc!, svd_compact, svd_full, svd_trunc, svd_vals, svd_vals!,
-    exp, exp!,
+    exp, exp!, exponential, exponential!,
     eigh_full!, eigh_full, eigh_trunc!, eigh_trunc, eigh_vals!, eigh_vals,
     eig_full!, eig_full, eig_trunc!, eig_trunc, eig_vals!, eig_vals,
     eigen, eigen!,

src/factorizations/matrixalgebrakit.jl

@@ -8,6 +8,7 @@ for f in
         :eig_full, :eig_vals, :eigh_full, :eigh_vals,
         :left_polar, :right_polar,
         :project_hermitian, :project_antihermitian, :project_isometric,
+        :exponential,
     ]
     f! = Symbol(f, :!)
     @eval function MAK.default_algorithm(::typeof($f!), ::Type{T}; kwargs...) where {T <: AbstractTensorMap}
@@ -18,6 +19,11 @@ for f in
     end
 end
 
+MAK.default_algorithm(::typeof(exponential!), ::Type{Tuple{E, T}}; kwargs...) where {E <: Number, T <: AbstractTensorMap} =
+    MAK.default_algorithm(exponential!, blocktype(T); kwargs...)
+MAK.copy_input(::typeof(exponential), (τ, t)::Tuple{E, T}) where {E <: Number, T <: AbstractTensorMap} =
+    (τ, copy_oftype(t, (E <: Complex ? complex : identity)(factorisation_scalartype(exponential, t))))
+
 _select_truncation(f, ::AbstractTensorMap, trunc::TruncationStrategy) = trunc
 function _select_truncation(::typeof(left_null!), ::AbstractTensorMap, trunc::NamedTuple)
     return MAK.null_truncation_strategy(; trunc...)
@@ -49,9 +55,10 @@ for f! in (
         :qr_null!, :lq_null!,
         :svd_vals!, :eig_vals!, :eigh_vals!,
         :project_hermitian!, :project_antihermitian!, :project_isometric!,
+        :exponential!,
     )
     @eval function MAK.$f!(t::AbstractTensorMap, N, alg::AbstractAlgorithm)
-        $(f! in (:eig_vals!, :eigh_vals!, :project_hermitian!, :project_antihermitian!) && :(LinearAlgebra.checksquare(t)))
+        $(f! in (:eig_vals!, :eigh_vals!, :project_hermitian!, :project_antihermitian!, :exponential!) && :(LinearAlgebra.checksquare(t)))
         foreachblock(t, N) do _, (tblock, Nblock)
             Nblock′ = $f!(tblock, Nblock, alg)
             # deal with the case where the output is not the same as the input
@@ -62,6 +69,19 @@ for f! in (
     end
 end
 
+# Exponential with Tuple
+function MAK.exponential!((τ, t)::Tuple{E, T}, N, alg::AbstractAlgorithm) where {E <: Number, T <: AbstractTensorMap}
+    LinearAlgebra.checksquare(t)
+    foreachblock(t, N) do _, (tblock, Nblock)
+        Nblock′ = exponential!((τ, tblock), Nblock, alg)
+        # deal with the case where the output is not the same as the input
+        Nblock === Nblock′ || copy!(Nblock, Nblock′)
+        return nothing
+    end
+    return N
+end
+
+
 MAK.zero!(t::AbstractTensorMap) = zerovector!(t)
 
 # Default algorithm
@@ -76,6 +96,7 @@ for f in [
         :left_polar, :right_polar,
         :left_orth, :right_orth, :left_null, :right_null,
         :project_hermitian, :project_antihermitian, :project_isometric,
+        :exponential,
     ]
     f! = Symbol(f, :!)
     @eval MAK.$f!(t::AbstractTensorMap, alg::DefaultAlgorithm) =
@@ -95,6 +116,12 @@ for f in [
         MAK.$f!(t, out, MAK.select_algorithm(MAK.$f!, t, nothing; alg.kwargs...))
 end
 
+# resolve `DefaultAlgorithm` for the tuple form at the tensor level, mirroring the loop below
+MAK.exponential!((τ, t)::Tuple{E, T}, alg::DefaultAlgorithm) where {E <: Number, T <: AbstractTensorMap} =
+    MAK.exponential!((τ, t), MAK.select_algorithm(exponential!, t, nothing; alg.kwargs...))
+MAK.exponential!((τ, t)::Tuple{E, T}, out, alg::DefaultAlgorithm) where {E <: Number, T <: AbstractTensorMap} =
+    MAK.exponential!((τ, t), out, MAK.select_algorithm(exponential!, t, nothing; alg.kwargs...))
+
 
 # Singular value decomposition
 # ----------------------------
@@ -221,3 +248,9 @@ MAK.initialize_output(::typeof(project_antihermitian!), tsrc::AbstractTensorMap,
     tsrc
 MAK.initialize_output(::typeof(project_isometric!), tsrc::AbstractTensorMap, ::AbstractAlgorithm) =
     similar(tsrc)
+
+# Exponential
+# ----------------
+MAK.initialize_output(::typeof(exponential!), t::AbstractTensorMap, ::AbstractAlgorithm) = t
+MAK.initialize_output(::typeof(exponential!), (τ, t)::Tuple{Number, AbstractTensorMap}, ::AbstractAlgorithm) = t
+MAK.initialize_output(::typeof(exponential!), (τ, t)::Tuple{T1, AbstractTensorMap{T2}}, ::AbstractAlgorithm) where {T1 <: Complex, T2 <: Real} = similar(t, complex(eltype(t)))

src/factorizations/utility.jl

@@ -5,7 +5,7 @@ end
 factorisation_scalartype(f, t) = factorisation_scalartype(t)
 
 function copy_oftype(t::AbstractTensorMap, T::Type{<:Number})
-    return copy!(similar(t, T, space(t)), t)
+    return copy!(similar(t, T), t)
 end
 
 function _reverse!(t::AbstractTensorMap; dims = :)

src/tensors/linalg.jl

@@ -41,7 +41,7 @@ function compose(A::AbstractTensorMap, B::AbstractTensorMap)
 end
 Base.:*(t1::AbstractTensorMap, t2::AbstractTensorMap) = compose(t1, t2)
 
-Base.exp(t::AbstractTensorMap) = exp!(copy(t))
+Base.exp(t::AbstractTensorMap) = exponential(t)
 function Base.:^(t::AbstractTensorMap, p::Integer)
     return p < 0 ? Base.power_by_squaring(inv(t), -p) : Base.power_by_squaring(t, p)
 end
@@ -416,15 +416,7 @@ function Base.:(/)(t1::AbstractTensorMap, t2::AbstractTensorMap)
     return t
 end
 
-# TensorMap exponentation:
-function exp!(t::TensorMap)
-    domain(t) == codomain(t) ||
-        error("Exponential of a tensor only exist when domain == codomain.")
-    for (c, b) in blocks(t)
-        copy!(b, LinearAlgebra.exp!(b))
-    end
-    return t
-end
+@deprecate exp!(t) exponential!(t)
 
 # Sylvester equation with TensorMap objects:
 function LinearAlgebra.sylvester(A::AbstractTensorMap, B::AbstractTensorMap, C::AbstractTensorMap)

test/tensors/exponential.jl

@@ -0,0 +1,91 @@
+using Test, TestExtras
+using TensorKit
+using MatrixAlgebraKit: DefaultAlgorithm, MatrixFunctionViaLA, MatrixFunctionViaEig,
+    MatrixFunctionViaEigh, MatrixFunctionViaTaylor
+using Random
+
+spacelist = default_spacelist(fast_tests)
+scalartypes = (Float32, Float64, ComplexF64)
+
+# algorithms that agree on Hermitian input
+hermitian_algs = (
+    MatrixFunctionViaLA(), MatrixFunctionViaEig(DefaultAlgorithm()),
+    MatrixFunctionViaEigh(DefaultAlgorithm()), MatrixFunctionViaTaylor(),
+)
+# algorithms valid for general (non-Hermitian) input
+general_algs = (
+    MatrixFunctionViaLA(), MatrixFunctionViaEig(DefaultAlgorithm()),
+    MatrixFunctionViaTaylor(),
+)
+
+for V in spacelist
+    I = sectortype(first(V))
+    Istr = TensorKit.type_repr(I)
+    println("---------------------------------------")
+    println("Matrix exponentials with symmetry: $Istr")
+    println("---------------------------------------")
+    @timedtestset "Matrix exponentials with symmetry: $Istr" verbose = true begin
+        V1, V2, V3, V4, V5 = V
+        W = V1 ⊗ V2 ⊗ V3
+        Vd = fuse(V1 ⊗ V2)
+
+        # explicit (non-diagonal) algorithms only apply to dense endomorphisms;
+        # `DiagonalTensorMap` inputs are exercised through the default algorithm below
+        @testset "exponential for Hermitian matrices" begin
+            for T1 in scalartypes, T2 in scalartypes,
+                    A in (randn(T1, V1, V1), randn(T1, W, W))
+
+                A = project_hermitian!(A)
+                τ = rand(T2)
+
+                expA = @constinferred exponential(A)
+                expτA = @constinferred exponential((τ, A))
+
+                for alg in hermitian_algs
+                    @test expA ≈ @constinferred exponential(A, alg)
+                    @test expτA ≈ @constinferred exponential((τ, A), alg)
+                end
+            end
+        end
+
+        @testset "exponential! for general matrices" begin
+            for T1 in scalartypes, T2 in scalartypes,
+                    A in (
+                        randn(T1, V1, V1), randn(T1, W, W),
+                        DiagonalTensorMap(randn(T1, reduceddim(Vd)), Vd),
+                    )
+
+                τ = rand(T2)
+
+                expA = @constinferred exponential(A)
+                if !(A isa DiagonalTensorMap) # diagonal only supports the default algorithm == DiagonalAlgorithm
+                    for alg in general_algs
+                        @test expA ≈ exponential(A, alg)
+                    end
+                end
+
+                @test exponential!(copy(A)) ≈ exponential!((1.0, copy(A)))
+
+                # exp(A)² == exp(2A)
+                @test expA * expA ≈ exponential((2, A))
+
+                # in-place semantics: aliases the input, unless a complex scalar forces
+                # a real tensor to widen to complex
+                Ain = copy(A)
+                A2 = @constinferred exponential!((τ, Ain))
+                A isa DiagonalTensorMap && @test A2 isa DiagonalTensorMap
+                if T1 <: Real && T2 <: Complex
+                    @test A2 !== Ain
+                else
+                    @test A2 === Ain
+                end
+
+                # exp(τA) and exp(-τA) are inverse
+                # the inverse roundtrip is ill-conditioned; only assert at full (Float64) precision
+                expτA = exponential!((τ, copy(A)))
+                expmτA = exponential!((-τ, copy(A)))
+                real(scalartype(expτA)) == Float64 && @test expτA * expmτA ≈ id(scalartype(expτA), domain(A))
+            end
+        end
+    end
+end