【Python】生存報告がてら昔書いたFFTのコードを晒す

お久しぶりです。ちゃんと生きてます。色々と立て込んでたり何だりでモチベーションが低下していました。このまま、フェードアウトしていくのもよろしくないと思いつつもネタがない・・・なんて思っていたところ、PCのデータを整理してたら、以前作成したFFT（高速フーリエ変換）のプログラムを発見。てことで、このソースコードを今回のネタにしようかなと思います。
なお、フーリエ変換やDFT（離散フーリエ変換）についての解説は割愛します。また、FFTについても多少は説明しますが、DFTからFFT アルゴリズムの導出過程など詳しい説明は省きます。また、データのサンプル数は2のべき乗個に限定します。

FFTのソースコード

ということで早速ソースコードです。まずは、モジュールのインポートとFFTの関数です。

import numpy as np
import matplotlib.pyplot as plt
import time

def fft(x):
    n = x.shape[0]
    X = np.zeros(n ,  dtype = 'complex')
    if n > 2:
        y = x[0:n:2] #(1)
        z = x[1:n:2]

        Y = fft(y) #(2)
        Z = fft(z)

        Wn = np.exp(-1j * (2 * np.pi * np.arange(0 , n/2))/float(n)) #(3)
        X[0:n/2] = Y + Wn * Z
        X[n/2:n] = Y - Wn * Z
        return X
    else:
        X[0] = x[0] + x[1] #(4)
        X[1] = x[0] - x[1]
        return X

モジュールはいつも通り、numpyをインポート。後は、後々の動作確認で使うモジュールです。
関数fftが今回の肝であるFFTを行う関数ですが、見ての通り、再帰を使って書いてます。通常、スタックオーバーフローとかを避けるために、ビットリバース処理なんかをして繰り返し構文で書くのですが、今回はそこまでやってません。こっちの方が簡単だし直感的なので再帰で書く方法を採用しています。
ソースコード内のコメントにあるように、ポイントは4つ。
まず、(1)で元の信号xから偶数番目だけ取り出した信号yと奇数番目だけ取り出した信号zを作ります。
(2)では、fft関数を再帰呼び出ししてyとzのFFTの結果Y,Zを得ています。このY,Zを利用すると元の信号xのFFTの結果Xは次の式で、表されます。
$\left\{ \begin{array}{l} X(i) = Y(i) + \exp \left( -j \frac{2 \pi i}{n} \right) Z(i) \: (0 \leq i < \frac{n}{2}) \\ X(\frac{n}{2} + i) = Y(i) - \exp \left( -j \frac{2 \pi i}{n} \right) Z(i) \: (0 \leq i < \frac{n}{2}) \end{array} \right.$

(3)はこの処理を行っています。n点のDFTをこの式のようにn/2点のDFTの和に書きなおすことができるってのが、FFTの勘所です。
そして、(1)、(2)の再帰呼び出し処理を繰り返していくと、最終的にn = 2に行きつきます（元のデータサンプル数が2のべき乗の場合）。ここが再帰の終着点で、2点のDFTは(4)のように凄く簡単な計算で求められます。
これだけです。numpy使えば簡単にFFTのプログラムが書けます。ただし、numpyには元々FFTを行うメソッドが用意されているので、普通はわざわざ書く必要はありません。処理時間もこっちの方が断然速いです。

動作確認

てことで、一応動作確認。

def main():
    #テスト用の信号生成
    n =2**10
    x = (3 * np.cos(250* 2*np.pi * np.arange(0,n)/n) +
         2 * np.sin(250* 2*np.pi * np.arange(0,n)/n) +
         np.cos(30* 2*np.pi * np.arange(0,n)/n) )

    tic = time.clock()
    X = fft(x) #FFTの計算
    toc = time.clock()
    print toc - tic #処理時間表示

    #以下、結果グラフ表示
    plt.figure()
    plt.plot(X.real , "r")
    plt.xlim((0, n))
    plt.xticks(np.arange(0, 1024 , 100))
    plt.hold(True)
    plt.plot(X.imag , "b")
    plt.show()

テスト用の信号は、データのサンプル数が1024個で、周波数250のcos波、sin波、周波数30のcos波をそれぞれ適当なゲインをかけて足し合わせたものを使います。で、結果はこちら。

赤が実数部、青が虚数部です。うん、それっぽい。numpyに用意されているfftメソッドも同様の結果を吐き出します。また処理時間は、1024点なら30[ms]程度で処理できます。しかし、既存のメソッドの方は0.2[ms]程度。速い・・・

まとめ

今回はFFTのプログラム（再帰版）について話しました。上述のように、FFTはループ処理で書くことの方が多いようで、僕が知っている参考書では、大抵そちらの方法のサンプルプログラムが記述されています。しかし、それらは全部C言語で書かれているものなので、他の言語のサンプルコードが載っているものだと、どうなのか気になるところです。
あと、今回はサンプル点数が2のべき乗個である場合という最も基本的なFFTのみ記述しましたが、サンプル点数が素数の場合のFFT アルゴリズムが提案されていたりと応用的なFFT アルゴリズムも数々あります。これらを調べてみるのも、面白いかもしれません。

関係ないですが、転職により愛知から大阪に引っ越ししてきました。新天地からの更新でした。