diff --git a/docs/src/user_interface/algorithms.md b/docs/src/user_interface/algorithms.md index e6030787..d90b7607 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 20564773..6153f2e6 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 7b9d5ef1..4d4e0084 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 00000000..64634446 --- /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 bb036916..0e324d20 100644 --- a/src/implementations/exponential.jl +++ b/src/implementations/exponential.jl @@ -114,3 +114,139 @@ 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, squarings = taylor_order_and_squarings(d, tol) + blocksize = ceil(Int, sqrt(order + 1)) + + # Rescale the (≤ p₀) powers we hold into powers of A/2ˢ + resize!(powers, min(p₀, blocksize)) + if squarings > 0 + f = inv(convert(T, 2)^squarings) + fk = one(T) + for k in 1:length(powers) + fk *= f + rmul!(powers[k], fk) # powers[k] ← (A/2ˢ)ᵏ + end + end + # Extend to `blocksize` powers + for p in (length(powers) + 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` and number of squarings `s` minimizing the Paterson–Stockmeyer +# matrix-multiplication count subject to the Taylor remainder bound `(θ/2ˢ)ᵐ⁺¹/(m+1)! ≤ tol`. +# +# The effective norm `θ` per candidate order is the sharpest Al-Mohy–Higham quantity +# `αₚ = max(dₚ, dₚ₊₁)` (with `dₚ = ‖Aᵖ‖^(1/p)`) that is valid for that order: since every +# `k ≥ p(p-1)` is a nonnegative combination of `p` and `p+1`, `‖Aᵏ‖ ≤ αₚᵏ` holds for `k ≥ p(p-1)`, +# so `αₚ` bounds the degree-`m` remainder whenever `m+1 ≥ p(p-1)`. The `length(d)` powers already +# formed count as free in the cost model. +function taylor_order_and_squarings(d::AbstractVector{<:Real}, tol::Real) + p₀ = length(d) + log2tol = log2(Float64(tol)) + log2factorial = 0.0 + best_order = 1 + best_squarings = 0 + best_cost = typemax(Int) + for order in 1:100 + log2factorial += log2(order + 1) + # sharpest valid αₚ = min over p with p(p-1) ≤ order+1 and p+1 ≤ p₀ + θ = Float64(d[1]) + for p in 1:(p₀ - 1) + p * (p - 1) ≤ order + 1 || break + θ = min(θ, max(Float64(d[p]), Float64(d[p + 1]))) + end + # `θ == 0` (a nilpotent Aᵖ) makes the remainder vanish exactly ⇒ no squarings needed; + # guard the `log2(0) = -Inf` before it reaches `ceil(Int, ⋅)`. + excess = log2(θ) - (log2tol + log2factorial) / (order + 1) + squarings = excess > 0 ? ceil(Int, excess) : 0 + blocksize = ceil(Int, sqrt(order + 1)) + cost = max(0, blocksize - p₀) + (cld(order + 1, blocksize) - 1) + squarings + if cost < best_cost + best_cost = cost + best_order = order + best_squarings = squarings + end + end + return best_order, best_squarings +end + +# Evaluate ∑ₖ₌₀ᵐ Aᵏ/k! via the Paterson–Stockmeyer scheme, returning a freshly allocated matrix. +# `powers` holds A, A², …, A^blocksize (already scaled), with blocksize = ceil(√(order+1)). +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 = fld(order, blocksize) + taylor_block!(result, powers, invfactorial, numblocks, blocksize, order) + for j in (numblocks - 1):-1:0 + mul!(block, result, powers[blocksize]) + taylor_block!(result, powers, invfactorial, j, blocksize, order) + result .+= block + end + return result +end + +# Sub-polynomial ∑ᵢ₌₀ᵇ⁻¹ c_{jb+i} Aⁱ of degree `blocksize - 1`, written into `out`. +function taylor_block!(out::AbstractMatrix, powers, invfactorial, j::Integer, blocksize::Integer, order::Integer) + base = j * blocksize + fill!(out, zero(eltype(out))) + diagview(out) .= invfactorial[base + 1] + for i in 1:(blocksize - 1) + k = base + i + k > order && break + out .+= invfactorial[k + 1] .* powers[i] + end + return out +end diff --git a/src/interface/exponential.jl b/src/interface/exponential.jl index bebcdb1f..b8d97ab3 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 8474f9bc..e1da0ecf 100644 --- a/src/interface/matrixfunctions.jl +++ b/src/interface/matrixfunctions.jl @@ -8,6 +8,20 @@ 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)`; these powers are reused by the Paterson–Stockmeyer evaluation. +As this algorithm requires no LAPACK support, it also applies at arbitrary precision. +""" +@algdef MatrixFunctionViaTaylor + """ MatrixFunctionViaEigh(eigh_alg) diff --git a/test/exponential.jl b/test/exponential.jl index ad7628c3..c71c1edc 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,60 @@ end @test expAτ ≈ expAτ2 @test A == Ac end + +@testset "exponential via Taylor for T = $T" for T in GenericFloats + rng = StableRNG(123) + m = 54 + + A = randn(rng, T, m, m) ./ (2 * m) + τ = randn(rng, T) + Ac = copy(A) + alg = MatrixFunctionViaTaylor() + reltol = max(1.0e-7, sqrt(eps(real(T)))) + + expA = @constinferred exponential(A, alg) + @test A == Ac + @test isapprox(expA * exponential(-A, alg), I; rtol = reltol) + @test isapprox(ComplexF64.(expA), exp(ComplexF64.(A)); rtol = reltol) + + expτA = @constinferred exponential((τ, A), alg) + @test isapprox(ComplexF64.(expτA), exp(ComplexF64(τ) .* ComplexF64.(A)); rtol = reltol) +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