airfRANS-model-exploration/src/airfrans_frontier/training/normalize.py

from __future__ import annotations

import json
from dataclasses import dataclass
from collections.abc import Sequence
from pathlib import Path

import numpy as np
from numpy.typing import NDArray

FloatArray = NDArray[np.float32]


@dataclass(frozen=True)
class NormalizationStats:
    feature_mean: FloatArray
    feature_std: FloatArray
    target_mean: FloatArray
    target_std: FloatArray
    feature_names: tuple[str, ...]
    target_names: tuple[str, ...]

    def to_dict(self) -> dict[str, object]:
        return {
            "feature_names": list(self.feature_names),
            "target_names": list(self.target_names),
            "feature_mean": self.feature_mean.tolist(),
            "feature_std": self.feature_std.tolist(),
            "target_mean": self.target_mean.tolist(),
            "target_std": self.target_std.tolist(),
        }

    @classmethod
    def from_dict(cls, data: dict[str, object]) -> NormalizationStats:
        return cls(
            feature_names=tuple(str(value) for value in _list(data, "feature_names")),
            target_names=tuple(str(value) for value in _list(data, "target_names")),
            feature_mean=np.asarray(_list(data, "feature_mean"), dtype=np.float32),
            feature_std=np.asarray(_list(data, "feature_std"), dtype=np.float32),
            target_mean=np.asarray(_list(data, "target_mean"), dtype=np.float32),
            target_std=np.asarray(_list(data, "target_std"), dtype=np.float32),
        )


def compute_normalization_stats(
    features: FloatArray,
    targets: FloatArray,
    *,
    feature_names: tuple[str, ...],
    target_names: tuple[str, ...],
    min_std: float = 1e-6,
) -> NormalizationStats:
    _validate_matrix(features, "features")
    _validate_matrix(targets, "targets")
    if features.shape[1] != len(feature_names):
        raise ValueError("Feature name count does not match feature matrix width")
    if targets.shape[1] != len(target_names):
        raise ValueError("Target name count does not match target matrix width")

    feature_mean = features.mean(axis=0, dtype=np.float64).astype(np.float32)
    feature_std = features.std(axis=0, dtype=np.float64).astype(np.float32)
    target_mean = targets.mean(axis=0, dtype=np.float64).astype(np.float32)
    target_std = targets.std(axis=0, dtype=np.float64).astype(np.float32)

    feature_std = _clamp_std(feature_std, min_std)
    target_std = _clamp_std(target_std, min_std)

    return NormalizationStats(
        feature_mean=feature_mean,
        feature_std=feature_std,
        target_mean=target_mean,
        target_std=target_std,
        feature_names=feature_names,
        target_names=target_names,
    )


def normalize_features(
    features: FloatArray,
    stats: NormalizationStats,
    *,
    raw_feature_names: Sequence[str] = (),
) -> FloatArray:
    _validate_matrix(features, "features")
    if features.shape[1] != stats.feature_mean.shape[0]:
        raise ValueError("Feature width does not match normalization stats")
    normalized = (features - stats.feature_mean) / stats.feature_std
    if raw_feature_names:
        name_to_index = {name: index for index, name in enumerate(stats.feature_names)}
        missing = [name for name in raw_feature_names if name not in name_to_index]
        if missing:
            raise ValueError(f"Raw feature names are missing from normalization stats: {missing}")
        for name in raw_feature_names:
            index = name_to_index[name]
            normalized[:, index] = features[:, index]
    return np.ascontiguousarray(normalized, dtype=np.float32)


def normalize_targets(targets: FloatArray, stats: NormalizationStats) -> FloatArray:
    _validate_matrix(targets, "targets")
    if targets.shape[1] != stats.target_mean.shape[0]:
        raise ValueError("Target width does not match normalization stats")
    return np.ascontiguousarray((targets - stats.target_mean) / stats.target_std, dtype=np.float32)


def denormalize_targets(targets: FloatArray, stats: NormalizationStats) -> FloatArray:
    _validate_matrix(targets, "targets")
    if targets.shape[1] != stats.target_mean.shape[0]:
        raise ValueError("Target width does not match normalization stats")
    return np.ascontiguousarray(targets * stats.target_std + stats.target_mean, dtype=np.float32)


def save_normalization_stats(stats: NormalizationStats, path: str | Path) -> None:
    Path(path).write_text(json.dumps(stats.to_dict(), indent=2, sort_keys=True) + "\n")


def load_normalization_stats(path: str | Path) -> NormalizationStats:
    data = json.loads(Path(path).read_text())
    if not isinstance(data, dict):
        raise ValueError(f"Invalid normalization stats file: {path}")
    return NormalizationStats.from_dict(data)


def _clamp_std(std: FloatArray, min_std: float) -> FloatArray:
    result = std.copy()
    result[result < min_std] = 1.0
    return result.astype(np.float32)


def _validate_matrix(values: FloatArray, name: str) -> None:
    if values.ndim != 2:
        raise ValueError(f"Expected {name} to be 2D; got shape {values.shape}")
    if not np.isfinite(values).all():
        raise ValueError(f"Expected finite values in {name}")


def _list(data: dict[str, object], key: str) -> list[object]:
    value = data.get(key)
    if not isinstance(value, list):
        raise ValueError(f"Expected list for normalization key: {key}")
    return value