workspace/cil/models/tangent_encoder.py

from __future__ import annotations

from dataclasses import dataclass

try:
    import torch
except ImportError:  # pragma: no cover
    torch = None


@dataclass
class TangentNormalizer:
    """Mean/std normalization for spline tangent codes."""

    mean: list[float]
    std: list[float]

    @classmethod
    def fit(cls, values: "torch.Tensor", *, eps: float = 1.0e-6) -> "TangentNormalizer":
        if torch is None:  # pragma: no cover
            raise ImportError("TangentNormalizer.fit requires torch")
        mean = values.mean(dim=0)
        std = values.std(dim=0).clamp_min(eps)
        return cls(mean=mean.detach().cpu().tolist(), std=std.detach().cpu().tolist())

    def transform(self, values: "torch.Tensor") -> "torch.Tensor":
        if torch is None:  # pragma: no cover
            raise ImportError("TangentNormalizer.transform requires torch")
        mean = torch.as_tensor(self.mean, dtype=values.dtype, device=values.device)
        std = torch.as_tensor(self.std, dtype=values.dtype, device=values.device)
        return (values - mean) / std

    def inverse_transform(self, values: "torch.Tensor") -> "torch.Tensor":
        if torch is None:  # pragma: no cover
            raise ImportError("TangentNormalizer.inverse_transform requires torch")
        mean = torch.as_tensor(self.mean, dtype=values.dtype, device=values.device)
        std = torch.as_tensor(self.std, dtype=values.dtype, device=values.device)
        return values * std + mean

    def to_dict(self) -> dict[str, list[float]]:
        return {"mean": self.mean, "std": self.std}

    @classmethod
    def from_dict(cls, payload: dict[str, list[float]]) -> "TangentNormalizer":
        return cls(mean=[float(v) for v in payload["mean"]], std=[float(v) for v in payload["std"]])