"""The Short-Time Fourier Transform: a sliding-window DFT and its inverse. A single DFT describes a whole signal with one spectrum, which only makes sense when the signal is stationary. Speech, music, machine vibration and radar returns are not: their spectra change over time. The STFT handles this by cutting the signal into short overlapping frames, windowing each one, and taking its DFT. The result is a time-frequency map, the data behind a spectrogram. This module implements the transform from first principles (frame, window, ``rfft``) rather than wrapping ``scipy.signal.stft``, so the steps are visible and testable. The forward and inverse form an exact analysis or synthesis pair: ``istft(stft(x)) == x`` on the covered interior for any window with positive overlap-add, because the inverse normalises by the actual overlap-added window. The ``cola`` helper checks the separate, stronger Constant-OverLap-Add condition that lets a real-time system overlap-add without that per-sample division, and that keeps resynthesis well behaved after the spectrogram is modified (filtering, denoising). Conventions: input is real, so the transform is one-sided (``rfft``); frames that fall fully inside the signal are used (no boundary padding); the time of a frame is the time of its centre sample. Reference: Oppenheim & Schafer, Discrete-Time Signal Processing, 3rd ed., sect. 10.3 (time-dependent Fourier transform); D. Gabor, "Theory of communication," J. IEE 93(26), 1946 (the time-frequency resolution bound). """ import numpy as np from scipy.signal import get_window def _frame_starts(n_samples, nperseg, hop): """Start index of every frame that fits fully inside the signal.""" if n_samples < nperseg: return np.empty(0, dtype=int) return np.arange(0, n_samples - nperseg + 1, hop, dtype=int) def stft(x, fs=1.0, window="hann", nperseg=256, noverlap=None): """Short-Time Fourier Transform of a real signal. Parameters ---------- x : array_like Real input signal. fs : float Sample rate in Hz. window : str or array_like Window name (passed to ``scipy.signal.get_window``) or an explicit length-``nperseg`` array. nperseg : int Frame length in samples. noverlap : int, optional Samples of overlap between adjacent frames. Defaults to ``nperseg // 2`` (50% overlap). Returns ------- freqs : ndarray One-sided frequency bins in Hz, shape ``(nperseg // 2 + 1,)``. times : ndarray Frame-centre times in seconds, shape ``(nframes,)``. Z : ndarray Complex STFT, shape ``(nfreq, nframes)``. """ x = np.asarray(x, dtype=float) if noverlap is None: noverlap = nperseg // 2 if not 0 <= noverlap < nperseg: raise ValueError("need 0 <= noverlap < nperseg") hop = nperseg - noverlap win = _resolve_window(window, nperseg) starts = _frame_starts(len(x), nperseg, hop) frames = np.stack([x[s:s + nperseg] * win for s in starts]) if len(starts) \ else np.empty((0, nperseg)) Z = np.fft.rfft(frames, axis=1).T # (nfreq, nframes) freqs = np.fft.rfftfreq(nperseg, 1.0 / fs) times = (starts + nperseg / 2.0) / fs return freqs, times, Z def istft(Z, fs=1.0, window="hann", nperseg=256, noverlap=None): """Invert an STFT by windowed overlap-add. Reconstructs the signal exactly on the interior covered by the frames. Each inverse frame is overlap-added and the result is divided by the overlap-added window, so the reconstruction is exact for any window with positive coverage, not only COLA-satisfying ones. Returns ------- t : ndarray Time axis in seconds. x : ndarray Reconstructed real signal. """ Z = np.asarray(Z) if noverlap is None: noverlap = nperseg // 2 hop = nperseg - noverlap win = _resolve_window(window, nperseg) nframes = Z.shape[1] frames_time = np.fft.irfft(Z.T, n=nperseg, axis=1) # (nframes, nperseg) n_out = (nframes - 1) * hop + nperseg if nframes else 0 x = np.zeros(n_out) wsum = np.zeros(n_out) for i in range(nframes): s = i * hop x[s:s + nperseg] += frames_time[i] wsum[s:s + nperseg] += win nonzero = wsum > 1e-12 x[nonzero] /= wsum[nonzero] t = np.arange(n_out) / fs return t, x def cola(window, nperseg, noverlap): """Check the Constant-OverLap-Add condition for a window and hop. A window satisfies COLA at a given hop if, on the steady-state interior, its shifted copies sum to a constant. When they do, overlap-add resynthesis needs no per-sample normalisation (the sum is just a fixed gain), and modifying the spectrogram before resynthesis does not colour the output. Returns ------- ok : bool True if the overlap-add sum is constant to a relative 1e-6. constant : float The steady-state overlap-add value (mean over the interior). """ win = _resolve_window(window, nperseg) hop = nperseg - noverlap nframes = 2 * (nperseg // hop) + 4 # enough to reach steady state n = (nframes - 1) * hop + nperseg s = np.zeros(n) for i in range(nframes): s[i * hop:i * hop + nperseg] += win interior = s[nperseg:-nperseg] # drop the ramp-up/ramp-down ends if interior.size == 0: return False, 0.0 constant = float(interior.mean()) ok = bool(np.ptp(interior) <= 1e-6 * max(abs(constant), 1e-12)) return ok, constant def spectrogram(x, fs=1.0, window="hann", nperseg=256, noverlap=None): """Power spectrogram: the squared magnitude of the STFT. Returns ``freqs, times, Sxx`` with ``Sxx = |Z|**2`` in shape ``(nfreq, nframes)``. """ freqs, times, Z = stft(x, fs, window, nperseg, noverlap) return freqs, times, np.abs(Z) ** 2 def _resolve_window(window, nperseg): """Return a length-nperseg window array from a name or an array.""" if isinstance(window, str): return get_window(window, nperseg, fftbins=True) win = np.asarray(window, dtype=float) if win.shape != (nperseg,): raise ValueError("explicit window must have length nperseg") return win