stft.py

import numpy as np


def get_hann_window(M):
    # M+1 evenly spaced numbers between -pi and pi
    fac = np.linspace(-np.pi, np.pi, M + 1, dtype=np.float32)
    w = 0.5 * (np.ones(M + 1, dtype=np.float32) + np.cos(fac))
    return w[:-1]


def pad_reflect(y, N):
    # Pads using reflected values.
    # E.g., array [0, 1, 2, 3, 4, 5, 6, 7, 8, 9], padded with N=3 results in
    # [3, 2, 1, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 8, 7, 6]
    start = y[N:0:-1]
    end = y[-2:-(N + 2):-1]
    return np.concatenate((start, y, end))


def stft(y, n_fft, hop_length):
    win_length = n_fft

    # Pad the time series so that frames are centered
    y = pad_reflect(y, n_fft // 2)

    # Pre-allocate the STFT matrix
    n_columns = (y.shape[0] - (win_length - hop_length)) // hop_length
    stft_matrix = np.empty((int(1 + n_fft // 2), n_columns), dtype=np.complex64)

    fft_window = get_hann_window(win_length)
    for i in range(n_columns):
        stft_matrix[:, i] = np.fft.rfft(fft_window * y[i * hop_length:i * hop_length + n_fft])

    return stft_matrix