| """Data, normalization, metrics and checkpoint contracts shared by the CLIs.""" |
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| import os |
| from pathlib import Path |
| import random |
| import sys |
|
|
| import numpy as np |
| from scipy.io import loadmat |
| import torch |
| import torch.nn.functional as F |
| import yaml |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT)) |
| from model.hno import HNO |
|
|
|
|
| def configuration(description): |
| parser = argparse.ArgumentParser(description=description) |
| parser.add_argument('--config', default='config/config.yaml') |
| parser.add_argument('--device', default='cuda', choices=['cuda', 'cpu']) |
| return parser |
|
|
|
|
| def read_config(path): |
| with open(path, encoding='utf-8') as f: |
| c = yaml.safe_load(f) |
| for section in ['model', 'data', 'training', 'output']: |
| if section not in c: |
| raise ValueError(f'Missing config section: {section}') |
| if c['model']['grid_size'] != (c['data']['raw_shape'][1] - 1) // c['data']['stride'] + 1: |
| raise ValueError('Configured grid and strided raw data disagree') |
| if c['training']['precision'] != 'float32': |
| raise ValueError('This reproduction uses FP32 model parameters') |
| return c |
|
|
|
|
| def set_seed(seed): |
| random.seed(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
| if torch.cuda.is_available(): |
| torch.cuda.manual_seed_all(seed) |
|
|
|
|
| def compute_device(name): |
| if name == 'cuda': |
| if not os.environ.get('SLURM_JOB_ID'): |
| raise RuntimeError('Accelerator computation requires a Slurm allocation') |
| if not torch.cuda.is_available(): |
| raise RuntimeError('No accelerator visible in this Slurm allocation') |
| return torch.device(name) |
|
|
|
|
| def read_split(c, split): |
| d = c['data'] |
| count = d['n_train' if split == 'train' else 'n_eval'] |
| path = Path(d['root']) / d['train_file' if split == 'train' else 'eval_file'] |
| print(f'Reading {split}: {path}, samples={count}', flush=True) |
| raw = loadmat(path, variable_names=['coeff', 'sol']) |
| arrays = [] |
| for key in ['coeff', 'sol']: |
| a = raw.pop(key) |
| if list(a.shape) != d['raw_shape']: |
| raise ValueError(f'{path}:{key} shape {a.shape}, expected {d["raw_shape"]}') |
| if not 0 < count <= len(a): |
| raise ValueError('Invalid sample count') |
| selected = np.array(a[:count, ::d['stride'], ::d['stride']], copy=True, order='C') |
| if not np.isfinite(selected).all(): |
| raise ValueError(f'Nonfinite {key} data') |
| tensor = torch.from_numpy(selected) |
| if key == 'coeff': |
| tensor = tensor.float() |
| arrays.append(tensor) |
| del a |
| if arrays[0].shape[1:] != (c['model']['grid_size'], c['model']['grid_size']): |
| raise ValueError('Unexpected downsampled grid') |
| return tuple(arrays) |
|
|
|
|
| def training_statistics(coeff, solution): |
| |
| return {name: value.reshape(1, 1, 1) for name, value in { |
| 'input_mean': coeff.mean(), 'input_std': coeff.std() + 1e-8, |
| 'target_mean': solution.mean(), 'target_std': solution.std() + 1e-8, |
| }.items()} |
|
|
|
|
| def encode_input(x, stats): |
| return ((x - stats['input_mean']) / stats['input_std']).unsqueeze(-1) |
|
|
|
|
| def decode_output(y, stats): |
| return y[..., 0] * stats['target_std'] + stats['target_mean'] |
|
|
|
|
| def relative_l2(prediction, target): |
| denominator = torch.linalg.vector_norm(target.flatten(1), dim=1) |
| return torch.linalg.vector_norm((prediction - target).flatten(1), dim=1) / denominator.clamp_min(1e-12) |
|
|
|
|
| def central_gradient(field, spacing): |
| extended = F.pad(field, (1, 1, 1, 1)) |
| dx = (extended[:, 1:-1, 2:] - extended[:, 1:-1, :-2]) / (2 * spacing) |
| dy = (extended[:, 2:, 1:-1] - extended[:, :-2, 1:-1]) / (2 * spacing) |
| return dx, dy |
|
|
|
|
| def objective(prediction, target, training): |
| primary = relative_l2(prediction, target).sum() |
| interior_prediction = F.pad(prediction[:, 1:-1, 1:-1], (1, 1, 1, 1)) |
| predicted_gradients = central_gradient(interior_prediction, training['derivative_spacing']) |
| target_gradients = central_gradient(target, training['derivative_spacing']) |
| auxiliary = sum(relative_l2(p, t).sum() for p, t in zip(predicted_gradients, target_gradients)) |
| return primary + training['derivative_weight'] * auxiliary, primary, auxiliary |
|
|
|
|
| def atomic_json(path, value): |
| path = Path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| temporary = path.with_suffix(path.suffix + '.tmp') |
| temporary.write_text(json.dumps(value, indent=2, ensure_ascii=False, allow_nan=False), encoding='utf-8') |
| temporary.replace(path) |
|
|
|
|
| def atomic_checkpoint(path, value): |
| path = Path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| temporary = path.with_suffix('.tmp') |
| torch.save(value, temporary) |
| temporary.replace(path) |
|
|
|
|
| def load_checkpoint(path, device): |
| checkpoint = torch.load(path, map_location='cpu', weights_only=True) |
| if checkpoint['format_version'] != 1: |
| raise ValueError('Unsupported checkpoint schema') |
| model = HNO(**checkpoint['model_config']).to(device) |
| model.load_state_dict(checkpoint['model_state'], strict=True) |
| stats = {k: v.to(device) for k, v in checkpoint['normalization'].items()} |
| return model, stats, checkpoint |
|
|
|
|
| def source_hashes(): |
| paths = [ROOT/'model/hno.py', *sorted((ROOT/'scripts').glob('*.py'))] |
| return {str(p.relative_to(ROOT)): hashlib.sha256(p.read_bytes()).hexdigest() for p in paths} |
|
|
