← Walkthroughmapa/preprocessing/stft.py
Source snapshot · source SHA-256 e6745c9d5a9d
1"""Whole-recording magnitude STFT, ported from the research cache producer."""23import torch456def _single_stft_raw_view(7 waveform: torch.Tensor,8 *,9 sample_rate: int,10 nperseg: int,11 hop_length: int,12 k0: int,13 k1: int,14 log_eps: float,15 apply_log: bool = False,16 cartesian: bool = False,17) -> torch.Tensor:18 """2STFT single-band raw |STFT| (§3a): ONE Hann STFT, keep the inclusive19 rfft-bin slice ``[k0, k1]``. ``cartesian=True`` (3STFT slow band, §4) instead20 returns the two real components stacked on the freq axis ([Re ++ Im], F →21 2·n_bins) to preserve phase; magnitude/log are bypassed.2223 Same STFT call as :func:`_multi_stft_raw_view` (``center=True``,24 ``normalized=False``, value axis = raw ``|STFT|``) — only the bin selection25 differs. Input ``(..., C, T_samples)`` → output ``(..., C, k1-k0+1, T_bin)``26 with ``T_bin = 1 + T_samples // hop_length``.27 """28 win = torch.hann_window(nperseg, device=waveform.device)29 wf = waveform30 if wf.shape[-1] < nperseg:31 # NeuralSet's prepare() probes with sub-second inputs — pad so32 # torch.stft(center=True) can reflect-pad without crashing. Real33 # 1 s windows never hit this branch.34 wf = torch.nn.functional.pad(wf, (0, nperseg - wf.shape[-1]))35 spec = torch.stft(36 wf,37 n_fft=nperseg,38 hop_length=hop_length,39 win_length=nperseg,40 window=win,41 return_complex=True,42 normalized=False,43 center=True,44 )45 band = spec[..., k0 : k1 + 1, :] # (..., n_bins, T) complex46 if cartesian:47 # Phase-bearing slow band (§4): two real channels stacked on the freq48 # axis as [Re of n_bins ++ Im of n_bins] → F doubles to 2·n_bins. Keeps49 # phase; apply_log is rejected for cartesian upstream so it never applies.50 return torch.cat([band.real, band.imag], dim=-2) # (..., 2·n_bins, T)51 mag = band.abs() # (..., n_bins, T)52 if apply_log:53 return torch.log(mag + log_eps)54 return mag555657def _single_stft_raw_view_chunked(58 waveform: torch.Tensor,59 *,60 sample_rate: int,61 nperseg: int,62 hop_length: int,63 k0: int,64 k1: int,65 log_eps: float,66 apply_log: bool = False,67 cartesian: bool = False,68 chunk_frames: int = 960,69) -> torch.Tensor:70 """Memory-bounded whole-recording single-band raw |STFT|, byte-identical to71 one un-chunked :func:`_single_stft_raw_view` over the full ``waveform``.7273 The single-band analog of :func:`_multi_stft_raw_view_chunked`: overlap-save74 in ``chunk_frames``-wide output blocks, each reading ``half = nperseg // 2``75 extra REAL samples of context per side, then keeping only the frames whose76 full analysis window was real (the ``center=True`` reflect frames survive77 only at the TRUE recording ends, exactly as the un-chunked call produces).7879 Correctness needs ``half`` to be a multiple of ``hop_length`` so every80 segment starts on the global hop grid — true for both 2STFT bands81 (low 256 = 1·256; high 64 = 1·64). Pinned by the acceptance test.8283 Input ``(..., C, T_samples)`` → ``(..., C, k1-k0+1, T_bin)`` with84 ``T_bin = 1 + T_samples // hop_length``.85 """86 if hop_length <= 0 or nperseg % hop_length != 0:87 raise ValueError(88 f"chunked single-band STFT needs nperseg ({nperseg}) divisible by "89 f"hop_length ({hop_length}) so segment starts land on the hop grid"90 )91 half = nperseg // 292 if half % hop_length != 0:93 raise ValueError(94 f"chunked single-band STFT needs half=nperseg//2 ({half}) divisible "95 f"by hop_length ({hop_length}) for segment→global frame alignment"96 )97 n_samples = int(waveform.shape[-1])98 total_frames = 1 + n_samples // hop_length99100 def _view(seg: torch.Tensor) -> torch.Tensor:101 return _single_stft_raw_view(102 seg,103 sample_rate=sample_rate,104 nperseg=nperseg,105 hop_length=hop_length,106 k0=k0,107 k1=k1,108 log_eps=log_eps,109 apply_log=apply_log,110 cartesian=cartesian,111 )112113 # Whole recording fits in one pass (or is too short to chunk): direct call.114 if total_frames <= chunk_frames or n_samples <= 2 * half + nperseg:115 return _view(waveform)116117 blocks: list[torch.Tensor] = []118 f = 0119 while f < total_frames:120 f_end = min(f + chunk_frames, total_frames)121 lo = max(0, (f * hop_length) - half)122 hi = min(n_samples, (f_end - 1) * hop_length + half + 1)123 if hi - lo < nperseg:124 lo = ((max(0, hi - nperseg)) // hop_length) * hop_length125 seg = waveform[..., lo:hi]126 spec_seg = _view(seg)127 g0 = f - lo // hop_length128 g1 = f_end - lo // hop_length129 blocks.append(spec_seg[..., g0:g1])130 f = f_end131 return torch.cat(blocks, dim=-1)