This is an automated email from the git hooks/post-receive script. yoh pushed a commit to tag 0.4 in repository python-mne.
commit 4676d831a4872599eaf3b905dbd48411187c46a1 Author: Alexandre Gramfort <[email protected]> Date: Thu Jul 12 17:06:21 2012 +0200 ENH : stft with tight frame and sine window --- mne/time_frequency/stft.py | 195 ++++++++++++++++++++++++++++++++++ mne/time_frequency/tests/test_stft.py | 38 +++++++ 2 files changed, 233 insertions(+) diff --git a/mne/time_frequency/stft.py b/mne/time_frequency/stft.py new file mode 100644 index 0000000..15b72cf --- /dev/null +++ b/mne/time_frequency/stft.py @@ -0,0 +1,195 @@ +from math import ceil +import numpy as np +from scipy.fftpack import fft, ifft, fftfreq + + +def stft(x, wsize, tstep=None, verbose=True): + """STFT Short-Term Fourier Transform using a sine window. + + The transformation is designed to be a tight frame that can be + perfectly inverted. It only returns the positive frequencies. + + Parameters + ---------- + x: 2d array of size n_signals x T + containing multi-channels signal + wsize: int + length of the STFT window in samples (must be a multiple of 4) + tstep: int + step between successive windows in samples (must be a multiple of 2, + a divider of wsize and smaller than wsize/2) (default: wsize/2) + verbose: bool + Verbose output or not. + + Returns + ------- + X: 3d array of shape [n_signals, wsize / 2 + 1, n_step] + STFT coefficients for positive frequencies with + n_step = ceil(T / tstep) + + Usage + ----- + X = stft(x, wsize) + X = stft(x, wsize, tstep) + + See also + -------- + istft + stftfreq + """ + if x.ndim == 1: + x = x[None, :] + + n_signals, T = x.shape + + ### Errors and warnings ### + if tstep is None: + tstep = wsize / 2 + + if wsize % 4: + raise ValueError('The window length must be a multiple of 4.') + + if (wsize % tstep) or (tstep % 2): + raise ValueError('The step size must be a multiple of 2 and a ' + 'divider of the window length.') + + if tstep > wsize / 2: + raise ValueError('The step size must be smaller than half the ' + 'window length.') + + n_step = int(ceil(T / tstep)) + n_freq = wsize / 2 + 1 + if verbose: + print "Number of frequencies: %d" % n_freq + print "Number of time steps: %d" % n_step + + X = np.zeros((n_signals, n_freq, n_step), dtype=np.complex) + + if n_signals == 0: + return X + + # Defining sine window + win = np.sin(np.arange(.5, wsize + .5) / wsize * np.pi) + win2 = win ** 2 + + swin = np.zeros((n_step - 1) * tstep + wsize) + for t in range(n_step): + swin[t * tstep:t * tstep + wsize] += win2 + swin = np.sqrt(wsize * swin) + + # Zero-padding and Pre-processing for edges + xp = np.zeros((n_signals, wsize - tstep + T + n_step * tstep - T), + dtype=x.dtype) + xp[:, (wsize - tstep) / 2: (wsize - tstep) / 2 + T] = x + x = xp + + for t in range(n_step): + # Framing + wwin = win / swin[t * tstep: t * tstep + wsize] + frame = np.conj(x[:, t * tstep: t * tstep + wsize]) * wwin[None, :] + # FFT + fframe = fft(frame) + X[:, :, t] = fframe[:, :n_freq] + + return X + + +def istft(X, tstep=None, Tx=None): + """ISTFT Inverse Short-Term Fourier Transform using a sine window + + Parameters + ---------- + X: 3d array of shape [n_signals, wsize / 2 + 1, n_step] + The STFT coefficients for positive frequencies + tstep: int + step between successive windows in samples (must be a multiple of 2, + a divider of wsize and smaller than wsize/2) (default: wsize/2) + Tx: int + Length of returned signal. If None Tx = n_step * tstep + + Returns + ------- + x: 1d array of length Tx + vector containing the inverse STFT signal + + Usage + ----- + x = istft(X) + x = istft(X, tstep) + + See also + -------- + stft. + """ + ### Errors and warnings ### + n_signals, n_win, n_step = X.shape + if (n_win % 2 == 0): + ValueError('The number of rows of the STFT matrix must be odd.') + + wsize = 2 * (n_win - 1) + if tstep is None: + tstep = wsize / 2 + + if wsize % tstep: + raise ValueError('The step size must be a divider of two times the ' + 'number of rows of the STFT matrix minus two.') + + if wsize % 2: + raise ValueError('The step size must be a multiple of 2.') + + if tstep > wsize / 2: + raise ValueError('The step size must be smaller than the number of ' + 'rows of the STFT matrix minus one.') + + if Tx is None: + Tx = n_step * tstep + + T = n_step * tstep + + x = np.zeros((n_signals, T + wsize - tstep), dtype=np.float) + + if n_signals == 0: + return x[:, :Tx] + + ### Computing inverse STFT signal ### + # Defining sine window + win = np.sin(np.arange(.5, wsize + .5) / wsize * np.pi) + # win = win / norm(win); + # Pre-processing for edges + swin = np.zeros(T + wsize - tstep, dtype=np.float) + for t in range(n_step): + swin[t * tstep:t * tstep + wsize] += win ** 2 + swin = np.sqrt(swin / wsize) + + fframe = np.empty((n_signals, n_win + wsize / 2 - 1), dtype=X.dtype) + for t in range(n_step): + # IFFT + fframe[:, :n_win] = X[:, :, t] + fframe[:, n_win:] = np.conj(X[:, wsize / 2 - 1: 0: -1, t]) + frame = ifft(fframe) + wwin = win / swin[t * tstep:t * tstep + wsize] + # Overlap-add + x[:, t * tstep: t * tstep + wsize] += np.real(np.conj(frame) * wwin) + + # Truncation + x = x[:, (wsize - tstep) / 2: (wsize - tstep) / 2 + T + 1][:, :Tx].copy() + return x + + +def stftfreq(wsize): + """Frequencies of stft transformation + + Parameters + ---------- + wsize : int + Size of stft window + + Returns + ------- + freqs : array + The positive frequencies returned by stft + """ + n_freq = wsize / 2 + 1 + freqs = fftfreq(wsize) + freqs = np.abs(freqs[:n_freq]) + return freqs diff --git a/mne/time_frequency/tests/test_stft.py b/mne/time_frequency/tests/test_stft.py new file mode 100644 index 0000000..4f99933 --- /dev/null +++ b/mne/time_frequency/tests/test_stft.py @@ -0,0 +1,38 @@ +import numpy as np +from scipy import linalg +from numpy.testing import assert_almost_equal, assert_array_almost_equal +from nose.tools import assert_true + +from ..stft import stft, istft, stftfreq + + +def test_stft(): + "Test stft and istft tight frame property" + for T in [253, 256]: # try with even and odd numbers + t = np.linspace(0, 20, T) + x = np.sin(30 * t) + x = np.array([x, x + 1.]) + wsize = 128 + tstep = 4 + X = stft(x, wsize, tstep) + xp = istft(X, tstep, Tx=T) + + freqs = stftfreq(wsize) + + assert_true(X.shape[1] == len(freqs)) + assert_true(np.all(freqs >= 0.)) + + assert_array_almost_equal(x, xp, decimal=6) + + # Symmetrize X to get also negative frequencies to guarantee + # norm conservation thanks to tight frame property + X = np.concatenate([X[:, 1:, :][:, ::-1, :], X], axis=1) + + assert_almost_equal(linalg.norm(X.ravel()), linalg.norm(x.ravel()), + decimal=2) + + # Try with empty array + x = np.zeros((0, T)) + X = stft(x, wsize, tstep) + xp = istft(X, tstep, T) + assert_true(xp.shape == x.shape) -- Alioth's /usr/local/bin/git-commit-notice on /srv/git.debian.org/git/debian-med/python-mne.git _______________________________________________ debian-med-commit mailing list [email protected] http://lists.alioth.debian.org/cgi-bin/mailman/listinfo/debian-med-commit
