Taylor series native exponential implementation (#243)

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

Author: Jutho

docs/src/user_interface/algorithms.md

@@ -102,6 +102,7 @@ The following algorithms for matrix functions are available.
 
 | Algorithm | Applicable matrix functions | Key keyword arguments |
 |:----------|:--------------------------|:----------------------|
+| [`MatrixFunctionViaTaylor`](@ref) | exponential | `tol`, `balance` |
 | [`MatrixFunctionViaLA`](@ref) | exponential | |
 | [`MatrixFunctionViaEig`](@ref) | exponential | `eig_alg` |
 | [`MatrixFunctionViaEigh`](@ref) | exponential | `eigh_alg` |

docs/src/user_interface/matrix_functions.md

@@ -23,14 +23,16 @@ Additionally, the `f!` method typically assumes that it is allowed to destroy th
 ## Exponential
 
 The [exponential](https://en.wikipedia.org/wiki/Matrix_exponential) of a square matrix `A` is used in many scientific applications, as it arises in the solution of an autonomous linear differential equation.
-An implementation for the matrix exponential based on a Padé approximation is available in `LinearAlgebra`, and can be accessed by the algorithm [`MatrixFunctionViaLA`](@ref).
-For more generic data types, the exponential can be calculated by first calculating the (hermitian) eigenvalue decomposition, and then computing
-the scalar exponential of the diagonal elements.
+The default algorithm [`MatrixFunctionViaTaylor`](@ref) is a pure-Julia scaling-and-squaring evaluation of the Taylor series.
+As it requires no LAPACK support, it also applies to generic data types at arbitrary precision.
+Alternatively, an implementation based on a Padé approximation is available in `LinearAlgebra`, and can be accessed by the algorithm [`MatrixFunctionViaLA`](@ref).
+The exponential can also be calculated by first calculating the (hermitian) eigenvalue decomposition, and then computing the scalar exponential of the diagonal elements.
 This strategy is implemented via the algorithms [`MatrixFunctionViaEig`](@ref) and [`MatrixFunctionViaEigh`](@ref), and call `eig_full` and `eigh_full`, respectively.
 Additionally, in order to calculate `exp(τ * A)`, the function `exponential` can be called with `(τ, A)`, using the same algorithms as before.
 
 ```@docs; canonical=false
 exponential
+MatrixAlgebraKit.MatrixFunctionViaTaylor
 MatrixAlgebraKit.MatrixFunctionViaLA
 MatrixAlgebraKit.MatrixFunctionViaEig
 MatrixAlgebraKit.MatrixFunctionViaEigh

src/MatrixAlgebraKit.jl

@@ -41,7 +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 MatrixFunctionViaLA, MatrixFunctionViaEig, MatrixFunctionViaEigh, MatrixFunctionViaTaylor
 export DefaultAlgorithm
 export DiagonalAlgorithm
 export NativeBlocked
@@ -94,6 +94,7 @@ include("common/safemethods.jl")
 include("common/view.jl")
 include("common/regularinv.jl")
 include("common/matrixproperties.jl")
+include("common/balancing.jl")
 include("common/utility.jl")
 
 include("yalapack.jl")

src/common/balancing.jl

@@ -0,0 +1,56 @@
+"""
+    balance!(A::AbstractMatrix; maxiter=100) -> A, scale
+
+Balance the square matrix `A` in place through a diagonal similarity `A ← D⁻¹ A D` that reduces its norm.
+Each sweep computes, for every index at once, the power-of-two factor that best equalizes the
+off-diagonal row and column norms (a simultaneous, Osborne-style variant of the Parlett–Reinsch scaling),
+and applies all factors together; sweeps repeat until no index changes or `maxiter` is reached.
+
+The scaling factors are integer powers of two, the machine radix, so that the transformation is exact even
+in floating point arithmetic. The row and column norms are simultaneously adapted to avoid scalar indexing
+and lots of kernel calls on GPUs.
+
+The returned `scale` holds the diagonal of `D`, so that a matrix function of the original
+input can be recovered from a matrix function `f` of the balanced matrix through
+`D f(D⁻¹AD) D⁻¹`, i.e. `expA[i, j] = scale[i] * f(B)[i, j] / scale[j]`.
+"""
+function balance!(A::AbstractMatrix{T}; maxiter::Integer = 100) where {T}
+    n = LinearAlgebra.checksquare(A)
+    R = real(T)
+    β = convert(R, 2)
+    logβ = log(β)
+    scale = fill!(similar(A, R, n), one(R))
+
+    colnorm = similar(A, R, n)
+    rownorm = similar(A, R, n)
+    f = similar(A, R, n)
+    colsum = reshape(colnorm, 1, n)
+    rowsum = reshape(rownorm, n, 1)
+    d = abs.(diagview(A))
+    fᵀ = transpose(f)
+
+    for _ in 1:maxiter
+        fill!(colsum, zero(R))
+        Base.mapreducedim!(abs, +, colsum, A)
+        fill!(rowsum, zero(R))
+        Base.mapreducedim!(abs, +, rowsum, A)
+        colnorm .-= d
+        rownorm .-= d
+        f .= _balance_factor.(colnorm, rownorm, β, logβ)
+        all(isone, f) && break
+        # apply Aᵢⱼ ← Aᵢⱼ fⱼ / fᵢ (i.e. column j scaled by fⱼ, row i by 1/fᵢ) and accumulate.
+        A .= A .* fᵀ ./ f
+        scale .*= f
+    end
+
+    return A, scale
+end
+
+# Nearest power-of-`β` factor `f` that equalizes the scaled off-diagonal norms
+# `colnorm·f` and `rownorm/f`, kept only when it reduces their sum (avoids oscillation and
+# leaves degenerate rows/columns untouched).
+@inline function _balance_factor(colnorm::R, rownorm::R, β::R, logβ::R) where {R}
+    (colnorm > 0 && rownorm > 0) || return one(R)
+    f = β^round(log(rownorm / colnorm) / (2 * logβ))
+    return (colnorm * f + rownorm / f) < convert(R, 19 // 20) * (colnorm + rownorm) ? f : one(R)
+end

src/implementations/exponential.jl

@@ -114,3 +114,154 @@ function exponential!((τ, A)::Tuple{Number, AbstractMatrix}, expA, alg::Diagona
     check_input(exponential!, (τ, A), expA, alg)
     return map_diagonal!(x -> exp(x * τ), expA, A)
 end
+
+# Taylor logic
+# ------------
+function exponential!(A::AbstractMatrix, expA, alg::MatrixFunctionViaTaylor)
+    check_input(exponential!, A, expA, alg)
+    m = LinearAlgebra.checksquare(A)
+    T = eltype(A)
+    R = real(T)
+    tol = convert(R, get(alg.kwargs, :tol, eps(R)))
+
+    dobalance = get(alg.kwargs, :balance, false)
+    scale = dobalance ? balance!(A)[2] : nothing
+
+    # Form a minimal set of powers of A up front and use them to sharpen the norm estimate
+    # through the Al-Mohy–Higham quantities dₚ = ‖Aᵖ‖^(1/p).
+    p₀ = min(convert(Int, get(alg.kwargs, :estimate_order, 4)), m)
+    powers = Vector{typeof(A)}(undef, p₀)
+    powers[1] = A
+    d = Vector{R}(undef, p₀)
+    d[1] = LinearAlgebra.opnorm(A, 1)
+    iszero(d[1]) && return one!(expA)
+    for p in 2:p₀
+        powers[p] = powers[p - 1] * A
+        d[p] = LinearAlgebra.opnorm(powers[p], 1)^(1 / p)
+    end
+
+    order, blocksize, squarings = taylor_order_and_squarings(d, tol)
+
+    # Rescale the p₀ powers we hold into powers of A/2ˢ
+    if squarings > 0
+        f = inv(convert(T, 2)^squarings)
+        fk = one(T)
+        for k in 1:p₀
+            fk *= f
+            rmul!(powers[k], fk) # powers[k] ← (A/2ˢ)ᵏ
+        end
+    end
+    # Extend from p₀ to `blocksize` powers (blocksize ≥ p₀ always, see taylor_order_and_squarings)
+    for p in (p₀ + 1):blocksize
+        push!(powers, powers[p - 1] * powers[1])
+    end
+
+
+    # Paterson–Stockmeyer evaluation, followed by squaring
+    X = taylor_polynomial(powers, order)
+    Y = expA # can reuse this memory
+    for _ in 1:squarings
+        mul!(Y, X, X)
+        X, Y = Y, X
+    end
+
+    if dobalance
+        expA .= scale .* X ./ transpose(scale)
+    else
+        X === expA || copyto!(expA, X)
+    end
+    return expA
+end
+
+# Truncation order `m`, blocksize `b` and number of squarings `s` minimizing the total
+# matrix-multiplication count needed to make the Taylor remainder bound `(θ/2ˢ)ᵐ⁺¹/(m+1)! ≤ tol`.
+function taylor_order_and_squarings(d::AbstractVector{<:Real}, tol::Real)
+    p₀ = length(d)
+    log2tol = Float64(log2(tol)) # log2 before conversion: high-precision `tol` may underflow Float64
+    log2factorial = 0.0 # log2((order + 1)!), accumulated incrementally across increasing orders
+    prev_order = best_order = best_blocksize = best_squarings = 0
+    best_cost = typemax(Int)
+
+    budget = 0 # multiplications spent on powers and Horner steps, excluding squarings
+    while budget < best_cost # total cost ≥ budget, so larger budgets cannot improve
+        # highest order = numblocks * blocksize reachable within `budget`: the sum
+        # numblocks + blocksize is fixed, so balance the split, keeping blocksize ≥ p₀
+        blocksize = max(p₀, (budget + p₀ + 1) ÷ 2)
+        numblocks = budget + p₀ + 1 - blocksize
+        order = numblocks * blocksize
+
+        # sharpest valid αₚ = min over p with p(p-1) ≤ order+1 and p+1 ≤ p₀
+        # Al-Mohy, A. H. and Higham, N. J., "A New Scaling and Squaring Algorithm for the Matrix Exponential",
+        # SIAM J. Matrix Anal. Appl., 31(3), 970–989, 2009, Lemma 4.1 and Theorem 4.2(a).
+        θ = d[1]
+        for p in 1:(p₀ - 1)
+            p * (p - 1) ≤ order + 1 || break
+            θ = min(θ, max(d[p], d[p + 1]))
+        end
+
+        # compute squarings needed to make `(θ/2ˢ)ᵐ⁺¹/(m+1)! ≤ tol`
+        for k in (prev_order + 2):(order + 1)
+            log2factorial += log2(k)
+        end
+        prev_order = order
+        excess = Float64(log2(θ)) - (log2tol + log2factorial) / (order + 1)
+        squarings = excess > 0 ? ceil(Int, excess) : 0 # be careful with θ = 0
+
+
+        cost = budget + squarings
+        if cost ≤ best_cost # ties → larger budget → fewer squarings (more accurate)
+            best_cost = cost
+            best_order = order
+            best_blocksize = blocksize
+            best_squarings = squarings
+        end
+        budget += 1
+    end
+
+    return best_order, best_blocksize, best_squarings
+end
+
+# Evaluate ∑ₖ₌₀ᵐ Aᵏ/k! via the Paterson–Stockmeyer scheme, returning a freshly allocated matrix.
+# `powers` holds A, A², …, A^blocksize (already scaled). The polynomial is split as
+# c₀ I + ∑ⱼ₌₁ᴷ A^((j-1)b) Bⱼ with blocks Bⱼ = ∑ᵢ₌₁ᵇ c_{(j-1)b+i} Aⁱ and K = cld(order, blocksize),
+# evaluated by Horner over A^b using K - 1 multiplications; the constant term is added to the
+# diagonal at the end.
+function taylor_polynomial(powers::AbstractVector{<:AbstractMatrix}, order::Integer)
+    A = powers[1]
+    T = eltype(A)
+    blocksize = length(powers)
+
+    invfactorial = Vector{T}(undef, order + 1)
+    invfactorial[1] = one(T)
+    for k in 1:order
+        invfactorial[k + 1] = invfactorial[k] / k
+    end
+
+    result = similar(A)
+    block = similar(A)
+    numblocks = cld(order, blocksize)
+    taylor_block!(result, powers, invfactorial, numblocks, blocksize, order)
+    for blockindex in (numblocks - 1):-1:1
+        mul!(block, result, powers[blocksize])
+        taylor_block!(result, powers, invfactorial, blockindex, blocksize, order)
+        result .+= block
+    end
+    diagview(result) .+= invfactorial[1]
+    return result
+end
+
+# Block Bⱼ = ∑ᵢ₌₁ᵇ c_{offset+i} Aⁱ with offset = (blockindex - 1) * blocksize, truncated at
+# total degree `order`, written into `out`.
+function taylor_block!(out::AbstractMatrix, powers, invfactorial, blockindex::Integer, blocksize::Integer, order::Integer)
+    offset = (blockindex - 1) * blocksize
+    for i in 1:blocksize
+        k = offset + i
+        k > order && break
+        if i == 1
+            out .= invfactorial[k + 1] .* powers[i]
+        else
+            out .+= invfactorial[k + 1] .* powers[i]
+        end
+    end
+    return out
+end

src/interface/exponential.jl

@@ -24,7 +24,7 @@ Compute the exponential of the square matrix `A` or `τ * A`,
 # -------------------
 default_exponential_algorithm(A; kwargs...) = default_exponential_algorithm(typeof(A); kwargs...)
 function default_exponential_algorithm(T::Type; kwargs...)
-    return MatrixFunctionViaLA(; kwargs...)
+    return MatrixFunctionViaTaylor(; kwargs...)
 end
 function default_exponential_algorithm(::Type{T}; kwargs...) where {T <: Diagonal}
     return DiagonalAlgorithm(; kwargs...)

src/interface/matrixfunctions.jl

@@ -8,6 +8,26 @@ Algorithm type to denote finding the exponential of `A` via the implementation o
 """
 @algdef MatrixFunctionViaLA
 
+"""
+    MatrixFunctionViaTaylor(; tol=eps, balance=true, estimate_order=4)
+
+Algorithm type to denote finding the exponential of `A` through a pure-Julia scaling-and-squaring
+evaluation of its Taylor series, following Fasi & Higham (2018).
+The truncation order and the number of squarings are chosen to reach a relative accuracy `tol`,
+and the Taylor polynomial is evaluated with the Paterson–Stockmeyer scheme.
+When `balance` is `true`, `A` is first balanced by a diagonal similarity.
+`estimate_order` sets how many powers of `A` are formed up front to sharpen the norm estimate via the
+Al-Mohy–Higham quantities `‖Aᵖ‖^(1/p)` (Al-Mohy & Higham, 2009); these powers are reused by the
+Paterson–Stockmeyer evaluation.
+As this algorithm requires no LAPACK support, it also applies at arbitrary precision.
+
+## References
+
+- A. H. Al-Mohy and N. J. Higham, "A New Scaling and Squaring Algorithm for the Matrix
+  Exponential", SIAM J. Matrix Anal. Appl., 31(3), 970–989, 2009.
+"""
+@algdef MatrixFunctionViaTaylor
+
 """
     MatrixFunctionViaEigh(eigh_alg)
 

test/exponential.jl

@@ -5,6 +5,7 @@ using StableRNGs
 using MatrixAlgebraKit: diagview
 using LinearAlgebra
 using LinearAlgebra: exp
+using CUDA, AMDGPU
 
 BLASFloats = (Float32, Float64, ComplexF32, ComplexF64)
 GenericFloats = (Float16, ComplexF16, BigFloat, Complex{BigFloat})
@@ -21,7 +22,7 @@ GenericFloats = (Float16, ComplexF16, BigFloat, Complex{BigFloat})
     @test expA ≈ expA2
     @test A == Ac
 
-    algs = (MatrixFunctionViaLA(), MatrixFunctionViaEig(LAPACK_Simple()))
+    algs = (MatrixFunctionViaLA(), MatrixFunctionViaEig(LAPACK_Simple()), MatrixFunctionViaTaylor())
     @testset "algorithm $alg" for alg in algs
         expA2 = @constinferred exponential(A, alg)
         @test expA ≈ expA2
@@ -46,7 +47,7 @@ end
     @test expAτ ≈ expAτ2
     @test A == Ac
 
-    algs = (MatrixFunctionViaLA(), MatrixFunctionViaEig(LAPACK_Simple()))
+    algs = (MatrixFunctionViaLA(), MatrixFunctionViaEig(LAPACK_Simple()), MatrixFunctionViaTaylor())
     @testset "algorithm $alg" for alg in algs
         expAτ2 = @constinferred exponential((τ, A), alg)
         @test expAτ ≈ expAτ2
@@ -86,3 +87,41 @@ end
     @test expAτ ≈ expAτ2
     @test A == Ac
 end
+
+# GPU tests
+# ---------
+# The Taylor exponential is backend-generic, so the same code runs on GPU. Compare device
+# results against the CPU reference, exercising both `balance` settings (fix 1 & 3) and the
+# scaled `(τ, A)` entrypoint. A badly-scaled matrix exercises the balancing path. If any step
+# fell back to scalar indexing these would error under GPUArrays' scalar-indexing guard.
+function test_exponential_gpu(ArrayT, T)
+    rng = StableRNG(123)
+    m = 54
+    A = randn(rng, T, m, m) ./ (2 * m)
+    τ = randn(rng, T)
+
+    # badly-scaled similarity transform Aᵢⱼ ← Aᵢⱼ sᵢ / sⱼ, to give balancing work to do
+    s = exp10.(range(-real(T)(3), real(T)(3), length = m))
+    Abad = A .* s ./ transpose(s)
+
+    for M in (A, Abad)
+        M_gpu = ArrayT(M)
+        for alg in (MatrixFunctionViaTaylor(), MatrixFunctionViaTaylor(; balance = false))
+            @test Array(exponential(M_gpu, alg)) ≈ exponential(M, alg)
+            @test Array(exponential((τ, M_gpu), alg)) ≈ exponential((τ, M), alg)
+        end
+    end
+    return nothing
+end
+
+if CUDA.functional()
+    @testset "exponential on CUDA for T = $T" for T in BLASFloats
+        test_exponential_gpu(CuArray, T)
+    end
+end
+
+if AMDGPU.functional()
+    @testset "exponential on AMDGPU for T = $T" for T in BLASFloats
+        test_exponential_gpu(ROCArray, T)
+    end
+end