tinycast/normalization.py

"""Per-context-window min-max normalization used by the encoder path."""
from __future__ import annotations

import torch
import torch.nn as nn


class WindowMinMax(nn.Module):
    """Per-context-window min-max normalization for the encoder path.

    Per-window statistics:

        x_min   = x.min(1, keepdim=True)[0].detach()
        x_max   = x.max(1, keepdim=True)[0].detach()
        x_range = (x_max - x_min).clamp(min=1e-5).detach()
        x_norm  = (x - x_min) / x_range            # → [0, 1]

    Stats are detached: gradients do NOT flow through normalization.
    Uses explicit ``transform`` / ``inverse_transform`` so the backbone can
    hold the (x_min, x_range) tuple across encoder + decoder and apply the
    inverse at loss / forecast time.
    """

    def __init__(self, eps_clamp: float = 1e-5) -> None:
        super().__init__()
        self.eps_clamp = eps_clamp

    def transform(
        self, x: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Normalize ``x: (B, L, 1)`` to [0, 1] per context window.

        Returns ``(x_normalized, x_min, x_range)`` with stats detached and
        ``x_range >= eps_clamp`` to avoid div-by-zero on constant series.

        Robust to NaN / +inf / -inf in ``x``: those positions are filled
        with 0 BEFORE computing min/max.
        """
        x_filled = torch.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0)
        x_min = x_filled.min(dim=1, keepdim=True).values.detach()
        x_max = x_filled.max(dim=1, keepdim=True).values.detach()
        x_range = (x_max - x_min).clamp(min=self.eps_clamp).detach()
        x_normalized = (x_filled - x_min) / x_range
        return x_normalized, x_min, x_range

    @staticmethod
    def inverse_transform(
        y_pred_normalized: torch.Tensor,
        x_min: torch.Tensor,
        x_range: torch.Tensor,
    ) -> torch.Tensor:
        """Un-normalize ``(B, p)`` predictions back to raw magnitude.

        ``x_min`` / ``x_range`` come from a prior ``transform`` call and are
        ``(B, 1, 1)``; the trailing dim is squeezed for broadcasting against
        ``(B, p)``.
        """
        return y_pred_normalized * x_range.squeeze(-1) + x_min.squeeze(-1)