平均二乗変位(mean square displacement:MSD)

　MSDはランダムウォークや分子動力学法などでよく出現するもので､粒子の動きの尺度として用いられます｡また､自己拡散係数の計算に用いられることもあります｡

　x方向のmsdは以下の式で表されます｡

\begin{align}
\mathrm{MSD}(\tau) &\equiv \langle |x(\tau+t) - x(t)|^2 \rangle \\
&= \sum_{\tau=0}^{n-1} \frac{1}{n-\tau} \sum_{t=1}^{n-\tau}\bigl(x(\tau+t)-x(t)\bigr)^2
\end{align}

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

　MSDは愚直に計算した場合､計算量は$O(N^2)$となります｡$O(N^2)$とは､データ数を$N$としたときに計算量(計算回数)が$N^2$に比例するという事を指します｡

　$x(1)$~$x(5)$のデータを例に考えてみます｡定義どおりのアルゴリズムは以下のような計算になるでしょう｡

$$\mathrm{MSD}(0) = 0$$

$$\mathrm{MSD}(1) = \frac{\bigl(x(2)-x(1)\bigr)^2 +
\bigl(x(3)-x(2)\bigr)^2 + \bigl(x(4)-x(3)\bigr)^2 + \bigl(x(5)-x(4)\bigr)^2}{4}$$

$$\mathrm{MSD}(2) = \frac{\bigl(x(3)-x(1)\bigr)^2 +
\bigl(x(4)-x(2)\bigr)^2 + \bigl(x(5)-x(3)\bigr)^2}{3}$$

$$\mathrm{MSD}(3) = \frac{\bigl(x(4)-x(1)\bigr)^2 + \bigl(x(5)-x(2)\bigr)^2}{2}$$

$$\mathrm{MSD}(4) = \bigl(x(5)-x(1)\bigr)^2$$

O(NlogN)アルゴリズムのための前提知識

　$O(N\log{N})$のアルゴリズムを使うためには二つの知識が必要です｡

　「累積和」と「FFTの線形畳み込みを用いた自己相関関数(autocorrelation function:ACF)」です｡
説明は長くなるので詳しく知りたい方はググってください｡ACFについてはグリーン久保関数のO(N log N)計算でより詳しく解説しています。

　累積和とは予め配列の和を計算しておくことで計算量を減少させる有名なテクニックです｡発想も単純で実装も容易であるため､知らず知らずのうちに見たり､使ったりしたことがある人も多いと思います｡

　自己相関関数は文字通り自己(データ自身)の相関(主に時間相関)を表すもので､N個のデータ点をN次元ベクトルとみなしたときの､自分自身とずれた自分との内積です｡

　強い言語には自己相関関数や相関関数､線形畳み込みの関数が用意されているので､それを用いると良いでしょう｡

#　O(NlogN)アルゴリズム

　ではやっていきましょう｡

　まずは、上式の展開です｡

　便宜上､$\mathrm{MSD}(0)$を以下の式で表します｡

$$\mathrm{MSD}(0) = \frac{(x(1)-x(1))^2 + (x(2)-x(2))^2 + (x(3)-x(3))^2 + (x(4)-x(4))^2 + (x(5)-x(5))^2}{5}$$

　はい｡それでは､$\mathrm{MSD}(0)$から$\mathrm{MSD}(4)$までを展開します｡
　しました｡そしたら各項ごとにΣでまとめましょう。

$$\mathrm{MSD}(0) = \frac{1}{5} \biggl(
\sum_{i=1}^5x(i)^2+\sum_{i=1}^5 x(i)^2-2\times\sum_{i=1}^5x(i) x(i) \biggr)$$

$$\mathrm{MSD}(1) = \frac{1}{4} \biggl(
\sum_{i=1}^4x(i)^2+\sum_{i=2}^5 x(i)^2-2\times\sum_{i=1}^4x(i)x(i+1) \biggr)$$

$$\mathrm{MSD}(2) = \frac{1}{3} \biggl(
\sum_{i=1}^3x(i)^2+\sum_{i=3}^5 x(i)^2-2\times\sum_{i=1}^3x(i)x(i+2) \biggr)$$

$$\mathrm{MSD}(3) = \frac{1}{2} \biggl(
\sum_{i=1}^2x(i)^2+\sum_{i=4}^5 x(i)^2-2\times\sum_{i=1}^2x(i)x(i+3) \biggr)$$

$$\mathrm{MSD}(4) = \frac{1}{1} \biggl(
\sum_{i=1}^1x(i)^2+\sum_{i=5}^5 x(i)^2-2\times\sum_{i=1}^1x(i)x(i+4) \biggr)$$

一般化すると､数列の長さを$n$とした場合､

$$\mathrm{MSD}(t) = \frac{1}{n-t} \biggl(
\sum_{i=1}^{n-t}x(i)^2+\sum_{i=t+1}^n x(i)^2-2\times\sum_{i=1}^{n-t}x(i)x(i+t) \biggr)$$

第二項は$\sum_{i=t+1}^n = \sum_{i=1}^n - \sum_{i=1}^t$と分解し、

$$\mathrm{MSD}(t) = \frac{1}{n-t} \biggl(
\sum_{i=1}^{n-t}x(i)^2+\biggl(\sum_{i=1}^n x(i)^2 -
\sum_{i=1}^{t} x(i)^2 \biggr)-2\times\sum_{i=1}^{n-t}x(i)x(i+t)\biggr)$$

$x^2$に対する累積和を$S_j = \sum_{i=1}^jx(i)^2$とすると､第一項、先程分解した項達は累積和で表わせ、

$$\mathrm{MSD}(t) = \frac{1}{n-t} \bigl( (S_{n-t} + S_{n}-S_{t}-2\times\sum_{i=1}^{n-t}x(i)x(i+t)\bigr)$$

となる。また､最後の項は畳み込みのため、$x$の時間間隔$t$での自己相関関数(線形畳み込みによる)を$\mathrm{ACF}(t)$としますと､

$$\mathrm{MSD}(t) = \frac{1}{n-t} \bigl(S_{n-t} + S_{n}-S_{t}-2\times \mathrm{ACF}(t) \bigr)$$

となり、無事平均二乗変位を累積和と畳み込みのみで表せました｡

　二乗と累積和は$O(N)$､自己相関関数はそのままだと$O(N^2)$となりますが､FFTと線形畳み込みを利用することで$O(N\log N)$となります｡よって$0 \leq t \leq N-1$の範囲における$\mathrm{MSD}(t)$の計算は$O(N\log N)$で行うことができます｡よかったですね。

どのくらい早くなるか

　実用する上で､例えば粒子数500, 三次元, 1粒子1次元あたりの軌跡のデータ点数を5000とした場合､愚直に計算すると$50035000^2 = 37500000000 > 10^{10}$となるのに対し､累積和とFFTを利用することで$50035000*\log_2 5000 = 92157842 < 10^8$ となり､約400倍高速化出来ました｡ベンチマークで比較するとより幸せになれます｡

まとめ

　平均自乗変位なんてメジャーな計算なので､もしかしたら当たり前のことかも知れませんが､事前サーチしても出てこなかったので一応記事にしてみました｡

　ちなみに､輸送係数を求める際に使われるGreen-kubo公式でよく出てくる､「自己相関関数の積分」､これも自己相関関数部分がボトルネックになり$O(N^2)$です｡しかも､Green-kubo公式はMSDに比べ多くの統計量が必要で､平気でデータ数が500000とかになります｡計算範囲を$M$として､$O(NM)$と頑張る方法もあったりしますが､自己相関関数は線形畳み込みによる計算で行えるため､実は$O(N \log N)$で計算可能です｡これについても記事を書きました。