[feat] squareroot matrix function

.buildkite/pipeline.yml

@@ -26,6 +26,7 @@ steps:
           - "rocm"
         group:
           - "decompositions"
+          - "matrixfunctions"
           - "mooncake"
       adjustments:
         - with:

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`.

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)
 

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
+```

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
 

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)