← Walkthroughmapa/data/normalize.py
Source snapshot · source SHA-256 b0bb952533a3
1"""Session-wise median/MAD normalization. No artifact detection or clipping."""23from __future__ import annotations45import torch67SCALE_TO_SIGMA = 1.48268910class SessionRobustZNormalizer:11 """Fit each contact/bin on a complete session, then reuse the frozen statistics.1213 This scaler does not identify artifacts or clip normalized values.14 """1516 def __init__(self, *, sigma_floor: float = 1e-6) -> None:17 self.sigma_floor = sigma_floor18 self.median: torch.Tensor | None = None19 self.sigma: torch.Tensor | None = None20 self._valid_bin_mask: torch.Tensor | None = None2122 @classmethod23 def from_stats(24 cls,25 *,26 median: torch.Tensor,27 sigma: torch.Tensor,28 sigma_floor: float = 1e-6,29 valid_bin_mask: torch.Tensor | None = None,30 ) -> SessionRobustZNormalizer:31 obj = cls(sigma_floor=sigma_floor)32 obj.median, obj.sigma = median, sigma33 obj._valid_bin_mask = valid_bin_mask34 return obj3536 def fit(37 self, frames: torch.Tensor, *, valid_bin_mask: torch.Tensor | None = None,38 ) -> SessionRobustZNormalizer:39 x = (40 frames41 if frames.dtype in (torch.float16, torch.float32, torch.float64)42 else frames.float()43 )44 # Time is the last axis: (C,F,T_session) -> statistics (C,F,1).45 # Each contact/bin gets its own median and scaled median absolute deviation.46 self.median = x.median(dim=-1, keepdim=True).values47 self.sigma = SCALE_TO_SIGMA * (x - self.median).abs().median(dim=-1, keepdim=True).values48 self._valid_bin_mask = valid_bin_mask49 return self5051 def transform(self, x: torch.Tensor) -> torch.Tensor:52 if self.median is None or self.sigma is None:53 raise RuntimeError("SessionRobustZNormalizer.transform called before fit()")54 x = x if x.dtype in (torch.float16, torch.float32, torch.float64) else x.float()55 # Broadcast stored session statistics across the clip's time axis.56 # Near-constant bins are set to zero below; clipping is done later in the stem.57 z = (x - self.median) / self.sigma.clamp(min=self.sigma_floor)58 z = torch.where(self.sigma >= self.sigma_floor, z, torch.zeros_like(z))59 if self._valid_bin_mask is not None:60 z = torch.where(self._valid_bin_mask.to(dtype=torch.bool), z, torch.zeros_like(z))61 return z6263 def __call__(self, x: torch.Tensor) -> torch.Tensor:64 return self.transform(x)