evalpoly for matrix polynomials

Author: stevengj

src/generic.jl

@@ -2091,3 +2091,83 @@ end
 function copytrito!(B::StridedMatrixStride1{T}, A::StridedMatrixStride1{T}, uplo::AbstractChar) where {T<:BlasFloat}
     LAPACK.lacpy!(B, A, uplo)
 end
+
+# non-inplace fallback for evalpoly(X, p)
+function _evalpoly(X::AbstractMatrix, p)
+    Base.require_one_based_indexing(p)
+    p0 = isempty(p) ? Base.reduce_empty_iter(+, p) : p[end]
+    Xone = one(X)
+    S = Base.promote_op(*, typeof(Xone), typeof(Xone))(Xone) * p0
+    for i = length(p)-1:-1:1
+        S = X * S + @inbounds(p[i] isa AbstractMatrix ? p[i] : p[i] * I)
+    end
+    return S
+end
+
+_scalarval(x::Number) = x
+_scalarval(x::UniformScaling) = x.λ
+
+"""
+    evalpoly!(Y::AbstractMatrix, X::AbstractMatrix, p)
+
+Evaluate the matrix polynomial ``Y = \\sum_k X^{k-1} p[k]``, storing the result
+in-place in `Y`, for the coefficients `p[k]` (a vector or tuple).  The coefficients
+can be scalars, matrices, or [`UniformScaling`](@ref).
+
+Similar to `evalpoly`, but may be more efficient by working more in-place.  (Some
+allocations may still be required, however.)
+"""
+function evalpoly!(Y::AbstractMatrix, X::AbstractMatrix, p::Union{AbstractVector,Tuple})
+    @boundscheck axes(Y,1) == axes(Y,2) == axes(X,1) == axes(X,2)
+    Base.require_one_based_indexing(p)
+
+    N = length(p)
+    pN = iszero(N) ? Base.reduce_empty_iter(+, p) : p[N]
+    if pN isa AbstractMatrix
+        Y .= pN
+    elseif N > 1 && p[N-1] isa Union{Number,UniformScaling}
+        # initialize Y to p[N-1] I + X p[N], in-place
+        Y .= X .* _scalarval(pN)
+        for i in axes(Y,1)
+            @inbounds Y[i,i] += p[N-1] * I
+        end
+        N -= 1
+    else
+        # initialize Y to one(Y) * pN in-place
+        for i in axes(Y,1)
+            for j in axes(Y,2)
+                @inbounds Y[i,j] = zero(Y[i,j])
+            end
+            @inbounds Y[i,i] += one(Y[i,i]) * pN
+        end
+    end
+    if N > 1
+        Z = similar(Y) # workspace for mul!
+        for i = N-1:-1:1
+            mul!(Z, X, Y)
+            if p[i] isa AbstractMatrix
+                Y .= p[i] .+ Z
+            else
+                # Y = p[i] * I + Z, in-place
+                Y .= Z
+                for j in axes(Y,1)
+                    @inbounds Y[j,j] += p[i] * I
+                end
+            end
+        end
+    end
+    return Y
+end
+
+Base.evalpoly(X::AbstractMatrix, p::Tuple) = _evalpoly(X, p)
+Base.evalpoly(X::AbstractMatrix, ::Tuple{}) = zero(one(X)) # dimensionless zero, i.e. 0 * x^0
+Base.evalpoly(X::AbstractMatrix, p::AbstractVector) = _evalpoly(X, p)
+
+Base.evalpoly(X::StridedMatrix{<:Number}, p::Tuple{Union{Number, UniformScaling}, Vararg{Union{Number, UniformScaling}}}) =
+    evalpoly!(similar(X, Base.promote_op(*, typeof(one(eltype(X))), typeof(_scalarval(p[begin])))), X, p)
+Base.evalpoly(X::StridedMatrix{<:Number}, p::Tuple{AbstractMatrix{<:Number}, Vararg{AbstractMatrix{<:Number}}}) =
+    evalpoly!(similar(X, Base.promote_op(*, typeof(one(eltype(X))), eltype(p[begin]))), X, p)
+Base.evalpoly(X::StridedMatrix{<:Number}, p::AbstractVector{<:Union{Number, UniformScaling}}) =
+    isempty(p) ? _evalpoly(X, p) : evalpoly!(similar(X, Base.promote_op(*, typeof(one(eltype(X))), typeof(_scalarval(p[begin])))), X, p)
+Base.evalpoly(X::StridedMatrix{<:Number}, p::AbstractVector{<:AbstractMatrix{<:Number}}) =
+    isempty(p) ? _evalpoly(X, p) : evalpoly!(similar(X, Base.promote_op(*, typeof(one(eltype(X))), eltype(p[begin]))), X, p)

test/generic.jl

@@ -837,4 +837,28 @@ end
     end
 end
 
+using LinearAlgebra: _evalpoly
+naive_evalpoly(X, p) = sum(X^(i-1) * p[i] for i=1:length(p))
+
+@testset "evalpoly" begin
+    for X in ([1 2 3;4 5 6;7 8 9], UpperTriangular([1 2 3;0 5 6;0 0 9]), SymTridiagonal([1,2,3],[4,5]))
+        @test @inferred(evalpoly(X, ())) == zero(X) == evalpoly(X, Int[])
+        @test @inferred(evalpoly(X, (17,))) == one(X) * 17
+        @test @inferred(_evalpoly(X, [1,2,3,4])) == @inferred(evalpoly(X, [1,2,3,4])) ==
+              @inferred(evalpoly(X, (1,2,3,4))) ==
+              naive_evalpoly(X, [1,2,3,4]) == 1*one(X) + 2*X + 3X^2 + 4X^3
+        @test typeof(evalpoly(X, [1,2,3])) == typeof(evalpoly(X, (1,2,3))) == typeof(_evalpoly(X, [1,2,3])) ==
+              typeof(X * X)
+
+        for N in (1,2,4), p in (rand(-10:10, N), UniformScaling.(rand(-10:10, N)), [rand(-5:5,3,3) for _ = 1:N])
+            @test _evalpoly(X, p) == evalpoly(X, p) == evalpoly(X, Tuple(p)) == naive_evalpoly(X, p)
+        end
+        for N in (1,2,4), p in (rand(N), UniformScaling.(rand(N)), [rand(3,3) for _ = 1:N])
+            @test _evalpoly(X, p) ≈ evalpoly(X, p) ≈ evalpoly(X, Tuple(p)) ≈ naive_evalpoly(X, p)
+        end
+
+        @test_throws MethodError evalpoly(X, [])
+    end
+end
+
 end # module TestGeneric