141 lines
5.3 KiB
Python
141 lines
5.3 KiB
Python
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
|