From d9d8df71e298a39fb653d023fc29e813a6217433 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Sat, 15 Aug 2026 05:15:40 -0400 Subject: [PATCH 1/6] test: fold the matrix-function tests into a `matrixfunctions` group The exponential tests lived in three places: `common/exponential.jl` for the BLAS floats and one file each in the `genericschur` and `genericlinearalgebra` groups for the generic ones. Move the assertions into the shared TestSuite and leave a single driver that covers every element type, which removes both generic groups. Co-Authored-By: Claude Opus 5 (1M context) --- .buildkite/pipeline.yml | 1 + test/README.md | 19 ++- test/common/exponential.jl | 145 ---------------- test/genericlinearalgebra/exponential.jl | 50 ------ test/genericschur/exponential.jl | 51 ------ test/matrixfunctions/exponential.jl | 108 ++++++++++++ test/testsuite/TestSuite.jl | 10 ++ test/testsuite/matrixfunctions/exponential.jl | 160 ++++++++++++++++++ 8 files changed, 293 insertions(+), 251 deletions(-) delete mode 100644 test/common/exponential.jl delete mode 100644 test/genericlinearalgebra/exponential.jl delete mode 100644 test/genericschur/exponential.jl create mode 100644 test/matrixfunctions/exponential.jl create mode 100644 test/testsuite/matrixfunctions/exponential.jl diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 162753ee1..f81b1d055 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -26,6 +26,7 @@ steps: - "rocm" group: - "decompositions" + - "matrixfunctions" - "mooncake" adjustments: - with: diff --git a/test/README.md b/test/README.md index 18eb02a1d..c428592cb 100644 --- a/test/README.md +++ b/test/README.md @@ -4,7 +4,7 @@ Tests are driven by [ParallelTestRunner.jl](https://github.com/JuliaTesting/Para Every `.jl` file under `test/` is auto-discovered and executed in its own worker process, so test files share no state and must be self-contained. -The test environment is a [Pkg workspace](https://pkgdocs.julialang.org/dev/workspaces/) member: +The test environment is a [Pkg workspace](https://pkgdocs.julialang.org/v1/toml-files/#Workspaces) member: `test/Project.toml` declares the test dependencies and resolves the parent package through `[sources]`, while the whole workspace shares a single `Manifest.toml` at the repository root. @@ -73,13 +73,21 @@ file placed directly in `test/` would belong to no group and never run in CI. | Group | Contents | |-------|----------| -| `common` | Algorithm selection and defaults, truncation strategies, projections, matrix exponential, Aqua code-quality checks | +| `common` | Algorithm selection and defaults, truncation strategies, projections, Aqua code-quality checks | | `decompositions` | `qr`, `lq`, `svd`, `eig`, `eigh`, `gen_eig`, `schur`, `polar`, `orthnull` on CPU and GPU array types | +| `matrixfunctions` | `exponential` on CPU and GPU array types | | `chainrules` | ChainRulesCore rules, exercised through ChainRulesTestUtils and Zygote | | `mooncake` | Mooncake AD rules | | `enzyme` | Enzyme AD rules, exercised through EnzymeTestUtils | -| `genericlinearalgebra` | `MatrixAlgebraKitGenericLinearAlgebraExt` | -| `genericschur` | `MatrixAlgebraKitGenericSchurExt` | + +The GPU runners select groups explicitly, through the `group` matrix of `.buildkite/pipeline.yml`, so +a new group needs to be listed there to get GPU coverage. + +The generic element types (`BigFloat` and friends) are not a group of their own: each driver covers +them alongside the BLAS floats, naming the `MatrixAlgebraKitGenericSchurExt` and +`MatrixAlgebraKitGenericLinearAlgebraExt` algorithms explicitly. Both extensions are loaded at once, +so a bare `QRIteration()` resolves its driver to GenericLinearAlgebra, which provides no `geev!`; +`eig`-based algorithms therefore have to spell out `QRIteration(; driver = GS())`. Two directories are *not* groups: `testsuite/` holds the shared implementation described below and is excluded both from discovery (in `runtests.jl`) and from CI (`exclude: '["testsuite"]'`). @@ -124,7 +132,8 @@ Supporting infrastructure in the module: - `precision(T)` — the default tolerance, `sqrt(eps(real(T)))`. - Predicates used throughout the assertions: `isleftnull`, `isrightnull`, `isleftcomplete`, `isrightcomplete`, `has_positive_diagonal`. -- `instantiate_unitary`, `instantiate_rank_deficient_matrix` — inputs with prescribed structure. +- `instantiate_unitary`, `instantiate_rank_deficient_matrix`, `instantiate_smallnorm_matrix` — inputs + with prescribed structure or spectrum. `test/linearmap.jl` is a further helper, defining a `LinearMap` wrapper that is deliberately *not* an `AbstractMatrix`, used to check the generic code paths. It is excluded from discovery and included diff --git a/test/common/exponential.jl b/test/common/exponential.jl deleted file mode 100644 index 981a0a265..000000000 --- a/test/common/exponential.jl +++ /dev/null @@ -1,145 +0,0 @@ -using MatrixAlgebraKit -using Test -using TestExtras -using StableRNGs -using MatrixAlgebraKit: diagview -using LinearAlgebra -using LinearAlgebra: exp -using CUDA, AMDGPU - -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()), MatrixFunctionViaTaylor()) - @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()), MatrixFunctionViaTaylor()) - @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 non-Matrix input $T" for T in BLASFloats - rng = StableRNG(123) - m = 12 - A = LinearAlgebra.normalize!(randn(rng, T, m, m)) - expA = LinearAlgebra.exp(A) - - wrappers = ( - ("view", B -> view(B, :, :)), - ("PermutedDimsArray", B -> PermutedDimsArray(permutedims(B), (2, 1))), - ("ReshapedArray", B -> reshape(view(vec(B), 1:(m * m)), m, m)), - ) - @testset "$name" for (name, wrap) in wrappers - W = wrap(copy(A)) - @test !(W isa Matrix) - @test exponential!(W) ≈ expA - end -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 - -# 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 diff --git a/test/genericlinearalgebra/exponential.jl b/test/genericlinearalgebra/exponential.jl deleted file mode 100644 index 5ffdaa10b..000000000 --- a/test/genericlinearalgebra/exponential.jl +++ /dev/null @@ -1,50 +0,0 @@ -using MatrixAlgebraKit -using Test -using TestExtras -using StableRNGs -using MatrixAlgebraKit: diagview -using LinearAlgebra -using GenericLinearAlgebra - -is_buildkite = get(ENV, "BUILDKITE", "false") == "true" - -if !is_buildkite - GenericFloats = (BigFloat, Complex{BigFloat}) - - @testset "exponential! for T = $T" for T in GenericFloats - rng = StableRNG(123) - m = 54 - - A = project_hermitian!(randn(rng, T, m, m)) - D, V = @constinferred eigh_full(A) - algs = (MatrixFunctionViaEigh(GLA_QRIteration()),) - @testset "algorithm $alg" for alg in algs - expA = @constinferred exponential!(copy(A); alg) - expA2 = @constinferred exponential(A; alg) - @test expA2 ≈ expA - - Dexp, Vexp = @constinferred eigh_full(expA) - @test diagview(Dexp) ≈ LinearAlgebra.exp.(diagview(D)) - end - end - - using GenericSchur - @testset "exponential! for T1 = $T1, T2 = $T2" for T1 in GenericFloats, T2 in GenericFloats - rng = StableRNG(123) - m = 54 - A = project_hermitian!(randn(rng, T1, m, m)) - τ = randn(rng, T2) - - D, V = @constinferred eigh_full(A) - algs = (MatrixFunctionViaEigh(GLA_QRIteration()),) - @testset "algorithm $alg" for alg in algs - expτA = @constinferred exponential!((τ, copy(A)); alg) - expτA2 = @constinferred exponential((τ, A); alg) - @test expτA2 ≈ expτA - - Dexp, Vexp = @constinferred eig_full(expτA) - - @test sort(diagview(Dexp); by = real) ≈ sort(LinearAlgebra.exp.(diagview(D) .* τ); by = real) - end - end -end diff --git a/test/genericschur/exponential.jl b/test/genericschur/exponential.jl deleted file mode 100644 index 7754e02c7..000000000 --- a/test/genericschur/exponential.jl +++ /dev/null @@ -1,51 +0,0 @@ -using MatrixAlgebraKit -using Test -using TestExtras -using StableRNGs -using MatrixAlgebraKit: diagview -using LinearAlgebra -using GenericSchur - -is_buildkite = get(ENV, "BUILDKITE", "false") == "true" - -if !is_buildkite - GenericFloats = (BigFloat, Complex{BigFloat}) - - @testset "exponential! for T = $T" for T in GenericFloats - rng = StableRNG(123) - m = 54 - - A = randn(rng, T, m, m) - D, V = @constinferred eig_full(A) - algs = (MatrixFunctionViaEig(GS_QRIteration()),) - expA_LA = @constinferred exponential(A) - @testset "algorithm $alg" for alg in algs - expA = @constinferred exponential!(copy(A)) - expA2 = @constinferred exponential(A; alg = alg) - @test expA ≈ expA_LA - @test expA2 ≈ expA - - Dexp, Vexp = @constinferred eig_full(expA) - @test sort(diagview(Dexp); by = imag) ≈ sort(LinearAlgebra.exp.(diagview(D)); by = imag) - end - end - - @testset "exponential! for T1 = $T1, T2 = $T2" for T1 in GenericFloats, T2 in GenericFloats - rng = StableRNG(123) - m = 54 - - A = randn(rng, T1, m, m) - τ = randn(rng, T2) - - D, V = @constinferred eig_full(A) - algs = (MatrixFunctionViaEig(GS_QRIteration()),) - @testset "algorithm $alg" for alg in algs - expτA = @constinferred exponential!((τ, copy(A))) - expτA2 = @constinferred exponential((τ, A); alg) - @test expτA2 ≈ expτA - - Dexp, Vexp = @constinferred eig_full(expτA) - @test sort(diagview(Dexp); by = x -> (imag(x), real(x))) ≈ sort(LinearAlgebra.exp.(diagview(D) .* τ); by = x -> (imag(x), real(x))) - end - end -end diff --git a/test/matrixfunctions/exponential.jl b/test/matrixfunctions/exponential.jl new file mode 100644 index 000000000..8829de6e8 --- /dev/null +++ b/test/matrixfunctions/exponential.jl @@ -0,0 +1,108 @@ +using MatrixAlgebraKit +using LinearAlgebra: Diagonal +using MatrixAlgebraKit: GLA, GS +using CUDA, AMDGPU +using GenericSchur, GenericLinearAlgebra + +if @isdefined(fast_tests) && fast_tests + BLASFloats = (Float64, ComplexF64) + GenericFloats = (BigFloat, Complex{BigFloat}) +else + BLASFloats = (Float32, Float64, ComplexF32, ComplexF64) + GenericFloats = (BigFloat, Complex{BigFloat}) +end +# only the `Diagonal` fast path applies to these, as they have no `eig`/`eigh` support +DiagonalOnlyFloats = (Float16, ComplexF16) + +@isdefined(TestSuite) || include("../testsuite/TestSuite.jl") +using .TestSuite + +is_buildkite = get(ENV, "BUILDKITE", "false") == "true" + +m = 54 + +# CPU tests +# --------- +if !is_buildkite + for T in BLASFloats + TestSuite.seed_rng!(123) + LAPACK_EIG_ALGS = ( + MatrixFunctionViaLA(), + MatrixFunctionViaEig(QRIteration()), + MatrixFunctionViaTaylor(), + ) + LAPACK_EIGH_ALGS = (MatrixFunctionViaEigh(QRIteration()), MatrixFunctionViaEigh(DivideAndConquer())) + TestSuite.test_exponential(T, (m, m)) + TestSuite.test_exponential_algs(T, (m, m), LAPACK_EIG_ALGS) + TestSuite.test_exponential_scaled(T, (m, m), LAPACK_EIG_ALGS) + TestSuite.test_exponential_hermitian(T, (m, m), LAPACK_EIGH_ALGS) + TestSuite.test_exponential_taylor(T, (m, m)) + TestSuite.test_exponential_reference(T, (m, m)) + TestSuite.test_exponential_wrappers(T, (12, 12)) + TestSuite.test_exponential_domain(T, (m, m), LAPACK_EIGH_ALGS) + end + + # `eig` comes from GenericSchur, `eigh` from GenericLinearAlgebra. Both are loaded here, so name + # the driver explicitly instead of relying on `default_driver`. + for T in GenericFloats + TestSuite.seed_rng!(123) + GS_ALGS = (MatrixFunctionViaEig(QRIteration(; driver = GS())), MatrixFunctionViaTaylor()) + GLA_ALGS = (MatrixFunctionViaEigh(QRIteration(; driver = GLA())),) + TestSuite.test_exponential_algs(T, (24, 24), GS_ALGS) + TestSuite.test_exponential_scaled(T, (24, 24), GS_ALGS) + TestSuite.test_exponential_hermitian(T, (24, 24), GLA_ALGS) + TestSuite.test_exponential_taylor(T, (24, 24)) + end + + for T in (BLASFloats..., GenericFloats..., DiagonalOnlyFloats...) + TestSuite.seed_rng!(123) + AT = Diagonal{T, Vector{T}} + test_spectrum = !(T in DiagonalOnlyFloats) + TestSuite.test_exponential(AT, m) + TestSuite.test_exponential_algs(AT, m, (DiagonalAlgorithm(),)) + TestSuite.test_exponential_scaled(AT, m, (DiagonalAlgorithm(),)) + TestSuite.test_exponential_hermitian(AT, m, (DiagonalAlgorithm(),); test_spectrum) + TestSuite.test_exponential_reference(AT, m; test_hermitian = !(T in GenericFloats)) + end +end + +# CUDA tests +# ---------- +# a general dense matrix is supported on device through the native `MatrixFunctionViaTaylor` +if CUDA.functional() + for T in BLASFloats + TestSuite.seed_rng!(123) + CUDA_EIGH_ALGS = (MatrixFunctionViaEigh(Jacobi()), MatrixFunctionViaEigh(DivideAndConquer())) + TestSuite.test_exponential_algs(CuMatrix{T}, (m, m), (MatrixFunctionViaTaylor(),)) + TestSuite.test_exponential_scaled(CuMatrix{T}, (m, m), (MatrixFunctionViaTaylor(),)) + TestSuite.test_exponential_taylor(CuMatrix{T}, (m, m)) + TestSuite.test_exponential_hermitian(CuMatrix{T}, (m, m), CUDA_EIGH_ALGS) + TestSuite.test_exponential_domain(CuMatrix{T}, (m, m), CUDA_EIGH_ALGS) + + AT = Diagonal{T, CuVector{T}} + TestSuite.test_exponential(AT, m) + TestSuite.test_exponential_algs(AT, m, (DiagonalAlgorithm(),)) + TestSuite.test_exponential_scaled(AT, m, (DiagonalAlgorithm(),)) + TestSuite.test_exponential_hermitian(AT, m, (DiagonalAlgorithm(),)) + end +end + +# AMDGPU tests +# ------------ +if AMDGPU.functional() + for T in BLASFloats + TestSuite.seed_rng!(123) + ROC_EIGH_ALGS = (MatrixFunctionViaEigh(Jacobi()), MatrixFunctionViaEigh(DivideAndConquer())) + TestSuite.test_exponential_algs(ROCMatrix{T}, (m, m), (MatrixFunctionViaTaylor(),)) + TestSuite.test_exponential_scaled(ROCMatrix{T}, (m, m), (MatrixFunctionViaTaylor(),)) + TestSuite.test_exponential_taylor(ROCMatrix{T}, (m, m)) + TestSuite.test_exponential_hermitian(ROCMatrix{T}, (m, m), ROC_EIGH_ALGS) + TestSuite.test_exponential_domain(ROCMatrix{T}, (m, m), ROC_EIGH_ALGS) + + AT = Diagonal{T, ROCVector{T}} + TestSuite.test_exponential(AT, m) + TestSuite.test_exponential_algs(AT, m, (DiagonalAlgorithm(),)) + TestSuite.test_exponential_scaled(AT, m, (DiagonalAlgorithm(),)) + TestSuite.test_exponential_hermitian(AT, m, (DiagonalAlgorithm(),)) + end +end diff --git a/test/testsuite/TestSuite.jl b/test/testsuite/TestSuite.jl index 36ae68304..bffbfe204 100644 --- a/test/testsuite/TestSuite.jl +++ b/test/testsuite/TestSuite.jl @@ -102,6 +102,12 @@ function instantiate_rank_deficient_matrix(::Type{T}, sz; trunc = truncrank(div( return Diagonal(diag(mul!(A, V, C))) end +# spectral radius at most one, keeping `exponential` well inside the principal branch +function instantiate_smallnorm_matrix(T, sz) + A = instantiate_matrix(T, sz) + return A / norm(A) +end + include("ad_utils.jl") include("projections.jl") @@ -117,6 +123,10 @@ include("decompositions/eigh.jl") include("decompositions/orthnull.jl") include("decompositions/svd.jl") +# Matrix functions +# ---------------- +include("matrixfunctions/exponential.jl") + # Mooncake # -------- include("mooncake/mooncake.jl") diff --git a/test/testsuite/matrixfunctions/exponential.jl b/test/testsuite/matrixfunctions/exponential.jl new file mode 100644 index 000000000..08e2654a1 --- /dev/null +++ b/test/testsuite/matrixfunctions/exponential.jl @@ -0,0 +1,160 @@ +using TestExtras +using LinearAlgebra: LinearAlgebra, I +using MatrixAlgebraKit: ishermitian + +# `exp(A) * exp(-A) ≈ I` holds for every algorithm and backend, since `A` and `-A` commute + +function test_exponential(T::Type, sz; kwargs...) + summary_str = testargs_summary(T, sz) + return @testset "exponential $summary_str" begin + A = instantiate_smallnorm_matrix(T, sz) + Ac = deepcopy(A) + + expA = @testinferred exponential(A) + @test eltype(expA) == eltype(A) + @test expA * exponential(-A) ≈ I + @test A == Ac + + # the in-place method may not be able to reuse the provided output + expA2 = @testinferred exponential!(deepcopy(A), deepcopy(expA)) + @test expA2 ≈ expA + end +end + +function test_exponential_algs(T::Type, sz, algs; kwargs...) + summary_str = testargs_summary(T, sz) + return @testset "exponential algorithm $alg $summary_str" for alg in algs + A = instantiate_smallnorm_matrix(T, sz) + Ac = deepcopy(A) + + expA = @testinferred exponential(A, alg) + @test eltype(expA) == eltype(A) + @test expA * exponential(-A, alg) ≈ I + @test A == Ac + end +end + +# the scaled entrypoint `exponential((τ, A))` computes `exp(τ * A)` +function test_exponential_scaled(T::Type, sz, algs; kwargs...) + R = real(eltype(T)) + summary_str = testargs_summary(T, sz) + return @testset "exponential scaled algorithm $alg $summary_str" for alg in algs + A = instantiate_smallnorm_matrix(T, sz) + Ac = deepcopy(A) + + # both a scalar of the matrix eltype and a real one, to exercise promotion + @testset "τ::$(typeof(τ))" for τ in (randn(rng, eltype(T)), randn(rng, R)) + expτA = @testinferred exponential((τ, A), alg) + @test eltype(expτA) == eltype(A) + @test expτA ≈ exponential(τ * A, alg) + @test expτA * exponential((-τ, A), alg) ≈ I + @test A == Ac + end + end +end + +# `exp(A)` of a real hermitian `A` is built as a symmetric product and is hermitian to the last bit; +# the complex case goes through `V exp(D) V'` and is hermitian only up to roundoff +function test_exponential_hermitian( + T::Type, sz, algs; + exact_hermiticity = eltype(T) <: Real, test_spectrum = true, kwargs... + ) + summary_str = testargs_summary(T, sz) + return @testset "exponential hermitian algorithm $alg $summary_str" for alg in algs + A = project_hermitian!(instantiate_smallnorm_matrix(T, sz)) + Ac = deepcopy(A) + + expA = @testinferred exponential(A, alg) + @test eltype(expA) == eltype(A) + @test expA * exponential(-A, alg) ≈ I + @test A == Ac + + if exact_hermiticity + @test ishermitian(expA) + # `eigh_vals` rejects anything but an exactly hermitian matrix + test_spectrum && @test eigh_vals(expA) ≈ exp.(eigh_vals(A)) + else + @test ishermitian(expA; rtol = precision(T)) + end + + τ = randn(rng, real(eltype(T))) + expτA = @testinferred exponential((τ, A), alg) + @test expτA ≈ exponential(τ * A, alg) + exact_hermiticity && @test ishermitian(expτA) + end +end + +# `MatrixFunctionViaTaylor` is the only algorithm that applies to a general matrix on device and at +# arbitrary precision. Its balancing step is exercised with a badly-scaled similarity transform +# `Aᵢⱼ ← Aᵢⱼ sᵢ / sⱼ`; dense input only, as that would densify a `Diagonal`. +function test_exponential_taylor(T::Type, sz; kwargs...) + R = real(eltype(T)) + summary_str = testargs_summary(T, sz) + return @testset "exponential Taylor $summary_str" begin + A = instantiate_smallnorm_matrix(T, sz) + n = size(A, 1) + s = similar(A, R, n) + copyto!(s, exp10.(range(-R(3), R(3), length = n))) + Abad = A .* s ./ transpose(s) + + @testset "balance = $balance" for balance in (true, false) + alg = MatrixFunctionViaTaylor(; balance) + expA = @testinferred exponential(A, alg) + @test eltype(expA) == eltype(A) + @test expA * exponential(-A, alg) ≈ I + + # `Abad` is too ill-conditioned for the inverse check, but balancing must not + # change the result + for M in (A, Abad) + expM = @testinferred exponential(M, alg) + @test eltype(expM) == eltype(M) + @test expM ≈ exponential(M, MatrixFunctionViaTaylor(; balance = !balance)) + end + end + end +end + +# `MatrixFunctionViaEigh` requires hermitian input, and says so rather than silently projecting +function test_exponential_domain(T::Type, sz, algs; kwargs...) + summary_str = testargs_summary(T, sz) + return @testset "exponential domain algorithm $alg $summary_str" for alg in algs + A = instantiate_smallnorm_matrix(T, sz) + @test_throws DomainError exponential(A, alg) + end +end + +# the kernels must not assume a strided layout +function test_exponential_wrappers(T::Type, sz; kwargs...) + summary_str = testargs_summary(T, sz) + return @testset "exponential non-Matrix input $summary_str" begin + A = instantiate_smallnorm_matrix(T, sz) + m = size(A, 1) + expA = exponential(A) + + wrappers = ( + ("view", B -> view(B, :, :)), + ("PermutedDimsArray", B -> PermutedDimsArray(permutedims(B), (2, 1))), + ("ReshapedArray", B -> reshape(view(vec(B), 1:(m * m)), m, m)), + ) + @testset "$name" for (name, wrap) in wrappers + W = wrap(deepcopy(A)) + @test !(W isa Matrix) + @test exponential!(W) ≈ expA + end + end +end + +# cross-check against `LinearAlgebra`, host arrays only. `LinearAlgebra` has no matrix functions +# for a `Hermitian` wrapper outside the BLAS floats, hence `test_hermitian`. +function test_exponential_reference(T::Type, sz; test_hermitian = true, kwargs...) + summary_str = testargs_summary(T, sz) + return @testset "exponential vs LinearAlgebra $summary_str" begin + A = instantiate_smallnorm_matrix(T, sz) + @test exponential(A) ≈ LinearAlgebra.exp(A) + + if test_hermitian + H = project_hermitian!(instantiate_smallnorm_matrix(T, sz)) + @test exponential(H) ≈ LinearAlgebra.exp(LinearAlgebra.Hermitian(H)) + end + end +end From b8ebf7601bdda0de687fe64b21ee587abd56f7b0 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Sat, 15 Aug 2026 05:19:52 -0400 Subject: [PATCH 2/6] feat: add the `squareroot` matrix function `squareroot` computes the principal square root, through `LinearAlgebra`, an (hermitian) eigenvalue decomposition, or the `Diagonal` fast path. The scalar type of the output matches the input, so a real matrix whose spectrum reaches the negative real axis raises a `DomainError`; eigenvalues that are negative only within `domain_atol` are clamped onto zero instead. `MatrixFunctionViaEig` and `MatrixFunctionViaEigh` gain a `domain_atol` field for this, and `src/implementations/matrixfunctions.jl` collects the input handling, reconstruction and domain helpers that the other matrix functions will share. Co-Authored-By: Claude Opus 5 (1M context) --- ext/MatrixAlgebraKitGenericSchurExt.jl | 7 ++ src/MatrixAlgebraKit.jl | 4 + src/common/defaults.jl | 21 ++++ src/implementations/matrixfunctions.jl | 162 +++++++++++++++++++++++++ src/implementations/squareroot.jl | 64 ++++++++++ src/interface/matrixfunctions.jl | 38 +++++- src/interface/squareroot.jl | 40 ++++++ 7 files changed, 331 insertions(+), 5 deletions(-) create mode 100644 src/implementations/matrixfunctions.jl create mode 100644 src/implementations/squareroot.jl create mode 100644 src/interface/squareroot.jl diff --git a/ext/MatrixAlgebraKitGenericSchurExt.jl b/ext/MatrixAlgebraKitGenericSchurExt.jl index c34f1abc4..8b7aeb100 100644 --- a/ext/MatrixAlgebraKitGenericSchurExt.jl +++ b/ext/MatrixAlgebraKitGenericSchurExt.jl @@ -21,6 +21,13 @@ function MatrixAlgebraKit.default_exponential_algorithm( return MatrixFunctionViaEig(eig_alg) end +function MatrixAlgebraKit.default_squareroot_algorithm( + type::Type{T}; domain_atol::Real = -1.0, kwargs... + ) where {T <: StridedMatrix{<:GSFloat}} + eig_alg = MatrixAlgebraKit.default_eig_algorithm(type; kwargs...) + return MatrixFunctionViaEig(eig_alg; domain_atol) +end + function geev!(::GS, A::AbstractMatrix, Dd::AbstractVector, V::AbstractMatrix; kwargs...) D, Vmat = GenericSchur.eigen!(A) copyto!(Dd, D) diff --git a/src/MatrixAlgebraKit.jl b/src/MatrixAlgebraKit.jl index 4d4e0084e..a11179f84 100644 --- a/src/MatrixAlgebraKit.jl +++ b/src/MatrixAlgebraKit.jl @@ -31,6 +31,7 @@ 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 squareroot, squareroot! export Householder, Native_HouseholderQR, Native_HouseholderLQ export DivideAndConquer, SafeDivideAndConquer, QRIteration, Bisection, Jacobi, SVDViaPolar @@ -115,6 +116,7 @@ include("interface/schur.jl") include("interface/polar.jl") include("interface/orthnull.jl") include("interface/exponential.jl") +include("interface/squareroot.jl") include("implementations/projections.jl") include("implementations/truncation.jl") @@ -128,6 +130,8 @@ include("implementations/schur.jl") include("implementations/polar.jl") include("implementations/orthnull.jl") include("implementations/exponential.jl") +include("implementations/matrixfunctions.jl") +include("implementations/squareroot.jl") include("common/gauge.jl") # needs to be defined after the functions are diff --git a/src/common/defaults.jl b/src/common/defaults.jl index c64c24ad4..17c20eb8e 100644 --- a/src/common/defaults.jl +++ b/src/common/defaults.jl @@ -43,6 +43,27 @@ Default tolerance for deciding to warn if the provided `A` is not hermitian. """ default_hermitian_tol(A) = eps(norm(A, Inf))^(3 / 4) +""" + default_domain_atol(λ, alg) + +Default absolute tolerance for deciding when the eigenvalues `λ` should be considered +to lie outside of the domain of a matrix function, e.g. on the negative real axis for +[`squareroot`](@ref) of a real matrix. + +The tolerance has to absorb the error with which `alg` obtained `λ`, and that error differs by +orders of magnitude between the algorithms, so the default is algorithm-dependent. +It is both the default clamping radius of [`squareroot`](@ref), which exposes it as `domain_atol`, +and the never user-settable tolerance with which a complex eigenvalue of a real matrix is decided to +lie *on* the negative real axis. +See [Domain considerations](@ref sec_matrixfunction_domain) for the resulting values. +""" +function default_domain_atol end + +# roundoff scales only with number of elements and spectrum - e.g. hermitian precision +# conditioning is less strict +_roundoff_domain_atol(λ) = length(λ) * eps(real(float(one(eltype(λ))))) * maximum(abs, λ; init = abs(zero(eltype(λ)))) +_conditioning_domain_atol(λ) = defaulttol(λ) * maximum(abs, λ; init = abs(zero(eltype(λ)))) + const DEFAULT_FIXGAUGE = Ref(true) diff --git a/src/implementations/matrixfunctions.jl b/src/implementations/matrixfunctions.jl new file mode 100644 index 000000000..6feddb691 --- /dev/null +++ b/src/implementations/matrixfunctions.jl @@ -0,0 +1,162 @@ +# Inputs +# ------ +_matrixfunction_copy_input(A::AbstractMatrix) = copy!(similar(A, float(eltype(A))), A) +_matrixfunction_copy_input(A::Diagonal) = map_diagonal(float, A) + +function _matrixfunction_check_input(A::AbstractMatrix, out, ::AbstractAlgorithm) + m = LinearAlgebra.checksquare(A) + @check_size(out, (m, m)) + @check_scalar(out, A) + return nothing +end + +function _matrixfunction_check_input(A::AbstractMatrix, out, ::DiagonalAlgorithm) + m = LinearAlgebra.checksquare(A) + @assert isdiag(A) + @assert out isa Diagonal + @check_size(out, (m, m)) + @check_scalar(out, A) + return nothing +end + +# Reconstruction from an eigenvalue decomposition +# ----------------------------------------------- +# both take the already-transformed eigenvalues `fD = f(D)` and rebuild `f(A)` + +# `f(A) = V f(D) V⁻¹`. A real matrix has a complex decomposition but a real `f(A)` whenever `f(D)` +# closes under conjugation, so the imaginary part is dropped after the solve; the callers are +# responsible for rejecting the inputs where it would not be. +function _apply_eig!(fA, V, fD) + if eltype(fA) <: Real + VfD = V * fD + fAc = rdiv!(VfD, LinearAlgebra.lu!(V)) + return fA .= real.(fAc) + else + fA .= V .* transpose(diagview(fD)) + return rdiv!(fA, LinearAlgebra.lu!(V)) + end +end + +# `f(A) = V f(D) V'` for a hermitian `A`; the product is hermitian only up to roundoff +_apply_eigh!(fA, V, fD) = project_hermitian!(mul!(fA, V * fD, V')) + +# Domain handling +# --------------- +# `atol` is the slack on the domain boundary, i.e. how far outside the domain an eigenvalue may +# stray before it counts as a genuine violation rather than a rounding artifact. `axis_atol` is the +# accuracy with which the algorithm resolves whether an eigenvalue lies *on* the negative real axis +# at all, which is a property of the eigensolver rather than a user choice. + +# each algorithm obtains its eigenvalues with a different accuracy, so each brings its own default +default_domain_atol(λ, ::DiagonalAlgorithm) = _roundoff_domain_atol(λ) +default_domain_atol(λ, ::MatrixFunctionViaEigh) = _roundoff_domain_atol(λ) +default_domain_atol(λ, ::MatrixFunctionViaEig) = _conditioning_domain_atol(λ) + +# a negative tolerance denotes the runtime default, which keeps the algorithm types concrete +_domain_atol(alg::Union{MatrixFunctionViaEig, MatrixFunctionViaEigh}) = alg.domain_atol +_domain_atol(alg::DiagonalAlgorithm) = get(alg.kwargs, :domain_atol, -1.0) + +# callers resolve the default before delegating to an inner `DiagonalAlgorithm`, so that the +# tolerance follows the algorithm that computed the eigenvalues +function _resolve_domain_atol(λ, alg) + atol = _domain_atol(alg) + R = real(float(eltype(λ))) + return atol < 0 ? convert(R, default_domain_atol(λ, alg)) : convert(R, atol) +end + +_axis_atol(λ, alg) = convert(real(float(eltype(λ))), default_domain_atol(λ, alg)) + +# the throwing branches live in `@noinline` helpers so that the reductions and broadcasts +# below stay free of error-path code, which keeps them GPU friendly +@noinline function throw_negative_eigenvalue(λmin, atol, what) + return throw( + DomainError( + λmin, + "The matrix has $what beyond `domain_atol = $atol` and the result of this matrix function is complex. " * + "Pass a complex matrix to obtain the principal value, or increase `domain_atol` if the eigenvalue is a rounding artifact." + ) + ) +end + +# real eigenvalues are on the real axis by construction, so `axis_atol` is unused here +function _check_domain_eigenvalues(λ::AbstractVector{<:Real}, atol::Real, axis_atol::Real = atol) + λmin = minimum(λ; init = zero(eltype(λ))) + λmin < -atol && throw_negative_eigenvalue(λmin, atol, "a negative real eigenvalue") + return λ +end + +# for the complex eigenvalues of a real matrix, only the ones (numerically) on the negative real +# axis obstruct a real result; complex-conjugate pairs do not +_onaxis(x, axis_atol) = abs(imag(x)) <= axis_atol && real(x) < 0 + +function _check_domain_eigenvalues(λ::AbstractVector{<:Complex}, atol::Real, axis_atol::Real = atol) + λmin = mapreduce(x -> _onaxis(x, axis_atol) ? real(x) : zero(real(x)), min, λ; init = zero(real(eltype(λ)))) + λmin < -atol && throw_negative_eigenvalue(λmin, atol, "an eigenvalue on the negative real axis") + return λ +end + +# move the eigenvalues that violate the domain within `atol` onto the boundary +function _clamp_domain_eigenvalues!(λ::AbstractVector{<:Real}, atol::Real, axis_atol::Real = atol) + _check_domain_eigenvalues(λ, atol, axis_atol) + λ .= max.(λ, zero(eltype(λ))) + return λ +end + +function _clamp_domain_eigenvalues!(λ::AbstractVector{<:Complex}, atol::Real, axis_atol::Real = atol) + _check_domain_eigenvalues(λ, atol, axis_atol) + λ .= ifelse.(_onaxis.(λ, axis_atol), zero(eltype(λ)), λ) + return λ +end + +# Domain handling for `MatrixFunctionViaLA` +# ----------------------------------------- +# `LinearAlgebra` never exposes the spectrum, so the check happens in result space and +# `domain_atol` bounds the imaginary part of `f(A)` instead + +@noinline function throw_la_kwargs(f, ks) + return throw( + ArgumentError("`MatrixFunctionViaLA` only accepts the `domain_atol` keyword argument for `$f`, got $ks") + ) +end + +# `MatrixFunctionViaLA` accepts generic keywords, so the kernels validate the ones they support +function _la_domain_atol(alg::MatrixFunctionViaLA, f) + ks = keys(alg.kwargs) + (isempty(ks) || ks == (:domain_atol,)) || throw_la_kwargs(f, ks) + return get(alg.kwargs, :domain_atol, -1.0) +end + +@noinline function throw_complex_result(f, atol, imagmax) + return throw( + DomainError( + f, + "The result of this matrix function applied to the given real matrix is complex (eigenvalues on the negative real axis): " * + "its imaginary part reaches $imagmax, beyond `domain_atol = $atol`. Pass a complex matrix to obtain the principal " * + "value, or increase `domain_atol` if the imaginary part is a rounding artifact." + ) + ) +end + +@noinline function throw_nonfinite_result(f) + return throw( + DomainError( + f, + "The result of this matrix function is not finite, which signals a (numerically) singular input for which it is undefined. " * + "Use `MatrixFunctionViaEig`/`MatrixFunctionViaEigh` to have the spectrum itself checked against `domain_atol`." + ) + ) +end + +# a rounding-level imaginary part is not a domain violation: `LinearAlgebra` casts back to real only +# when the imaginary part vanishes identically +function _la_project_real!(fA, fAc, domain_atol::Real, f) + all(isfinite, fAc) || throw_nonfinite_result(f) + R = real(eltype(fA)) + # the working precision is that of the output: `LinearAlgebra` computes in complex arithmetic + # throughout, so e.g. a `Float32` input promotes all the way to `ComplexF64` + atol = domain_atol < 0 ? defaulttol(fA) * convert(R, norm(fAc, Inf)) : convert(R, domain_atol) + imagmax = convert(R, maximum(abs ∘ imag, fAc; init = zero(real(eltype(fAc))))) + imagmax <= atol || throw_complex_result(f, atol, imagmax) + fA .= real.(fAc) + return fA +end diff --git a/src/implementations/squareroot.jl b/src/implementations/squareroot.jl new file mode 100644 index 000000000..6cbb2e84d --- /dev/null +++ b/src/implementations/squareroot.jl @@ -0,0 +1,64 @@ +# Inputs +# ------ +copy_input(::typeof(squareroot), A::AbstractMatrix) = _matrixfunction_copy_input(A) + +function check_input(::typeof(squareroot!), A::AbstractMatrix, sqrtA, alg::AbstractAlgorithm) + return _matrixfunction_check_input(A, sqrtA, alg) +end + +# Algorithm selection +# ------------------- +squareroot!(A::AbstractMatrix, alg::DefaultAlgorithm) = squareroot!(A, select_algorithm(squareroot!, A, nothing; alg.kwargs...)) +squareroot!(A::AbstractMatrix, out, alg::DefaultAlgorithm) = squareroot!(A, out, select_algorithm(squareroot!, A, nothing; alg.kwargs...)) + +# Outputs +# ------- +initialize_output(::typeof(squareroot!), A::AbstractMatrix, ::AbstractAlgorithm) = A + +# Implementation +# -------------- +function squareroot!(A::AbstractMatrix, sqrtA, alg::MatrixFunctionViaLA) + check_input(squareroot!, A, sqrtA, alg) + domain_atol = _la_domain_atol(alg, squareroot!) + # `LinearAlgebra.sqrt` of a real matrix is real whenever the principal square root is + sqrtAc = LinearAlgebra.sqrt(A) + if eltype(sqrtAc) <: Complex && !(eltype(sqrtA) <: Complex) + _la_project_real!(sqrtA, sqrtAc, domain_atol, squareroot!) + else + copy!(sqrtA, sqrtAc) + end + return sqrtA +end + +function squareroot!(A::AbstractMatrix, sqrtA, alg::MatrixFunctionViaEigh) + check_input(squareroot!, A, sqrtA, alg) + D, V = eigh_full!(A, alg.eigh_alg) + λ = diagview(D) + _clamp_domain_eigenvalues!(λ, _resolve_domain_atol(λ, alg)) + λ .= sqrt.(λ) + return _apply_eigh!(sqrtA, V, D) +end + +function squareroot!(A::AbstractMatrix, sqrtA, alg::MatrixFunctionViaEig) + check_input(squareroot!, A, sqrtA, alg) + D, V = eig_full!(A, alg.eig_alg) + λ = diagview(D) + atol = _resolve_domain_atol(λ, alg) + # a real result requires the spectrum to stay off the negative real axis; whether an eigenvalue + # sits *on* that axis keeps the algorithm default however `domain_atol` was set + eltype(A) <: Real && _clamp_domain_eigenvalues!(λ, atol, _axis_atol(λ, alg)) + diag_alg = DiagonalAlgorithm(; domain_atol = atol) + return _apply_eig!(sqrtA, V, squareroot!(D, D, diag_alg)) +end + +# Diagonal logic +# -------------- +function squareroot!(A::AbstractMatrix, sqrtA, alg::DiagonalAlgorithm) + check_input(squareroot!, A, sqrtA, alg) + λ = diagview(sqrtA) + copy!(λ, diagview(A)) + # `sqrt(0) = 0`, so the domain includes its boundary + eltype(λ) <: Real && _clamp_domain_eigenvalues!(λ, _resolve_domain_atol(λ, alg)) + λ .= sqrt.(λ) + return sqrtA +end diff --git a/src/interface/matrixfunctions.jl b/src/interface/matrixfunctions.jl index b0ea57da8..08886345e 100644 --- a/src/interface/matrixfunctions.jl +++ b/src/interface/matrixfunctions.jl @@ -1,10 +1,16 @@ # ================================ -# EXPONENTIAL ALGORITHMS +# MATRIX FUNCTION ALGORITHMS # ================================ """ - MatrixFunctionViaLA() + MatrixFunctionViaLA(; domain_atol=-1) -Algorithm type to denote finding the exponential of `A` via the implementation of `LinearAlgebra`. +Algorithm type to denote computing a function of a matrix `A` via the implementation of `LinearAlgebra`. +For a matrix function with a restricted domain, i.e. [`squareroot`](@ref), `domain_atol` specifies the +absolute tolerance on the imaginary part of the result below which a complex result is attributed to +rounding rather than to a domain violation, with a negative value denoting the default tolerance. +Note that this measures a different quantity than the eigenvalue tolerance of +[`MatrixFunctionViaEig`](@ref) and [`MatrixFunctionViaEigh`](@ref), as `LinearAlgebra` does not +expose the spectrum; see [Domain considerations](@ref sec_matrixfunction_domain). """ @algdef MatrixFunctionViaLA @@ -29,31 +35,53 @@ As this algorithm requires no LAPACK support, it also applies at arbitrary preci @algdef MatrixFunctionViaTaylor """ - MatrixFunctionViaEigh(eigh_alg) + MatrixFunctionViaEigh(eigh_alg; domain_atol=-1) 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. +`domain_atol` applies to [`squareroot`](@ref): it is the absolute tolerance within which negative +eigenvalues are treated as rounding artifacts and clamped onto zero, with a negative value denoting +the default tolerance [`default_domain_atol`](@ref). Raising it accepts more matrices; +see [Domain considerations](@ref sec_matrixfunction_domain). """ struct MatrixFunctionViaEigh{A <: AbstractAlgorithm} <: AbstractAlgorithm eigh_alg::A + domain_atol::Float64 # negative value for runtime defaults end +MatrixFunctionViaEigh(eigh_alg::AbstractAlgorithm; domain_atol::Real = -1.0) = + MatrixFunctionViaEigh(eigh_alg, Float64(domain_atol)) function Base.show(io::IO, alg::MatrixFunctionViaEigh) print(io, "MatrixFunctionViaEigh(") _show_alg(io, alg.eigh_alg) + alg.domain_atol < 0 || print(io, "; domain_atol=", alg.domain_atol) return print(io, ")") end """ - MatrixFunctionViaEig(eig_alg) + MatrixFunctionViaEig(eig_alg; domain_atol=-1) 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. +`domain_atol` applies to [`squareroot`](@ref): it is the absolute tolerance within which eigenvalues +on the negative real axis are treated as rounding artifacts and clamped onto zero, with a negative +value denoting the default tolerance [`default_domain_atol`](@ref). Raising it accepts more matrices; +see [Domain considerations](@ref sec_matrixfunction_domain). + +!!! warning + This algorithm presumes a well-conditioned eigenbasis. For a defective or nearly defective + matrix both its result and its domain verdict are unreliable, since the eigenvalues themselves + are resolved only to `eps^(1/k)` for a Jordan block of size `k`. Prefer + [`MatrixFunctionViaLA`](@ref), which is Schur-based, for such matrices. """ struct MatrixFunctionViaEig{A <: AbstractAlgorithm} <: AbstractAlgorithm eig_alg::A + domain_atol::Float64 # negative value for runtime defaults end +MatrixFunctionViaEig(eig_alg::AbstractAlgorithm; domain_atol::Real = -1.0) = + MatrixFunctionViaEig(eig_alg, Float64(domain_atol)) function Base.show(io::IO, alg::MatrixFunctionViaEig) print(io, "MatrixFunctionViaEig(") _show_alg(io, alg.eig_alg) + alg.domain_atol < 0 || print(io, "; domain_atol=", alg.domain_atol) return print(io, ")") end diff --git a/src/interface/squareroot.jl b/src/interface/squareroot.jl new file mode 100644 index 000000000..e7be7d28b --- /dev/null +++ b/src/interface/squareroot.jl @@ -0,0 +1,40 @@ +# Square root +# ----------- + +""" + squareroot(A; kwargs...) -> sqrtA + squareroot(A, alg::AbstractAlgorithm) -> sqrtA + squareroot!(A, [sqrtA]; kwargs...) -> sqrtA + squareroot!(A, [sqrtA], alg::AbstractAlgorithm) -> sqrtA + +Compute the principal square root `sqrtA` of the square matrix `A`, i.e. the square root +whose eigenvalues have nonnegative real part. + +The scalar type of the output matches that of the input. +As a consequence, a real matrix with eigenvalues on the negative real axis, for which +the principal square root is complex, leads to a `DomainError`; pass a complex matrix +to obtain the principal value. +Real eigenvalues that are negative within a tolerance `domain_atol` are treated as rounding +artifacts and clamped to zero, so that raising `domain_atol` accepts more matrices. It defaults to +[`default_domain_atol`](@ref); see [Domain considerations](@ref sec_matrixfunction_domain). + +!!! note + The bang method `squareroot!` 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 `sqrtA` as output. +""" +@functiondef squareroot + +# Algorithm selection +# ------------------- +default_squareroot_algorithm(A; kwargs...) = default_squareroot_algorithm(typeof(A); kwargs...) +function default_squareroot_algorithm(T::Type; kwargs...) + return MatrixFunctionViaLA(; kwargs...) +end +function default_squareroot_algorithm(::Type{T}; kwargs...) where {T <: Diagonal} + return DiagonalAlgorithm(; kwargs...) +end + +function default_algorithm(::typeof(squareroot!), ::Type{A}; kwargs...) where {A} + return default_squareroot_algorithm(A; kwargs...) +end From f0a60958e78a8274e71c7dfb81c99c61becbb160 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Sat, 15 Aug 2026 05:24:37 -0400 Subject: [PATCH 3/6] test: add tests for `squareroot` MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The assertions are invariants — `sqrt(A)^2 ≈ A`, hermiticity, the elementwise action on the spectrum — so that they also apply on GPU and to downstream array types, with a host-only cross-check against `LinearAlgebra`. The domain tests prescribe the full spectrum, and probe a genuinely negative eigenvalue, one that is negative by roundoff, and one in between that only an explicit `domain_atol` admits. Co-Authored-By: Claude Opus 5 (1M context) --- test/README.md | 9 +- test/matrixfunctions/squareroot.jl | 109 +++++++++++++++++ test/testsuite/TestSuite.jl | 39 ++++++ test/testsuite/matrixfunctions/squareroot.jl | 122 +++++++++++++++++++ 4 files changed, 276 insertions(+), 3 deletions(-) create mode 100644 test/matrixfunctions/squareroot.jl create mode 100644 test/testsuite/matrixfunctions/squareroot.jl diff --git a/test/README.md b/test/README.md index c428592cb..400f1fc88 100644 --- a/test/README.md +++ b/test/README.md @@ -75,7 +75,7 @@ file placed directly in `test/` would belong to no group and never run in CI. |-------|----------| | `common` | Algorithm selection and defaults, truncation strategies, projections, Aqua code-quality checks | | `decompositions` | `qr`, `lq`, `svd`, `eig`, `eigh`, `gen_eig`, `schur`, `polar`, `orthnull` on CPU and GPU array types | -| `matrixfunctions` | `exponential` on CPU and GPU array types | +| `matrixfunctions` | `exponential`, `squareroot` on CPU and GPU array types | | `chainrules` | ChainRulesCore rules, exercised through ChainRulesTestUtils and Zygote | | `mooncake` | Mooncake AD rules | | `enzyme` | Enzyme AD rules, exercised through EnzymeTestUtils | @@ -132,8 +132,11 @@ Supporting infrastructure in the module: - `precision(T)` — the default tolerance, `sqrt(eps(real(T)))`. - Predicates used throughout the assertions: `isleftnull`, `isrightnull`, `isleftcomplete`, `isrightcomplete`, `has_positive_diagonal`. -- `instantiate_unitary`, `instantiate_rank_deficient_matrix`, `instantiate_smallnorm_matrix` — inputs - with prescribed structure or spectrum. +- `instantiate_unitary`, `instantiate_rank_deficient_matrix` — inputs with prescribed structure. +- `instantiate_smallnorm_matrix`, `instantiate_offaxis_matrix`, `instantiate_posdef_matrix`, + `instantiate_hermitian_spectrum` — inputs with a prescribed *spectrum*, which the matrix functions + need: `squareroot` is only defined away from the negative real axis, so the plain `randn` of + `instantiate_matrix` will not do. `test/linearmap.jl` is a further helper, defining a `LinearMap` wrapper that is deliberately *not* an `AbstractMatrix`, used to check the generic code paths. It is excluded from discovery and included diff --git a/test/matrixfunctions/squareroot.jl b/test/matrixfunctions/squareroot.jl new file mode 100644 index 000000000..f9d75f41d --- /dev/null +++ b/test/matrixfunctions/squareroot.jl @@ -0,0 +1,109 @@ +using MatrixAlgebraKit +using LinearAlgebra: Diagonal +using MatrixAlgebraKit: GLA, GS +using CUDA, AMDGPU +using GenericSchur, GenericLinearAlgebra + +if @isdefined(fast_tests) && fast_tests + BLASFloats = (Float64, ComplexF64) + GenericFloats = (BigFloat, Complex{BigFloat}) +else + BLASFloats = (Float32, Float64, ComplexF32, ComplexF64) + GenericFloats = (BigFloat, Complex{BigFloat}) +end +# only the `Diagonal` fast path applies to these, as they have no `eig`/`eigh` support +DiagonalOnlyFloats = (Float16, ComplexF16) + +@isdefined(TestSuite) || include("../testsuite/TestSuite.jl") +using .TestSuite + +is_buildkite = get(ENV, "BUILDKITE", "false") == "true" + +m = 54 +md = 4 # the domain tests prescribe the full spectrum, so keep them small + +# CPU tests +# --------- +if !is_buildkite + for T in BLASFloats + TestSuite.seed_rng!(123) + LAPACK_EIG_ALGS = (MatrixFunctionViaLA(), MatrixFunctionViaEig(QRIteration())) + LAPACK_EIGH_ALGS = ( + MatrixFunctionViaEigh(QRIteration()), + MatrixFunctionViaEigh(DivideAndConquer()), + ) + TestSuite.test_squareroot(T, (m, m)) + TestSuite.test_squareroot_algs(T, (m, m), LAPACK_EIG_ALGS) + TestSuite.test_squareroot_hermitian(T, (m, m), LAPACK_EIG_ALGS; exact_hermiticity = false) + TestSuite.test_squareroot_hermitian(T, (m, m), LAPACK_EIGH_ALGS) + TestSuite.test_squareroot_reference(T, (m, m)) + TestSuite.test_squareroot_domain(T, (md, md), LAPACK_EIG_ALGS) + TestSuite.test_squareroot_domain(T, (md, md), LAPACK_EIGH_ALGS; hermitian_output = true) + end + + # `eig` comes from GenericSchur, `eigh` from GenericLinearAlgebra. Both are loaded here, so name + # the driver explicitly instead of relying on `default_driver`. + for T in GenericFloats + TestSuite.seed_rng!(123) + GS_ALGS = (MatrixFunctionViaEig(QRIteration(; driver = GS())),) + GLA_ALGS = (MatrixFunctionViaEigh(QRIteration(; driver = GLA())),) + TestSuite.test_squareroot_algs(T, (24, 24), GS_ALGS) + TestSuite.test_squareroot_hermitian(T, (24, 24), GS_ALGS; exact_hermiticity = false) + TestSuite.test_squareroot_hermitian(T, (24, 24), GLA_ALGS) + TestSuite.test_squareroot_domain(T, (md, md), GS_ALGS) + TestSuite.test_squareroot_domain(T, (md, md), GLA_ALGS; hermitian_output = true) + end + + for T in (BLASFloats..., GenericFloats..., DiagonalOnlyFloats...) + TestSuite.seed_rng!(123) + AT = Diagonal{T, Vector{T}} + test_spectrum = !(T in DiagonalOnlyFloats) + TestSuite.test_squareroot(AT, m) + TestSuite.test_squareroot_algs(AT, m, (DiagonalAlgorithm(),)) + TestSuite.test_squareroot_hermitian(AT, m, (DiagonalAlgorithm(),); test_spectrum) + TestSuite.test_squareroot_reference(AT, m; test_hermitian = !(T in GenericFloats)) + TestSuite.test_squareroot_domain(AT, md, (DiagonalAlgorithm(),)) + end +end + +# CUDA tests +# ---------- +# general dense matrices are not supported on device: `MatrixFunctionViaLA` would call LAPACK on +# device memory, and `MatrixFunctionViaEig` scalar-indexes in its `lu!`-based solve +if CUDA.functional() + for T in BLASFloats + TestSuite.seed_rng!(123) + CUDA_EIGH_ALGS = ( + MatrixFunctionViaEigh(Jacobi()), + MatrixFunctionViaEigh(DivideAndConquer()), + ) + TestSuite.test_squareroot_hermitian(CuMatrix{T}, (m, m), CUDA_EIGH_ALGS) + TestSuite.test_squareroot_domain(CuMatrix{T}, (md, md), CUDA_EIGH_ALGS; hermitian_output = true) + + AT = Diagonal{T, CuVector{T}} + TestSuite.test_squareroot(AT, m) + TestSuite.test_squareroot_algs(AT, m, (DiagonalAlgorithm(),)) + TestSuite.test_squareroot_hermitian(AT, m, (DiagonalAlgorithm(),)) + TestSuite.test_squareroot_domain(AT, md, (DiagonalAlgorithm(),)) + end +end + +# AMDGPU tests +# ------------ +if AMDGPU.functional() + for T in BLASFloats + TestSuite.seed_rng!(123) + ROC_EIGH_ALGS = ( + MatrixFunctionViaEigh(Jacobi()), + MatrixFunctionViaEigh(DivideAndConquer()), + ) + TestSuite.test_squareroot_hermitian(ROCMatrix{T}, (m, m), ROC_EIGH_ALGS) + TestSuite.test_squareroot_domain(ROCMatrix{T}, (md, md), ROC_EIGH_ALGS; hermitian_output = true) + + AT = Diagonal{T, ROCVector{T}} + TestSuite.test_squareroot(AT, m) + TestSuite.test_squareroot_algs(AT, m, (DiagonalAlgorithm(),)) + TestSuite.test_squareroot_hermitian(AT, m, (DiagonalAlgorithm(),)) + TestSuite.test_squareroot_domain(AT, md, (DiagonalAlgorithm(),)) + end +end diff --git a/test/testsuite/TestSuite.jl b/test/testsuite/TestSuite.jl index bffbfe204..251608243 100644 --- a/test/testsuite/TestSuite.jl +++ b/test/testsuite/TestSuite.jl @@ -102,12 +102,50 @@ function instantiate_rank_deficient_matrix(::Type{T}, sz; trunc = truncrank(div( return Diagonal(diag(mul!(A, V, C))) end +# `squareroot` is only defined away from the negative real axis, so the matrix functions need +# matrices with a controlled spectrum rather than the plain `randn` of `instantiate_matrix` + # spectral radius at most one, keeping `exponential` well inside the principal branch function instantiate_smallnorm_matrix(T, sz) A = instantiate_matrix(T, sz) return A / norm(A) end +# spectrum inside the unit disk around 1, i.e. clear of the negative real axis and of zero +instantiate_offaxis_matrix(T, sz) = instantiate_smallnorm_matrix(T, sz) + I + +# hermitian positive definite, so that the `eigh`-based algorithms apply +function instantiate_posdef_matrix(T, sz) + A = instantiate_matrix(T, sz) + return project_hermitian!(A * A') + I +end + +# hermitian with the prescribed (real) spectrum `λ`, for probing the domain boundary +function instantiate_hermitian_spectrum(T, sz, λ) + A = instantiate_matrix(T, sz) + n = size(A, 1) + @assert length(λ) == n + dv = A isa Diagonal ? diagview(A) : A + Ddiag = similar(dv, eltype(A), n) + # convert on the host first: `copyto!` onto a device array does not convert eltypes + copyto!(Ddiag, convert(Vector{eltype(A)}, collect(λ))) + A isa Diagonal && return Diagonal(Ddiag) + V = instantiate_unitary(T, A, n) + return project_hermitian!(V * Diagonal(Ddiag) * V') +end + +# rebuild `alg` with an explicit `domain_atol`, so that the domain tests need not spell out the +# inner decomposition algorithm a second time +with_domain_atol(alg::MatrixFunctionViaEig, atol) = MatrixFunctionViaEig(alg.eig_alg; domain_atol = atol) +with_domain_atol(alg::MatrixFunctionViaEigh, atol) = MatrixFunctionViaEigh(alg.eigh_alg; domain_atol = atol) +with_domain_atol(::MatrixAlgebraKit.DiagonalAlgorithm, atol) = DiagonalAlgorithm(; domain_atol = atol) +with_domain_atol(::MatrixFunctionViaLA, atol) = MatrixFunctionViaLA(; domain_atol = atol) + +# a tolerance generous enough to admit an eigenvalue at `-√eps`. For `MatrixFunctionViaLA` it bounds +# the imaginary part of the result rather than the spectrum, which sits on a coarser scale. +domain_test_atol(::MatrixAlgebraKit.AbstractAlgorithm, R) = cbrt(eps(R)) +domain_test_atol(::MatrixFunctionViaLA, R) = one(R) / 2 + include("ad_utils.jl") include("projections.jl") @@ -126,6 +164,7 @@ include("decompositions/svd.jl") # Matrix functions # ---------------- include("matrixfunctions/exponential.jl") +include("matrixfunctions/squareroot.jl") # Mooncake # -------- diff --git a/test/testsuite/matrixfunctions/squareroot.jl b/test/testsuite/matrixfunctions/squareroot.jl new file mode 100644 index 000000000..1dc6d106a --- /dev/null +++ b/test/testsuite/matrixfunctions/squareroot.jl @@ -0,0 +1,122 @@ +using TestExtras +using LinearAlgebra: LinearAlgebra, I +using MatrixAlgebraKit: ishermitian + +# the assertions are invariants rather than comparisons against a reference implementation, so that +# the same bodies apply on GPU and to downstream array types. `test_squareroot_reference` is the +# exception and is host-only by construction. + +function test_squareroot(T::Type, sz; kwargs...) + summary_str = testargs_summary(T, sz) + return @testset "squareroot $summary_str" begin + A = instantiate_offaxis_matrix(T, sz) + Ac = deepcopy(A) + + sqrtA = @testinferred squareroot(A) + @test eltype(sqrtA) == eltype(A) + @test sqrtA * sqrtA ≈ A + @test A == Ac + + # the in-place method may not be able to reuse the provided output + sqrtA2 = @testinferred squareroot!(deepcopy(A), deepcopy(sqrtA)) + @test sqrtA2 ≈ sqrtA + end +end + +function test_squareroot_algs(T::Type, sz, algs; kwargs...) + summary_str = testargs_summary(T, sz) + return @testset "squareroot algorithm $alg $summary_str" for alg in algs + A = instantiate_offaxis_matrix(T, sz) + Ac = deepcopy(A) + + sqrtA = @testinferred squareroot(A, alg) + @test eltype(sqrtA) == eltype(A) + @test sqrtA * sqrtA ≈ A + @test A == Ac + end +end + +# the `eigh`-based kernels project the result, so it is hermitian to the last bit; pass +# `exact_hermiticity = false` for algorithms that route through the general `eig` path instead. +# The elementwise spectrum check needs exact hermiticity, since `eigh_vals` rejects anything else. +function test_squareroot_hermitian( + T::Type, sz, algs; + exact_hermiticity = true, test_spectrum = true, kwargs... + ) + summary_str = testargs_summary(T, sz) + return @testset "squareroot hermitian algorithm $alg $summary_str" for alg in algs + A = instantiate_posdef_matrix(T, sz) + Ac = deepcopy(A) + + sqrtA = @testinferred squareroot(A, alg) + @test eltype(sqrtA) == eltype(A) + @test sqrtA * sqrtA ≈ A + @test A == Ac + + if exact_hermiticity + @test ishermitian(sqrtA) + test_spectrum && @test eigh_vals(sqrtA) ≈ sqrt.(eigh_vals(A)) + else + @test ishermitian(sqrtA; rtol = precision(T)) + end + end +end + +# a matrix whose spectrum reaches the negative real axis has a complex principal square root, which +# a type-stable real output cannot represent. Pass `hermitian_output = true` for the algorithms that +# promise a hermitian result: those must reject a negative eigenvalue whatever the scalar type. +function test_squareroot_domain(T::Type, sz, algs; hermitian_output = false, kwargs...) + R = real(eltype(T)) + n = sz isa Tuple ? first(sz) : sz + summary_str = testargs_summary(T, sz) + return @testset "squareroot domain algorithm $alg $summary_str" for alg in algs + # genuinely negative eigenvalue + λ = collect(R, 1:n) + λ[1] = -one(R) + A = instantiate_hermitian_spectrum(T, sz, λ) + + if eltype(T) <: Real || hermitian_output + @test_throws DomainError squareroot(A, alg) + else + sqrtA = @testinferred squareroot(A, alg) + @test sqrtA * sqrtA ≈ A + end + + # roundoff-scale negative eigenvalue: clamped onto the boundary rather than rejected + λclamp = collect(R, 1:n) + λclamp[1] = -eps(R) + Aclamp = instantiate_hermitian_spectrum(T, sz, λclamp) + sqrtAclamp = @testinferred squareroot(Aclamp, alg) + @test eltype(sqrtAclamp) == eltype(Aclamp) + @test sqrtAclamp * sqrtAclamp ≈ Aclamp atol = sqrt(eps(R)) + + # an eigenvalue beyond every default tolerance is out of domain, while an explicit + # `domain_atol` admits it after all + λwide = collect(R, 1:n) + λwide[1] = -sqrt(eps(R)) + Awide = instantiate_hermitian_spectrum(T, sz, λwide) + if eltype(T) <: Real || hermitian_output + @test_throws DomainError squareroot(Awide, alg) + end + wide_alg = with_domain_atol(alg, domain_test_atol(alg, R)) + sqrtAwide = @testinferred squareroot(Awide, wide_alg) + @test eltype(sqrtAwide) == eltype(Awide) + # accepting is backward stable, but only to the size of the eigenvalue that was discarded + @test sqrtAwide * sqrtAwide ≈ Awide atol = sqrt(sqrt(eps(R))) + end +end + +# cross-check against `LinearAlgebra`, host arrays only. `LinearAlgebra` has no matrix functions +# for a `Hermitian` wrapper outside the BLAS floats, hence `test_hermitian`. +function test_squareroot_reference(T::Type, sz; test_hermitian = true, kwargs...) + summary_str = testargs_summary(T, sz) + return @testset "squareroot vs LinearAlgebra $summary_str" begin + A = instantiate_offaxis_matrix(T, sz) + @test squareroot(A) ≈ LinearAlgebra.sqrt(A) + + if test_hermitian + H = instantiate_posdef_matrix(T, sz) + @test squareroot(H) ≈ LinearAlgebra.sqrt(LinearAlgebra.Hermitian(H)) + end + end +end From a8ddbb119b0e306d141c5118fc8afdb3ed539e37 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Sat, 15 Aug 2026 05:26:33 -0400 Subject: [PATCH 4/6] docs: document `squareroot` and the domain handling Describe the shared matrix function algorithms once, and add a section on the domain of `squareroot`: what `domain_atol` clamps, why an accepted result can be off by more than the tolerance, and why the tolerance of `MatrixFunctionViaLA` measures something else entirely. Co-Authored-By: Claude Opus 5 (1M context) --- docs/src/changelog.md | 6 ++ docs/src/user_interface/algorithms.md | 13 ++-- docs/src/user_interface/matrix_functions.md | 81 ++++++++++++++++++--- 3 files changed, 86 insertions(+), 14 deletions(-) diff --git a/docs/src/changelog.md b/docs/src/changelog.md index cc2946b68..20d81a369 100644 --- a/docs/src/changelog.md +++ b/docs/src/changelog.md @@ -22,8 +22,14 @@ When releasing a new version, move the "Unreleased" changes to a new version sec ### Added +- New matrix function `squareroot`, computing the principal square root, supporting the + `MatrixFunctionViaLA`, `MatrixFunctionViaEig`, `MatrixFunctionViaEigh` and `DiagonalAlgorithm` + algorithms ([#261](https://github.com/QuantumKitHub/MatrixAlgebraKit.jl/pull/261)). + ### Changed +- `MatrixFunctionViaEig` and `MatrixFunctionViaEigh` gained a `domain_atol` field, settable through + the keyword argument of the same name. ([#261](https://github.com/QuantumKitHub/MatrixAlgebraKit.jl/pull/261)). - `qr_compact!`, `qr_full!`, `lq_compact!` and `lq_full!` now extract `R` (or `L`) before constructing `Q`, so that an inplace `Q` (supplying `A` itself as output for `Q`) can be combined with computing `R` (or `L`) and with `positive = true`. diff --git a/docs/src/user_interface/algorithms.md b/docs/src/user_interface/algorithms.md index d90b76074..feaff7738 100644 --- a/docs/src/user_interface/algorithms.md +++ b/docs/src/user_interface/algorithms.md @@ -98,16 +98,19 @@ The following algorithms for matrix decompositions are available. | [`PolarViaSVD`](@ref) | polar | positional `svd_alg` argument | | [`PolarNewton`](@ref) | polar | `maxiter`, `tol` | +For full docstring details on each algorithm type, see the corresponding section in [Decompositions](@ref). + 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` | +| [`MatrixFunctionViaTaylor`](@ref) | exponential | `tol`, `balance`, `estimate_order` | +| [`MatrixFunctionViaLA`](@ref) | exponential, squareroot | `domain_atol` (squareroot) | +| [`MatrixFunctionViaEig`](@ref) | exponential, squareroot | positional `eig_alg` argument, `domain_atol` (squareroot) | +| [`MatrixFunctionViaEigh`](@ref) | exponential, squareroot | positional `eigh_alg` argument, `domain_atol` (squareroot) | -For full docstring details on each algorithm type, see the corresponding section in [Decompositions](@ref). +Note that `domain_atol` means something different for `MatrixFunctionViaLA`, which bounds the imaginary part of the result rather than the distance of an eigenvalue to the domain boundary. +For full docstring details on each algorithm type, and for how the tolerance is meant to be used, see [Matrix functions](@ref) and in particular [Domain considerations](@ref sec_matrixfunction_domain). ## [Driver Selection](@id sec_driverselection) diff --git a/docs/src/user_interface/matrix_functions.md b/docs/src/user_interface/matrix_functions.md index 6153f2e6a..26b82fc99 100644 --- a/docs/src/user_interface/matrix_functions.md +++ b/docs/src/user_interface/matrix_functions.md @@ -20,20 +20,83 @@ For a full description of how to select and configure algorithms, see [Algorithm Importantly, for generic code patterns it is recommended to always use the output `F` explicitly, since some implementations may not be able to reuse the provided memory. Additionally, the `f!` method typically assumes that it is allowed to destroy the input `A`, and making use of the contents of `A` afterwards should be deemed as undefined behavior. -## Exponential +## Algorithms -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. -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. +The matrix functions share a common set of algorithms, which differ in how they reduce the problem to a scalar function of the eigenvalues: + +- [`MatrixFunctionViaLA`](@ref) defers to the implementation of `LinearAlgebra`, which is Schur-based for [`squareroot`](@ref) and a Padé approximation for [`exponential`](@ref). +- [`MatrixFunctionViaEig`](@ref) and [`MatrixFunctionViaEigh`](@ref) first compute an eigenvalue decomposition, through `eig_full` and `eigh_full` respectively, and then apply the scalar function to the eigenvalues. The latter requires a hermitian input, and in return its result is hermitian by construction. +- [`MatrixFunctionViaTaylor`](@ref) applies to [`exponential`](@ref) only, and evaluates its Taylor series through scaling and squaring. As it requires no LAPACK support, it also applies to generic data types at arbitrary precision. +- [`DiagonalAlgorithm`](@ref) is the fast path for a `Diagonal` input, and simply maps the scalar function over the diagonal. ```@docs; canonical=false -exponential MatrixAlgebraKit.MatrixFunctionViaTaylor MatrixAlgebraKit.MatrixFunctionViaLA MatrixAlgebraKit.MatrixFunctionViaEig MatrixAlgebraKit.MatrixFunctionViaEigh ``` + +## [Domain considerations](@id sec_matrixfunction_domain) + +Not every matrix function is defined for every square matrix: [`squareroot`](@ref) requires the eigenvalues to avoid the negative real axis, and its principal value is complex whenever eigenvalues on that axis are present. +In MatrixAlgebraKit, we aim to keep type stability, and thus the scalar type of the output always matches that of the input. +As such, a real matrix with eigenvalues on the negative real axis leads to a `DomainError`. +You should pass a complex matrix instead to obtain the complex principal value. + +The hard part is that eigenvalues are *computed*, not given, so deciding whether a matrix is in domain means comparing a number against a tolerance, which the keyword `domain_atol` controls. +Since `sqrt(0) = 0`, the domain boundary belongs to the domain, and an eigenvalue that is negative only by roundoff can be clamped onto zero and the computation continues. +Raising `domain_atol` therefore *accepts* more matrices. + +Clamping is not as cheap as its tolerance suggests, however. +It is backward stable, but only to the size of the eigenvalue that was discarded, and the forward error it incurs is *not* of that size: clamping an eigenvalue at `-δ` perturbs the square root by `O(√δ)`, so an accepted result computed at the default tolerance can differ from the exact principal value by considerably more than the tolerance itself. + +Additionally, not all algorithms have acces to the spectrum. +[`MatrixFunctionViaLA`](@ref) defers to `LinearAlgebra`, which decides internally whether a real result exists and hands back a complex matrix when it does not. +There are no eigenvalues to compare against anything, so the only quantity available is the imaginary part of the *result* — a different quantity on a different scale, since an eigenvalue at `-δ` shows up as an imaginary part of order `√δ`. +The two tolerances are not comparable, and a value tuned for one algorithm should not be carried over to the other. +A tolerance is not optional there either, as `LinearAlgebra` casts a result back to a real matrix only when its imaginary part vanishes identically. + +| Algorithm | What `domain_atol` does | Default | +|:----------|:------------------------|:--------| +| [`DiagonalAlgorithm`](@ref) | eigenvalues negative within the tolerance are clamped to zero | `n * eps * maximum(abs, λ)` | +| [`MatrixFunctionViaEigh`](@ref) | as above, on the eigenvalues from `eigh_full` | `n * eps * maximum(abs, λ)` | +| [`MatrixFunctionViaEig`](@ref) | as above, on the eigenvalues from `eig_full` | `defaulttol(λ) * maximum(abs, λ)` | +| [`MatrixFunctionViaLA`](@ref) | bounds `maximum(abs ∘ imag, sqrt(A))` instead | `defaulttol(A) * norm(sqrt(A), Inf)` | + +The first two use the same rule as `LinearAlgebra.sqrt(::Hermitian; rtol = eps(T) * size(A, 1))`, so for hermitian input MatrixAlgebraKit and `LinearAlgebra` accept and reject the same matrices. +`MatrixFunctionViaEig` is deliberately looser, at `eps^(2/3)` rather than `eps`, since the accuracy of its eigenvalues is additionally limited by the conditioning of the eigenvectors. + +One tolerance is not a choice about the domain, and so it is not user-settable. +When a real matrix is diagonalized by [`MatrixFunctionViaEig`](@ref), its eigenvalues come back complex, and deciding whether one lies *on* the negative real axis — as opposed to being one of a complex-conjugate pair, which obstructs nothing — is a question about the accuracy of the eigensolver rather than about the domain. +That test always uses [`default_domain_atol`](@ref), however you set `domain_atol`. +The other algorithms face no such question, since their eigenvalues are real to begin with. + +!!! warning "`MatrixFunctionViaEig` and defective matrices" + The eigenvalues of a Jordan block of size `k` are resolved only to `eps^(1/k)`, which exceeds every tolerance on this page. + A real matrix with a defective negative eigenvalue can therefore have its spectrum reported as a complex-conjugate pair well off the axis, be judged in domain, and yield a result whose imaginary part is silently discarded. + This is not specific to the domain test: `MatrixFunctionViaEig` reconstructs `f(A)` by inverting the eigenvector matrix, so for a defective or nearly defective matrix its result is unreliable whether the input is real or complex. + Use the Schur-based [`MatrixFunctionViaLA`](@ref) for such matrices. + +```@docs; canonical=false +MatrixAlgebraKit.default_domain_atol +``` + +## 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. +It is defined for every square matrix, so the domain considerations above do not apply to it. +The default algorithm is [`MatrixFunctionViaTaylor`](@ref), which is the only one that also covers generic data types at arbitrary precision. +Additionally, in order to calculate `exp(τ * A)`, the function `exponential` can be called with `(τ, A)`, using the same algorithms. + +```@docs; canonical=false +exponential +``` + +## Square root + +The principal [square root](https://en.wikipedia.org/wiki/Square_root_of_a_matrix) of a square matrix `A` is the unique square root whose eigenvalues have nonnegative real part. +It is computed by the function [`squareroot`](@ref), with [`MatrixFunctionViaLA`](@ref) as the default algorithm, and is subject to the domain considerations above. + +```@docs; canonical=false +squareroot +``` From cb057fce004dfdeb3a8d549b1569e224f9a0a9f5 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Wed, 26 Aug 2026 09:23:46 -0400 Subject: [PATCH 5/6] simplify further --- docs/src/changelog.md | 5 +- docs/src/user_interface/algorithms.md | 8 +- docs/src/user_interface/matrix_functions.md | 84 ++++++++------------ ext/MatrixAlgebraKitGenericSchurExt.jl | 6 +- src/common/defaults.jl | 16 ++-- src/implementations/exponential.jl | 4 +- src/implementations/matrixfunctions.jl | 60 ++++++-------- src/implementations/squareroot.jl | 14 ++-- src/interface/matrixfunctions.jl | 67 +++++++--------- src/interface/squareroot.jl | 8 +- test/matrixfunctions/squareroot.jl | 4 +- test/testsuite/TestSuite.jl | 12 +-- test/testsuite/matrixfunctions/squareroot.jl | 25 ++++-- 13 files changed, 140 insertions(+), 173 deletions(-) diff --git a/docs/src/changelog.md b/docs/src/changelog.md index 20d81a369..ac4bc68dd 100644 --- a/docs/src/changelog.md +++ b/docs/src/changelog.md @@ -28,8 +28,9 @@ When releasing a new version, move the "Unreleased" changes to a new version sec ### Changed -- `MatrixFunctionViaEig` and `MatrixFunctionViaEigh` gained a `domain_atol` field, settable through - the keyword argument of the same name. ([#261](https://github.com/QuantumKitHub/MatrixAlgebraKit.jl/pull/261)). +- `MatrixFunctionViaEig` and `MatrixFunctionViaEigh` are now defined through `@algdef`, so that both + the wrapped decomposition algorithm (`eig_alg` / `eigh_alg`, still accepted positionally) and the + new `domain_atol` are optional keyword arguments. ([#261](https://github.com/QuantumKitHub/MatrixAlgebraKit.jl/pull/261)). - `qr_compact!`, `qr_full!`, `lq_compact!` and `lq_full!` now extract `R` (or `L`) before constructing `Q`, so that an inplace `Q` (supplying `A` itself as output for `Q`) can be combined with computing `R` (or `L`) and with `positive = true`. diff --git a/docs/src/user_interface/algorithms.md b/docs/src/user_interface/algorithms.md index feaff7738..3e348697d 100644 --- a/docs/src/user_interface/algorithms.md +++ b/docs/src/user_interface/algorithms.md @@ -105,11 +105,11 @@ The following algorithms for matrix functions are available. | Algorithm | Applicable matrix functions | Key keyword arguments | |:----------|:--------------------------|:----------------------| | [`MatrixFunctionViaTaylor`](@ref) | exponential | `tol`, `balance`, `estimate_order` | -| [`MatrixFunctionViaLA`](@ref) | exponential, squareroot | `domain_atol` (squareroot) | -| [`MatrixFunctionViaEig`](@ref) | exponential, squareroot | positional `eig_alg` argument, `domain_atol` (squareroot) | -| [`MatrixFunctionViaEigh`](@ref) | exponential, squareroot | positional `eigh_alg` argument, `domain_atol` (squareroot) | +| [`MatrixFunctionViaLA`](@ref) | exponential, squareroot | — | +| [`MatrixFunctionViaEig`](@ref) | exponential, squareroot | `eig_alg` (also positional), `domain_atol` (squareroot) | +| [`MatrixFunctionViaEigh`](@ref) | exponential, squareroot | `eigh_alg` (also positional), `domain_atol` (squareroot) | -Note that `domain_atol` means something different for `MatrixFunctionViaLA`, which bounds the imaginary part of the result rather than the distance of an eigenvalue to the domain boundary. +Note that [`MatrixFunctionViaLA`](@ref) accepts no keyword arguments, since it has no access to the spectrum and thus cannot honor a `domain_atol`. For full docstring details on each algorithm type, and for how the tolerance is meant to be used, see [Matrix functions](@ref) and in particular [Domain considerations](@ref sec_matrixfunction_domain). ## [Driver Selection](@id sec_driverselection) diff --git a/docs/src/user_interface/matrix_functions.md b/docs/src/user_interface/matrix_functions.md index 26b82fc99..367adc71a 100644 --- a/docs/src/user_interface/matrix_functions.md +++ b/docs/src/user_interface/matrix_functions.md @@ -17,12 +17,12 @@ f!(A, [F]; kwargs...) -> F... Here, the input matrix is always the first argument, and optionally the output can be provided as well. The keywords are algorithm-specific, and can be used to influence the behavior of the algorithms. For a full description of how to select and configure algorithms, see [Algorithm Selection](@ref sec_algorithmselection). -Importantly, for generic code patterns it is recommended to always use the output `F` explicitly, since some implementations may not be able to reuse the provided memory. -Additionally, the `f!` method typically assumes that it is allowed to destroy the input `A`, and making use of the contents of `A` afterwards should be deemed as undefined behavior. +Importantly, for generic code patterns it is recommended to always use the output `F` explicitly, rather than relying on the in-place functionality, since some implementations may not be able to reuse the provided memory. +Additionally, the `f!` method typically assumes that it is allowed to destroy the input `A`, and making use of the contents of `A` afterwards is undefined behavior. ## Algorithms -The matrix functions share a common set of algorithms, which differ in how they reduce the problem to a scalar function of the eigenvalues: +The matrix functions share a common set of algorithms, which differ in how they reduce the problem to a scalar function of the eigenvalues, along with more specialized implementations for specific functions: - [`MatrixFunctionViaLA`](@ref) defers to the implementation of `LinearAlgebra`, which is Schur-based for [`squareroot`](@ref) and a Padé approximation for [`exponential`](@ref). - [`MatrixFunctionViaEig`](@ref) and [`MatrixFunctionViaEigh`](@ref) first compute an eigenvalue decomposition, through `eig_full` and `eigh_full` respectively, and then apply the scalar function to the eigenvalues. The latter requires a hermitian input, and in return its result is hermitian by construction. @@ -36,55 +36,10 @@ MatrixAlgebraKit.MatrixFunctionViaEig MatrixAlgebraKit.MatrixFunctionViaEigh ``` -## [Domain considerations](@id sec_matrixfunction_domain) - -Not every matrix function is defined for every square matrix: [`squareroot`](@ref) requires the eigenvalues to avoid the negative real axis, and its principal value is complex whenever eigenvalues on that axis are present. -In MatrixAlgebraKit, we aim to keep type stability, and thus the scalar type of the output always matches that of the input. -As such, a real matrix with eigenvalues on the negative real axis leads to a `DomainError`. -You should pass a complex matrix instead to obtain the complex principal value. - -The hard part is that eigenvalues are *computed*, not given, so deciding whether a matrix is in domain means comparing a number against a tolerance, which the keyword `domain_atol` controls. -Since `sqrt(0) = 0`, the domain boundary belongs to the domain, and an eigenvalue that is negative only by roundoff can be clamped onto zero and the computation continues. -Raising `domain_atol` therefore *accepts* more matrices. - -Clamping is not as cheap as its tolerance suggests, however. -It is backward stable, but only to the size of the eigenvalue that was discarded, and the forward error it incurs is *not* of that size: clamping an eigenvalue at `-δ` perturbs the square root by `O(√δ)`, so an accepted result computed at the default tolerance can differ from the exact principal value by considerably more than the tolerance itself. - -Additionally, not all algorithms have acces to the spectrum. -[`MatrixFunctionViaLA`](@ref) defers to `LinearAlgebra`, which decides internally whether a real result exists and hands back a complex matrix when it does not. -There are no eigenvalues to compare against anything, so the only quantity available is the imaginary part of the *result* — a different quantity on a different scale, since an eigenvalue at `-δ` shows up as an imaginary part of order `√δ`. -The two tolerances are not comparable, and a value tuned for one algorithm should not be carried over to the other. -A tolerance is not optional there either, as `LinearAlgebra` casts a result back to a real matrix only when its imaginary part vanishes identically. - -| Algorithm | What `domain_atol` does | Default | -|:----------|:------------------------|:--------| -| [`DiagonalAlgorithm`](@ref) | eigenvalues negative within the tolerance are clamped to zero | `n * eps * maximum(abs, λ)` | -| [`MatrixFunctionViaEigh`](@ref) | as above, on the eigenvalues from `eigh_full` | `n * eps * maximum(abs, λ)` | -| [`MatrixFunctionViaEig`](@ref) | as above, on the eigenvalues from `eig_full` | `defaulttol(λ) * maximum(abs, λ)` | -| [`MatrixFunctionViaLA`](@ref) | bounds `maximum(abs ∘ imag, sqrt(A))` instead | `defaulttol(A) * norm(sqrt(A), Inf)` | - -The first two use the same rule as `LinearAlgebra.sqrt(::Hermitian; rtol = eps(T) * size(A, 1))`, so for hermitian input MatrixAlgebraKit and `LinearAlgebra` accept and reject the same matrices. -`MatrixFunctionViaEig` is deliberately looser, at `eps^(2/3)` rather than `eps`, since the accuracy of its eigenvalues is additionally limited by the conditioning of the eigenvectors. - -One tolerance is not a choice about the domain, and so it is not user-settable. -When a real matrix is diagonalized by [`MatrixFunctionViaEig`](@ref), its eigenvalues come back complex, and deciding whether one lies *on* the negative real axis — as opposed to being one of a complex-conjugate pair, which obstructs nothing — is a question about the accuracy of the eigensolver rather than about the domain. -That test always uses [`default_domain_atol`](@ref), however you set `domain_atol`. -The other algorithms face no such question, since their eigenvalues are real to begin with. - -!!! warning "`MatrixFunctionViaEig` and defective matrices" - The eigenvalues of a Jordan block of size `k` are resolved only to `eps^(1/k)`, which exceeds every tolerance on this page. - A real matrix with a defective negative eigenvalue can therefore have its spectrum reported as a complex-conjugate pair well off the axis, be judged in domain, and yield a result whose imaginary part is silently discarded. - This is not specific to the domain test: `MatrixFunctionViaEig` reconstructs `f(A)` by inverting the eigenvector matrix, so for a defective or nearly defective matrix its result is unreliable whether the input is real or complex. - Use the Schur-based [`MatrixFunctionViaLA`](@ref) for such matrices. - -```@docs; canonical=false -MatrixAlgebraKit.default_domain_atol -``` - ## 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. -It is defined for every square matrix, so the domain considerations above do not apply to it. +It is defined for every square matrix, so the [domain considerations](@ref sec_matrixfunction_domain) below do not apply to it. The default algorithm is [`MatrixFunctionViaTaylor`](@ref), which is the only one that also covers generic data types at arbitrary precision. Additionally, in order to calculate `exp(τ * A)`, the function `exponential` can be called with `(τ, A)`, using the same algorithms. @@ -95,8 +50,37 @@ exponential ## Square root The principal [square root](https://en.wikipedia.org/wiki/Square_root_of_a_matrix) of a square matrix `A` is the unique square root whose eigenvalues have nonnegative real part. -It is computed by the function [`squareroot`](@ref), with [`MatrixFunctionViaLA`](@ref) as the default algorithm, and is subject to the domain considerations above. +It is computed by the function [`squareroot`](@ref), with [`MatrixFunctionViaLA`](@ref) as the default algorithm, and is subject to the [domain considerations](@ref sec_matrixfunction_domain) below. ```@docs; canonical=false squareroot ``` + +## [Domain considerations](@id sec_matrixfunction_domain) + +Not every matrix function is defined for every square matrix, for example a real [`squareroot`](@ref) requires the eigenvalues to avoid the negative real axis, and its principal value is complex whenever eigenvalues on that axis are present. +In MatrixAlgebraKit, we aim to keep type stability, and thus the scalar type of the output always matches that of the input. +As such, a real matrix with eigenvalues on the negative real axis leads to a `DomainError`, and a complex matrix should be passed instead. + +The hard part is that eigenvalues are *computed*, and thus contain some inaccuracy from the method used to compute them. +Typically it can be beneficial to introduce some tolerance to compare the domain with, which is controlled by the `domain_atol` keyword. +Clamping these values does come at a cost, as e.g. an eigenvalue at `-δ` perturbs the square root by `O(√δ)`, so an accepted result computed at the default tolerance can differ from the exact principal value by considerably more than the tolerance itself. + +`domain_atol` defaults to [`default_domain_atol`](@ref), i.e. `n * eps * maximum(abs, λ)`, which is the accumulated roundoff of a spectrum computed to hermitian accuracy. +This is the same rule as `LinearAlgebra.sqrt(::Hermitian; rtol = eps(T) * size(A, 1))`, so for hermitian input MatrixAlgebraKit and `LinearAlgebra` accept and reject the same matrices. +The eigenvalues of [`MatrixFunctionViaEig`](@ref) are additionally limited by the conditioning of the eigenvectors, so for a poorly conditioned eigenbasis a larger `domain_atol` may have to be set explicitly. +The same default is used for the never user-settable tolerance with which a complex eigenvalue of a real matrix is decided to lie *on* the negative real axis, as that is a question about the accuracy of the eigensolver rather than about the domain. + +Additionally, not all algorithms have acces to the spectrum, so not all methods are suitable for eigenvalues close to the domain edges. +For example, [`MatrixFunctionViaLA`](@ref) defers to `LinearAlgebra`, which decides internally whether a real result exists and hands back a complex matrix when it does not. +There are no eigenvalues to compare against anything, so it rejects a complex result for a real input outright, and passing it `domain_atol` is an error rather than a silent no-op. + +!!! warning "`MatrixFunctionViaEig` and defective matrices" + The eigenvalues of a Jordan block of size `k` are resolved only to `eps^(1/k)`, which exceeds every tolerance on this page. + A real matrix with a defective negative eigenvalue can therefore have its spectrum reported as a complex-conjugate pair well off the axis, be judged in domain, and yield a result whose imaginary part is silently discarded. + This is not specific to the domain test: `MatrixFunctionViaEig` reconstructs `f(A)` by inverting the eigenvector matrix, so for a defective or nearly defective matrix its result is unreliable whether the input is real or complex. + Use the Schur-based [`MatrixFunctionViaLA`](@ref) for such matrices. + +```@docs; canonical=false +MatrixAlgebraKit.default_domain_atol +``` diff --git a/ext/MatrixAlgebraKitGenericSchurExt.jl b/ext/MatrixAlgebraKitGenericSchurExt.jl index 8b7aeb100..561cf4b28 100644 --- a/ext/MatrixAlgebraKitGenericSchurExt.jl +++ b/ext/MatrixAlgebraKitGenericSchurExt.jl @@ -22,10 +22,12 @@ function MatrixAlgebraKit.default_exponential_algorithm( end function MatrixAlgebraKit.default_squareroot_algorithm( - type::Type{T}; domain_atol::Real = -1.0, kwargs... + type::Type{T}; domain_atol = nothing, kwargs... ) where {T <: StridedMatrix{<:GSFloat}} + # the remaining keywords configure the eigensolver, `domain_atol` the domain check eig_alg = MatrixAlgebraKit.default_eig_algorithm(type; kwargs...) - return MatrixFunctionViaEig(eig_alg; domain_atol) + return isnothing(domain_atol) ? MatrixFunctionViaEig(eig_alg) : + MatrixFunctionViaEig(eig_alg; domain_atol) end function geev!(::GS, A::AbstractMatrix, Dd::AbstractVector, V::AbstractMatrix; kwargs...) diff --git a/src/common/defaults.jl b/src/common/defaults.jl index 17c20eb8e..4383646fc 100644 --- a/src/common/defaults.jl +++ b/src/common/defaults.jl @@ -44,25 +44,21 @@ Default tolerance for deciding to warn if the provided `A` is not hermitian. default_hermitian_tol(A) = eps(norm(A, Inf))^(3 / 4) """ - default_domain_atol(λ, alg) + default_domain_atol(λ) Default absolute tolerance for deciding when the eigenvalues `λ` should be considered to lie outside of the domain of a matrix function, e.g. on the negative real axis for [`squareroot`](@ref) of a real matrix. -The tolerance has to absorb the error with which `alg` obtained `λ`, and that error differs by -orders of magnitude between the algorithms, so the default is algorithm-dependent. +It is set to `length(λ) * eps * maximum(abs, λ)`, i.e. the accumulated roundoff of the spectrum, +which is the same rule as `LinearAlgebra.sqrt(::Hermitian; rtol = eps(T) * size(A, 1))`. It is both the default clamping radius of [`squareroot`](@ref), which exposes it as `domain_atol`, and the never user-settable tolerance with which a complex eigenvalue of a real matrix is decided to lie *on* the negative real axis. -See [Domain considerations](@ref sec_matrixfunction_domain) for the resulting values. +See [Domain considerations](@ref sec_matrixfunction_domain). """ -function default_domain_atol end - -# roundoff scales only with number of elements and spectrum - e.g. hermitian precision -# conditioning is less strict -_roundoff_domain_atol(λ) = length(λ) * eps(real(float(one(eltype(λ))))) * maximum(abs, λ; init = abs(zero(eltype(λ)))) -_conditioning_domain_atol(λ) = defaulttol(λ) * maximum(abs, λ; init = abs(zero(eltype(λ)))) +default_domain_atol(λ) = + length(λ) * eps(real(float(one(eltype(λ))))) * maximum(abs, λ; init = abs(zero(eltype(λ)))) const DEFAULT_FIXGAUGE = Ref(true) diff --git a/src/implementations/exponential.jl b/src/implementations/exponential.jl index 3e9c92412..c03ee8459 100644 --- a/src/implementations/exponential.jl +++ b/src/implementations/exponential.jl @@ -72,7 +72,7 @@ end function exponential!((τ, A)::Tuple{Number, AbstractMatrix}, expA, alg::MatrixFunctionViaEigh) check_input(exponential!, (τ, A), expA, alg) - D, V = eigh_full!(A, alg.eigh_alg) + D, V = eigh_full!(A, select_algorithm(eigh_full!, A, _eigh_alg(alg))) if eltype(A) <: Real if eltype(τ) <: Real VexpD = rmul!(V, exponential!((τ / 2, D), D)) @@ -92,7 +92,7 @@ end function exponential!((τ, A)::Tuple{Number, AbstractMatrix}, expA, alg::MatrixFunctionViaEig) check_input(exponential!, (τ, A), expA, alg) - D, V = eig_full!(A, alg.eig_alg) + D, V = eig_full!(A, select_algorithm(eig_full!, A, _eig_alg(alg))) if eltype(A) <: Real && eltype(τ) <: Real VexpD = V * exponential!((τ, D), D) expAc = rdiv!(VexpD, LinearAlgebra.lu!(V)) diff --git a/src/implementations/matrixfunctions.jl b/src/implementations/matrixfunctions.jl index 6feddb691..a0a149738 100644 --- a/src/implementations/matrixfunctions.jl +++ b/src/implementations/matrixfunctions.jl @@ -47,24 +47,22 @@ _apply_eigh!(fA, V, fD) = project_hermitian!(mul!(fA, V * fD, V')) # accuracy with which the algorithm resolves whether an eigenvalue lies *on* the negative real axis # at all, which is a property of the eigensolver rather than a user choice. -# each algorithm obtains its eigenvalues with a different accuracy, so each brings its own default -default_domain_atol(λ, ::DiagonalAlgorithm) = _roundoff_domain_atol(λ) -default_domain_atol(λ, ::MatrixFunctionViaEigh) = _roundoff_domain_atol(λ) -default_domain_atol(λ, ::MatrixFunctionViaEig) = _conditioning_domain_atol(λ) - -# a negative tolerance denotes the runtime default, which keeps the algorithm types concrete -_domain_atol(alg::Union{MatrixFunctionViaEig, MatrixFunctionViaEigh}) = alg.domain_atol -_domain_atol(alg::DiagonalAlgorithm) = get(alg.kwargs, :domain_atol, -1.0) +# an unset keyword denotes the runtime default +_domain_atol(alg) = get(alg.kwargs, :domain_atol, nothing) # callers resolve the default before delegating to an inner `DiagonalAlgorithm`, so that the # tolerance follows the algorithm that computed the eigenvalues function _resolve_domain_atol(λ, alg) - atol = _domain_atol(alg) R = real(float(eltype(λ))) - return atol < 0 ? convert(R, default_domain_atol(λ, alg)) : convert(R, atol) + atol = _domain_atol(alg) + return convert(R, isnothing(atol) ? default_domain_atol(λ) : atol) end -_axis_atol(λ, alg) = convert(real(float(eltype(λ))), default_domain_atol(λ, alg)) +_axis_atol(λ) = convert(real(float(eltype(λ))), default_domain_atol(λ)) + +# the wrapped decomposition algorithm is optional as well +_eig_alg(alg) = get(alg.kwargs, :eig_alg, nothing) +_eigh_alg(alg) = get(alg.kwargs, :eigh_alg, nothing) # the throwing branches live in `@noinline` helpers so that the reductions and broadcasts # below stay free of error-path code, which keeps them GPU friendly @@ -110,29 +108,33 @@ end # Domain handling for `MatrixFunctionViaLA` # ----------------------------------------- -# `LinearAlgebra` never exposes the spectrum, so the check happens in result space and -# `domain_atol` bounds the imaginary part of `f(A)` instead +# `LinearAlgebra` never exposes the spectrum, so there is nothing to compare against a tolerance: +# a complex result for a real input is a domain violation, full stop @noinline function throw_la_kwargs(f, ks) return throw( - ArgumentError("`MatrixFunctionViaLA` only accepts the `domain_atol` keyword argument for `$f`, got $ks") + ArgumentError( + "`MatrixFunctionViaLA` accepts no keyword arguments for `$f`, got $ks. In particular " * + "`domain_atol` is not supported, as `LinearAlgebra` does not expose the spectrum; " * + "use `MatrixFunctionViaEig` or `MatrixFunctionViaEigh` instead." + ) ) end -# `MatrixFunctionViaLA` accepts generic keywords, so the kernels validate the ones they support -function _la_domain_atol(alg::MatrixFunctionViaLA, f) +# `MatrixFunctionViaLA` accepts generic keywords, so the kernels reject the ones they cannot honor +function _check_la_kwargs(alg::MatrixFunctionViaLA, f) ks = keys(alg.kwargs) - (isempty(ks) || ks == (:domain_atol,)) || throw_la_kwargs(f, ks) - return get(alg.kwargs, :domain_atol, -1.0) + isempty(ks) || throw_la_kwargs(f, ks) + return nothing end -@noinline function throw_complex_result(f, atol, imagmax) +@noinline function throw_complex_result(f) return throw( DomainError( f, - "The result of this matrix function applied to the given real matrix is complex (eigenvalues on the negative real axis): " * - "its imaginary part reaches $imagmax, beyond `domain_atol = $atol`. Pass a complex matrix to obtain the principal " * - "value, or increase `domain_atol` if the imaginary part is a rounding artifact." + "The result of this matrix function applied to the given real matrix is complex (eigenvalues on the negative real axis). " * + "Pass a complex matrix to obtain the principal value, or use `MatrixFunctionViaEig`/`MatrixFunctionViaEigh` to have " * + "the spectrum itself checked against `domain_atol`." ) ) end @@ -146,17 +148,3 @@ end ) ) end - -# a rounding-level imaginary part is not a domain violation: `LinearAlgebra` casts back to real only -# when the imaginary part vanishes identically -function _la_project_real!(fA, fAc, domain_atol::Real, f) - all(isfinite, fAc) || throw_nonfinite_result(f) - R = real(eltype(fA)) - # the working precision is that of the output: `LinearAlgebra` computes in complex arithmetic - # throughout, so e.g. a `Float32` input promotes all the way to `ComplexF64` - atol = domain_atol < 0 ? defaulttol(fA) * convert(R, norm(fAc, Inf)) : convert(R, domain_atol) - imagmax = convert(R, maximum(abs ∘ imag, fAc; init = zero(real(eltype(fAc))))) - imagmax <= atol || throw_complex_result(f, atol, imagmax) - fA .= real.(fAc) - return fA -end diff --git a/src/implementations/squareroot.jl b/src/implementations/squareroot.jl index 6cbb2e84d..5512c1e78 100644 --- a/src/implementations/squareroot.jl +++ b/src/implementations/squareroot.jl @@ -19,20 +19,20 @@ initialize_output(::typeof(squareroot!), A::AbstractMatrix, ::AbstractAlgorithm) # -------------- function squareroot!(A::AbstractMatrix, sqrtA, alg::MatrixFunctionViaLA) check_input(squareroot!, A, sqrtA, alg) - domain_atol = _la_domain_atol(alg, squareroot!) + _check_la_kwargs(alg, squareroot!) # `LinearAlgebra.sqrt` of a real matrix is real whenever the principal square root is sqrtAc = LinearAlgebra.sqrt(A) if eltype(sqrtAc) <: Complex && !(eltype(sqrtA) <: Complex) - _la_project_real!(sqrtA, sqrtAc, domain_atol, squareroot!) - else - copy!(sqrtA, sqrtAc) + all(isfinite, sqrtAc) || throw_nonfinite_result(squareroot!) + throw_complex_result(squareroot!) end + copy!(sqrtA, sqrtAc) return sqrtA end function squareroot!(A::AbstractMatrix, sqrtA, alg::MatrixFunctionViaEigh) check_input(squareroot!, A, sqrtA, alg) - D, V = eigh_full!(A, alg.eigh_alg) + D, V = eigh_full!(A, select_algorithm(eigh_full!, A, _eigh_alg(alg))) λ = diagview(D) _clamp_domain_eigenvalues!(λ, _resolve_domain_atol(λ, alg)) λ .= sqrt.(λ) @@ -41,12 +41,12 @@ end function squareroot!(A::AbstractMatrix, sqrtA, alg::MatrixFunctionViaEig) check_input(squareroot!, A, sqrtA, alg) - D, V = eig_full!(A, alg.eig_alg) + D, V = eig_full!(A, select_algorithm(eig_full!, A, _eig_alg(alg))) λ = diagview(D) atol = _resolve_domain_atol(λ, alg) # a real result requires the spectrum to stay off the negative real axis; whether an eigenvalue # sits *on* that axis keeps the algorithm default however `domain_atol` was set - eltype(A) <: Real && _clamp_domain_eigenvalues!(λ, atol, _axis_atol(λ, alg)) + eltype(A) <: Real && _clamp_domain_eigenvalues!(λ, atol, _axis_atol(λ)) diag_alg = DiagonalAlgorithm(; domain_atol = atol) return _apply_eig!(sqrtA, V, squareroot!(D, D, diag_alg)) end diff --git a/src/interface/matrixfunctions.jl b/src/interface/matrixfunctions.jl index 08886345e..bd3435482 100644 --- a/src/interface/matrixfunctions.jl +++ b/src/interface/matrixfunctions.jl @@ -2,15 +2,12 @@ # MATRIX FUNCTION ALGORITHMS # ================================ """ - MatrixFunctionViaLA(; domain_atol=-1) + MatrixFunctionViaLA() Algorithm type to denote computing a function of a matrix `A` via the implementation of `LinearAlgebra`. -For a matrix function with a restricted domain, i.e. [`squareroot`](@ref), `domain_atol` specifies the -absolute tolerance on the imaginary part of the result below which a complex result is attributed to -rounding rather than to a domain violation, with a negative value denoting the default tolerance. -Note that this measures a different quantity than the eigenvalue tolerance of -[`MatrixFunctionViaEig`](@ref) and [`MatrixFunctionViaEigh`](@ref), as `LinearAlgebra` does not -expose the spectrum; see [Domain considerations](@ref sec_matrixfunction_domain). +In order to retain type stability, complex results for real inputs are rejected with a `DomainError`. +Use [`MatrixFunctionViaEig`](@ref) or [`MatrixFunctionViaEigh`](@ref) to check the spectrum itself against a tolerance. +See also [Domain considerations](@ref sec_matrixfunction_domain). """ @algdef MatrixFunctionViaLA @@ -35,36 +32,34 @@ As this algorithm requires no LAPACK support, it also applies at arbitrary preci @algdef MatrixFunctionViaTaylor """ - MatrixFunctionViaEigh(eigh_alg; domain_atol=-1) + MatrixFunctionViaEigh(eigh_alg; domain_atol = default_domain_atol(λ)) + MatrixFunctionViaEigh(; eigh_alg, domain_atol = default_domain_atol(λ)) 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. +The optional `eigh_alg` specifies which hermitian eigendecomposition implementation to use, either +positionally or as a keyword, and defaults to the one selected for the input. `domain_atol` applies to [`squareroot`](@ref): it is the absolute tolerance within which negative -eigenvalues are treated as rounding artifacts and clamped onto zero, with a negative value denoting -the default tolerance [`default_domain_atol`](@ref). Raising it accepts more matrices; +eigenvalues are treated as rounding artifacts and clamped onto zero, and defaults to +[`default_domain_atol`](@ref). Raising it accepts more matrices; see [Domain considerations](@ref sec_matrixfunction_domain). """ -struct MatrixFunctionViaEigh{A <: AbstractAlgorithm} <: AbstractAlgorithm - eigh_alg::A - domain_atol::Float64 # negative value for runtime defaults -end -MatrixFunctionViaEigh(eigh_alg::AbstractAlgorithm; domain_atol::Real = -1.0) = - MatrixFunctionViaEigh(eigh_alg, Float64(domain_atol)) -function Base.show(io::IO, alg::MatrixFunctionViaEigh) - print(io, "MatrixFunctionViaEigh(") - _show_alg(io, alg.eigh_alg) - alg.domain_atol < 0 || print(io, "; domain_atol=", alg.domain_atol) - return print(io, ")") -end +@algdef MatrixFunctionViaEigh + +@deprecate( + MatrixFunctionViaEigh(eigh_alg::AbstractAlgorithm; kwargs...), + MatrixFunctionViaEigh(; eigh_alg, kwargs...) +) """ - MatrixFunctionViaEig(eig_alg; domain_atol=-1) + MatrixFunctionViaEig(eig_alg; domain_atol = default_domain_atol(λ)) + MatrixFunctionViaEig(; eig_alg, domain_atol = default_domain_atol(λ)) 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. +The optional `eig_alg` specifies which eigendecomposition implementation to use, either positionally +or as a keyword, and defaults to the one selected for the input. `domain_atol` applies to [`squareroot`](@ref): it is the absolute tolerance within which eigenvalues -on the negative real axis are treated as rounding artifacts and clamped onto zero, with a negative -value denoting the default tolerance [`default_domain_atol`](@ref). Raising it accepts more matrices; +on the negative real axis are treated as rounding artifacts and clamped onto zero, and defaults to +[`default_domain_atol`](@ref). Raising it accepts more matrices; see [Domain considerations](@ref sec_matrixfunction_domain). !!! warning @@ -73,15 +68,9 @@ see [Domain considerations](@ref sec_matrixfunction_domain). are resolved only to `eps^(1/k)` for a Jordan block of size `k`. Prefer [`MatrixFunctionViaLA`](@ref), which is Schur-based, for such matrices. """ -struct MatrixFunctionViaEig{A <: AbstractAlgorithm} <: AbstractAlgorithm - eig_alg::A - domain_atol::Float64 # negative value for runtime defaults -end -MatrixFunctionViaEig(eig_alg::AbstractAlgorithm; domain_atol::Real = -1.0) = - MatrixFunctionViaEig(eig_alg, Float64(domain_atol)) -function Base.show(io::IO, alg::MatrixFunctionViaEig) - print(io, "MatrixFunctionViaEig(") - _show_alg(io, alg.eig_alg) - alg.domain_atol < 0 || print(io, "; domain_atol=", alg.domain_atol) - return print(io, ")") -end +@algdef MatrixFunctionViaEig + +@deprecate( + MatrixFunctionViaEig(eig_alg::AbstractAlgorithm; kwargs...), + MatrixFunctionViaEig(; eig_alg, kwargs...) +) diff --git a/src/interface/squareroot.jl b/src/interface/squareroot.jl index e7be7d28b..97efb03c6 100644 --- a/src/interface/squareroot.jl +++ b/src/interface/squareroot.jl @@ -14,9 +14,11 @@ The scalar type of the output matches that of the input. As a consequence, a real matrix with eigenvalues on the negative real axis, for which the principal square root is complex, leads to a `DomainError`; pass a complex matrix to obtain the principal value. -Real eigenvalues that are negative within a tolerance `domain_atol` are treated as rounding -artifacts and clamped to zero, so that raising `domain_atol` accepts more matrices. It defaults to -[`default_domain_atol`](@ref); see [Domain considerations](@ref sec_matrixfunction_domain). +For the algorithms that have access to the spectrum, real eigenvalues that are negative within a +tolerance `domain_atol` are treated as rounding artifacts and clamped to zero, so that raising +`domain_atol` accepts more matrices. It defaults to [`default_domain_atol`](@ref), and is not +supported by [`MatrixFunctionViaLA`](@ref); +see [Domain considerations](@ref sec_matrixfunction_domain). !!! note The bang method `squareroot!` optionally accepts the output structure and diff --git a/test/matrixfunctions/squareroot.jl b/test/matrixfunctions/squareroot.jl index f9d75f41d..d4bfc3446 100644 --- a/test/matrixfunctions/squareroot.jl +++ b/test/matrixfunctions/squareroot.jl @@ -37,7 +37,9 @@ if !is_buildkite TestSuite.test_squareroot_hermitian(T, (m, m), LAPACK_EIG_ALGS; exact_hermiticity = false) TestSuite.test_squareroot_hermitian(T, (m, m), LAPACK_EIGH_ALGS) TestSuite.test_squareroot_reference(T, (m, m)) - TestSuite.test_squareroot_domain(T, (md, md), LAPACK_EIG_ALGS) + # `MatrixFunctionViaLA` has no access to the spectrum, and thus no `domain_atol` + TestSuite.test_squareroot_domain(T, (md, md), (MatrixFunctionViaLA(),); test_domain_atol = false) + TestSuite.test_squareroot_domain(T, (md, md), (MatrixFunctionViaEig(QRIteration()),)) TestSuite.test_squareroot_domain(T, (md, md), LAPACK_EIGH_ALGS; hermitian_output = true) end diff --git a/test/testsuite/TestSuite.jl b/test/testsuite/TestSuite.jl index 251608243..32b12c57f 100644 --- a/test/testsuite/TestSuite.jl +++ b/test/testsuite/TestSuite.jl @@ -136,15 +136,9 @@ end # rebuild `alg` with an explicit `domain_atol`, so that the domain tests need not spell out the # inner decomposition algorithm a second time -with_domain_atol(alg::MatrixFunctionViaEig, atol) = MatrixFunctionViaEig(alg.eig_alg; domain_atol = atol) -with_domain_atol(alg::MatrixFunctionViaEigh, atol) = MatrixFunctionViaEigh(alg.eigh_alg; domain_atol = atol) -with_domain_atol(::MatrixAlgebraKit.DiagonalAlgorithm, atol) = DiagonalAlgorithm(; domain_atol = atol) -with_domain_atol(::MatrixFunctionViaLA, atol) = MatrixFunctionViaLA(; domain_atol = atol) - -# a tolerance generous enough to admit an eigenvalue at `-√eps`. For `MatrixFunctionViaLA` it bounds -# the imaginary part of the result rather than the spectrum, which sits on a coarser scale. -domain_test_atol(::MatrixAlgebraKit.AbstractAlgorithm, R) = cbrt(eps(R)) -domain_test_atol(::MatrixFunctionViaLA, R) = one(R) / 2 +function with_domain_atol(alg::MatrixAlgebraKit.Algorithm, atol) + return MatrixAlgebraKit.Algorithm{MatrixAlgebraKit.name(alg)}(; alg.kwargs..., domain_atol = atol) +end include("ad_utils.jl") diff --git a/test/testsuite/matrixfunctions/squareroot.jl b/test/testsuite/matrixfunctions/squareroot.jl index 1dc6d106a..32bea5c2d 100644 --- a/test/testsuite/matrixfunctions/squareroot.jl +++ b/test/testsuite/matrixfunctions/squareroot.jl @@ -64,8 +64,12 @@ end # a matrix whose spectrum reaches the negative real axis has a complex principal square root, which # a type-stable real output cannot represent. Pass `hermitian_output = true` for the algorithms that -# promise a hermitian result: those must reject a negative eigenvalue whatever the scalar type. -function test_squareroot_domain(T::Type, sz, algs; hermitian_output = false, kwargs...) +# promise a hermitian result: those must reject a negative eigenvalue whatever the scalar type, and +# `test_domain_atol = false` for the ones without access to the spectrum, which have no tolerance. +function test_squareroot_domain( + T::Type, sz, algs; + hermitian_output = false, test_domain_atol = true, kwargs... + ) R = real(eltype(T)) n = sz isa Tuple ? first(sz) : sz summary_str = testargs_summary(T, sz) @@ -90,19 +94,24 @@ function test_squareroot_domain(T::Type, sz, algs; hermitian_output = false, kwa @test eltype(sqrtAclamp) == eltype(Aclamp) @test sqrtAclamp * sqrtAclamp ≈ Aclamp atol = sqrt(eps(R)) - # an eigenvalue beyond every default tolerance is out of domain, while an explicit + # an eigenvalue beyond the default tolerance is out of domain, while an explicit # `domain_atol` admits it after all λwide = collect(R, 1:n) λwide[1] = -sqrt(eps(R)) Awide = instantiate_hermitian_spectrum(T, sz, λwide) if eltype(T) <: Real || hermitian_output @test_throws DomainError squareroot(Awide, alg) + else + sqrtAwide = @testinferred squareroot(Awide, alg) + @test sqrtAwide * sqrtAwide ≈ Awide + end + if test_domain_atol + wide_alg = with_domain_atol(alg, cbrt(eps(R))) + sqrtAwide = @testinferred squareroot(Awide, wide_alg) + @test eltype(sqrtAwide) == eltype(Awide) + # accepting is backward stable, but only to the size of the eigenvalue that was discarded + @test sqrtAwide * sqrtAwide ≈ Awide atol = sqrt(sqrt(eps(R))) end - wide_alg = with_domain_atol(alg, domain_test_atol(alg, R)) - sqrtAwide = @testinferred squareroot(Awide, wide_alg) - @test eltype(sqrtAwide) == eltype(Awide) - # accepting is backward stable, but only to the size of the eigenvalue that was discarded - @test sqrtAwide * sqrtAwide ≈ Awide atol = sqrt(sqrt(eps(R))) end end From efe738a40e33e09917cc86e131b778695a16dbf9 Mon Sep 17 00:00:00 2001 From: lkdvos Date: Wed, 26 Aug 2026 12:45:34 -0400 Subject: [PATCH 6/6] more square root algorithms --- docs/src/changelog.md | 15 +- docs/src/user_interface/algorithms.md | 1 + docs/src/user_interface/matrix_functions.md | 12 +- ext/MatrixAlgebraKitGenericSchurExt.jl | 9 -- src/MatrixAlgebraKit.jl | 4 +- src/common/quasitriangular.jl | 159 +++++++++++++++++++ src/implementations/matrixfunctions.jl | 17 +- src/implementations/squareroot.jl | 127 ++++++++++++++- src/interface/matrixfunctions.jl | 41 ++++- src/interface/squareroot.jl | 8 +- test/matrixfunctions/squareroot.jl | 41 ++++- test/testsuite/TestSuite.jl | 19 +++ test/testsuite/matrixfunctions/squareroot.jl | 59 +++++-- 13 files changed, 473 insertions(+), 39 deletions(-) create mode 100644 src/common/quasitriangular.jl diff --git a/docs/src/changelog.md b/docs/src/changelog.md index ac4bc68dd..9745974b2 100644 --- a/docs/src/changelog.md +++ b/docs/src/changelog.md @@ -23,11 +23,22 @@ When releasing a new version, move the "Unreleased" changes to a new version sec ### Added - New matrix function `squareroot`, computing the principal square root, supporting the - `MatrixFunctionViaLA`, `MatrixFunctionViaEig`, `MatrixFunctionViaEigh` and `DiagonalAlgorithm` - algorithms ([#261](https://github.com/QuantumKitHub/MatrixAlgebraKit.jl/pull/261)). + `MatrixFunctionViaSchur`, `MatrixFunctionViaLA`, `MatrixFunctionViaEig`, `MatrixFunctionViaEigh` + and `DiagonalAlgorithm` algorithms + ([#261](https://github.com/QuantumKitHub/MatrixAlgebraKit.jl/pull/261)). +- New algorithm `MatrixFunctionViaSchur`, a native implementation of the Schur method for + `squareroot`: the recursion of Björck & Hammarling on the (quasi-)triangular Schur factor, in the + real quasi-triangular variant of Higham so that a real input stays in real arithmetic, with the + recursive blocking of Deadman, Higham & Ralha. Unlike `MatrixFunctionViaLA` it is backward stable + for every input, honors `domain_atol`, accepts a `schur_alg` and a `blocksize`, and computes a + real square root of a real matrix at arbitrary precision, including half precision. ### Changed +- `MatrixFunctionViaSchur` is now the default algorithm for `squareroot` of a dense matrix, replacing + `MatrixFunctionViaLA`. As a consequence `domain_atol` is supported by default, a defective matrix + no longer requires selecting an algorithm by hand, and the `GenericSchur` extension no longer + overrides the default for `Float16`/`BigFloat` and friends. - `MatrixFunctionViaEig` and `MatrixFunctionViaEigh` are now defined through `@algdef`, so that both the wrapped decomposition algorithm (`eig_alg` / `eigh_alg`, still accepted positionally) and the new `domain_atol` are optional keyword arguments. ([#261](https://github.com/QuantumKitHub/MatrixAlgebraKit.jl/pull/261)). diff --git a/docs/src/user_interface/algorithms.md b/docs/src/user_interface/algorithms.md index 3e348697d..142f34e43 100644 --- a/docs/src/user_interface/algorithms.md +++ b/docs/src/user_interface/algorithms.md @@ -105,6 +105,7 @@ The following algorithms for matrix functions are available. | Algorithm | Applicable matrix functions | Key keyword arguments | |:----------|:--------------------------|:----------------------| | [`MatrixFunctionViaTaylor`](@ref) | exponential | `tol`, `balance`, `estimate_order` | +| [`MatrixFunctionViaSchur`](@ref) | squareroot | `schur_alg`, `blocksize`, `domain_atol` | | [`MatrixFunctionViaLA`](@ref) | exponential, squareroot | — | | [`MatrixFunctionViaEig`](@ref) | exponential, squareroot | `eig_alg` (also positional), `domain_atol` (squareroot) | | [`MatrixFunctionViaEigh`](@ref) | exponential, squareroot | `eigh_alg` (also positional), `domain_atol` (squareroot) | diff --git a/docs/src/user_interface/matrix_functions.md b/docs/src/user_interface/matrix_functions.md index 367adc71a..68f4025fe 100644 --- a/docs/src/user_interface/matrix_functions.md +++ b/docs/src/user_interface/matrix_functions.md @@ -24,13 +24,15 @@ Additionally, the `f!` method typically assumes that it is allowed to destroy th The matrix functions share a common set of algorithms, which differ in how they reduce the problem to a scalar function of the eigenvalues, along with more specialized implementations for specific functions: -- [`MatrixFunctionViaLA`](@ref) defers to the implementation of `LinearAlgebra`, which is Schur-based for [`squareroot`](@ref) and a Padé approximation for [`exponential`](@ref). +- [`MatrixFunctionViaSchur`](@ref) applies to [`squareroot`](@ref) only, and evaluates the function on the (quasi-)triangular factor of a Schur decomposition computed through `schur_full`. It is backward stable, independent of the conditioning of the eigenbasis, and applies to generic data types at arbitrary precision, but it indexes individual entries and is therefore unsuited to GPU arrays. - [`MatrixFunctionViaEig`](@ref) and [`MatrixFunctionViaEigh`](@ref) first compute an eigenvalue decomposition, through `eig_full` and `eigh_full` respectively, and then apply the scalar function to the eigenvalues. The latter requires a hermitian input, and in return its result is hermitian by construction. - [`MatrixFunctionViaTaylor`](@ref) applies to [`exponential`](@ref) only, and evaluates its Taylor series through scaling and squaring. As it requires no LAPACK support, it also applies to generic data types at arbitrary precision. +- [`MatrixFunctionViaLA`](@ref) defers to the implementation of `LinearAlgebra`, which is a Padé approximation for [`exponential`](@ref). For [`squareroot`](@ref) it dispatches on the structure of the input rather than on a fixed strategy, and is Schur-based only for a general matrix: a hermitian one is routed through an eigenvalue decomposition instead, with a tolerance of its own that `domain_atol` cannot reach, and a real matrix at other than BLAS precision is promoted to a complex Schur form, so that a real result is not recovered even when it exists. - [`DiagonalAlgorithm`](@ref) is the fast path for a `Diagonal` input, and simply maps the scalar function over the diagonal. ```@docs; canonical=false MatrixAlgebraKit.MatrixFunctionViaTaylor +MatrixAlgebraKit.MatrixFunctionViaSchur MatrixAlgebraKit.MatrixFunctionViaLA MatrixAlgebraKit.MatrixFunctionViaEig MatrixAlgebraKit.MatrixFunctionViaEigh @@ -50,7 +52,7 @@ exponential ## Square root The principal [square root](https://en.wikipedia.org/wiki/Square_root_of_a_matrix) of a square matrix `A` is the unique square root whose eigenvalues have nonnegative real part. -It is computed by the function [`squareroot`](@ref), with [`MatrixFunctionViaLA`](@ref) as the default algorithm, and is subject to the [domain considerations](@ref sec_matrixfunction_domain) below. +It is computed by the function [`squareroot`](@ref), with [`MatrixFunctionViaSchur`](@ref) as the default algorithm, and is subject to the [domain considerations](@ref sec_matrixfunction_domain) below. ```@docs; canonical=false squareroot @@ -71,6 +73,9 @@ This is the same rule as `LinearAlgebra.sqrt(::Hermitian; rtol = eps(T) * size(A The eigenvalues of [`MatrixFunctionViaEig`](@ref) are additionally limited by the conditioning of the eigenvectors, so for a poorly conditioned eigenbasis a larger `domain_atol` may have to be set explicitly. The same default is used for the never user-settable tolerance with which a complex eigenvalue of a real matrix is decided to lie *on* the negative real axis, as that is a question about the accuracy of the eigensolver rather than about the domain. +[`MatrixFunctionViaSchur`](@ref) needs no such tolerance at all: a real eigenvalue of a real matrix is a `1×1` block of its real Schur form, whereas a `2×2` block is a genuine complex-conjugate pair, which lies off the negative real axis and always has a real square root. +The block structure of the decomposition therefore answers exactly the question that the other algorithms have to settle numerically. + Additionally, not all algorithms have acces to the spectrum, so not all methods are suitable for eigenvalues close to the domain edges. For example, [`MatrixFunctionViaLA`](@ref) defers to `LinearAlgebra`, which decides internally whether a real result exists and hands back a complex matrix when it does not. There are no eigenvalues to compare against anything, so it rejects a complex result for a real input outright, and passing it `domain_atol` is an error rather than a silent no-op. @@ -79,7 +84,8 @@ There are no eigenvalues to compare against anything, so it rejects a complex re The eigenvalues of a Jordan block of size `k` are resolved only to `eps^(1/k)`, which exceeds every tolerance on this page. A real matrix with a defective negative eigenvalue can therefore have its spectrum reported as a complex-conjugate pair well off the axis, be judged in domain, and yield a result whose imaginary part is silently discarded. This is not specific to the domain test: `MatrixFunctionViaEig` reconstructs `f(A)` by inverting the eigenvector matrix, so for a defective or nearly defective matrix its result is unreliable whether the input is real or complex. - Use the Schur-based [`MatrixFunctionViaLA`](@ref) for such matrices. + Use the Schur-based [`MatrixFunctionViaSchur`](@ref), the default, for such matrices. + Note that a defective eigenvalue sitting *on* the negative real axis is ill-conditioned for every algorithm, since the computed copies of it straddle the branch cut of the scalar square root. ```@docs; canonical=false MatrixAlgebraKit.default_domain_atol diff --git a/ext/MatrixAlgebraKitGenericSchurExt.jl b/ext/MatrixAlgebraKitGenericSchurExt.jl index 561cf4b28..c34f1abc4 100644 --- a/ext/MatrixAlgebraKitGenericSchurExt.jl +++ b/ext/MatrixAlgebraKitGenericSchurExt.jl @@ -21,15 +21,6 @@ function MatrixAlgebraKit.default_exponential_algorithm( return MatrixFunctionViaEig(eig_alg) end -function MatrixAlgebraKit.default_squareroot_algorithm( - type::Type{T}; domain_atol = nothing, kwargs... - ) where {T <: StridedMatrix{<:GSFloat}} - # the remaining keywords configure the eigensolver, `domain_atol` the domain check - eig_alg = MatrixAlgebraKit.default_eig_algorithm(type; kwargs...) - return isnothing(domain_atol) ? MatrixFunctionViaEig(eig_alg) : - MatrixFunctionViaEig(eig_alg; domain_atol) -end - function geev!(::GS, A::AbstractMatrix, Dd::AbstractVector, V::AbstractMatrix; kwargs...) D, Vmat = GenericSchur.eigen!(A) copyto!(Dd, D) diff --git a/src/MatrixAlgebraKit.jl b/src/MatrixAlgebraKit.jl index a11179f84..ba9ac383c 100644 --- a/src/MatrixAlgebraKit.jl +++ b/src/MatrixAlgebraKit.jl @@ -42,7 +42,8 @@ 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, MatrixFunctionViaTaylor +export MatrixFunctionViaLA, MatrixFunctionViaEig, MatrixFunctionViaEigh, MatrixFunctionViaTaylor, + MatrixFunctionViaSchur export DefaultAlgorithm export DiagonalAlgorithm export NativeBlocked @@ -94,6 +95,7 @@ include("common/pullbacks.jl") include("common/safemethods.jl") include("common/view.jl") include("common/regularinv.jl") +include("common/quasitriangular.jl") include("common/matrixproperties.jl") include("common/balancing.jl") include("common/utility.jl") diff --git a/src/common/quasitriangular.jl b/src/common/quasitriangular.jl new file mode 100644 index 000000000..d23e2b5c8 --- /dev/null +++ b/src/common/quasitriangular.jl @@ -0,0 +1,159 @@ +# Methods for upper quasi-triangular matrices, i.e. the real Schur form: block upper triangular with +# diagonal blocks of size 1 for a real eigenvalue and 2 for a complex-conjugate pair. A complex +# Schur form is the special case where every block has size 1. + +# `T[i+1, i]` is the only subdiagonal entry that can be nonzero, so it decides the size of the block +# starting at `i`, while `T[i, i-1]` decides the size of the one ending there +@inline _quasitriu_blocksize(T, i, n) = @inbounds ((i < n && !iszero(T[i + 1, i])) ? 2 : 1) +@inline _quasitriu_blocksize_end(T, i) = @inbounds ((i > 1 && !iszero(T[i, i - 1])) ? 2 : 1) + +# bisection point that does not cut a 2x2 block, or `0` for a matrix that is a single block +function _quasitriu_split(T, n) + n <= 1 && return 0 + s = n ÷ 2 + @inbounds iszero(T[s + 1, s]) || (s += 1) + return s < n ? s : 0 +end + +# scratch space for the 4x4 system of the 2x2-by-2x2 Sylvester solve, allocated once per call +_quasitriu_workspace(A) = (similar(A, (4, 4)), similar(A, (4,))) + +# Sylvester equations +# ------------------- + +# `A * X + X * B = C` for upper quasi-triangular `A` and `B`, in place in `C`; note the opposite sign +# convention to `LinearAlgebra.sylvester`. A singular system surfaces as a non-finite result. +function _quasitriu_sylvester!(A, B, C, ws, blocksize::Int) + m, n = size(C) + # split the larger dimension; an indivisible one is a single 2x2 block, and then so is the other + sa = (m >= n && m > blocksize) ? _quasitriu_split(A, m) : 0 + sb = (iszero(sa) && n > blocksize) ? _quasitriu_split(B, n) : 0 + if !iszero(sa) + # `A = [A11 A12; 0 A22]` and `X = [X1; X2]`, so `X2` is determined on its own + r1, r2 = 1:sa, (sa + 1):m + A11, A12, A22 = view(A, r1, r1), view(A, r1, r2), view(A, r2, r2) + C1, C2 = view(C, r1, :), view(C, r2, :) + _quasitriu_sylvester!(A22, B, C2, ws, blocksize) + mul!(C1, A12, C2, -1, 1) + _quasitriu_sylvester!(A11, B, C1, ws, blocksize) + elseif !iszero(sb) + # `B = [B11 B12; 0 B22]` and `X = [X1 X2]`, so `X1` is determined on its own + c1, c2 = 1:sb, (sb + 1):n + B11, B12, B22 = view(B, c1, c1), view(B, c1, c2), view(B, c2, c2) + C1, C2 = view(C, :, c1), view(C, :, c2) + _quasitriu_sylvester!(A, B11, C1, ws, blocksize) + mul!(C2, C1, B12, -1, 1) + _quasitriu_sylvester!(A, B22, C2, ws, blocksize) + else + _quasitriu_sylvester_point!(A, B, C, ws) + end + return C +end + +# Bartels-Stewart: block columns of `B` from the left and block rows of `A` from the bottom, so that +# the corrections only involve blocks that are already solved +function _quasitriu_sylvester_point!(A, B, C, ws) + m, n = size(C) + j = 1 + @inbounds while j <= n + q = _quasitriu_blocksize(B, j, n) + J = j:(j + q - 1) + i = m + while i >= 1 + p = _quasitriu_blocksize_end(A, i) + I = (i - p + 1):i + Cij = view(C, I, J) + i < m && mul!(Cij, view(A, I, (i + 1):m), view(C, (i + 1):m, J), -1, 1) + j > 1 && mul!(Cij, view(C, I, 1:(j - 1)), view(B, 1:(j - 1), J), -1, 1) + _quasitriu_sylvester_block!(view(A, I, I), view(B, J, J), Cij, ws) + i -= p + end + j += q + end + return C +end + +# `A * X + X * B = C` for blocks of size 1 or 2, in place in `C` +Base.@propagate_inbounds function _quasitriu_sylvester_block!(A, B, C, ws) + p, q = size(C) + if p == 1 && q == 1 + d = A[1, 1] + B[1, 1] + # a vanishing right hand side is consistent with a singular equation, and selects the root + # that vanishes along with it + C[1, 1] = (iszero(d) && iszero(C[1, 1])) ? d : C[1, 1] / d + elseif p == 2 && q == 1 + b = B[1, 1] + _solve_adjugate2!(A[1, 1] + b, A[1, 2], A[2, 1], A[2, 2] + b, C) + elseif p == 1 && q == 2 + # `a * x + x * B = c` transposes into a system for the row vector `x` + a = A[1, 1] + _solve_adjugate2!(B[1, 1] + a, B[2, 1], B[1, 2], B[2, 2] + a, C) + else + _quasitriu_sylvester_2x2!(A, B, C, ws) + end + return C +end + +# the 2x2-by-2x2 case is the 4x4 system `(I ⊗ A + Bᵀ ⊗ I) * vec(X) = vec(C)` +function _quasitriu_sylvester_2x2!(A, B, C, (M, v)) + z = zero(eltype(M)) + @inbounds for j in 1:2, i in 1:2 + r = i + 2 * (j - 1) + for l in 1:2, k in 1:2 + M[r, k + 2 * (l - 1)] = (j == l ? A[i, k] : z) + (i == k ? B[l, j] : z) + end + v[r] = C[i, j] + end + _solve_gauss4!(M, v) + @inbounds for j in 1:2, i in 1:2 + C[i, j] = v[i + 2 * (j - 1)] + end + return C +end + +# Small dense solvers +# ------------------- + +# `[m11 m12; m21 m22] * x = c` in place in `c`, through the adjugate +Base.@propagate_inbounds function _solve_adjugate2!(m11, m12, m21, m22, c) + d = m11 * m22 - m12 * m21 + c1, c2 = c[1], c[2] + c[1] = (m22 * c1 - m12 * c2) / d + c[2] = (m11 * c2 - m21 * c1) / d + return c +end + +# 4x4 Gaussian elimination with partial pivoting, in place in `M` and `v` +function _solve_gauss4!(M, v) + n = 4 + @inbounds for k in 1:n + p, amax = k, abs(M[k, k]) + for i in (k + 1):n + a = abs(M[i, k]) + a > amax && ((p, amax) = (i, a)) + end + if p != k + for j in k:n + M[k, j], M[p, j] = M[p, j], M[k, j] + end + v[k], v[p] = v[p], v[k] + end + piv = M[k, k] + for i in (k + 1):n + f = M[i, k] / piv + iszero(f) && continue + for j in (k + 1):n + M[i, j] -= f * M[k, j] + end + v[i] -= f * v[k] + end + end + @inbounds for k in n:-1:1 + s = v[k] + for j in (k + 1):n + s -= M[k, j] * v[j] + end + v[k] = s / M[k, k] + end + return v +end diff --git a/src/implementations/matrixfunctions.jl b/src/implementations/matrixfunctions.jl index a0a149738..a95abf22f 100644 --- a/src/implementations/matrixfunctions.jl +++ b/src/implementations/matrixfunctions.jl @@ -63,6 +63,10 @@ _axis_atol(λ) = convert(real(float(eltype(λ))), default_domain_atol(λ)) # the wrapped decomposition algorithm is optional as well _eig_alg(alg) = get(alg.kwargs, :eig_alg, nothing) _eigh_alg(alg) = get(alg.kwargs, :eigh_alg, nothing) +_schur_alg(alg) = get(alg.kwargs, :schur_alg, nothing) + +# `0` denotes the algorithm default, `1` the unblocked algorithm, following `qr_householder!` +_blocksize(alg) = get(alg.kwargs, :blocksize, 0) # the throwing branches live in `@noinline` helpers so that the reductions and broadcasts # below stay free of error-path code, which keeps them GPU friendly @@ -116,7 +120,7 @@ end ArgumentError( "`MatrixFunctionViaLA` accepts no keyword arguments for `$f`, got $ks. In particular " * "`domain_atol` is not supported, as `LinearAlgebra` does not expose the spectrum; " * - "use `MatrixFunctionViaEig` or `MatrixFunctionViaEigh` instead." + "use `MatrixFunctionViaSchur`, `MatrixFunctionViaEig` or `MatrixFunctionViaEigh` instead." ) ) end @@ -139,12 +143,19 @@ end ) end -@noinline function throw_nonfinite_result(f) +@noinline function throw_nonfinite_result(f, advice) return throw( DomainError( f, "The result of this matrix function is not finite, which signals a (numerically) singular input for which it is undefined. " * - "Use `MatrixFunctionViaEig`/`MatrixFunctionViaEigh` to have the spectrum itself checked against `domain_atol`." + advice ) ) end + +const _NONFINITE_LA_ADVICE = "Use `MatrixFunctionViaEig`/`MatrixFunctionViaEigh` to have the spectrum itself checked against `domain_atol`." + +# the Schur form does expose the spectrum, so a singular result there is a property of the input +# rather than of the tolerance: it means two of the eigenvalue square roots cancel, as a repeated +# zero eigenvalue does, leaving one of the equations for the off-diagonal entries unsolvable +const _NONFINITE_SCHUR_ADVICE = "Two of the eigenvalue square roots sum to zero, which leaves the result undefined whatever the tolerance." diff --git a/src/implementations/squareroot.jl b/src/implementations/squareroot.jl index 5512c1e78..92b0fe2b1 100644 --- a/src/implementations/squareroot.jl +++ b/src/implementations/squareroot.jl @@ -23,7 +23,7 @@ function squareroot!(A::AbstractMatrix, sqrtA, alg::MatrixFunctionViaLA) # `LinearAlgebra.sqrt` of a real matrix is real whenever the principal square root is sqrtAc = LinearAlgebra.sqrt(A) if eltype(sqrtAc) <: Complex && !(eltype(sqrtA) <: Complex) - all(isfinite, sqrtAc) || throw_nonfinite_result(squareroot!) + all(isfinite, sqrtAc) || throw_nonfinite_result(squareroot!, _NONFINITE_LA_ADVICE) throw_complex_result(squareroot!) end copy!(sqrtA, sqrtAc) @@ -62,3 +62,128 @@ function squareroot!(A::AbstractMatrix, sqrtA, alg::DiagonalAlgorithm) λ .= sqrt.(λ) return sqrtA end + +# Schur logic +# ----------- +function squareroot!(A::AbstractMatrix, sqrtA, alg::MatrixFunctionViaSchur) + check_input(squareroot!, A, sqrtA, alg) + T, Z, vals = schur_full!(A, select_algorithm(schur_full!, A, _schur_alg(alg))) + # the (quasi-)diagonal of `T` is the spectrum, and it is the block structure rather than a + # tolerance that decides whether an eigenvalue lies on the negative real axis + eltype(T) <: Real && _clamp_domain_quasitriu!(T, _resolve_domain_atol(vals, alg)) + R = _squareroot_quasitriu!(zero!(similar(T)), T, _squareroot_blocksize(T, alg)) + all(isfinite, R) || throw_nonfinite_result(squareroot!, _NONFINITE_SCHUR_ADVICE) + ZR = mul!(T, Z, R) # `T` shares its storage with `A` and is no longer needed + has_equal_storage(sqrtA, T) || return mul!(sqrtA, ZR, Z') + return copy!(sqrtA, mul!(R, ZR, Z')) +end + +# How deep it pays to recurse follows whether the multiplications reach a level-3 BLAS: measured +# over `n = 500` to `2000` a BLAS float is fastest at `2`-`4` and loses up to `1.7x` at `64`, while +# `BigFloat` prefers the entrywise algorithm outright and loses `2.2x` at `1`. +function _squareroot_blocksize(T, alg) + blocksize = _blocksize(alg) + blocksize > 0 && return blocksize + return eltype(T) <: BlasFloat ? 4 : 64 +end + +# A real result requires the real eigenvalues to be nonnegative, and those are exactly the `1x1` +# blocks: a `2x2` block has `T[i, i+1] * T[i+1, i] < 0`, hence a conjugate pair off the negative +# real axis, which always has a real square root. Hence no `axis_atol` here. +function _clamp_domain_quasitriu!(T::AbstractMatrix{<:Real}, atol::Real) + n = size(T, 1) + tmin = zero(eltype(T)) + @inbounds begin + i = 1 + while i <= n + p = _quasitriu_blocksize(T, i, n) + p == 1 && (tmin = min(tmin, T[i, i])) + i += p + end + tmin < -atol && throw_negative_eigenvalue(tmin, atol, "a negative real eigenvalue") + i = 1 + while i <= n + p = _quasitriu_blocksize(T, i, n) + p == 1 && T[i, i] < 0 && (T[i, i] = zero(eltype(T))) + i += p + end + end + return T +end + +# `R = [R11 R12; 0 R22]` turns `R^2 = T` into two smaller square roots and the Sylvester equation +# `R11 * R12 + R12 * R22 = T12`, whose corrections are matrix multiplications +function _squareroot_quasitriu!(R, T, blocksize::Int, ws = _quasitriu_workspace(R)) + n = size(T, 1) + s = n > blocksize ? _quasitriu_split(T, n) : 0 + iszero(s) && return _squareroot_quasitriu_point!(R, T, ws) + r1, r2 = 1:s, (s + 1):n + R11, R22 = view(R, r1, r1), view(R, r2, r2) + _squareroot_quasitriu!(R11, view(T, r1, r1), blocksize, ws) + _squareroot_quasitriu!(R22, view(T, r2, r2), blocksize, ws) + R12 = copy!(view(R, r1, r2), view(T, r1, r2)) + _quasitriu_sylvester!(R11, R22, R12, ws, blocksize) + return R +end + +# Björck-Hammarling: the diagonal blocks are scalar or `2x2` square roots, the off-diagonal blocks +# solve `R_ii * R_ij + R_ij * R_jj = T_ij - sum(R_ik * R_kj for i < k < j)`, by block column so that +# the sum only involves blocks that are already known +function _squareroot_quasitriu_point!(R, T, ws) + n = size(T, 1) + @inbounds begin + i = 1 + while i <= n + p = _quasitriu_blocksize(T, i, n) + if p == 1 + R[i, i] = _squareroot_diag(T[i, i]) + else + I = i:(i + 1) + _squareroot_2x2!(view(R, I, I), view(T, I, I)) + end + i += p + end + j = 1 + while j <= n + q = _quasitriu_blocksize(T, j, n) + J = j:(j + q - 1) + i = j - 1 + while i >= 1 + p = _quasitriu_blocksize_end(T, i) + I = (i - p + 1):i + Rij = copy!(view(R, I, J), view(T, I, J)) + K = (i + 1):(j - 1) + isempty(K) || mul!(Rij, view(R, I, K), view(R, K, J), -1, 1) + _quasitriu_sylvester_block!(view(R, I, I), view(R, J, J), Rij, ws) + i -= p + end + j += q + end + end + return R +end + +# `sqrt` is discontinuous across the negative real axis, where the sign of a computed zero imaginary +# part is noise; pinning the branch keeps two copies of a negative eigenvalue from receiving +# opposite roots, which would leave the off-diagonal equations singular. +_squareroot_diag(t::Real) = sqrt(t) +function _squareroot_diag(t::Complex) + imt = imag(t) + return sqrt(iszero(imt) ? complex(real(t), abs(imt)) : t) +end + +# Real square root of a `2x2` block with eigenvalues `θ ± im * μ`, following Higham (2008), Alg. 6.5 +# and eqs. (6.8)-(6.9); the standardized form produced by LAPACK `?gees` and `GenericSchur.gschur!` +# has equal diagonal entries. +Base.@propagate_inbounds function _squareroot_2x2!(R, T) + θ, b, c = T[1, 1], T[1, 2], T[2, 1] + μ = sqrt(abs(b)) * sqrt(abs(c)) + # the real part of `sqrt(θ + im * μ)`, in the form that avoids cancellation for `θ < 0` + t = sqrt((abs(θ) + hypot(θ, μ)) / 2) + α = θ >= zero(θ) ? t : μ / (2 * t) + R[1, 1] = α + R[2, 2] = α + R[1, 2] = b / (2 * α) + R[2, 1] = c / (2 * α) + return R +end diff --git a/src/interface/matrixfunctions.jl b/src/interface/matrixfunctions.jl index bd3435482..1ef8dad90 100644 --- a/src/interface/matrixfunctions.jl +++ b/src/interface/matrixfunctions.jl @@ -66,7 +66,7 @@ see [Domain considerations](@ref sec_matrixfunction_domain). This algorithm presumes a well-conditioned eigenbasis. For a defective or nearly defective matrix both its result and its domain verdict are unreliable, since the eigenvalues themselves are resolved only to `eps^(1/k)` for a Jordan block of size `k`. Prefer - [`MatrixFunctionViaLA`](@ref), which is Schur-based, for such matrices. + [`MatrixFunctionViaSchur`](@ref), which is Schur-based, for such matrices. """ @algdef MatrixFunctionViaEig @@ -74,3 +74,42 @@ see [Domain considerations](@ref sec_matrixfunction_domain). MatrixFunctionViaEig(eig_alg::AbstractAlgorithm; kwargs...), MatrixFunctionViaEig(; eig_alg, kwargs...) ) + +""" + MatrixFunctionViaSchur(; schur_alg, blocksize = 0, domain_atol = default_domain_atol(λ)) + +Algorithm type for computing a function of a matrix from its Schur decomposition, by evaluating the +function on the (quasi-)triangular Schur factor and transforming back. +It applies to [`squareroot`](@ref) only, where the triangular factor is obtained through the +recursion of Björck & Hammarling (1983), in the real quasi-triangular variant of Higham (1987) so +that a real input is treated in real arithmetic, and with the recursive blocking of Deadman, Higham +& Ralha (2013) to move the bulk of the work into matrix multiplications. +As this algorithm requires no LAPACK support beyond the Schur decomposition itself, it also applies +at arbitrary precision; it is however not suitable for GPU arrays, as the recursion indexes +individual entries. + +The optional `schur_alg` specifies which Schur decomposition implementation to use, and defaults to +the one selected for the input. +`blocksize` is the size at which the recursion switches over to the entrywise algorithm, where `1` +recurses as far as it can and any value at or above the matrix size skips the recursion entirely. +It defaults to `0`, which selects a small threshold for the scalar types that have a level-3 BLAS +to spend the work in, and a large one for those that do not. +`domain_atol` is the absolute tolerance within which negative eigenvalues are treated as rounding +artifacts and clamped onto zero, and defaults to [`default_domain_atol`](@ref). Raising it accepts +more matrices; see [Domain considerations](@ref sec_matrixfunction_domain). + +Unlike [`MatrixFunctionViaEig`](@ref), this algorithm is backward stable and does not rely on the +conditioning of the eigenbasis, so it is the appropriate choice for a defective or nearly defective +matrix. + +## References + +- Å. Björck and S. Hammarling, "A Schur method for the square root of a matrix", + Linear Algebra Appl., 52/53, 127–140, 1983. +- N. J. Higham, "Computing real square roots of a real matrix", Linear Algebra Appl., 88/89, + 405–430, 1987. +- E. Deadman, N. J. Higham and R. Ralha, "Blocked Schur Algorithms for Computing the Matrix Square + Root", Applied Parallel and Scientific Computing, Lecture Notes in Computer Science 7782, + 171–182, 2013. +""" +@algdef MatrixFunctionViaSchur diff --git a/src/interface/squareroot.jl b/src/interface/squareroot.jl index 97efb03c6..4bac37ee0 100644 --- a/src/interface/squareroot.jl +++ b/src/interface/squareroot.jl @@ -16,9 +16,9 @@ the principal square root is complex, leads to a `DomainError`; pass a complex m to obtain the principal value. For the algorithms that have access to the spectrum, real eigenvalues that are negative within a tolerance `domain_atol` are treated as rounding artifacts and clamped to zero, so that raising -`domain_atol` accepts more matrices. It defaults to [`default_domain_atol`](@ref), and is not -supported by [`MatrixFunctionViaLA`](@ref); -see [Domain considerations](@ref sec_matrixfunction_domain). +`domain_atol` accepts more matrices. It defaults to [`default_domain_atol`](@ref), and is +supported by every algorithm except [`MatrixFunctionViaLA`](@ref), which has no access to the +spectrum; see [Domain considerations](@ref sec_matrixfunction_domain). !!! note The bang method `squareroot!` optionally accepts the output structure and @@ -31,7 +31,7 @@ see [Domain considerations](@ref sec_matrixfunction_domain). # ------------------- default_squareroot_algorithm(A; kwargs...) = default_squareroot_algorithm(typeof(A); kwargs...) function default_squareroot_algorithm(T::Type; kwargs...) - return MatrixFunctionViaLA(; kwargs...) + return MatrixFunctionViaSchur(; kwargs...) end function default_squareroot_algorithm(::Type{T}; kwargs...) where {T <: Diagonal} return DiagonalAlgorithm(; kwargs...) diff --git a/test/matrixfunctions/squareroot.jl b/test/matrixfunctions/squareroot.jl index d4bfc3446..42e24fb2d 100644 --- a/test/matrixfunctions/squareroot.jl +++ b/test/matrixfunctions/squareroot.jl @@ -11,8 +11,7 @@ else BLASFloats = (Float32, Float64, ComplexF32, ComplexF64) GenericFloats = (BigFloat, Complex{BigFloat}) end -# only the `Diagonal` fast path applies to these, as they have no `eig`/`eigh` support -DiagonalOnlyFloats = (Float16, ComplexF16) +HalfFloats = (Float16, ComplexF16) @isdefined(TestSuite) || include("../testsuite/TestSuite.jl") using .TestSuite @@ -21,6 +20,11 @@ is_buildkite = get(ENV, "BUILDKITE", "false") == "true" m = 54 md = 4 # the domain tests prescribe the full spectrum, so keep them small +mh = 8 # half precision only holds up on a small matrix + +# What half precision can be asked of `squareroot` is set by `schur_full`, not by the kernel, whose +# own residual stays below `eps(Float16)` here +rtol16 = 0.1 # CPU tests # --------- @@ -32,14 +36,23 @@ if !is_buildkite MatrixFunctionViaEigh(QRIteration()), MatrixFunctionViaEigh(DivideAndConquer()), ) + SCHUR_ALGS = ( + MatrixFunctionViaSchur(), + MatrixFunctionViaSchur(; schur_alg = QRIteration(; expert = true)), + ) TestSuite.test_squareroot(T, (m, m)) TestSuite.test_squareroot_algs(T, (m, m), LAPACK_EIG_ALGS) + TestSuite.test_squareroot_algs(T, (m, m), SCHUR_ALGS) TestSuite.test_squareroot_hermitian(T, (m, m), LAPACK_EIG_ALGS; exact_hermiticity = false) + TestSuite.test_squareroot_hermitian(T, (m, m), SCHUR_ALGS; exact_hermiticity = false) TestSuite.test_squareroot_hermitian(T, (m, m), LAPACK_EIGH_ALGS) TestSuite.test_squareroot_reference(T, (m, m)) + TestSuite.test_squareroot_blocked(T, (m, m), (MatrixFunctionViaSchur(),)) + TestSuite.test_squareroot_defective(T, (6, 6), (MatrixFunctionViaSchur(),)) # `MatrixFunctionViaLA` has no access to the spectrum, and thus no `domain_atol` TestSuite.test_squareroot_domain(T, (md, md), (MatrixFunctionViaLA(),); test_domain_atol = false) TestSuite.test_squareroot_domain(T, (md, md), (MatrixFunctionViaEig(QRIteration()),)) + TestSuite.test_squareroot_domain(T, (md, md), (MatrixFunctionViaSchur(),)) TestSuite.test_squareroot_domain(T, (md, md), LAPACK_EIGH_ALGS; hermitian_output = true) end @@ -49,17 +62,37 @@ if !is_buildkite TestSuite.seed_rng!(123) GS_ALGS = (MatrixFunctionViaEig(QRIteration(; driver = GS())),) GLA_ALGS = (MatrixFunctionViaEigh(QRIteration(; driver = GLA())),) + GS_SCHUR_ALGS = (MatrixFunctionViaSchur(; schur_alg = QRIteration(; driver = GS())),) + # `LinearAlgebra.sqrt` promotes a real generic matrix with complex eigenvalues to a complex + # Schur form, so the Schur route is the only one here that is both real and backward stable + TestSuite.test_squareroot(T, (24, 24)) TestSuite.test_squareroot_algs(T, (24, 24), GS_ALGS) + TestSuite.test_squareroot_algs(T, (24, 24), GS_SCHUR_ALGS) TestSuite.test_squareroot_hermitian(T, (24, 24), GS_ALGS; exact_hermiticity = false) + TestSuite.test_squareroot_hermitian(T, (24, 24), GS_SCHUR_ALGS; exact_hermiticity = false) TestSuite.test_squareroot_hermitian(T, (24, 24), GLA_ALGS) + TestSuite.test_squareroot_blocked(T, (12, 12), GS_SCHUR_ALGS) TestSuite.test_squareroot_domain(T, (md, md), GS_ALGS) + TestSuite.test_squareroot_domain(T, (md, md), GS_SCHUR_ALGS) TestSuite.test_squareroot_domain(T, (md, md), GLA_ALGS; hermitian_output = true) end - for T in (BLASFloats..., GenericFloats..., DiagonalOnlyFloats...) + # `eigh` is unavailable in half precision, so the Schur route is what covers these + for T in HalfFloats + TestSuite.seed_rng!(123) + HALF_SCHUR_ALGS = (MatrixFunctionViaSchur(; schur_alg = QRIteration(; driver = GS())),) + TestSuite.test_squareroot(T, (mh, mh); rtol = rtol16) + TestSuite.test_squareroot_algs(T, (mh, mh), HALF_SCHUR_ALGS; rtol = rtol16) + TestSuite.test_squareroot_domain(T, (md, md), HALF_SCHUR_ALGS; atol = rtol16) + # the hermitian generator `A * A' + I` is out of reach: on it `GenericSchur` returns a + # half-precision decomposition wrong by `0.4` relative (`6e-7` in `Float32`), so nothing + # downstream of it can be asserted + end + + for T in (BLASFloats..., GenericFloats..., HalfFloats...) TestSuite.seed_rng!(123) AT = Diagonal{T, Vector{T}} - test_spectrum = !(T in DiagonalOnlyFloats) + test_spectrum = !(T in HalfFloats) TestSuite.test_squareroot(AT, m) TestSuite.test_squareroot_algs(AT, m, (DiagonalAlgorithm(),)) TestSuite.test_squareroot_hermitian(AT, m, (DiagonalAlgorithm(),); test_spectrum) diff --git a/test/testsuite/TestSuite.jl b/test/testsuite/TestSuite.jl index 32b12c57f..141777a82 100644 --- a/test/testsuite/TestSuite.jl +++ b/test/testsuite/TestSuite.jl @@ -134,12 +134,31 @@ function instantiate_hermitian_spectrum(T, sz, λ) return project_hermitian!(V * Diagonal(Ddiag) * V') end +# a matrix with a defective eigenvalue `λ`, i.e. a 2x2 Jordan block, and simple eigenvalues +# elsewhere. The similarity transformation is unitary, so the conditioning of the problem comes +# from the Jordan block alone rather than from the transformation. +function instantiate_defective_matrix(T, sz, λ) + n = sz isa Tuple ? first(sz) : sz + @assert n >= 2 + A = instantiate_matrix(T, (n, n)) + J = fill!(similar(A), zero(eltype(A))) + copyto!(diagview(J), convert(Vector{eltype(A)}, [λ, λ, collect(3:n)...])) + J[1, 2] = one(eltype(A)) + V = instantiate_unitary(T, A, n) + return V * J * V' +end + # rebuild `alg` with an explicit `domain_atol`, so that the domain tests need not spell out the # inner decomposition algorithm a second time function with_domain_atol(alg::MatrixAlgebraKit.Algorithm, atol) return MatrixAlgebraKit.Algorithm{MatrixAlgebraKit.name(alg)}(; alg.kwargs..., domain_atol = atol) end +# likewise for the blocking threshold of `MatrixFunctionViaSchur` +function with_blocksize(alg::MatrixAlgebraKit.Algorithm, blocksize) + return MatrixAlgebraKit.Algorithm{MatrixAlgebraKit.name(alg)}(; alg.kwargs..., blocksize = blocksize) +end + include("ad_utils.jl") include("projections.jl") diff --git a/test/testsuite/matrixfunctions/squareroot.jl b/test/testsuite/matrixfunctions/squareroot.jl index 32bea5c2d..a83b3fcf4 100644 --- a/test/testsuite/matrixfunctions/squareroot.jl +++ b/test/testsuite/matrixfunctions/squareroot.jl @@ -6,7 +6,10 @@ using MatrixAlgebraKit: ishermitian # the same bodies apply on GPU and to downstream array types. `test_squareroot_reference` is the # exception and is host-only by construction. -function test_squareroot(T::Type, sz; kwargs...) +# `rtol`, and `atol` in `test_squareroot_domain`, is the tolerance on the residual `sqrtA^2 ≈ A`. +# Raise it for a scalar type whose decomposition is resolved to less, as in half precision. + +function test_squareroot(T::Type, sz; rtol = precision(T), kwargs...) summary_str = testargs_summary(T, sz) return @testset "squareroot $summary_str" begin A = instantiate_offaxis_matrix(T, sz) @@ -14,7 +17,7 @@ function test_squareroot(T::Type, sz; kwargs...) sqrtA = @testinferred squareroot(A) @test eltype(sqrtA) == eltype(A) - @test sqrtA * sqrtA ≈ A + @test isapprox(sqrtA * sqrtA, A; rtol) @test A == Ac # the in-place method may not be able to reuse the provided output @@ -23,7 +26,7 @@ function test_squareroot(T::Type, sz; kwargs...) end end -function test_squareroot_algs(T::Type, sz, algs; kwargs...) +function test_squareroot_algs(T::Type, sz, algs; rtol = precision(T), kwargs...) summary_str = testargs_summary(T, sz) return @testset "squareroot algorithm $alg $summary_str" for alg in algs A = instantiate_offaxis_matrix(T, sz) @@ -31,7 +34,7 @@ function test_squareroot_algs(T::Type, sz, algs; kwargs...) sqrtA = @testinferred squareroot(A, alg) @test eltype(sqrtA) == eltype(A) - @test sqrtA * sqrtA ≈ A + @test isapprox(sqrtA * sqrtA, A; rtol) @test A == Ac end end @@ -41,7 +44,7 @@ end # The elementwise spectrum check needs exact hermiticity, since `eigh_vals` rejects anything else. function test_squareroot_hermitian( T::Type, sz, algs; - exact_hermiticity = true, test_spectrum = true, kwargs... + exact_hermiticity = true, test_spectrum = true, rtol = precision(T), kwargs... ) summary_str = testargs_summary(T, sz) return @testset "squareroot hermitian algorithm $alg $summary_str" for alg in algs @@ -50,14 +53,14 @@ function test_squareroot_hermitian( sqrtA = @testinferred squareroot(A, alg) @test eltype(sqrtA) == eltype(A) - @test sqrtA * sqrtA ≈ A + @test isapprox(sqrtA * sqrtA, A; rtol) @test A == Ac if exact_hermiticity @test ishermitian(sqrtA) - test_spectrum && @test eigh_vals(sqrtA) ≈ sqrt.(eigh_vals(A)) + test_spectrum && @test isapprox(eigh_vals(sqrtA), sqrt.(eigh_vals(A)); rtol) else - @test ishermitian(sqrtA; rtol = precision(T)) + @test ishermitian(sqrtA; rtol) end end end @@ -68,7 +71,8 @@ end # `test_domain_atol = false` for the ones without access to the spectrum, which have no tolerance. function test_squareroot_domain( T::Type, sz, algs; - hermitian_output = false, test_domain_atol = true, kwargs... + hermitian_output = false, test_domain_atol = true, + atol = sqrt(eps(real(eltype(T)))), kwargs... ) R = real(eltype(T)) n = sz isa Tuple ? first(sz) : sz @@ -92,7 +96,7 @@ function test_squareroot_domain( Aclamp = instantiate_hermitian_spectrum(T, sz, λclamp) sqrtAclamp = @testinferred squareroot(Aclamp, alg) @test eltype(sqrtAclamp) == eltype(Aclamp) - @test sqrtAclamp * sqrtAclamp ≈ Aclamp atol = sqrt(eps(R)) + @test isapprox(sqrtAclamp * sqrtAclamp, Aclamp; atol) # an eigenvalue beyond the default tolerance is out of domain, while an explicit # `domain_atol` admits it after all @@ -110,7 +114,7 @@ function test_squareroot_domain( sqrtAwide = @testinferred squareroot(Awide, wide_alg) @test eltype(sqrtAwide) == eltype(Awide) # accepting is backward stable, but only to the size of the eigenvalue that was discarded - @test sqrtAwide * sqrtAwide ≈ Awide atol = sqrt(sqrt(eps(R))) + @test isapprox(sqrtAwide * sqrtAwide, Awide; atol = sqrt(atol)) end end end @@ -129,3 +133,36 @@ function test_squareroot_reference(T::Type, sz; test_hermitian = true, kwargs... end end end + +# A general matrix carries 2x2 blocks in its real Schur form while a hermitian one carries none, and +# the two structures drive the recursion differently. Every block size shares one decomposition of +# the same input and differs only in the kernel, hence the comparison against each other. +function test_squareroot_blocked(T::Type, sz, algs; rtol = precision(T), kwargs...) + summary_str = testargs_summary(T, sz) + return @testset "squareroot blocking $alg $summary_str" for alg in algs + for A in (instantiate_offaxis_matrix(T, sz), instantiate_posdef_matrix(T, sz)) + sqrtA = squareroot(A, with_blocksize(alg, 1)) + @test isapprox(sqrtA * sqrtA, A; rtol) + for blocksize in (2, 3, 8) + @test squareroot(A, with_blocksize(alg, blocksize)) ≈ sqrtA + end + end + end +end + +# A Schur-based algorithm never inverts an eigenvector matrix, so it stays backward stable where +# `MatrixFunctionViaEig` resolves the eigenvalues to no better than `sqrt(eps)`; hence the residual +# is held to roundoff rather than to the default `≈`, which the eigenvector route would still meet. +# A defective eigenvalue *on* the negative real axis is left out: whether the perturbed pair +# surfaces as two real eigenvalues or as a conjugate pair is up to the eigensolver. +function test_squareroot_defective(T::Type, sz, algs; kwargs...) + R = real(eltype(T)) + n = sz isa Tuple ? first(sz) : sz + summary_str = testargs_summary(T, sz) + return @testset "squareroot defective algorithm $alg $summary_str" for alg in algs + A = instantiate_defective_matrix(T, sz, one(R) + one(R)) + sqrtA = @testinferred squareroot(A, alg) + @test eltype(sqrtA) == eltype(A) + @test norm(sqrtA * sqrtA - A) <= 100 * n * eps(R) * norm(A) + end +end