File size: 6,048 Bytes
b887635 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 | """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")
|