special case power!(A, 0/1)

Author: lkdvos

src/implementations/power.jl

@@ -30,11 +30,14 @@ power!(A::AbstractMatrix, p::Real, out, alg::DefaultAlgorithm) = power!(A, p, ou
 # -------
 initialize_output(::typeof(power!), A::AbstractMatrix, p::Real, ::AbstractAlgorithm) = A
 
+
 # Implementation
 # --------------
 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`"))
+    iszero(p) && return one!(powA)
+    isone(p) && ((powA === A || copy!(powA, A)); return powA)
     powAc = A^p
     if eltype(powAc) <: Complex && !(eltype(powA) <: Complex)
         # `LinearAlgebra` computes fractional powers of real matrices in complex
@@ -53,6 +56,8 @@ end
 
 function power!(A::AbstractMatrix, p::Real, powA, alg::MatrixFunctionViaEigh)
     check_input(power!, A, p, powA, alg)
+    iszero(p) && return one!(powA)
+    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)
@@ -68,6 +73,8 @@ 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)
     diag_alg = DiagonalAlgorithm(; domain_atol = alg.domain_atol)
     if eltype(A) <: Real
@@ -85,6 +92,8 @@ end
 # --------------
 function power!(A::AbstractMatrix, p::Real, powA, alg::DiagonalAlgorithm)
     check_input(power!, A, p, powA, alg)
+    iszero(p) && return one!(powA)
+    isone(p) && ((powA === A || copy!(powA, A)); return powA)
     λ = diagview(powA)
     copyto!(λ, diagview(A))
     if isinteger(p)

src/interface/power.jl

@@ -11,6 +11,9 @@ Compute the matrix power `powA = A^p` of the square matrix `A`.
 For integer `p` this is defined for any square matrix (invertible if `p < 0`);
 for fractional `p` the principal power `exp(p * log(A))` is computed.
 
+The exponents `p = 0` and `p = 1` are resolved directly as `I` and `A`, without computing any
+decomposition, so they apply to every square matrix regardless of the algorithm selected.
+
 The scalar type of the output matches that of the input.
 As a consequence, for fractional `p`, a real matrix with eigenvalues on the negative
 real axis, for which the principal power is complex, leads to a `DomainError`; pass a

test/matrixfunctions/power.jl

@@ -35,6 +35,8 @@ if !is_buildkite
         )
         TestSuite.test_power(T, (m, m))
         TestSuite.test_power_algs(T, (m, m), LAPACK_EIG_ALGS)
+        TestSuite.test_power_trivial(T, (m, m), LAPACK_EIG_ALGS)
+        TestSuite.test_power_trivial(T, (m, m), LAPACK_EIGH_ALGS; test_rejected_input = true)
         TestSuite.test_power_hermitian(T, (m, m), LAPACK_EIG_ALGS; exact_hermiticity = false)
         TestSuite.test_power_hermitian(T, (m, m), LAPACK_EIGH_ALGS)
         TestSuite.test_power_reference(T, (m, m))
@@ -50,6 +52,8 @@ if !is_buildkite
         GS_ALGS = (MatrixFunctionViaEig(QRIteration(; driver = GS())),)
         GLA_ALGS = (MatrixFunctionViaEigh(QRIteration(; driver = GLA())),)
         TestSuite.test_power_algs(T, (24, 24), GS_ALGS)
+        TestSuite.test_power_trivial(T, (24, 24), GS_ALGS)
+        TestSuite.test_power_trivial(T, (24, 24), GLA_ALGS; test_rejected_input = true)
         TestSuite.test_power_hermitian(T, (24, 24), GS_ALGS; exact_hermiticity = false)
         TestSuite.test_power_hermitian(T, (24, 24), GLA_ALGS)
         TestSuite.test_power_domain(T, (md, md), GS_ALGS)
@@ -63,6 +67,7 @@ if !is_buildkite
         test_spectrum = !(T in DiagonalOnlyFloats)
         TestSuite.test_power(AT, m)
         TestSuite.test_power_algs(AT, m, (DiagonalAlgorithm(),))
+        TestSuite.test_power_trivial(AT, m, (DiagonalAlgorithm(),))
         TestSuite.test_power_hermitian(AT, m, (DiagonalAlgorithm(),); test_spectrum)
         TestSuite.test_power_reference(AT, m; test_hermitian = !(T in GenericFloats))
         TestSuite.test_power_domain(AT, md, (DiagonalAlgorithm(),); test_singular = true)
@@ -82,11 +87,13 @@ if CUDA.functional()
             MatrixFunctionViaEigh(DivideAndConquer()),
         )
         TestSuite.test_power_hermitian(CuMatrix{T}, (m, m), CUDA_EIGH_ALGS)
+        TestSuite.test_power_trivial(CuMatrix{T}, (m, m), CUDA_EIGH_ALGS; test_rejected_input = true)
         TestSuite.test_power_domain(CuMatrix{T}, (md, md), CUDA_EIGH_ALGS; hermitian_output = true)
 
         AT = Diagonal{T, CuVector{T}}
         TestSuite.test_power(AT, m)
         TestSuite.test_power_algs(AT, m, (DiagonalAlgorithm(),))
+        TestSuite.test_power_trivial(AT, m, (DiagonalAlgorithm(),))
         TestSuite.test_power_hermitian(AT, m, (DiagonalAlgorithm(),))
         TestSuite.test_power_domain(AT, md, (DiagonalAlgorithm(),); test_singular = true)
     end
@@ -102,11 +109,13 @@ if AMDGPU.functional()
             MatrixFunctionViaEigh(DivideAndConquer()),
         )
         TestSuite.test_power_hermitian(ROCMatrix{T}, (m, m), ROC_EIGH_ALGS)
+        TestSuite.test_power_trivial(ROCMatrix{T}, (m, m), ROC_EIGH_ALGS; test_rejected_input = true)
         TestSuite.test_power_domain(ROCMatrix{T}, (md, md), ROC_EIGH_ALGS; hermitian_output = true)
 
         AT = Diagonal{T, ROCVector{T}}
         TestSuite.test_power(AT, m)
         TestSuite.test_power_algs(AT, m, (DiagonalAlgorithm(),))
+        TestSuite.test_power_trivial(AT, m, (DiagonalAlgorithm(),))
         TestSuite.test_power_hermitian(AT, m, (DiagonalAlgorithm(),))
         TestSuite.test_power_domain(AT, md, (DiagonalAlgorithm(),); test_singular = true)
     end

test/testsuite/matrixfunctions/power.jl

@@ -1,6 +1,6 @@
 using TestExtras
 using LinearAlgebra: LinearAlgebra, I, SingularException
-using MatrixAlgebraKit: ishermitian
+using MatrixAlgebraKit: ishermitian, one!
 
 # `power` takes the exponent as a second positional argument, and splits into an integer branch
 # (repeated multiplication, defined for any square matrix) and a fractional branch (the principal
@@ -71,6 +71,42 @@ function test_power_algs(T::Type, sz, algs; ps = (2, -1), qs = (1 // 2, -1 // 4)
     end
 end
 
+# `A^0 = I` and `A^1 = A` hold for every square matrix, and both are short-circuited before any
+# decomposition is computed, so the results are exact rather than merely approximate.
+#
+# `test_rejected_input = true` for algorithms that cannot handle a general matrix at all (the
+# `eigh`-based ones): feeding them input they would reject proves the decomposition really is
+# skipped, rather than being computed and happening to give the right answer.
+function test_power_trivial(T::Type, sz, algs; test_rejected_input = false, kwargs...)
+    R = real(eltype(T))
+    summary_str = testargs_summary(T, sz)
+    return @testset "power trivial exponents algorithm $alg $summary_str" for alg in algs
+        A = instantiate_offaxis_matrix(T, sz)
+        Id = one!(deepcopy(A))
+
+        @testset "integer p = $p" for p in (0, 1)
+            powA = @testinferred power(A, p, alg)
+            @test powA == (iszero(p) ? Id : A)
+            # in-place, where the output aliases the input
+            @test power!(deepcopy(A), p, alg) == (iszero(p) ? Id : A)
+        end
+        @testset "float p = $p" for p in (zero(R), one(R))
+            powA = @testinferred power(A, p, alg)
+            @test powA == (iszero(p) ? Id : A)
+        end
+
+        if test_rejected_input
+            B = instantiate_smallnorm_matrix(T, sz)
+            IdB = one!(deepcopy(B))
+            # confirm the premise: this algorithm cannot decompose `B`
+            @test_throws DomainError power(B, 2, alg)
+            # yet the trivial exponents never reach the decomposition
+            @test power(B, 0, alg) == IdB
+            @test power(B, 1, alg) == B
+        end
+    end
+end
+
 # See the corresponding comment in `squareroot.jl` for `exact_hermiticity`.
 function test_power_hermitian(
         T::Type, sz, algs;