refactor: share input handling and reconstruction between matrix functions

`squareroot`, `logarithm` and `power` all map a square matrix to a matrix
of the same size and scalar type, so their `copy_input`, `check_input` and
`initialize_output` contracts were identical modulo the function name --
six methods repeated three times over. Fold the bodies into
`_matrixfunction_copy_input` and `_matrixfunction_check_input`, whose own
dispatch absorbs the `Diagonal` and `DiagonalAlgorithm` specializations, so
each function keeps just the one line that names it.

The three `MatrixFunctionViaEig` kernels shared an identical skeleton,
differing only in which scalar function is applied to the eigenvalues:
solve `V f(D) V⁻¹` and, for real input, discard the imaginary part left by
the complex decomposition. That becomes `_apply_eig!`, taking the already
transformed eigenvalues. Similarly `_apply_eigh!` captures the
`V f(D) V'`-then-project shape shared by `logarithm` and integer `power`,
and pairs with the existing `_mul_herm!` for the cases where a symmetric
product makes the result hermitian by construction instead.

Deliberately left alone: the `MatrixFunctionViaLA` kernels, whose bodies
look alike but must not be merged -- `_copy_result!` did exactly that and
was removed in 2698e26, because `squareroot`/`logarithm` need a strict
complex-result check while `power` has to tolerate rounding-level imaginary
components. `exponential` also stays out, as its real/complex branch keys
on the scalar `τ` as well as the matrix eltype.

Behaviour is unchanged: the matrix-function tests report the same 4701
passing as before the refactor.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: lkdvos

src/implementations/logarithm.jl

@@ -1,24 +1,9 @@
 # Inputs
 # ------
-function copy_input(::typeof(logarithm), A::AbstractMatrix)
-    return copy!(similar(A, float(eltype(A))), A)
-end
-copy_input(::typeof(logarithm), A::Diagonal) = map_diagonal(float, A)
+copy_input(::typeof(logarithm), A::AbstractMatrix) = _matrixfunction_copy_input(A)
 
 function check_input(::typeof(logarithm!), A::AbstractMatrix, logA, alg::AbstractAlgorithm)
-    m = LinearAlgebra.checksquare(A)
-    @check_size(logA, (m, m))
-    @check_scalar(logA, A)
-    return nothing
-end
-
-function check_input(::typeof(logarithm!), A::AbstractMatrix, logA, ::DiagonalAlgorithm)
-    m = LinearAlgebra.checksquare(A)
-    @assert isdiag(A)
-    @assert logA isa Diagonal
-    @check_size(logA, (m, m))
-    @check_scalar(logA, A)
-    return nothing
+    return _matrixfunction_check_input(A, logA, alg)
 end
 
 # Algorithm selection
@@ -48,24 +33,17 @@ end
 function logarithm!(A::AbstractMatrix, logA, alg::MatrixFunctionViaEigh)
     check_input(logarithm!, A, logA, alg)
     D, V = eigh_full!(A, alg.eigh_alg)
-    VD = V * logarithm!(D, D, DiagonalAlgorithm(; domain_atol = alg.domain_atol))
-    mul!(logA, VD, V')
-    return project_hermitian!(logA)
+    diag_alg = DiagonalAlgorithm(; domain_atol = alg.domain_atol)
+    return _apply_eigh!(logA, V, logarithm!(D, D, diag_alg))
 end
 
 function logarithm!(A::AbstractMatrix, logA, alg::MatrixFunctionViaEig)
     check_input(logarithm!, A, logA, alg)
     D, V = eig_full!(A, alg.eig_alg)
+    # a real result requires the spectrum to stay off the negative real axis
+    eltype(A) <: Real && _clamp_domain_eigenvalues!(D, alg.domain_atol)
     diag_alg = DiagonalAlgorithm(; domain_atol = alg.domain_atol)
-    if eltype(A) <: Real
-        _clamp_domain_eigenvalues!(D, alg.domain_atol)
-        VlogD = V * logarithm!(D, D, diag_alg)
-        logAc = rdiv!(VlogD, LinearAlgebra.lu!(V))
-        return logA .= real.(logAc)
-    else
-        logA .= V .* transpose(diagview(logarithm!(D, D, diag_alg)))
-        return rdiv!(logA, LinearAlgebra.lu!(V))
-    end
+    return _apply_eig!(logA, V, logarithm!(D, D, diag_alg))
 end
 
 # Diagonal logic

src/implementations/matrixfunctions.jl

@@ -1,3 +1,52 @@
+# Shared input and output handling
+# --------------------------------
+# `squareroot`, `logarithm` and `power` all map a square matrix to a matrix of the same size and
+# scalar type, so they agree on `copy_input`, `check_input` and `initialize_output`; only their
+# kernels differ. Each implementation file forwards to the helpers below, which keeps the
+# per-function definitions to the one line that names the function.
+
+_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
+
+# Shared reconstruction from an eigenvalue decomposition
+# ------------------------------------------------------
+# Both take the already-transformed eigenvalues `fD = f(D)` and rebuild `f(A)`.
+
+# `f(A) = V f(D) V⁻¹` for a general `A`. 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, so it is
+# projected afterwards; where `f` admits a square root, prefer building `f(A)` as a symmetric
+# product via `_mul_herm!`, which is hermitian by construction.
+_apply_eigh!(fA, V, fD) = project_hermitian!(mul!(fA, V * fD, V'))
+
 # Shared helpers for matrix functions with a restricted domain
 # -------------------------------------------------------------
 

src/implementations/power.jl

@@ -1,24 +1,9 @@
 # Inputs
 # ------
-function copy_input(::typeof(power), A::AbstractMatrix, p::Real)
-    return copy!(similar(A, float(eltype(A))), A), p
-end
-copy_input(::typeof(power), A::Diagonal, p::Real) = map_diagonal(float, A), p
+copy_input(::typeof(power), A::AbstractMatrix, p::Real) = _matrixfunction_copy_input(A), p
 
 function check_input(::typeof(power!), A::AbstractMatrix, p::Real, powA, alg::AbstractAlgorithm)
-    m = LinearAlgebra.checksquare(A)
-    @check_size(powA, (m, m))
-    @check_scalar(powA, A)
-    return nothing
-end
-
-function check_input(::typeof(power!), A::AbstractMatrix, p::Real, powA, ::DiagonalAlgorithm)
-    m = LinearAlgebra.checksquare(A)
-    @assert isdiag(A)
-    @assert powA isa Diagonal
-    @check_size(powA, (m, m))
-    @check_scalar(powA, A)
-    return nothing
+    return _matrixfunction_check_input(A, powA, alg)
 end
 
 # Algorithm selection
@@ -60,32 +45,20 @@ function power!(A::AbstractMatrix, p::Real, powA, alg::MatrixFunctionViaEigh)
     isone(p) && ((powA === A || copy!(powA, A)); return powA)
     D, V = eigh_full!(A, alg.eigh_alg)
     diag_alg = DiagonalAlgorithm(; domain_atol = alg.domain_atol)
-    if isinteger(p)
-        VD = V * power!(D, p, D, diag_alg)
-        mul!(powA, VD, V')
-        return project_hermitian!(powA)
-    else
-        # `A^p = (V * D^(p/2)) * (V * D^(p/2))'` is hermitian by construction
-        Vs = rmul!(V, power!(D, p / 2, D, diag_alg))
-        return _mul_herm!(powA, Vs)
-    end
+    isinteger(p) && return _apply_eigh!(powA, V, power!(D, p, D, diag_alg))
+    # `A^p = (V * D^(p/2)) * (V * D^(p/2))'` is hermitian by construction
+    return _mul_herm!(powA, rmul!(V, power!(D, p / 2, D, diag_alg)))
 end
 
 function power!(A::AbstractMatrix, p::Real, powA, alg::MatrixFunctionViaEig)
     check_input(power!, A, p, powA, alg)
     iszero(p) && return one!(powA)
     isone(p) && ((powA === A || copy!(powA, A)); return powA)
     D, V = eig_full!(A, alg.eig_alg)
+    # only a fractional power of a real matrix needs the spectrum off the negative real axis
+    eltype(A) <: Real && !isinteger(p) && _clamp_domain_eigenvalues!(D, alg.domain_atol)
     diag_alg = DiagonalAlgorithm(; domain_atol = alg.domain_atol)
-    if eltype(A) <: Real
-        isinteger(p) || _clamp_domain_eigenvalues!(D, alg.domain_atol)
-        VpD = V * power!(D, p, D, diag_alg)
-        powAc = rdiv!(VpD, LinearAlgebra.lu!(V))
-        return powA .= real.(powAc)
-    else
-        powA .= V .* transpose(diagview(power!(D, p, D, diag_alg)))
-        return rdiv!(powA, LinearAlgebra.lu!(V))
-    end
+    return _apply_eig!(powA, V, power!(D, p, D, diag_alg))
 end
 
 # Diagonal logic

src/implementations/squareroot.jl

@@ -1,24 +1,9 @@
 # Inputs
 # ------
-function copy_input(::typeof(squareroot), A::AbstractMatrix)
-    return copy!(similar(A, float(eltype(A))), A)
-end
-copy_input(::typeof(squareroot), A::Diagonal) = map_diagonal(float, A)
+copy_input(::typeof(squareroot), A::AbstractMatrix) = _matrixfunction_copy_input(A)
 
 function check_input(::typeof(squareroot!), A::AbstractMatrix, sqrtA, alg::AbstractAlgorithm)
-    m = LinearAlgebra.checksquare(A)
-    @check_size(sqrtA, (m, m))
-    @check_scalar(sqrtA, A)
-    return nothing
-end
-
-function check_input(::typeof(squareroot!), A::AbstractMatrix, sqrtA, ::DiagonalAlgorithm)
-    m = LinearAlgebra.checksquare(A)
-    @assert isdiag(A)
-    @assert sqrtA isa Diagonal
-    @check_size(sqrtA, (m, m))
-    @check_scalar(sqrtA, A)
-    return nothing
+    return _matrixfunction_check_input(A, sqrtA, alg)
 end
 
 # Algorithm selection
@@ -56,16 +41,10 @@ end
 function squareroot!(A::AbstractMatrix, sqrtA, alg::MatrixFunctionViaEig)
     check_input(squareroot!, A, sqrtA, alg)
     D, V = eig_full!(A, alg.eig_alg)
+    # a real result requires the spectrum to stay off the negative real axis
+    eltype(A) <: Real && _clamp_domain_eigenvalues!(D, alg.domain_atol)
     diag_alg = DiagonalAlgorithm(; domain_atol = alg.domain_atol)
-    if eltype(A) <: Real
-        _clamp_domain_eigenvalues!(D, alg.domain_atol)
-        VsqrtD = V * squareroot!(D, D, diag_alg)
-        sqrtAc = rdiv!(VsqrtD, LinearAlgebra.lu!(V))
-        return sqrtA .= real.(sqrtAc)
-    else
-        sqrtA .= V .* transpose(diagview(squareroot!(D, D, diag_alg)))
-        return rdiv!(sqrtA, LinearAlgebra.lu!(V))
-    end
+    return _apply_eig!(sqrtA, V, squareroot!(D, D, diag_alg))
 end
 
 # Diagonal logic