misleep 0.3.0__py3-none-any.whl

misleep/analysis/transformer/inference.py

@@ -0,0 +1,328 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from importlib.resources import files
+from pathlib import Path
+from typing import Iterable, List, Optional, Sequence, Tuple
+
+import numpy as np
+
+from .configs import EEGTransformerConfig, PreprocessConfig
+from .models import build_model
+from .preprocessing import (
+    apply_bandpass,
+    build_channel_filters,
+    epoch_signal,
+    exponential_moving_standardize,
+)
+
+
+def default_checkpoint_path() -> Path:
+    """Return the path of the packaged CausalTransformer checkpoint."""
+    ref = files("misleep.analysis.transformer.checkpoints")
+    return Path(str(ref)) / "CausalTransformer_best.pt"
+
+
+@dataclass()
+class AutoStageConfig:
+    """
+    Configuration for offline auto staging with the transformer models.
+
+    Most users can leave this at the defaults and only set ``model_path``
+    and ``sf``.
+    """
+
+    model_name: str = "CausalTransformer"
+    model_path: Optional[Path] = None
+    sf: float = 305.1758
+    epoch_seconds: float = 5.0
+    filter_order: int = 4
+    eeg_band: Tuple[float, float] = (0.5, 30.0)
+    emg_band: Tuple[float, float] = (15.0, 100.0)
+    context_window: Optional[int] = None
+    finetune_frac: float = 1.0
+    finetune_epochs: int = 5
+    finetune_lr: float = 5e-5
+    finetune_batch: int = 16
+    label_stride_seconds: float = 5.0
+    label_mode: str = "stage_id"
+    output_stride_seconds: float = 5.0
+    output_label_mode: str = "stage_id"
+    device: Optional[str] = None
+    override_model_sample_rate: Optional[float] = None
+
+
+def _validate_1d(arr: np.ndarray, name: str) -> np.ndarray:
+    arr = np.asarray(arr).squeeze()
+    if arr.ndim != 1:
+        raise ValueError(f"{name} must be 1D after squeeze; got shape {arr.shape}")
+    return arr.astype(np.float32)
+
+
+def _map_label_value(value: int, mode: str) -> int:
+    if mode == "stage_id":
+        if value in (1, 2, 3):
+            return value - 1
+        return -1
+    if mode == "class":
+        if value in (0, 1, 2):
+            return value
+        return -1
+    raise ValueError("label_mode must be 'stage_id' or 'class'")
+
+
+def _labels_to_sample_labels(
+    labels: Sequence[int],
+    total_samples: int,
+    sf: float,
+    label_stride_seconds: float,
+    label_mode: str,
+) -> np.ndarray:
+    sample_labels = np.full(total_samples, -1, dtype=np.int16)
+    if not labels:
+        return sample_labels
+    samples_per_unit = int(round(sf * label_stride_seconds))
+    if samples_per_unit <= 0:
+        return sample_labels
+    for idx, raw_label in enumerate(labels):
+        mapped = _map_label_value(int(raw_label), label_mode)
+        if mapped < 0:
+            continue
+        start = idx * samples_per_unit
+        end = min(start + samples_per_unit, total_samples)
+        if start >= total_samples:
+            break
+        sample_labels[start:end] = mapped
+    return sample_labels
+
+
+def _compute_epoch_labels(
+    sample_labels: np.ndarray,
+    cfg: PreprocessConfig,
+    min_valid_ratio: float = 0.8,
+) -> np.ndarray:
+    samples_per_epoch = cfg.samples_per_epoch()
+    hop = cfg.hop_length_samples()
+    total = sample_labels.shape[0]
+    if total < samples_per_epoch:
+        return np.zeros((0,), dtype=np.int64)
+    epoch_labels: List[int] = []
+    start = 0
+    while start + samples_per_epoch <= total:
+        end = start + samples_per_epoch
+        label_slice = sample_labels[start:end]
+        valid = label_slice[label_slice >= 0]
+        if valid.size / samples_per_epoch < min_valid_ratio:
+            epoch_labels.append(-1)
+        else:
+            epoch_labels.append(int(np.bincount(valid, minlength=3).argmax()))
+        start += hop
+    return np.asarray(epoch_labels, dtype=np.int64)
+
+
+def _prepare_sequences(
+    epochs: np.ndarray,
+    labels: np.ndarray,
+    context_window: int,
+) -> Tuple[np.ndarray, np.ndarray]:
+    sequences: list[np.ndarray] = []
+    seq_labels: list[int] = []
+    buffer: list[np.ndarray] = []
+    epoch_shape = epochs.shape[1:] if epochs.ndim >= 3 else ()
+    for idx, epoch in enumerate(epochs):
+        buffer.append(epoch)
+        if len(buffer) > context_window:
+            buffer.pop(0)
+        seq = list(buffer)
+        if len(seq) < context_window:
+            seq = [seq[0]] * (context_window - len(seq)) + seq
+        arr = np.stack(seq, axis=0)
+        sequences.append(arr.astype(np.float32))
+        seq_labels.append(int(labels[idx]))
+    if not sequences:
+        return np.empty((0, context_window, *epoch_shape), dtype=np.float32), \
+            np.empty((0,), dtype=np.int64)
+    return np.stack(sequences, axis=0), np.asarray(seq_labels, dtype=np.int64)
+
+
+def _finetune_on_labeled(
+    model,
+    sequences: np.ndarray,
+    labels: np.ndarray,
+    device,
+    epochs_steps: int,
+    batch_size: int,
+    lr: float,
+    frac: float,
+) -> None:
+    import torch
+
+    if sequences.size == 0 or labels.size == 0 or frac <= 0 or epochs_steps <= 0:
+        return
+    total = len(labels)
+    use_count = max(1, min(int(round(total * frac)), total))
+    sequences = sequences[:use_count]
+    labels = labels[:use_count]
+    dataset = torch.utils.data.TensorDataset(
+        torch.from_numpy(sequences),
+        torch.from_numpy(labels.astype(np.int64)),
+    )
+    loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True, drop_last=False)
+    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
+    criterion = torch.nn.CrossEntropyLoss()
+    model.train()
+    for _ in range(epochs_steps):
+        for batch_x, batch_y in loader:
+            batch_x = batch_x.to(device)
+            batch_y = batch_y.to(device)
+            optimizer.zero_grad(set_to_none=True)
+            logits = model(batch_x)
+            loss = criterion(logits, batch_y)
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            optimizer.step()
+    model.eval()
+
+
+def _expand_labels(labels: Sequence[int], epoch_seconds: float, output_stride_seconds: float) -> List[int]:
+    if output_stride_seconds <= 0:
+        raise ValueError("output_stride_seconds must be > 0")
+    if abs(epoch_seconds / output_stride_seconds - round(epoch_seconds / output_stride_seconds)) > 1e-6:
+        raise ValueError("epoch_seconds must be an integer multiple of output_stride_seconds")
+    repeat = int(round(epoch_seconds / output_stride_seconds))
+    return [label for label in labels for _ in range(repeat)]
+
+
+def auto_stage_llm(
+    EEG: Sequence[float],
+    EMG: Sequence[float],
+    label: Optional[Sequence[int]] = None,
+    *,
+    config: Optional[AutoStageConfig] = None,
+) -> List[int]:
+    """
+    Auto stage a recording using the offline transformer models.
+
+    Parameters
+    ----------
+    EEG : array-like
+        EEG 1-D signal.
+    EMG : array-like
+        EMG 1-D signal.
+    label : list[int] | None
+        Optional reference labels used for on-the-fly finetuning. Use 1-3
+        for ``stage_id`` mode (NREM/REM/Wake) or 0-2 for ``class`` mode.
+        Provide per-second or per-epoch labels depending on
+        ``config.label_stride_seconds``.
+    config : AutoStageConfig | None
+        Configuration for preprocessing, finetuning, and output formatting.
+
+    Returns
+    -------
+    list of int
+        Predicted state codes. Resolution is controlled by
+        ``config.output_stride_seconds``.
+
+    Raises
+    ------
+    ImportError
+        When PyTorch is not installed.
+    """
+    try:
+        import torch
+    except ImportError as e:  # pragma: no cover
+        raise ImportError(
+            "The transformer auto-staging requires PyTorch. "
+            "Install it with: pip install 'misleep[transformer]' "
+            "or pip install torch") from e
+
+    cfg = config or AutoStageConfig()
+    model_path = cfg.model_path or default_checkpoint_path()
+    if not Path(model_path).exists():
+        raise FileNotFoundError(f"Checkpoint not found: {model_path}")
+
+    eeg = _validate_1d(np.asarray(EEG), "EEG")
+    emg = _validate_1d(np.asarray(EMG), "EMG")
+    if eeg.shape[0] != emg.shape[0]:
+        raise ValueError("EEG and EMG must have the same length.")
+
+    device = torch.device(cfg.device or ("cuda" if torch.cuda.is_available() else "cpu"))
+
+    payload = torch.load(model_path, map_location=device, weights_only=False)
+    model_cfg: EEGTransformerConfig = payload["config"]
+    if cfg.override_model_sample_rate is not None:
+        model_cfg.stft.sample_rate = float(cfg.override_model_sample_rate)
+
+    preprocess_cfg = PreprocessConfig(
+        sample_rate=cfg.sf,
+        epoch_seconds=cfg.epoch_seconds,
+        filter_order=cfg.filter_order,
+        eeg_band=cfg.eeg_band,
+        emg_band=cfg.emg_band,
+    )
+    context_window = cfg.context_window or model_cfg.context_window
+
+    signal = np.stack([eeg, emg], axis=0)
+    sos = build_channel_filters(preprocess_cfg)
+    filtered = apply_bandpass(signal, sos)
+    standardized = exponential_moving_standardize(filtered, init_window=preprocess_cfg.samples_per_epoch())
+    epochs = epoch_signal(standardized, preprocess_cfg)
+    if epochs.size == 0:
+        return []
+
+    labels_in = label or []
+    sample_labels = _labels_to_sample_labels(
+        labels_in,
+        total_samples=signal.shape[1],
+        sf=cfg.sf,
+        label_stride_seconds=cfg.label_stride_seconds,
+        label_mode=cfg.label_mode,
+    )
+    epoch_labels = _compute_epoch_labels(sample_labels, preprocess_cfg)
+    if epochs.shape[0] != epoch_labels.shape[0]:
+        limit = min(epochs.shape[0], epoch_labels.shape[0])
+        epochs = epochs[:limit]
+        epoch_labels = epoch_labels[:limit]
+
+    sequences, seq_labels = _prepare_sequences(epochs, epoch_labels, context_window)
+    model = build_model(cfg.model_name, model_cfg).to(device)
+    model.load_state_dict(payload["model_state"])
+    model.eval()
+
+    labeled_mask = seq_labels >= 0
+    if labeled_mask.any():
+        _finetune_on_labeled(
+            model,
+            sequences[labeled_mask],
+            seq_labels[labeled_mask],
+            device,
+            epochs_steps=cfg.finetune_epochs,
+            batch_size=cfg.finetune_batch,
+            lr=cfg.finetune_lr,
+            frac=cfg.finetune_frac,
+        )
+
+    preds: List[int] = []
+    loader = torch.utils.data.DataLoader(
+        torch.utils.data.TensorDataset(torch.from_numpy(sequences)),
+        batch_size=64,
+        shuffle=False,
+        num_workers=0,
+        drop_last=False,
+    )
+    with torch.no_grad():
+        for (batch_x,) in loader:
+            batch_x = batch_x.to(device)
+            logits = model(batch_x)
+            pred = logits.argmax(dim=-1).cpu().numpy()
+            preds.extend(pred.tolist())
+
+    if cfg.output_label_mode == "stage_id":
+        preds = [p + 1 for p in preds]
+    elif cfg.output_label_mode != "class":
+        raise ValueError("output_label_mode must be 'stage_id' or 'class'")
+
+    if cfg.output_stride_seconds != cfg.epoch_seconds:
+        preds = _expand_labels(preds, cfg.epoch_seconds, cfg.output_stride_seconds)
+
+    return preds

misleep/analysis/transformer/models.py

@@ -0,0 +1,220 @@
+from __future__ import annotations
+
+import math
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .configs import EEGTransformerConfig, STFTConfig
+
+
+class ConvFeatureExtractor(nn.Module):
+    def __init__(self, in_channels: int, base_filters: int, kernel_size: int,
+                 conv_blocks: int, embed_dim: int) -> None:
+        super().__init__()
+        padding = kernel_size // 2
+        layers = []
+        current_channels = in_channels
+        filters = base_filters
+        for _ in range(conv_blocks):
+            layers.append(nn.Conv1d(current_channels, filters, kernel_size, padding=padding))
+            layers.append(nn.BatchNorm1d(filters))
+            layers.append(nn.ReLU(inplace=True))
+            current_channels = filters
+            filters *= 2
+        self.conv = nn.Sequential(*layers)
+        self.pool = nn.AdaptiveAvgPool1d(1)
+        self.proj = nn.Linear(current_channels, embed_dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        features = self.conv(x)
+        pooled = self.pool(features).squeeze(-1)
+        return F.relu(self.proj(pooled))
+
+
+def infer_stft_feature_dim(cfg: STFTConfig) -> int:
+    nyquist = cfg.sample_rate / 2.0
+    max_freq = min(cfg.f_max, nyquist)
+    freq_resolution = cfg.sample_rate / cfg.n_fft
+    max_bin = max(int(math.floor(max_freq / max(freq_resolution, 1e-6))), 0)
+    return min(cfg.n_fft // 2 + 1, max_bin + 1)
+
+
+def compute_stft_features(x: torch.Tensor, cfg: STFTConfig) -> torch.Tensor:
+    bsz, window, channels, samples = x.shape
+    if channels == 0:
+        raise ValueError("Input tensor must contain at least one channel for STFT features.")
+    device = x.device
+    window_fn = torch.hann_window(cfg.win_length, device=device)
+    eeg = x[:, :, 0, :]  # use the first channel for spectral features
+    flat = eeg.reshape(bsz * window, samples)
+    spec = torch.stft(
+        flat,
+        n_fft=cfg.n_fft,
+        hop_length=cfg.hop_length,
+        win_length=cfg.win_length,
+        window=window_fn,
+        return_complex=True,
+        onesided=True,
+        normalized=False,
+    )
+    freqs = torch.fft.rfftfreq(cfg.n_fft, d=1.0 / cfg.sample_rate).to(device)
+    mask = freqs <= cfg.f_max
+    spec = spec[:, mask, :]
+    mag = torch.log1p(spec.abs())
+    mag = mag.mean(dim=-1)  # (B*W, F_keep)
+    return mag.view(bsz, window, -1)
+
+
+class ConvEpochEncoder(nn.Module):
+    def __init__(self, in_channels: int, hidden_dim: int) -> None:
+        super().__init__()
+        layers = []
+        current = in_channels
+        for out in (hidden_dim // 2, hidden_dim):
+            layers.append(nn.Conv1d(current, out, kernel_size=7, padding=3, stride=2))
+            layers.append(nn.BatchNorm1d(out))
+            layers.append(nn.GELU())
+            current = out
+        layers.append(nn.AdaptiveAvgPool1d(1))
+        self.net = nn.Sequential(*layers)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        bsz, window, channels, samples = x.shape
+        out = self.net(x.view(bsz * window, channels, samples)).squeeze(-1)
+        return out.view(bsz, window, -1)
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim: int, hidden: int, dropout: float) -> None:
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(dim, hidden),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden, dim),
+            nn.Dropout(dropout),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.net(x)
+
+
+class CausalBlock(nn.Module):
+    def __init__(self, dim: int, heads: int, dropout: float) -> None:
+        super().__init__()
+        self.attn = nn.MultiheadAttention(dim, heads, dropout=dropout, batch_first=True)
+        self.ln1 = nn.LayerNorm(dim)
+        self.ln2 = nn.LayerNorm(dim)
+        self.ff = FeedForward(dim, dim * 4, dropout)
+
+    def forward(self, x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+        attn_out, _ = self.attn(self.ln1(x), self.ln1(x), self.ln1(x), attn_mask=mask)
+        x = x + attn_out
+        x = x + self.ff(self.ln2(x))
+        return x
+
+
+class CausalTransformerClassifier(nn.Module):
+    def __init__(self, cfg: EEGTransformerConfig) -> None:
+        super().__init__()
+        self.cfg = cfg
+        self.context_window = max(1, cfg.context_window)
+        self.local_encoder = ConvEpochEncoder(2, cfg.d_model)
+        spec_dim = infer_stft_feature_dim(cfg.stft)
+        self.spec_linear = nn.Sequential(
+            nn.Linear(spec_dim, cfg.d_model),
+            nn.GELU(),
+            nn.LayerNorm(cfg.d_model),
+        )
+        self.fuse = nn.Sequential(
+            nn.Linear(cfg.d_model * 2, cfg.d_model),
+            nn.GELU(),
+        )
+        self.pos_embed = nn.Parameter(
+            torch.randn(1, self.context_window, cfg.d_model) / math.sqrt(cfg.d_model))
+        self.blocks = nn.ModuleList(
+            [CausalBlock(cfg.d_model, cfg.num_heads, cfg.dropout) for _ in range(cfg.num_layers)])
+        self.norm = nn.LayerNorm(cfg.d_model)
+        self.head = nn.Linear(cfg.d_model, cfg.num_classes)
+
+    def forward(self, x: torch.Tensor, return_features: bool = False,
+                ablate_spec: bool = False):
+        if x.dim() == 3:
+            x = x.unsqueeze(1)
+        bsz, window, channels, samples = x.shape
+        assert window == self.context_window, f"Expected {self.context_window} epochs, got {window}"
+        local = self.local_encoder(x)
+        spec = compute_stft_features(x, self.cfg.stft)
+        if ablate_spec:
+            spec = torch.zeros_like(spec)
+        if spec.size(-1) != self.spec_linear[0].in_features:
+            target = self.spec_linear[0].in_features
+            if spec.size(-1) > target:
+                spec = spec[..., :target]
+            else:
+                spec = F.pad(spec, (0, target - spec.size(-1)))
+        spec = self.spec_linear(spec)
+        tokens = self.fuse(torch.cat([local, spec], dim=-1)) + self.pos_embed
+        mask = torch.triu(torch.ones(window, window, device=x.device), diagonal=1).bool()
+        for block in self.blocks:
+            tokens = block(tokens, mask)
+        features = self.norm(tokens[:, -1, :])
+        logits = self.head(features)
+        if return_features:
+            return logits, features
+        return logits
+
+
+class ConvTransformerClassifier(nn.Module):
+    def __init__(self, cfg: EEGTransformerConfig) -> None:
+        super().__init__()
+        self.cfg = cfg
+        self.context_window = max(1, cfg.context_window)
+        self.feature_extractor = ConvFeatureExtractor(
+            in_channels=2,
+            base_filters=cfg.conv_base_filters,
+            kernel_size=cfg.conv_kernel_size,
+            conv_blocks=cfg.conv_blocks,
+            embed_dim=cfg.d_model,
+        )
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=cfg.d_model,
+            nhead=cfg.num_heads,
+            dim_feedforward=cfg.dim_feedforward,
+            dropout=cfg.dropout,
+            batch_first=True,
+        )
+        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=cfg.num_layers)
+        self.pos_embed = nn.Parameter(torch.zeros(1, self.context_window, cfg.d_model))
+        self.norm = nn.LayerNorm(cfg.d_model)
+        self.head = nn.Sequential(
+            nn.Linear(cfg.d_model, cfg.dim_feedforward),
+            nn.ReLU(inplace=True),
+            nn.Dropout(cfg.dropout),
+            nn.Linear(cfg.dim_feedforward, cfg.num_classes),
+        )
+
+    def forward(self, x: torch.Tensor, return_features: bool = False):
+        if x.dim() == 3:
+            x = x.unsqueeze(1)
+        bsz, window, channels, samples = x.shape
+        assert window == self.context_window, f"Expected {self.context_window} epochs, got {window}"
+        feats = self.feature_extractor(x.view(bsz * window, channels, samples)).view(bsz, window, -1)
+        tokens = feats + self.pos_embed[:, :window, :]
+        encoded = self.transformer(tokens)
+        pooled = self.norm(encoded.mean(dim=1))
+        logits = self.head(pooled)
+        if return_features:
+            return logits, pooled
+        return logits
+
+
+def build_model(name: str, config: EEGTransformerConfig) -> nn.Module:
+    if name == "CausalTransformer":
+        return CausalTransformerClassifier(config)
+    if name == "ConvTransformer":
+        return ConvTransformerClassifier(config)
+    raise ValueError("Unknown model. Choices: CausalTransformer, ConvTransformer.")

misleep/analysis/transformer/preprocessing.py

@@ -0,0 +1,105 @@
+from __future__ import annotations
+
+from typing import Optional, Sequence, Tuple
+
+import numpy as np
+from scipy import signal
+
+from .configs import PreprocessConfig
+
+
+def design_bandpass_sos(low_cut: float, high_cut: float, sample_rate: float, order: int) -> np.ndarray:
+    return signal.butter(order, [low_cut, high_cut], btype="bandpass",
+                         fs=sample_rate, output="sos")
+
+
+def build_channel_filters(cfg: PreprocessConfig) -> Tuple[np.ndarray, np.ndarray]:
+    eeg_low, eeg_high = cfg.eeg_band
+    emg_low, emg_high = cfg.emg_band
+    eeg_sos = design_bandpass_sos(eeg_low, eeg_high, cfg.sample_rate, cfg.filter_order)
+    emg_sos = design_bandpass_sos(emg_low, emg_high, cfg.sample_rate, cfg.filter_order)
+    return eeg_sos, emg_sos
+
+
+def apply_bandpass(signal_arr: np.ndarray, sos_filters: Sequence[np.ndarray],
+                   real_time: bool = False) -> np.ndarray:
+    """
+    Band-pass filter a multi-channel signal.
+
+    Parameters
+    ----------
+    signal_arr : ndarray
+        Shape ``(channels, samples)``.
+    sos_filters : sequence of ndarray
+        One SOS filter per channel (the last filter is reused if fewer are given).
+    real_time : bool
+        When True, applies causal filtering (``sosfilt``) to emulate online
+        processing; otherwise zero-phase ``sosfiltfilt`` is used.
+    """
+    if len(sos_filters) == 0:
+        raise ValueError("At least one SOS filter must be provided.")
+    filtered = np.empty_like(signal_arr, dtype=np.float32)
+    filter_fn = signal.sosfilt if real_time else signal.sosfiltfilt
+    for ch in range(signal_arr.shape[0]):
+        sos = sos_filters[min(ch, len(sos_filters) - 1)]
+        filtered[ch] = filter_fn(sos, signal_arr[ch].astype(np.float32))
+    return filtered
+
+
+def exponential_moving_standardize(
+    signal_arr: np.ndarray,
+    alpha: float = 0.999,
+    init_window: Optional[int] = None,
+) -> np.ndarray:
+    """
+    Exponential moving standardization applied per channel.
+
+    ``x'_k = (x_k - mu_k) / sqrt(sigma_k^2)`` with
+    ``mu_k = (1 - alpha) * x_k + alpha * mu_{k-1}`` and
+    ``sigma_k^2 = (1 - alpha) * (x_k - mu_k)^2 + alpha * sigma_{k-1}^2``.
+    """
+    if signal_arr.ndim != 2:
+        raise ValueError("Signal array must have shape (channels, samples) for standardization.")
+    channels, samples = signal_arr.shape
+    standardized = np.empty_like(signal_arr, dtype=np.float32)
+    window = samples if init_window is None else max(1, min(samples, int(init_window)))
+    init_slice = signal_arr[:, :window]
+    mu = init_slice.mean(axis=1, keepdims=True).astype(np.float32)
+    var = init_slice.var(axis=1, keepdims=True).astype(np.float32)
+    var = np.clip(var, 1e-6, None)
+    for idx in range(samples):
+        x = signal_arr[:, idx: idx + 1].astype(np.float32)
+        mu = (1.0 - alpha) * x + alpha * mu
+        var = (1.0 - alpha) * (x - mu) ** 2 + alpha * var
+        std = np.sqrt(np.maximum(var, 1e-6))
+        standardized[:, idx] = ((x - mu) / std)[:, 0]
+    return standardized
+
+
+def epoch_signal(signal_arr: np.ndarray, cfg: PreprocessConfig) -> np.ndarray:
+    """Slice a signal into non-overlapping epochs of ``cfg.epoch_seconds``."""
+    c, total = signal_arr.shape
+    samples_per_epoch = cfg.samples_per_epoch()
+    hop = cfg.hop_length_samples()
+    if total < samples_per_epoch:
+        return np.zeros((0, c, samples_per_epoch), dtype=np.float32)
+    epochs = []
+    start = 0
+    while start + samples_per_epoch <= total:
+        epochs.append(signal_arr[:, start: start + samples_per_epoch])
+        start += hop
+    if not epochs:
+        return np.zeros((0, c, samples_per_epoch), dtype=np.float32)
+    return np.stack(epochs, axis=0).astype(np.float32)
+
+
+def augment_epoch(
+    epoch: np.ndarray,
+    scale_range: Tuple[float, float] = (0.9, 1.1),
+    noise_std_ratio: float = 0.02,
+) -> np.ndarray:
+    """Light data augmentation (scaling + Gaussian noise) for one epoch."""
+    scaled = epoch * np.random.uniform(*scale_range)
+    noise_std = noise_std_ratio * np.std(scaled, axis=-1, keepdims=True)
+    noise = np.random.randn(*scaled.shape).astype(np.float32) * noise_std
+    return scaled + noise