Pythonで離散フーリエ変換

gistfile1.py

#!/usr/bin/python
# coding: utf-8

import sys
import os
import wave
import numpy as np
import pylab as pl

def dft_test(filename):
    """Discrete fourier transform

    options:
       -f=<str>, --filename=<str>    filename    
    """
    if not os.path.isfile(filename):
        print "File not found."
        sys.exit()

    data,fs = _read_wav(filename)

    # hopsize
    N = 512
    # start posision
    start_pos = 0
    cut_signal = data[start_pos:start_pos+N]
    X = _dft(cut_signal, N)
    # amplitude
    X = [abs(x) for x in X]

    cut_signal = _normalize(cut_signal)
    X = _normalize(X)

    _plot(cut_signal, X, fs)

def _normalize(x):
    return [float(c)/max(x) for c in x]

def _plot(x, X, fs):
    timeList = [float(c)/fs for c in range(len(x))]
    freqList = [float(c)*fs/len(X) for c in range(len(X))]

    pl.subplot(2,1,1)
    pl.plot(timeList, x, lw=1.5)
    pl.xlim(0, float(len(x))/fs)
    pl.xlabel("Time [second]")

    pl.subplot(2,1,2)
    pl.plot(freqList, X, lw=1.5)
    pl.xlabel("Frequency [Hz]")
    pl.xlim([0, fs/2])

    pl.savefig("a.png", format="png", transparent=True)
    pl.show()

def _read_wav(filename):
    wf = wave.open(filename, "rb")
    data = wf.readframes(wf.getnframes())
    data = pl.frombuffer(data, dtype="int16")
    fs = wf.getframerate()

    return data, fs

# Discrete fourier transform
def _dft(x, N):
    return _dft_inner(x, N, 1)

# Inverse-discrete fourier transform
def _idft(X, N):
    return _dft_inner(x, N, -1)

# implementation
def _dft_inner(x, N, flag):
    assert len(x) == N
    X = [(0 + 0j)] * N
    for k in range(N):
        tmp = 0 + 0j
        for n in range(N):
            tmp += x[n] * np.exp(-flag*(0+1j)*2.0*np.pi*n*k/N)
        X[k] = tmp
    return X

if __name__=='__main__':
    import clime.now