Juliaで「対角化された行列型」を作った

2022/1/20追記
こちらに非常に有益な情報&コードの改善案がありますので併せてご覧ください。
・黒木玄さんのTwitterでの指摘
・コード改善案
追記終

対角化とは。何が嬉しいのか。

対角化とは、行列$M$を、正則行列$P$、対角行列$D$を用いて

M = PDP^{-1}

と分解することです。これ以降、$M$は対角化可能な行列であるとします。
対角化の何が嬉しいかというと、そのご利益の一つは、$M^n$が、

\begin{align}
    M^n &= (PDP^{-1})(PDP^{-1}) \cdots (PDP^{-1})\\
        &= PD(P^{-1}P)D(P^{-1} \cdots P)DP^{-1}\\
        &= PDD\cdots P^{-1}\\
        &= PD^nP^{-1}
\end{align}

となり、少ない計算量で$M^n$が求められます。（$D^n$は単に$D$の対角要素を$n$乗すれば良い。）その系として、

\exp(M) = P\exp(D)P^{-1}

が従います。（$\exp(D)$は$D$の対角要素の指数関数。）
対角化は結構重い処理ですが、対角化された行列のスカラー倍は対角化し直す必要はありません（$M = PDP^{-1}$を見れば明らかだと思います）。なので行列$M$は固定で、
・いろいろなスカラー$s$に対し$\exp(Ms)$を求める
・いろいろなベクトル$V$に対して$\exp(M)V$を求める
時には対角化は有効です。

Juliaで対角化

using LinearAlgebra
M = rand(100,100)#対角化したい行列
E, P = eigen(M)

さすがJulia、超簡単だ。Eは上節で言うところの$D$の対角成分を要素に持つベクトルで、Pは固有ベクトルを並べてできた行列で、上節での$P$に対応します。
当然ですが、

exp(eigen(M))

は動きません。そこで、既存の関数にそのまま突っ込める「対角化された行列型」を作りました。

対角化された行列型

2022/1/20追記
以下のコードは型安定性の観点から、望ましいものではありません。
詳しくは、
・黒木玄さんのTwitterでの指摘
・公式ドキュメント > Performance Tips > Avoid fields with abstract type
をご覧ください。
追記終

module EigenDecompression

export EigenDecompose, eigDecomp
using LinearAlgebra
import Base.*, Base./

#対角化された行列型
struct EigenDecompose{T<:Number} <: AbstractMatrix{T}
    P::AbstractMatrix{T}
    D::Diagonal{T}
    invP::AbstractMatrix{T}
end

#普通のMatrixを対角化する
function eigDecomp(mat::AbstractMatrix)
    E, P = eigen(mat)
    EigenDecompose(P, Diagonal(E), inv(P))
end

#EigenDecompose型に対する関数
Base.exp(eig::EigenDecompose) = EigenDecompose(eig.P, exp(eig.D), eig.invP)
*(eig::EigenDecompose, vec::AbstractVector) = eig.P * eig.D * eig.invP * vec
*(eig::EigenDecompose, sc::Number) = EigenDecompose(eig.P, eig.D*sc, eig.invP)
/(eig::EigenDecompose, sc::Number) = EigenDecompose(eig.P, eig.D/sc, eig.invP)

#普通のMatrixに戻す
Base.Array(eig::EigenDecompose) = eig.P * eig.D * eig.invP

end

関数は私が使うものしか実装していないので、そのままREPLに突っ込むとエラーが出るかもしれません。お好みでお好きな関数を追加してください。

テスト

2022/1/17追記：テストコードを修正

using Main.EigenDecompression

M = rand(100, 100)
eM = eigDecomp(M)
for i in 1:100
    v = rand(100)
    rnd = rand()
    @assert exp(M*rnd)*v ≈ exp(eM*rnd)*v
end

OK.

ベンチマーク

2022/1/17追記：ベンチマークを修正

using BenchmarkTools

M = rand(100, 100);

#普通な方
function bench1(M)
    for i in 1:100
        v = rand(100)
        exp(M*rand())*v
    end
end

#今回実装した方
function bench2(M)
    eM = eigDecomp(M)
    for i in 1:100
        v = rand(100)
        exp(eM*rand())*v
    end
end

結果

はやい。
（でもREPLよりJupyterの方が10倍くらい速いのはなぜ・・・？でもふだんJupyterしか使わないから見なかったことにしよう）