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/docs/src/changelog.md b/docs/src/changelog.md index cc2946b68..9745974b2 100644 --- a/docs/src/changelog.md +++ b/docs/src/changelog.md @@ -22,8 +22,26 @@ 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 + `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)). - `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..142f34e43 100644 --- a/docs/src/user_interface/algorithms.md +++ b/docs/src/user_interface/algorithms.md @@ -98,16 +98,20 @@ 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` | - -For full docstring details on each algorithm type, see the corresponding section in [Decompositions](@ref). +| [`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) | + +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 6153f2e6a..68f4025fe 100644 --- a/docs/src/user_interface/matrix_functions.md +++ b/docs/src/user_interface/matrix_functions.md @@ -17,23 +17,76 @@ 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. -## 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, along with more specialized implementations for specific functions: + +- [`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 -exponential MatrixAlgebraKit.MatrixFunctionViaTaylor +MatrixAlgebraKit.MatrixFunctionViaSchur MatrixAlgebraKit.MatrixFunctionViaLA MatrixAlgebraKit.MatrixFunctionViaEig MatrixAlgebraKit.MatrixFunctionViaEigh ``` + +## 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](@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. + +```@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 [`MatrixFunctionViaSchur`](@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. + +[`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. + +!!! 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 [`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/src/MatrixAlgebraKit.jl b/src/MatrixAlgebraKit.jl index 4d4e0084e..ba9ac383c 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 @@ -41,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 @@ -93,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") @@ -115,6 +118,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 +132,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..4383646fc 100644 --- a/src/common/defaults.jl +++ b/src/common/defaults.jl @@ -43,6 +43,23 @@ 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(λ) + +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. + +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). +""" +default_domain_atol(λ) = + length(λ) * eps(real(float(one(eltype(λ))))) * maximum(abs, λ; init = abs(zero(eltype(λ)))) + const DEFAULT_FIXGAUGE = Ref(true) 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/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 new file mode 100644 index 000000000..a95abf22f --- /dev/null +++ b/src/implementations/matrixfunctions.jl @@ -0,0 +1,161 @@ +# 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. + +# 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) + R = real(float(eltype(λ))) + atol = _domain_atol(alg) + return convert(R, isnothing(atol) ? default_domain_atol(λ) : atol) +end + +_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 +@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 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` accepts no keyword arguments for `$f`, got $ks. In particular " * + "`domain_atol` is not supported, as `LinearAlgebra` does not expose the spectrum; " * + "use `MatrixFunctionViaSchur`, `MatrixFunctionViaEig` or `MatrixFunctionViaEigh` instead." + ) + ) +end + +# `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) || throw_la_kwargs(f, ks) + return nothing +end + +@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). " * + "Pass a complex matrix to obtain the principal value, or use `MatrixFunctionViaEig`/`MatrixFunctionViaEigh` to have " * + "the spectrum itself checked against `domain_atol`." + ) + ) +end + +@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. " * + 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 new file mode 100644 index 000000000..92b0fe2b1 --- /dev/null +++ b/src/implementations/squareroot.jl @@ -0,0 +1,189 @@ +# 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) + _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) + all(isfinite, sqrtAc) || throw_nonfinite_result(squareroot!, _NONFINITE_LA_ADVICE) + 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, select_algorithm(eigh_full!, A, _eigh_alg(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, 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(λ)) + 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 + +# 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 b0ea57da8..1ef8dad90 100644 --- a/src/interface/matrixfunctions.jl +++ b/src/interface/matrixfunctions.jl @@ -1,10 +1,13 @@ # ================================ -# EXPONENTIAL ALGORITHMS +# MATRIX FUNCTION ALGORITHMS # ================================ """ MatrixFunctionViaLA() -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`. +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 @@ -29,31 +32,84 @@ As this algorithm requires no LAPACK support, it also applies at arbitrary preci @algdef MatrixFunctionViaTaylor """ - MatrixFunctionViaEigh(eigh_alg) + 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, 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 -end -function Base.show(io::IO, alg::MatrixFunctionViaEigh) - print(io, "MatrixFunctionViaEigh(") - _show_alg(io, alg.eigh_alg) - return print(io, ")") -end +@algdef MatrixFunctionViaEigh + +@deprecate( + MatrixFunctionViaEigh(eigh_alg::AbstractAlgorithm; kwargs...), + MatrixFunctionViaEigh(; eigh_alg, kwargs...) +) """ - MatrixFunctionViaEig(eig_alg) + 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. -""" -struct MatrixFunctionViaEig{A <: AbstractAlgorithm} <: AbstractAlgorithm - eig_alg::A -end -function Base.show(io::IO, alg::MatrixFunctionViaEig) - print(io, "MatrixFunctionViaEig(") - _show_alg(io, alg.eig_alg) - return print(io, ")") -end +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, and defaults to +[`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 + [`MatrixFunctionViaSchur`](@ref), which is Schur-based, for such matrices. +""" +@algdef MatrixFunctionViaEig + +@deprecate( + 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 new file mode 100644 index 000000000..4bac37ee0 --- /dev/null +++ b/src/interface/squareroot.jl @@ -0,0 +1,42 @@ +# 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. +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 +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 + 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 MatrixFunctionViaSchur(; 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 diff --git a/test/README.md b/test/README.md index 18eb02a1d..400f1fc88 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`, `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 | -| `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"]'`). @@ -125,6 +133,10 @@ Supporting infrastructure in the module: - Predicates used throughout the assertions: `isleftnull`, `isrightnull`, `isleftcomplete`, `isrightcomplete`, `has_positive_diagonal`. - `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/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/matrixfunctions/squareroot.jl b/test/matrixfunctions/squareroot.jl new file mode 100644 index 000000000..42e24fb2d --- /dev/null +++ b/test/matrixfunctions/squareroot.jl @@ -0,0 +1,144 @@ +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 +HalfFloats = (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 +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 +# --------- +if !is_buildkite + for T in BLASFloats + TestSuite.seed_rng!(123) + LAPACK_EIG_ALGS = (MatrixFunctionViaLA(), MatrixFunctionViaEig(QRIteration())) + LAPACK_EIGH_ALGS = ( + 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 + + # `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())),) + 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 + + # `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 HalfFloats) + 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 36ae68304..141777a82 100644 --- a/test/testsuite/TestSuite.jl +++ b/test/testsuite/TestSuite.jl @@ -102,6 +102,63 @@ 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 + +# 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") @@ -117,6 +174,11 @@ include("decompositions/eigh.jl") include("decompositions/orthnull.jl") include("decompositions/svd.jl") +# Matrix functions +# ---------------- +include("matrixfunctions/exponential.jl") +include("matrixfunctions/squareroot.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 diff --git a/test/testsuite/matrixfunctions/squareroot.jl b/test/testsuite/matrixfunctions/squareroot.jl new file mode 100644 index 000000000..a83b3fcf4 --- /dev/null +++ b/test/testsuite/matrixfunctions/squareroot.jl @@ -0,0 +1,168 @@ +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. + +# `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) + Ac = deepcopy(A) + + sqrtA = @testinferred squareroot(A) + @test eltype(sqrtA) == eltype(A) + @test isapprox(sqrtA * sqrtA, A; rtol) + @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; 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) + Ac = deepcopy(A) + + sqrtA = @testinferred squareroot(A, alg) + @test eltype(sqrtA) == eltype(A) + @test isapprox(sqrtA * sqrtA, A; rtol) + @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, rtol = precision(T), 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 isapprox(sqrtA * sqrtA, A; rtol) + @test A == Ac + + if exact_hermiticity + @test ishermitian(sqrtA) + test_spectrum && @test isapprox(eigh_vals(sqrtA), sqrt.(eigh_vals(A)); rtol) + else + @test ishermitian(sqrtA; rtol) + 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, 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, + atol = sqrt(eps(real(eltype(T)))), 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 isapprox(sqrtAclamp * sqrtAclamp, Aclamp; atol) + + # 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 isapprox(sqrtAwide * sqrtAwide, Awide; atol = sqrt(atol)) + 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_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 + +# 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