目次
- Cooley-Tukey FFT
- FFT による巡回畳み込み
- Bluestein's FFT ← 今回の話
本題
Cooley-Tukey FFT では配列長が素数などの場合は計算量を削減することができなかった。
Bluestein's FFT は Cooley-Tukey FFT より数倍遅いものの、どのような配列長であっても計算量 $O(N \log N)$ でフーリエ変換することができる。
配列長 $N$ の配列 $x$ があったとき、その DFT (離散フーリエ変換) の結果 $X$ は
\begin{aligned}
& X[k] =
\sum_{n=0}^{N-1} x[n]
\exp\left(-2\pi i \frac{n k}{N}\right) &
& (0 \le k < N) &
\end{aligned}
で定義される。
ここで $nk$ を
nk =
\frac{-(k-n)^2}{2} +
\frac{n^2}{2} + \frac{k^2}{2}
のように展開し整理すると
\begin{aligned}
& X[k] &=&
\sum_{n=0}^{N-1}
x[n]
\exp\left(-2\pi i \frac{n k}{N} \right) & \\
& &=&
\exp\left(-\pi i \frac{k^2}{N} \right)
\sum_{n=0}^{N-1}
x[n]
\exp\left(-\pi i \frac{n^2}{N} \right)
\exp\left(\pi i \frac{(k-n)^2}{N} \right) & \\
\end{aligned}
となる。さらに
\begin{aligned}
& a[n] &=&
x[n]
\exp\left(-\pi i \frac{n^2}{N} \right) & \\
& b[n] &=&
\exp\left(\pi i \frac{n^2}{N} \right) &
\end{aligned}
と置き換えると
\begin{aligned}
& X[k] =
b[k]^* \times
\left(
\sum_{n=0}^{N-1}
a[n]
b[k-n]
\right) &
& (b[k]^*\ は\ b[k]\ の複素共役) &
\end{aligned}
となり、これは畳み込みと同じ形となる。
つまり、これは「FFT による巡回畳み込み」で紹介した手法と似た要領で計算量 $O(N \log N)$ で解くことができる。
まずは $a, b$ の長さが $2N-1$ 以上になるよう $0$ 埋めで調節する。
また $b$ はマイナスインデックスの範囲を末尾に移動する。
\begin{aligned}
& a &= [&&&&&a_{0}&, &a_{1}&, &a_{2}&] \\
& b &= [&b_{-2}&, &b_{-1}&, &b_{0}&, &b_{1}&, &b_{2}&] \\
&&&&&&& \downarrow \\
& a' &= [&&&&&a_{0}&, &a_{1}&, &a_{2}&, &0&, &...&, &0&, &0&, &0&] \\
& b' &= [&b_{-2}&, &b_{-1}&, &b_{0}&, &b_{1}&, &b_{2}&, &0&, &...&, &0&] \\
&&= [&&&&&b_{0}&, &b_{1}&, &b_{2}&, &0&, &...&, &0&, &b_{-2}&, &b_{-1}&] \\
&&= [&&&&&b_{0}&, &b_{1}&, &b_{2}&, &0&, &...&, &0&, &b_{2}&, &b_{1}&]
\end{aligned}
このような配置にすると巡回畳み込みと同じ形にできる。
\begin{aligned}
& X[k] =
b'[k]^* \times
\left(
\sum_{n=0}^{N'-1}
a'[n]
b'[(k-n) \mod {N'}]
\right) &
\end{aligned}
あとは
X[k] = b'[k]^* \times \mathcal{F}^{-1}(\mathcal{F}(a') \cdot \mathcal{F}(b'))[k]
を計算すれば良いこととなる。
サンプルコード
最後にサンプルコードを示す。
このコードは 0BSD ライセンス (出典明記不要で自由に利用可能) の元に公開する。
fn main() {
// Cooley-Tukey FFT と Bluestein's FFT の計算時間を比較
println!("len\tfft_cooley_tukey (sec)\tfft_bluestein (sec)");
let mut rand = random();
let len_patterns: Vec<usize> = (0..=10).map(|x| 1 << x).collect();
for len in len_patterns {
let x: Vec<(f64, f64)> = (0..len).map(|_| (rand(), rand())).collect();
let start = std::time::Instant::now();
let y1 = fft_cooley_tukey(&x, false);
let cooley_tukey_dur = start.elapsed().as_secs_f64();
let start = std::time::Instant::now();
let y2 = fft_bluestein(&x, false);
let bluestein_dur = start.elapsed().as_secs_f64();
std::hint::black_box((y1, y2));
println!("{}\t{:.9}\t{:.9}", len, cooley_tukey_dur, bluestein_dur);
}
}
fn fft_bluestein(x: &Vec<(f64, f64)>, inverse: bool) -> Vec<(f64, f64)> {
// 事前に w を計算しておく
let sign = if inverse { 1.0 } else { -1.0 };
let w: Vec<(f64, f64)> = (0..x.len())
.map(|n| {
let theta = sign * std::f64::consts::PI * (n * n) as f64 / x.len() as f64;
(theta.cos(), theta.sin())
})
.collect();
// Bluestein's FFT
let len = (x.len() * 2 - 1).next_power_of_two();
let a: Vec<(f64, f64)> = (0..len)
.map(|n| match n {
_ if n < x.len() => (
x[n].0 * w[n].0 - x[n].1 * w[n].1,
x[n].0 * w[n].1 + x[n].1 * w[n].0,
),
_ => (0.0, 0.0),
})
.collect();
let b: Vec<(f64, f64)> = (0..len)
.map(|n| match n {
_ if n < x.len() => (w[n].0, -w[n].1),
_ if len - n < x.len() => (w[len - n].0, -w[len - n].1),
_ => (0.0, 0.0),
})
.collect();
let a_fft = fft_cooley_tukey(&a, false);
let b_fft = fft_cooley_tukey(&b, false);
let y_fft: Vec<(f64, f64)> = (a_fft.iter().zip(b_fft.iter()))
.map(|(a, b)| (a.0 * b.0 - a.1 * b.1, a.0 * b.1 + a.1 * b.0))
.collect();
let mut y = fft_cooley_tukey(&y_fft, true);
y.truncate(x.len());
for k in 0..y.len() {
y[k] = (
y[k].0 * b[k].0 - y[k].1 * -b[k].1,
y[k].0 * -b[k].1 + y[k].1 * b[k].0,
);
if inverse {
y[k].0 /= y.len() as f64;
y[k].1 /= y.len() as f64;
}
}
return y;
}
fn fft_cooley_tukey(x: &Vec<(f64, f64)>, inverse: bool) -> Vec<(f64, f64)> {
// 事前に w を計算しておく
let sign = if inverse { 1.0 } else { -1.0 };
let w: Vec<(f64, f64)> = (0..x.len())
.map(|n| {
let theta = sign * std::f64::consts::TAU * n as f64 / x.len() as f64;
(theta.cos(), theta.sin())
})
.collect();
// 出力用 (y1 と y2 を交互に使い回す)
let mut y1 = x.clone();
let mut y2 = vec![(0f64, 0f64); x.len()];
// FFT
let mut len = y1.len();
let mut stride = 1;
while len >= 2 {
// len を 2 つの整数の積に分解
let len1 = (2..len).find(|&n| len % n == 0).unwrap_or(len);
let len2 = len / len1;
// バタフライ演算
y2.iter_mut().for_each(|y| *y = (0.0, 0.0));
for k1 in 0..len1 {
for n1 in 0..len1 {
for n2 in 0..len2 {
let k = len1 * n2 + k1;
let n = len2 * n1 + n2;
let w = w[(stride * n * k1) % w.len()];
for offset in 0..stride {
let k = stride * k + offset;
let n = stride * n + offset;
y2[k].0 += y1[n].0 * w.0 - y1[n].1 * w.1;
y2[k].1 += y1[n].0 * w.1 + y1[n].1 * w.0;
}
}
}
}
// y1 と y2 を入れ替えて次のステップへ
std::mem::swap(&mut y1, &mut y2);
len = len2;
stride *= len1;
}
// 逆変換の場合は正規化
if inverse {
for k in 0..y1.len() {
y1[k].0 /= y1.len() as f64;
y1[k].1 /= y1.len() as f64;
}
}
return y1;
}
// 疑似乱数生成 (xorshift32)
fn random() -> impl FnMut() -> f64 {
let mut state = 1u32;
move || {
state ^= state << 13;
state ^= state >> 17;
state ^= state << 5;
return (state - 1) as f64 / std::u32::MAX as f64;
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_fft_bluestein() {
// 乱数列を用意
let mut rand = random();
let len = 2 * 3 * 5 * 7;
let x: Vec<(f64, f64)> = (0..len).map(|_| (rand(), rand())).collect();
// Bluestein's FFT の結果が Cooley-Tukey FFT と一致することを確認する
let expect_fft = fft_cooley_tukey(&x, false);
// FFT
let actual_fft = fft_bluestein(&x, false);
assert_eq!(actual_fft.len(), expect_fft.len());
for (a, b) in expect_fft.iter().zip(actual_fft.iter()) {
assert!((a.0 - b.0).abs() < 1e-10, "{} != {}", a.0, b.0);
assert!((a.1 - b.1).abs() < 1e-10, "{} != {}", a.1, b.1);
}
// iFFT
let actual_ifft = fft_bluestein(&actual_fft, true);
assert_eq!(actual_ifft.len(), x.len());
for (a, b) in x.iter().zip(actual_ifft.iter()) {
assert!((a.0 - b.0).abs() < 1e-10, "{} != {}", a.0, b.0);
assert!((a.1 - b.1).abs() < 1e-10, "{} != {}", a.1, b.1);
}
}
}