from math import ceil import numpy as np from scipy.fftpack import fft, ifft, fftfreq from ..utils import logger, verbose @verbose def stft(x, wsize, tstep=None, verbose=None): """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, str, int, or None If not None, override default verbose level (see :func:`mne.verbose` and :ref:`Logging documentation ` for more). Returns ------- X : 3d array of shape [n_signals, wsize / 2 + 1, n_step] STFT coefficients for positive frequencies with n_step = ceil(T / tstep) Examples -------- X = stft(x, wsize) X = stft(x, wsize, tstep) See Also -------- istft stftfreq """ if not np.isrealobj(x): raise ValueError("x is not a real valued array") if x.ndim == 1: x = x[None, :] n_signals, T = x.shape wsize = int(wsize) # Errors and warnings if wsize % 4: raise ValueError('The window length must be a multiple of 4.') if tstep is None: tstep = wsize / 2 tstep = int(tstep) 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 / float(tstep))) n_freq = wsize // 2 + 1 logger.info("Number of frequencies: %d" % n_freq) logger.info("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 + (n_step - 1) * tstep), 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 = 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 Examples -------- 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] # 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, sfreq=None): # noqa: D401 """Frequencies of stft transformation. Parameters ---------- wsize : int Size of stft window sfreq : float Sampling frequency. If None the frequencies are given between 0 and pi otherwise it's given in Hz. Returns ------- freqs : array The positive frequencies returned by stft See Also -------- stft istft """ n_freq = wsize // 2 + 1 freqs = fftfreq(wsize) freqs = np.abs(freqs[:n_freq]) if sfreq is not None: freqs *= float(sfreq) return freqs def stft_norm2(X): """Compute L2 norm of STFT transform. It takes into account that stft only return positive frequencies. As we use tight frame this quantity is conserved by the stft. Parameters ---------- X : 3D complex array The STFT transforms Returns ------- norms2 : array The squared L2 norm of every row of X. """ X2 = (X * X.conj()).real # compute all L2 coefs and remove first and last frequency once. norms2 = (2. * X2.sum(axis=2).sum(axis=1) - np.sum(X2[:, 0, :], axis=1) - np.sum(X2[:, -1, :], axis=1)) return norms2 def stft_norm1(X): """Compute L1 norm of STFT transform. It takes into account that stft only return positive frequencies. Parameters ---------- X : 3D complex array The STFT transforms Returns ------- norms : array The L1 norm of every row of X. """ X_abs = np.abs(X) # compute all L1 coefs and remove first and last frequency once. norms = (2. * X_abs.sum(axis=(1, 2)) - np.sum(X_abs[:, 0, :], axis=1) - np.sum(X_abs[:, -1, :], axis=1)) return norms