refactor: reuse existing code patterns in matrix-function kernels

Move `_mul_herm!` from `implementations/polar.jl` to
`implementations/projections.jl` and use it for the symmetric-product
kernels of `squareroot` and fractional `power` via
`MatrixFunctionViaEigh`, so hermitian output is guaranteed through the
same mechanism as the polar decomposition (BLAS `herk`/`syrk`, explicit
projection otherwise, GPU overrides).

Rewrite the `MatrixFunctionViaLA` kernels in the style of the
`exponential` implementation: an explicit eltype branch with a
`real.(...)` assignment instead of the bespoke `_copy_result!` helper.
`squareroot`/`logarithm` use a strict complex-result check (LinearAlgebra
returns real results whenever the principal value is real), while
`power` tolerates rounding-level imaginary components since fractional
powers are computed in complex arithmetic upstream.

Author: lkdvos

src/implementations/logarithm.jl

@@ -35,8 +35,13 @@ initialize_output(::typeof(logarithm!), A::AbstractMatrix, ::AbstractAlgorithm)
 function logarithm!(A::AbstractMatrix, logA, alg::MatrixFunctionViaLA)
     check_input(logarithm!, A, logA, alg)
     isempty(alg.kwargs) || throw(ArgumentError("`MatrixFunctionViaLA` does not accept keyword arguments for `logarithm`"))
-    result = LinearAlgebra.log(A)
-    _copy_result!(logarithm!, logA, result)
+    # `LinearAlgebra.log` of a real matrix is real whenever the principal logarithm is,
+    # so a complex result with a real output signals a genuine domain violation
+    logAc = LinearAlgebra.log(A)
+    if eltype(logAc) <: Complex && !(eltype(logA) <: Complex)
+        throw(_realness_domainerror(logarithm!))
+    end
+    copy!(logA, logAc)
     return logA
 end
 

src/implementations/matrixfunctions.jl

@@ -61,31 +61,12 @@ function _check_nonzero_eigenvalues(λ, atol::Real)
 end
 
 # For `MatrixFunctionViaLA`, domain violations surface as a complex result from
-# `LinearAlgebra`; detect this after the fact to preserve type-stable output.
-# `LinearAlgebra` sometimes computes a real result in complex arithmetic (e.g. fractional
-# powers via a complex Schur decomposition), so rounding-level imaginary components do
-# not signal a domain violation and are simply discarded.
-function _copy_result!(f, out::AbstractMatrix, result::AbstractMatrix)
-    if eltype(result) <: Complex && !(eltype(out) <: Complex)
-        # base the tolerance on the working precision, which is the lower of the two:
-        # `LinearAlgebra` may internally promote, e.g. `Float32` fractional powers are
-        # computed with `Float64`-eltype results but `Float32`-level accuracy
-        atol = max(defaulttol(out), defaulttol(result)) * norm(result, Inf)
-        for x in result
-            if abs(imag(x)) > atol
-                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 `MatrixFunctionViaEigh`/`MatrixFunctionViaEig` with a suitable " *
-                            "`domain_atol` if the offending eigenvalues are rounding artifacts."
-                    )
-                )
-            end
-        end
-        out .= real.(result)
-    else
-        copy!(out, result)
-    end
-    return out
+# `LinearAlgebra` while the output should remain real.
+function _realness_domainerror(f)
+    return 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 `MatrixFunctionViaEigh`/`MatrixFunctionViaEig` with a suitable " *
+            "`domain_atol` if the offending eigenvalues are rounding artifacts."
+    )
 end

src/implementations/polar.jl

@@ -81,19 +81,6 @@ function right_polar!(A::AbstractMatrix, PWᴴ, alg::PolarViaSVD)
     return (P, Wᴴ)
 end
 
-# Implement `mul!(C, A', A)` and guarantee the result is hermitian.
-# For BLAS calls that dispatch to `syrk` or `herk` this works automatically
-# for GPU this currently does not seem to be guaranteed so we manually project
-function _mul_herm!(C, A)
-    mul!(C, A, A')
-    project_hermitian!(C)
-    return C
-end
-function _mul_herm!(C::YALAPACK.BlasMat{T}, A::YALAPACK.BlasMat{T}) where {T <: YALAPACK.BlasFloat}
-    mul!(C, A, A')
-    return C
-end
-
 # Implementation via Newton
 # --------------------------
 function left_polar!(A::AbstractMatrix, WP, alg::PolarNewton)

src/implementations/power.jl

@@ -35,8 +35,19 @@ initialize_output(::typeof(power!), A::AbstractMatrix, p::Real, ::AbstractAlgori
 function power!(A::AbstractMatrix, p::Real, powA, alg::MatrixFunctionViaLA)
     check_input(power!, A, p, powA, alg)
     isempty(alg.kwargs) || throw(ArgumentError("`MatrixFunctionViaLA` does not accept keyword arguments for `power`"))
-    result = A^p
-    _copy_result!(power!, powA, result)
+    powAc = A^p
+    if eltype(powAc) <: Complex && !(eltype(powA) <: Complex)
+        # `LinearAlgebra` computes fractional powers of real matrices in complex
+        # arithmetic and only casts back to real when the result is exactly real,
+        # so rounding-level imaginary components do not signal a domain violation.
+        # The tolerance is based on the working precision, which may be lower than
+        # the result eltype suggests (e.g. `Float32` input promotes to `ComplexF64`).
+        atol = defaulttol(powA) * norm(powAc, Inf)
+        all(x -> abs(imag(x)) <= atol, powAc) || throw(_realness_domainerror(power!))
+        powA .= real.(powAc)
+        return powA
+    end
+    copy!(powA, powAc)
     return powA
 end
 
@@ -49,16 +60,16 @@ function power!(A::AbstractMatrix, p::Real, powA, alg::MatrixFunctionViaEigh)
         λ .= λ .^ p
         VD = V * D
         mul!(powA, VD, V')
+        return project_hermitian!(powA)
     else
         atol = something(alg.domain_atol, default_domain_atol(λ))
         p < 0 && _check_nonzero_eigenvalues(λ, atol)
         _clamp_domain_eigenvalues!(λ, atol)
         # `A^p = (V * D^(p/2)) * (V * D^(p/2))'` is hermitian by construction
         λ .= λ .^ (p / 2)
         Vs = rmul!(V, D)
-        mul!(powA, Vs, Vs')
+        return _mul_herm!(powA, Vs)
     end
-    return project_hermitian!(powA)
 end
 
 function power!(A::AbstractMatrix, p::Real, powA, alg::MatrixFunctionViaEig)

src/implementations/projections.jl

@@ -140,3 +140,16 @@ end
 
 _imimag(x::Real) = zero(x)
 _imimag(x::Complex) = im * imag(x)
+
+# Implement `mul!(C, A, A')` and guarantee the result is hermitian.
+# For BLAS calls that dispatch to `syrk` or `herk` this works automatically
+# for GPU this currently does not seem to be guaranteed so we manually project
+function _mul_herm!(C, A)
+    mul!(C, A, A')
+    project_hermitian!(C)
+    return C
+end
+function _mul_herm!(C::YALAPACK.BlasMat{T}, A::YALAPACK.BlasMat{T}) where {T <: YALAPACK.BlasFloat}
+    mul!(C, A, A')
+    return C
+end

src/implementations/squareroot.jl

@@ -35,8 +35,13 @@ initialize_output(::typeof(squareroot!), A::AbstractMatrix, ::AbstractAlgorithm)
 function squareroot!(A::AbstractMatrix, sqrtA, alg::MatrixFunctionViaLA)
     check_input(squareroot!, A, sqrtA, alg)
     isempty(alg.kwargs) || throw(ArgumentError("`MatrixFunctionViaLA` does not accept keyword arguments for `squareroot`"))
-    result = LinearAlgebra.sqrt(A)
-    _copy_result!(squareroot!, sqrtA, result)
+    # `LinearAlgebra.sqrt` of a real matrix is real whenever the principal square root is,
+    # so a complex result with a real output signals a genuine domain violation
+    sqrtAc = LinearAlgebra.sqrt(A)
+    if eltype(sqrtAc) <: Complex && !(eltype(sqrtA) <: Complex)
+        throw(_realness_domainerror(squareroot!))
+    end
+    copy!(sqrtA, sqrtAc)
     return sqrtA
 end
 
@@ -49,8 +54,7 @@ function squareroot!(A::AbstractMatrix, sqrtA, alg::MatrixFunctionViaEigh)
     # `sqrt(A) = (V * D^(1/4)) * (V * D^(1/4))'` is hermitian by construction
     λ .= sqrt.(sqrt.(λ))
     Vs = rmul!(V, D)
-    mul!(sqrtA, Vs, Vs')
-    return project_hermitian!(sqrtA)
+    return _mul_herm!(sqrtA, Vs)
 end
 
 function squareroot!(A::AbstractMatrix, sqrtA, alg::MatrixFunctionViaEig)