TL;DR
動機
みなさん、フーリエ変換していますか？自分は今年、大学のあるときは少なくとも週１でフーリエ変換の式を見ていたんじゃないかというくらいフーリエ変換してました。そんななか先日、Google先生がICZTの$O(n\log{n})$なアルゴリズムが発表されたという記事を持ってきてくれたので、読んで実装してみました。元の紹介記事はこちら(IOWA州大学)で、アルゴリズムの記事はこちら(nature)です。
なお、natureの方はわかりやすい英語でかつ非常にわかりやすく書いてくださっているので、理論的な部分はそちらを読むことを強くおすすめします。
CZT/ICZT
(離散)フーリエ変換((D)FT)は有名だと思うのですが、その一般化であるCZTを自分は知らなかったのでICZTのアルゴリズムの前に軽く説明してみたいと思います。なお、ICZTは逆CZTなのでCZTだけ説明します。さらに後述しますが、CZTは計算量の観点だと$O(n\log{n})$にも関わらずFFTのような入力サイズの制約がないため、任意サイズのFFTとして使われることもあります。
CZTとは
CZTとはChirp $z$-Transformの略です。そしてこの言葉はChirpと$z$-Transformに分割できます。Chirpとはチャープ信号(wiki)のことで、ざっくりいうと周波数が時変な信号のことです。このような信号はフーリエ変換では扱えません。また、$z$-Transformは離散な信号に対するラプラス変換のようなものです。すなわち、CZTとは離散でかつ周波数が時変な信号にも対応したラプラス変換(フーリエ変換)といえます。
CZTの式
CZTの式は入力として$N$次元ベクトル$\mathbf{x}=(x_0,x_1,\dots,x_{N-1})^T$、出力として$M$次元ベクトル$\mathbf{X}=(X_0,X_1,\dots,X_{M-1})^T$をとり、各$k=0,1,2,\dots,M-1$に対して次のように表されます。(記事と記号を合わせ、虚数単位を$i$で表すためにインデックスは$j$で表しています。)
$$
X_k=\sum_{j=0}^{N-1}x_jA^{-j}W^{jk}
$$
ここで、$A,W\in\mathbb{C}\setminus\{0\}$のスカラーです。なお、一般に$\mathbf{x,X}$は複素ベクトルであり、$N$と$M$は同じ値とは限りません。ちなみに、$A=1,W=\exp{(-i\frac{2\pi}{M})}$とするとDFTと一致します。
このまま計算すると計算量は$O(NM)$になりますが、ここから$\mathbf{A}=\mathrm{diag}\left(A^{-0},A^{-1},\dots,A^{-(N-1)}\right)$、$\mathbf{W}$を$W^{jk}$を順に並べた$M{\times}N$行列とすると$\mathbf{A,W}$は特殊な行列となるので、FFTとIFFTを用いることで$O(n\log{n})$となります。このときの式をICZTで使うので途中を省略して式だけ残します。
$$
\mathbf{X}=\mathbf{P\hat{W}QAx}
$$
なお、$\mathbf{P,Q}$は$\mathbf{W}$から得られる対角行列で、$\mathbf{\hat{W}}$は$\mathbf{W}=\mathbf{P\hat{W}Q}$を満たすToeplitz行列です。
ICZT
CZTは上記の通り$O(n\log{n})$で計算可能なのですが、ICZTはこれまで$O(n^2)$程度だったそうです。これを$O(n\log{n})$で可能にしたのが今回の発表だそうです。
CZTの最後の式を見ると、$\mathbf{\hat{W}}$以外は対角行列なので無条件で逆行列が存在し計算も$O(n)$ですが、$\mathbf{\hat{W}}^{-1}$は$N=M$、すなわち入出力の次元が変わらないときのみ存在します。これまた細かいところは元を参照していただくとして、逆行列が存在するとき、その逆行列はあるベクトル$\mathbf{u}=(u_0,u_1,\dots,u_{N-1})$が存在して、そのベクトルにより作られる下三角行列$\mathcal{A}$と上三角行列$\mathcal{D}$から次のように表せます。
$$
\mathbf{\hat{W}}^{-1}=\frac{1}{u_0}(\mathcal{AA}^T-\mathcal{D}^T\mathcal{D})
$$
肝心のベクトル$\mathbf{u}$ですが、やや込み入った形で、CZTのときの$W$を用いて各要素$u_k$は次式で表されます。
$$
u_k=(-1)^k\frac{W^{\frac{2k^2-(2N-1)k+N(N-1)}{2}}}{\displaystyle\prod_{s=1}^{N-k-1}{(W^s-1)}\prod_{s=1}^k{(W^s-1)}}
$$
よってこれを実装すれば$O(n\log{n})$のICZTが実装でます。
実装
natureの記事によると、CZT/ICZTの実装は2通り考えられるのですが、今回はデフォルトの方で実装していきます。違いは行列演算の仕方程度なので結果は変わりません。また、全体を見たい方はGitHubを参照してください。ここではハマったポイントをメインに書いていきます。
RustFFTでのハマリポイント
CZT/ICZT共に途中の行列演算でFFTを行うため今回はRustFFTを用いてFFTを行ったのですが、RustFFTでのIFFTはサンプル数で割っていないため、$(\mathrm{IFFT}\circ\mathrm{FFT})(x){\neq}x$です。なので、正しく逆変換をする場合は以下のようになります。
let fft = FFTplanner::new(false).plan_fft(n);
let mut out = vec![Complex::zero(); n()];
fft.process(x, &mut out);
let mut out = out
.into_iter()
.map(|c| c / T::from_usize(n).unwrap())
.collect::<Vec<_>>();
let ifft = FFTplanner::new(true).plan_fft(n);
ifft.process(&mut out, inv);
ここで、T::from_usize(n).unwrap()はnをusizeから浮動小数点としてf32とf64の両対応をするために導入した型Tに変換するためのコードです。なお上記コードは、上記GitHubでのsrc/lib.rs中のcirculant_multiplyから一部抜粋・改変したものです。
ICZT自体のハマリポイント
これは実装時ではなく、テストコードを書いていたときのハマリポイントなのですが、ICZTが存在する条件は$N=M$だけでなく、$s=1,2,\dots,n-1$に対して$W^s\neq1$もあり、一般にDFTとしてCZTを扱うときの$W=\exp{(-i\frac{2\pi}{M})}$を用いるとICZTが存在しなくなってしまいます。もっとも、位相をずらせば良いだけなのでそこまで問題ではないと思いますが、これでだいぶ時間を取られました...
ベンチマーク
cargo +nightly benchで簡易的にベンチマークを取ってみました。単位はナノ秒/回で、誤差は全体的におおよそ$5%$程度です。
| FFT(RustFFT) | CZT | ICZT | |
|---|---|---|---|
| $N=10$ | $46$ | $2,419$ | $7,308$ |
| $N=100$ | $1,161$ | $28,027$ | $92,122$ |
| $N=1000$ | $12,777$ | $324,400$ | $92,122$ |
| $N=10000$ | $258,424$ | $4,946,920$ | $16,163,651$ |
| $N=10007$ | $2,247,835$ | $4,860,720$ | $16,120,731$ |
CZTもICZTも内部で$2$の冪乗サイズのFFTを使っているのでFFTと比べると遅いですが、サンプル数に対する上昇はおおよそ$O(n\log{n})$っぽいとして良いと思います。一方、FFTはサンプル数が素数になった途端、一気に遅くなっています。これはRustFFTでサンプル数によって最適なフーリエ変換手法を選んでいるため、$N=10007$でそれまでの手法と異なるものを選んだからだと思われます。
感想
比較的簡単に$O(n\log{n})$なCZT/ICZTを実装できました。ただ、どう頑張ってもFFTより理論上3倍以上遅くなってしまったり、逆変換が$N=M$でも常に存在するとは限らなかったりと、FFTと同じ感覚で使うわけにはいかなさそうです。ただそれでも、FFTでは表現できないようなものを解析できますし、フーリエ変換自体が幅広い分野で使われているので、少しでも速くなる意味はどこかしらであるのだろうと思います。