scripts/common.py

"""Shared data, configuration, and metric utilities."""

from __future__ import annotations

import json
import os
import random
from pathlib import Path
from typing import Any

import numpy as np
import torch
import yaml


REQUIRED_KEYS = ("init_value", "last_value", "sparse_mask")


def load_config(path: str | Path) -> dict[str, Any]:
    config_path = Path(path).resolve()
    with config_path.open("r", encoding="utf-8") as handle:
        config = yaml.safe_load(handle)
    if not isinstance(config, dict):
        raise ValueError(f"Config must be a mapping: {config_path}")
    return config


def atomic_write_json(path: str | Path, payload: object) -> None:
    destination = Path(path)
    destination.parent.mkdir(parents=True, exist_ok=True)
    temporary = destination.with_suffix(destination.suffix + ".tmp")
    temporary.write_text(json.dumps(payload, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
    os.replace(temporary, destination)


def seed_everything(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed % (2**32))
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def load_split(
    data_root: str | Path,
    split: str,
    expected_count: int,
    expected_points: int = 4096,
) -> tuple[list[str], np.ndarray, np.ndarray, np.ndarray]:
    split_dir = Path(data_root) / split
    paths = sorted(split_dir.glob("*.npz"))
    if len(paths) != expected_count:
        raise ValueError(f"{split}: expected {expected_count} NPZ files, found {len(paths)}")
    names: list[str] = []
    inputs: list[np.ndarray] = []
    targets: list[np.ndarray] = []
    masks: list[np.ndarray] = []
    for path in paths:
        with np.load(path, allow_pickle=False) as sample:
            missing = [key for key in REQUIRED_KEYS if key not in sample.files]
            if missing:
                raise ValueError(f"{path}: missing arrays {missing}")
            values = [np.asarray(sample[key]) for key in REQUIRED_KEYS]
            for key, value in zip(REQUIRED_KEYS, values):
                if value.shape != (expected_points,) or value.dtype != np.float32:
                    raise ValueError(
                        f"{path}:{key} expected ({expected_points},) float32, got {value.shape} {value.dtype}"
                    )
                if not np.isfinite(value).all():
                    raise ValueError(f"{path}:{key} contains non-finite values")
            names.append(path.name)
            inputs.append(values[0])
            targets.append(values[1])
            masks.append(values[2])
    return names, np.stack(inputs), np.stack(targets), np.stack(masks)


def make_coordinate_features(grid_size: int, fourier_order: int) -> np.ndarray:
    axis = np.linspace(0.0, 1.0, grid_size, endpoint=True, dtype=np.float64)
    x_grid, y_grid = np.meshgrid(axis, axis, indexing="ij")
    raw = np.stack((x_grid.reshape(-1), y_grid.reshape(-1)), axis=1)
    feature_blocks = [raw]
    for frequency in range(1, fourier_order + 1):
        feature_blocks.append(np.sin(2.0 * np.pi * frequency * raw))
    for frequency in range(1, fourier_order + 1):
        feature_blocks.append(np.cos(2.0 * np.pi * frequency * raw))
    return np.concatenate(feature_blocks, axis=1).astype(np.float32)


def fit_target_transform(targets: np.ndarray) -> tuple[np.ndarray, float]:
    values = np.asarray(targets, dtype=np.float64)
    mean = values.mean(axis=0)
    centered = values - mean
    standard_deviation = float(centered.std(ddof=0))
    if not np.isfinite(standard_deviation) or standard_deviation <= 0:
        raise ValueError(f"Invalid target standard deviation: {standard_deviation}")
    return mean.astype(np.float32), standard_deviation


def balanced_train_queries(
    number_of_fields: int,
    points_per_field: int,
    total_queries: int,
    schedule_seed: int,
    zero_based_epoch: int,
) -> tuple[np.ndarray, np.ndarray]:
    base, remainder = divmod(total_queries, number_of_fields)
    if base + (1 if remainder else 0) > points_per_field:
        raise ValueError("balanced query request exceeds per-field points")
    rng = np.random.default_rng(schedule_seed + zero_based_epoch)
    field_order = rng.permutation(number_of_fields)
    counts = np.full(number_of_fields, base, dtype=np.int64)
    counts[field_order[:remainder]] += 1
    fields: list[np.ndarray] = []
    coordinates: list[np.ndarray] = []
    for field_index in range(number_of_fields):
        selected = np.sort(rng.choice(points_per_field, size=int(counts[field_index]), replace=False))
        fields.append(np.full(selected.size, field_index, dtype=np.int64))
        coordinates.append(selected.astype(np.int64, copy=False))
    field_indices = np.concatenate(fields)
    coordinate_indices = np.concatenate(coordinates)
    if field_indices.size != total_queries:
        raise RuntimeError("balanced sampler did not produce the requested query count")
    return field_indices, coordinate_indices


def fixed_tune_queries(
    number_of_fields: int,
    points_per_field: int,
    total_queries: int,
    seed: int = 0,
) -> tuple[np.ndarray, np.ndarray]:
    rng = np.random.default_rng(seed)
    global_indices = np.sort(
        rng.choice(number_of_fields * points_per_field, size=total_queries, replace=False)
    )
    return global_indices // points_per_field, global_indices % points_per_field


def relative_l2(target: np.ndarray, prediction: np.ndarray) -> float:
    target_values = np.asarray(target, dtype=np.float64).reshape(-1)
    prediction_values = np.asarray(prediction, dtype=np.float64).reshape(-1)
    return float(
        np.linalg.norm(target_values - prediction_values)
        / (np.linalg.norm(target_values) + 1.0e-12)
    )


def select_device(require_accelerator: bool) -> torch.device:
    if torch.cuda.is_available():
        return torch.device("cuda:0")
    if require_accelerator:
        raise RuntimeError("Formal training requires a visible DCU accelerator via torch.cuda")
    return torch.device("cpu")