JuliaのKrylov部分空間法ライブラリ、KrylovKit.jlの紹介

Qiitaに書いた記事をこちらにも載せておきます。内容は全く同じです。

Krylov部分空間法を使って線形方程式や固有値問題等を解きたいときに便利なパッケージ、 KrylovKit.jlを紹介します。

数値線形代数とKrylov部分空間法

ベクトル$\vec{b}$と行列$A$に対して、$\vec{b},A\vec{b},A^2\vec{b},\dots, A^{d-1}\vec{b}$によって張られる線形部分空間のことをKrylov部分空間と呼びます。

このKrylov部分空間の性質を利用した数値線形代数のアルゴリズムが数多く存在し、まとめてKrylov部分空間法と呼ばれます。有名な共役勾配法やLanczos法などもKrylov部分空間法の仲間です。

KrylovKit.jlでは、Krylov部分空間法を用いて線形方程式$A\vec{x}=\vec{b}$、固有値問題$A\vec{x} = \lambda \vec{x}$、一般固有値問題$A\vec{x}=\lambda B\vec{x}$、特異値分解$A=U\Sigma V^*$、行列の指数関数との積の計算$\exp(tA)\vec{x}$、等を行うことができます。

実際に利用可能なアルゴリズムの一覧は ドキュメントを参照してください。

使い方

ドキュメントがちゃんとしているので特に説明の必要は無い気がしますが

インストール

(v1.2) pkg> add KrylovKit

固有値問題

A = randn(1000,1000)
# Arnoldi法でAの固有値と固有ベクトルを絶対値が大きい方から10個求める
vals, vecs, info = eigsolve(A, 10, :LM)
# vecsは2次元配列ではなくVectorのVector
# infoにはiterationの回数等の情報が入っている

# 対称化
A += A'
# Lanczos法で最小固有値を計算
(val,), (vec,), info = eigsolve(A, 1, :SR; issymmetric=true)

線形方程式

b = randn(1000)
# CG法で A * x = b を解く
x, info = linsolve(A, b)

指数関数

# exp(0.1 * A)bを計算
x, info = exponentiate(A, 0.1, b) 

キーワード引数によって、各種パラメータや使用するアルゴリズムの指定が可能です。

KrylovKit.jlの特徴

Juliaには元々反復法パッケージが沢山存在します(例えば IterativeSolvers.jl)。 以下では、KrylovKit.jlがこれらの類似パッケージと比較してどのような点で優れているかについて、僕個人の意見を挙げていきます。 (実は ドキュメント内でも類似パッケージとの比較は行われていて、ほぼ同じ内容で向こうのほうが全然詳しいのですが…)

Matrix-freeな計算のサポート

Krylov部分空間法においては行列ベクトル積の計算さえできれば良いので、行列の成分を直接持つ必要はありません。 そのため、行列ベクトル積を関数として直接実装してしまうことによって、メモリ使用量を大幅に削減することができます。このようなやり方はMatrix-freeな計算と呼ばれています。

既存の反復法パッケージにおいてMatrix-freeな計算をしたいときには、 LinearMaps.jlのようなパッケージを組み合わせる必要がありました。LinearMaps.jlを使うと、関数からLinearMap型のオブジェクトを作成することができます。LinearMapAbstractMatrixと同じように振る舞うので、これをIterativeSolvers.jl等の引数として与えることでMatrix-freeな計算を行うことができます。

しかし、これはどうもまどろっこしいですね。Juliaでは関数を引数として与えることができるのですから、最初からソルバーに関数を渡してしまいたくなります。

KrylovKit.jlでは、最初から行列だけでなく関数を引数に取ることができます。

例えば固有値問題$A\vec{x}=\lambda \vec{x}$をMatrix-freeに解きたいときには、まず$f(\vec{x}) = A\vec{x}$を満たす関数fを実装し

values, vectors, info = eigsolve(f, inivec; kwargs...) # inivecは反復の初期ベクトル

のようにすればOKです。

fがパラメータを持つ場合は無名関数x -> f(x, args...)を作ることになりますが、juliaのdo構文を使えば

values, vectors, info = eigsolve(inivec; kwargs...) do x
    f(x, args...)
end

のようにすっきりと記述することができます。LinearMaps.jlなどを使ったやり方では、このようなパラメータをきれいに扱うのが面倒だった記憶があります。

ユーザー定義型や多次元配列の取り扱い

上の例を見てもらえれば分かる通り、Matrix-freeに計算する場合、fの引数の型はAbstractVectorである必要はありません。AbstractVectorのように振る舞うオブジェクトであれば十分です。

ここで言うAbstractVectorのように振る舞うとは、2つのオブジェクトx, yに対して、a * x + b * ya, bはスカラー)やdot(x, y)等が計算できる、ということを指します。 ドキュメントに従っていくつかの関数をオーバーロードすることで、好きな型のオブジェクトに対して直接Krylov部分空間法を使うことができます。 Juliaの特徴の一つである多重ディスパッチのおかげで、オーバーロードは直感的、簡潔に行うことができます。

例えば、僕が研究で使っているテンソルネットワークと呼ばれる計算手法においては、「テンソルからテンソルへの線形写像について固有値問題や線形方程式を解く」というケースが頻繁に登場します。このようなときに、テンソルを一旦ベクトルに直したりすることなくそのまま扱えるのは非常に楽です。 (実はこのライブラリの作者であるJutho Haegemanがテンソルネットワークを用いて研究を行っている理論物理学者なので、これは当たり前といえば当たり前なのかもしれません)

まとめ

JuliaのKrylov部分空間法パッケージ、KrylovKit.jlについて紹介しました。

まだ若いライブラリなので、BiCGStabが無かったりJacobi-Davidsonが無かったりと、アルゴリズムの種類について不足している部分は否めませんが、現時点でも非常に使いやすいパッケージです。

皆様ぜひお試しください。

Avatar
Hayate Nakano
Ph.D. student

My research interests include tensor-network methods, quantum many-body physics and non-equilibrium dynamics.