src/normalization.py

"""
Field normalization utilities compatible with The Well style (per-channel mean/std or RMS).
"""
from __future__ import annotations
import torch
from typing import Dict, Optional, Tuple


class FieldNormalizer:
    def __init__(self, mode: str = "zscore", eps: float = 1e-6):
        assert mode in ("zscore", "rms", "none")
        self.mode = mode
        self.eps = eps
        self.mean: Optional[torch.Tensor] = None   # (C,)
        self.std: Optional[torch.Tensor] = None    # (C,)
        self.rms: Optional[torch.Tensor] = None    # (C,)
        self.fitted = False

    @torch.no_grad()
    def fit(self, data: torch.Tensor):
        if isinstance(data, (list, tuple)):
            data = torch.cat([d.reshape(-1, d.shape[-3], d.shape[-2], d.shape[-1]) for d in data], dim=0)
        # Flatten to (..., C, H, W)
        if data.dim() == 5:          # (N, T, C, H, W)
            data = data.permute(0, 1, 3, 4, 2).reshape(-1, data.shape[2])
        elif data.dim() == 4:        # (N, C, H, W)
            data = data.permute(0, 2, 3, 1).reshape(-1, data.shape[1])
        else:
            raise ValueError(f"Unsupported shape {data.shape}")

        self.mean = data.mean(dim=0)
        self.std = data.std(dim=0).clamp_min(self.eps)
        self.rms = torch.sqrt((data ** 2).mean(dim=0)).clamp_min(self.eps)
        self.fitted = True
        return self

    def transform(self, x: torch.Tensor) -> torch.Tensor:
        """x: (..., C, H, W)"""
        if self.mode == "none" or not self.fitted:
            return x
        actual_c = x.shape[-3]
        expected_c = self.mean.shape[0]
        if actual_c != expected_c:
            raise ValueError(
                f"FieldNormalizer was fit on {expected_c} channels but received "
                f"data with {actual_c} channels. A normalizer fit on one domain "
                f"cannot be reused for a domain with a different channel count "
                f"-- fit a new FieldNormalizer for this domain instead of "
                f"reusing the previous one. (This commonly happens in continual "
                f"learning when successive domains have different channel "
                f"counts; MultiScaleEncoder itself now handles varying channel "
                f"counts, but normalization statistics are still per-domain.)"
            )
        # Broadcast stats to (C, 1, 1)
        if self.mode == "zscore":
            mean = self.mean.view(-1, 1, 1).to(x.device, x.dtype)
            std = self.std.view(-1, 1, 1).to(x.device, x.dtype)
            return (x - mean) / std
        else:  # rms
            rms = self.rms.view(-1, 1, 1).to(x.device, x.dtype)
            return x / rms

    def inverse(self, x: torch.Tensor) -> torch.Tensor:
        if self.mode == "none" or not self.fitted:
            return x
        if self.mode == "zscore":
            mean = self.mean.view(-1, 1, 1).to(x.device, x.dtype)
            std = self.std.view(-1, 1, 1).to(x.device, x.dtype)
            return x * std + mean
        else:
            rms = self.rms.view(-1, 1, 1).to(x.device, x.dtype)
            return x * rms

    def state_dict(self) -> Dict:
        return {
            "mode": self.mode,
            "mean": self.mean,
            "std": self.std,
            "rms": self.rms,
            "fitted": self.fitted,
        }

    def load_state_dict(self, d: Dict):
        self.mode = d["mode"]
        self.mean = d["mean"]
        self.std = d["std"]
        self.rms = d["rms"]
        self.fitted = d["fitted"]