🚀

Juliaでの対称行列の生成を高速化する

2025/01/06に公開

量子物理学・量子化学では行列固有値問題によく遭遇します. 固有値問題を解く部分については情報・数学系の研究者が速いソルバをたくさん開発してくれていますが, そもそもの解くための行列を作るところは我々が自らプログラムを書かなくてはいけません. TwoBody.jlの開発にあたって対称行列の生成を高速化した際の知見を備忘録としてまとめておきます.

対称行列を計算する例

実関数を基底に使ったRayleigh-Ritz法の計算では実対称行列を生成することになります. 下記の教科書より4つのガウス型軌道を用いた水素原子のSchrödinger方程式に対する計算例を紹介します.

入力, 水素原子の計算例
using LinearAlgebra
α = [13.00773, 1.962079, 0.444529, 0.1219492]
nₘₐₓ = length(α)
S = [(π/(α[i]+α[j]))^(3/2) for i=1:nₘₐₓ, j=1:nₘₐₓ]
T = [3*π^(3/2)*α[i]*α[j]/(α[i]+α[j])^(5/2) for i=1:nₘₐₓ, j=1:nₘₐₓ]
V = [-2*π/(α[i]+α[j]) for i=1:nₘₐₓ, j=1:nₘₐₓ]
H = T + V
E, C = eigen(Symmetric(H),Symmetric(S))

出力, 固有値のみ抜き出し
values:
4-element Vector{Float64}:
 -0.49927840566748605
  0.11321392045798521
  2.5922995719598148
 21.1443651901225

水素原子の基底状態 $E_\mathrm{n=1} = -0.5~E_\mathrm{h}$ に近い値が計算できています. ここで計算した対称行列 H と S を高速に生成する方法を考えます.

準備

下記の2つのパッケージを読み込み, 100×100サイズの行列を計算するためのパラメータを用意します.

パッケージの読み込み, 行列サイズの指定
using LinearAlgebra
using BenchmarkTools
α = [0.1*i for i in 1: 100] # [13.00773, 1.962079, 0.444529, 0.1219492]

※ この指数では明らかに悪条件問題なので実用上はそもそも解けませんが, 計算コストの計測のために使用します.

遅いけどシンプルな例

内包表記を使って計算した例です. 遅いですがシンプルに書ける点に魅力があります.

遅い例
function S_slow(α)
    nₘₐₓ = length(α)
    S = [(π/(α[i]+α[j]))^(3/2) for i=1:nₘₐₓ, j=1:nₘₐₓ]
    return LinearAlgebra.Symmetric(S)
end

@btime S_slow(α)
# 167.800 μs (4 allocations: 78.23 KiB)

シングルスレッドの例

対称行列なので右上の半分だけ計算します. 行列サイズが大きいとメモリ確保の時間などが無視され, 計算コストはだいたい半分になります.

シングルスレッドの例
function S_fast(α)
    nₘₐₓ = length(α)
    S = Array{Float64}(undef, nₘₐₓ, nₘₐₓ)
    for j in 1:nₘₐₓ
        for i in 1:j
            @inbounds S[i,j] = (π/(α[i]+α[j]))^(3/2)
        end
    end
    return LinearAlgebra.Symmetric(S)
end

@btime S_fast(α)
# 84.200 μs (4 allocations: 78.23 KiB)

マルチスレッドの例

こちらの通りJuliaを julia -t 4 や julia -t auto で起動すればスレッド並列が使えます. IJulia.jlのJupyter Notebook/Labではこちらの通り using IJulia; installkernel("Julia (Multi-threads)", "--threads=auto") でマルチスレッドのカーネルが追加でき, using IJulia; jupyterlab(detached=true) で起動できます.　起動後はマルチスレッドになっていることを確認します.

スレッド数の確認
Threads.nthreads()
# 8

外側の for に Threads.@threads を付けると速くなります.

マルチスレッドの例
function S_parallel(α)
    nₘₐₓ = length(α)
    S = Array{Float64}(undef, nₘₐₓ, nₘₐₓ)
    Threads.@threads for j in 1:nₘₐₓ
        for i in 1:j
            @inbounds S[i,j] = (π/(α[i]+α[j]))^(3/2)
        end
    end
    return LinearAlgebra.Symmetric(S)
end

@btime S_parallel(α)
# 29.900 μs (46 allocations: 83.23 KiB)

まとめ

対称行列であることを利用すると計算時間が概ね半分（84.2 / 167.8 = 0.5017）になりました. 端数はメモリの確保や対称行列への変換, キャッシュミスなどに起因する時間だと思われます. マルチスレッドの場合, スレッド数が8なので8倍速くなってほしいですが, そうもいきません. 少なくとも倍以上は速くなる（29.9 / 84.2 = 0.355）ということは確かめられました. 各スレッドに割り振られる計算を均等になるようにするなど工夫すればもっと速くなるかもしれません. ひとまずThreads.@threadsを追記するだけでも簡単に倍以上速くなるという点は強調しておきます.

計算時間まとめ
800 μs (4 allocations: 78.23 KiB)
200 μs (4 allocations: 78.23 KiB)
900 μs (46 allocations: 83.23 KiB)

ノートブック

Discussion

ログインするとコメントできます