FFT書かないといけないんだけど

世の中にはCooley-Tukeyの解説しかないんですよ。並列の時代なのでStockhamとか知りたいんですが…

で、難しい英語の解説とか読まなくてもサンプルコードがあれば実装できるので別に問題ないんだけど、やっぱりちゃんと理解したい。

FFTの解説ってΣをごにょごにょする系とクロネッカー積をごにょごにょする系しかないんですよ。

Σだらけの式とか読みたくないし、世の中全部行列で説明できるなんてものきもい。高階のテンソルなんかももちろん行列ですよ。行列の行列とか行列の分割とかそういった世界ですよわけがわかりません。

最初はもっと数学数学したテンソル代数を使うものではないかとも考えたのですが、それは間違いで、抽象的に考えると計算が効率的になったのかなんてさっぱり分からないし、テンソルのメモリ上の表現なんかきれいさっぱりなくなってこれはもうなにをやってるかわからない。

でも、行列だけにしぼるとメモリ上の配置とか計算量だとか見やすいかというとまったくそうでない。

テンソル的なものを表現するのはアインシュタイン先生がわりと記法を編み出してくれてるんだけど、あれはみんな4次元の話だからインデックスには何次元かの型情報がいるんですよ。あと普通の内積しか使わないからインデックスの上下なんて対して意味がないんです。

というわけで、テンソルはこう書くことにしよう

x[i:32][j:32]
y[k:32]
z

xは32x32の行列だ。yは32次元のベクトルだ。zはスカラーだ。インデックスは縮約するときに使うぞ。

Cのプログラムに変換する方法はまったく自明だ

complex x[32][32];
complex y[32];
complex z;

あとはテンソル積と縮約だ。ただ掛け算するのをテンソル積だとする。

w[k:32][i:32][j:32] = y[k:32]*x[i:32][j:32]

テンソル積の結果はただたんにインデックスを並べるだけだ。同じ名前のインデックスがあると縮約する

v[i:32] = y[j:32]*x[i:32][j:32]

これはベクトルと行列の積だ。この二つをプログラムで書くと

for (k) for (i) for (j) w[k][i][j] = 0;
for (k) for (i) for (j) w[k][i][j] += y[k]*x[i][j];

for (i) v[i] = 0;
for (i) for (j) v[i] += y[j]*x[i][j];

というわけで、積は機械的に総和の計算になる。

とか書いてると全部のインデックスに型をつけるのはめんどくさいので、一度型を書いたインデックスは後では書かなくてもいいことにしよう。あと先に型を宣言しておいてもいいことにしよう。まぎらわしくないときはコンマでつなげてもいいことにしよう。

w[k:32][i:32][j:32] = y[k]*x[i,j]

とか

i,j,k:32

v[i] = y[j]*x[i,j]

テンソルの要素にアクセスしたいときもあるので、それは丸括弧を使うことにしよう

omega[j:32][k:32]

を

omega(j,k) := exp(-2*pi*j*k*sqrt(-1)/32)

で定義するみたいに使う。

グループ化を行うテンソルが重要で、(N*M)要素ベクトルからNxM行列を作るような操作Gをテンソルで書く。

i:N, j:M, k:N*M
y[i,j] = G[i,j,k]*x[k]

のように使うとグループ分けができる。定義は

G(i,j,k) := if k = i*M+j then 1 else 0

グループ分けの解除も同じGでできる

x[k] = G[i,j,k]*y[i,j]

完全シャッフルテンソルPは一度グループ分けして逆方向に解除するとできる。トランプのシャッフルのこと。山を二つに分けて完全に交互に挿入するシャッフル。FFTで重要。

i,j:2*N
y[i] = P[i,j]*x[j]
where
m:N, n:2
P[i][j] := G[m,n,i]*G[n,m,j]

当たり前だが、Gを同じ方向に２回すると元に戻る。恒等テンソルをidで書くことにすると

n,s:N, m,t:M, j,k:N*M
G[n,m,j]*G[n,m,k] = id[j,k]
G[n,m,j]*G[s,t,j] = id[n,s]*id[m,t]

こんな感じで0と1しかないテンソルは代入だと理解したほうがいいかも？

n:N, m:M, j:N*M
G(n,m)[j]*T[j] = T(n*M+m)
G[n,m](j)*U[n,m] = U(j div M, j mod M)

とかやってればP(i,j)とかも計算できるかも

i,j:2*N
P(i,j) = G[m:N][n:2](i)*G[n][m](j)
= G(i mod 2, i div 2, j)
= id((i mod 2)*N+(i div 2), j)
= if iが偶数 then id(i div 2, j) else id(i div 2 + N, j)
= if j < N then id(i, j*2) else id(i, (j-N)*2+1)

というわけで、P(i)[j]はiの偶奇によりjに(i div 2)か(i div 2 + N)を代入し、P[i](j)はjとNの大小で、iに((j-N)*2+1)か(j*2)を代入するテンソルだと分かりました。

離散フーリエ変換を書くとどうなるかというと、まずomega[i:N][j:N]を

omega(i,j) := exp(-2pi*i*j*sqrt(-1)/N)

で定義して、ベクトルx[i:N]のフーリエ変換y[i:N]は

y[i] = omega[i,j]*x[j]

と書ける。サイズが偶数のときを考えて、omegaをインデックスの偶奇と大小で2分割する。iを分割するかjを分割するかの流儀がある。
iを偶奇にjを大小に分割するほうが周波数間引きで、
jを偶奇にiを大小に分割するほうが時間間引きというのでたぶんあってる。

iを偶奇で分割するときは

omega'[i':N][i0:2][j0:2][j':N]=G[i', i0, i]*G[j0, j', j]*omega[i,j]

これを計算すると

omega'(i', 0, j0, j')
=G(i', 0)[i]*G(j0, j')[j]*omega[i,j]
=omega(i'*2, j0*N+j')
=omega(i', j')
=omega(i')[k:N]*id[k](j')

omega'(i', 1, j0, j')
=G(i', 1)[i]*G(j0, j')[j]*omega[i, j]
=omega(i'*2+1, j0*N+j')
=omega(i', j')*omega(1, j')*(-1)^j0
=omega(i')[k:N]*id[k](j')*omega(1, j')*(-1)^j0

というわけで、バタフライ演算B[i0:2][j0:2][i:N][j:N]を

B(0, j0, i, j) := id(i, j)
B(1, j0, i, j) := id(i, j)*omega(1, j)*(-1)^j0

で定めると

omega'[i':N][i0:2][j0:2][j':N] = omega[i':N][k:N]*B[i0:2][j0:2][k:N][j':N]

と書けました。両辺にGを２個掛けてomega[i:2*N][j:2*N]に戻すと

i0,j0:2, i',j':N

omega[i, j]
= G[i', i0, i]*G[j0, j', j]*omega'[i', i0, j0, j']
= G[i', i0, i] * omega[i', k] * B[i0, j0, k, j'] * G[j0, j', j]

というわけで、インデックスがjのベクトルをフーリエ変換するには、jを大小で2分割して新しいインデックスを j0, j' としてバタフライ演算を行い、インデックスが i0, k になって、各 i0 に小さいフーリエ変換を行い、インデックスが i0, i' になり、最後に偶奇分割を解除すればできることがわかりました。めでたしめでたし。計算は機械的にできるようになったけどインデックスつけるのがめんどくせえ。

iに関する分割解除をする前に大小で分割解除して、さらに大小に分割するようにすると

s0:2, s':N, s:2*N
omega[i, j]
= G[s', s0, i] * G[s0, s', s] * G[i0, i', s] * omega[i', k] * B[i0, j0, k, j'] * G[j0, j', j]
= P[i,s] * G[i0, i', s] * omega[i', k] * B[i0, j0, k, j'] * G[j0, j', j]
= P[i,s] * F[s,j]

とかいった感じで進めてればCooley-Tukeyとか簡単に理解できるのではないだろうか、続くかも？