グリーン久保公式(GK公式)のO(NlogN)計算

前回の平均自乗変位のアルゴリズム改善の記事に少し書いた、グリーン久保公式の高速化について簡単に紹介します。

グリーン久保公式(Green-Kubo relation:GK公式)

GK公式は平衡状態でのゆらぎと非平衡状態での動的性質の関係を表したものであり、分子動力学法では、平衡MDでのカレントのゆらぎから、輸送係数を導出するために用いられます。以下の式で求められ、輸送係数$\gamma$を時間相関関数の積分で厳密に数学的に表現しています。

$$
\gamma = \int^\infty_0{\dot{A}(t)\dot{A}(0)dt}
$$

エルゴート性があり、$\dot{A}$が長さ$N$の配列$A$で表すことができる場合、配列$A$の$t$番目の要素を$A_t$として、以下のように書き換えられます。

\begin{align}
\gamma &= \int^\infty_0{ <\dot{A}(t)\dot{A}(0)>dt} \\
&= \frac{1}{N-\tau}\sum_{\tau=0}^{N-1}\sum_{t=1}^{N-\tau}{A_t A_{t+\tau}} \\
&= \frac{1}{N-\tau}ACF(A) \\
\end{align}

$ACF(A)$というのは直線状畳み込みによる自己相関関数で、上の二重の$\sum$以降の部分の別名です。
GK公式の計算を行うときは分子動力学法ではこの式を一般に用います。
このGK公式は愚直に計算した場合、計算量は$O(N^2)$となります。$O(N^2)$とは、ランダヴの記号と言われているもので、ざっくりというと、データ数を$N$としたときに計算量(計算回数)からが$N$に比例するということを指します。
$N$個の実数$A_1,...,A_N$が離散的・データ数が有限・周期性がない関数として与えられた時のGK公式を具体的に考えてみます。

O(N^2)アルゴリズム

ボトルネックになっている部分は$A$の自己相関関数$ACF(A)$の部分です。

$$
ACF(A) = \sum_{\tau=0}^{N-1}\sum_{t=1}^{N-\tau}{A_t A_{t+\tau}}
$$

この式は愚直に実装すると$O(N^2)$です。Fortranで愚直な実装を書くと、

acf.f90

    function acf(a,n)
        use,intrinsic :: iso_fortran_env
        implicit none
        integer(int32),intent(in):: n
        real(real64),intent(in):: a(n)
        integer(int32):: tau, t
        real(real64):: acf(0:n-1)

        do tau=0,n-1
            do t=1,n-tau
                acf(tau) = a(t)*a(t+tau)
            end do
        end do
    end function

非常に単純な2重ループで計算ができることがわかります。

O(NlogN)アルゴリズム

自己相関関数$ACF$を計算する概略図です。

少し解説をすると、まず2べきのフーリエ変換を行うために配列長より長い最小の2べきの長さの配列に移し替えます(２べき以外でもFFTが行える場合は省略可能です)。ループ処理で$2^i$を探索で$O(\log{N})$、ビット演算により最上位ビットを求める場合は$O(1)$です。
次に配列長をさらに2倍します。これは、「直線状畳み込みをしたいが、FFT-IFFTによる畳み込みが円状畳み込みのため、配列長を2倍にして円状畳み込みを行い、その前半分を取り出すことで直線状畳み込みになる」ということを利用しています。

円状畳み込みによる自己相関関数$
ACF(A)
$は、以下のように変換できます。

\begin{align}
ACF(A) &= A_t \star A_t \\
&= \mathcal{F}^{-1}(\mathcal{F}(A_t \star A_t)) \\
&= \mathcal{F}^{-1}(\mathcal{F}(A_t * \overline{A_{-t}})) \\
&= \mathcal{F}^{-1}(\mathcal{F}(A_t) \circ \mathcal{F}(\overline{A_{-t}})) \\
&= \mathcal{F}^{-1}(\mathcal{F}(A_t) \circ \overline{\mathcal{F}(A_{-t}))} 
\end{align}

ここで記号$\star$は円状相関関数(correlation)、$*$は円状畳み込み(convolution)、$\circ$は配列の要素同士の積（アダマール積）を表しています。そして、$\overline{A}$は$A$の複素共役を示します。
上記の式変形により、円状畳み込みによる$ACF$は$O(N\log{N})$で求めることができます。
先にも言った通り、直線状畳み込みは、元の配列の2倍以上の長さのゼロ埋めした配列に移して円状畳み込みを行うことで計算ができるため、これにて直線状畳み込みも$O(N\log{N})$で計算することができました。よかったですね。与えられた関数$A$に周期性があると考える場合、円状畳込みが行え、ウィーナー・ヒンチンの公式(上記のものとほとんど同じ)を用いることでTCFを求めることができます。

自己相関関数の積分

自己相関関数を高速で計算できれば、あとはそれを任意の積分法で積分するだけです。積分は愚直に行っても$O(N)$で行えるため、特に高速化する必要はありません。ここでは省略します。

どのくらい早くなるか

ランダヴの記号を用いて計算量改善を示しても、わかりにくい方もいると思うので、どのくらい早くなるのか具体的な数字を用いて軽く解説します。
例えば、3次元、カレントのデータが1次元あたり$500000$個とした場合、愚直に計算すると、$(3\times 500000) ^ 2 > 10^{12}$となり、数十分程かかるだろうと予測できます。一方FFTを利用すると、$3\times 500000 \times \log_2{(3\times 500000)} < 10^8$となるため、数秒で終わります。理論上$\frac{3\times 500000}{\log_2{(3\times 500000)}}\approx 7000$ 倍高速化出来ます。
実際どれくらい高速化できるかは理論上の値よりは劣りますが、例えば私の研究室ではGK公式を用いた解析を1日がかりで行っていましたが、FFTを用いた実装に変えたところ、相関時間を今までの100倍ほど長く取ってもなお15分程度で終了するほどに高速化ができました。
この手法によって、既存の方法では緩和時間がとても長くて計算不可能であったような系に対しても解析が行えるようになります。