diff --git a/docs/src/user_interface/algorithms.md b/docs/src/user_interface/algorithms.md index e6030787d..d90b76074 100644 --- a/docs/src/user_interface/algorithms.md +++ b/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` | diff --git a/docs/src/user_interface/matrix_functions.md b/docs/src/user_interface/matrix_functions.md index 20564773c..6153f2e6a 100644 --- a/docs/src/user_interface/matrix_functions.md +++ b/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 diff --git a/src/MatrixAlgebraKit.jl b/src/MatrixAlgebraKit.jl index 7b9d5ef10..4d4e0084e 100644 --- a/src/MatrixAlgebraKit.jl +++ b/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") diff --git a/src/common/balancing.jl b/src/common/balancing.jl new file mode 100644 index 000000000..646344461 --- /dev/null +++ b/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 diff --git a/src/implementations/exponential.jl b/src/implementations/exponential.jl index bb036916c..a8727a93b 100644 --- a/src/implementations/exponential.jl +++ b/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 diff --git a/src/interface/exponential.jl b/src/interface/exponential.jl index bebcdb1f4..b8d97ab31 100644 --- a/src/interface/exponential.jl +++ b/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...) diff --git a/src/interface/matrixfunctions.jl b/src/interface/matrixfunctions.jl index 8474f9bcc..b0ea57da8 100644 --- a/src/interface/matrixfunctions.jl +++ b/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) diff --git a/test/exponential.jl b/test/exponential.jl index ad7628c3a..5d83f275c 100644 --- a/test/exponential.jl +++ b/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