1"""A short-time Fourier transform spectrogram in decibels."""
3from __future__ import annotations
5from dataclasses import dataclass
7import numpy as np
8import scipy.fft
12class Spectrogram:
13 times: np.ndarray
14 """Center time of each column [s]."""
15 freqs: np.ndarray
16 """Frequency of each row [Hz]."""
17 db: np.ndarray
18 """Magnitude in dB relative to the peak, ``(len(freqs), len(times))``."""
20 def plot(self, ax=None, **kwargs):
21 """Draws the spectrogram with matplotlib; see ``cymsim.plot.spectrogram``."""
22 from .plot import spectrogram
24 return spectrogram(self, ax=ax, **kwargs)
27def spectrogram(signal: np.ndarray, fs: float, window: int = 2048, hop: int = 256) -> Spectrogram:
28 """The STFT of ``signal`` with a periodic Hann window of ``window`` samples
29 and a hop of ``hop`` samples, in dB relative to the largest magnitude."""
30 x = np.asarray(signal, dtype=float)
31 if len(x) < window:
32 x = np.pad(x, (0, window - len(x)))
33 n_cols = (len(x) - window) // hop + 1
34 idx = np.arange(window)[None, :] + hop * np.arange(n_cols)[:, None]
35 win = 0.5 - 0.5 * np.cos(2 * np.pi * np.arange(window) / window)
36 X = scipy.fft.rfft(x[idx] * win[None, :], axis=1)
37 p = X.real**2 + X.imag**2
38 db = 10 * np.log10(p + 1e-30)
39 db -= db.max()
40 times = (hop * np.arange(n_cols) + window / 2) / fs
41 freqs = np.arange(window // 2 + 1) * fs / window
42 return Spectrogram(times, freqs, db.T)