noise

計算機科学や各種設定のメモ

Stockham アルゴリズムについて

アルゴリズム OCaml 計算機科学

目的

Fast Fourier Transform (FFT) のアルゴリズムとして Cooley-Turky 法が紹介されることが多いと思いますが、あまり触れられることのない、もうひとつのアルゴリズムである Stockham 法について書きたいと思います。

Cooley-Tukey との違い

簡潔に言うと in place か out of place の違いになります。
in place では一つの固定された領域を計算全体において使い回すことが可能です。一方、 out of place では少なくとも入力と同じサイズの記憶領域を余分に用意する必要があります。ですので空間効率では in-place 型が有利です。
Cooley-Tukey は in place なので入力となる配列の記憶領域のみで結果を求めることが可能です。Stockham は out of place なので余分に記憶領域を必要とします。しかし、Cooley-Tukey は配列要素の順序を計算の最初か最後で入れ替える必要があるのでその分の時間的コストがあります。(追記:このコストをスケジューリングによって隠蔽できるようです。)
Stockham では bit reversal で順序を入れ替える必要はありません。このような out of place 型アルゴリズムはベクトル型計算機やSIMD計算において効果を発揮するそうです。

式

Let n=2^m, and let \alpha = e^{\pm 2 \pi i / p} , where p = 2^t.
\begin{array} Y(j2^t+k) & = & X(j2^{t-1}+k)+\alpha^kX(j2^{t-1}+n/2+k) \\ Y(j2^t+2^{t-1}+k) & = & X(j2^{t-1}+k)-\alpha^kX(j2^{t-1}+n/2+k) \end{array}
\begin{array} 0 \le j < 2^{m-t} & , & 0 \le k < 2^{t-1} \end{array}
これを t について 1 から m まで繰り返す。

計算のようす

m = 3, n = 8 の場合。
オレンジ色が上式の和をとっている部分、緑色が差をとっている部分に対応しています。

実装

OCaml でのコード

以下に 2^n タイプの Stockham 法の実装を示します。

let stockham flag m arr =
  let    cadd  = Complex.add
  in let csub  = Complex.sub
  in let cmul  = Complex.mul
  in let cpow  = Complex.pow
  in let czero = Complex.zero
  in let ci    = Complex.i
  in let f2c   = Complex.of_float
  in
  let    n   = 1 lsl m (* total length of arr : 2^m *)
  in let arr_ref0 = ref [||]
  in let arr_ref1 = ref (Array.create n czero)
  in let result_ref = ref [||]
  in
  let normalize a = Array.modify (n |> float |> sqrt |> (/.) 1. |> f2c |> cmul) a
  in
  let stockham_1 t src dst =
    let index = 1 lsl (t-1) (* 2^(t-1) *)
    in let step = index lsl 1 (* 2^t *)
    in let phase = (if flag then -1. else 1.) *. 2. *. pi /. (index lsl 1 |> float) in
    for j = 0 to 1 lsl (m - t) - 1 do
      for k = 0 to index - 1 do
        let a = cadd (phase |> cos |> f2c) (phase |> sin |> f2c |> (cmul ci)) in
        let c = k |> float |> f2c |> cpow a in
        let x0 = src.(j*index + k) in
        let x1 = src.(j*index + (n/2)+ k) in
        dst.(j * step + k)         <- cadd x0 (cmul c x1);
        dst.(j * step + index + k) <- csub x0 (cmul c x1);
      done
    done;
  in
  arr_ref0 := arr;
  for t=1 to m do
    if t land 1 = 0 then
      stockham_1 t !arr_ref1 !arr_ref0
    else
      stockham_1 t !arr_ref0 !arr_ref1
  done;
  result_ref := if m land 1 = 0 then !arr_ref0 else !arr_ref1;
  normalize !result_ref;
  !result_ref;;

let fft m arr = stockham true m (Array.map Complex.of_float arr);;
let ifft m arr = Array.map Complex.to_float (stockham false m arr);;

実行

実行結果をみるために、まず入力データを用意します。
窓領域？の全体を周期とする振幅 1.0 の波を入力とします。

# let wave n i = let pos = (i |> float) /. (n |> float) in let phase1 = 2. *. pi *. pos in sin phase1;;
val wave : int -> int -> float = <fun>

# let input = Array.init 64 (wave 64);;
val input : float array =
  [|0.; 0.0980171403295606; 0.195090322016128248;
    (略)
    -0.0980171403295605|]

振幅は一般に 2 * (複素数のノルム / 入力要素数) で求まるそうです。しかし、今回は順FFT時に正規化係数を掛けてしまっているので調節します。この振幅を求める関数を FFT して得られた複素配列に対し適用します。

# let amplitude n c = (Complex.norm c) *. (float n |> sqrt) *. 2. /. (float n);;
val amplitude : int -> Batteries.Complex.t -> float = <fun>

各周期ごとの振幅を調べてみます。

# Array.map (amplitude 64) (fft 6 input);;
- : float Batteries.Array.mappable =
[|2.46435765674166141e-17; 1.; 5.86612358578035e-17; 4.9343576547353052e-17;
  (略)
  3.03942247316122e-17;
  4.9052532219224359e-17; 5.86612358578035151e-17; 1.|]

先頭の値は直流成分です。今回入力した波の情報はその次の値です。この値は 1. となっているので正しく波の振幅が求められていることがわかります。

備考

今回紹介したアルゴリズムは Stockham の2基底(2-radix)のバージョンとなります。他に4基底のものがありますしさまざまな改良が施されたものが実用的な数値計算ライブラリの中では用いられているようです。

参考