| """Shared configuration, raw Burgers data and metric contracts.""" |
| import argparse |
| import hashlib |
| import json |
| import os |
| from pathlib import Path |
| import random |
| import socket |
| import sys |
|
|
| import numpy as np |
| from scipy.io import loadmat |
| import torch |
| import yaml |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| if str(ROOT) not in sys.path: |
| sys.path.insert(0, str(ROOT)) |
| from model.ronom import RONOM |
|
|
|
|
| def parser(description): |
| p = argparse.ArgumentParser(description=description) |
| p.add_argument('--config', default='config/config.yaml') |
| return p |
|
|
|
|
| def load_config(filename): |
| with open(filename, encoding='utf-8') as handle: |
| cfg = yaml.safe_load(handle) |
| return cfg |
|
|
|
|
| def resolve(path): |
| path = Path(path).expanduser() |
| return path if path.is_absolute() else ROOT / path |
|
|
|
|
| def save_json(path, value): |
| path = Path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| tmp = path.with_suffix(path.suffix + '.tmp') |
| tmp.write_text(json.dumps(value, indent=2, ensure_ascii=False, allow_nan=False) + '\n') |
| tmp.replace(path) |
|
|
|
|
| def update_run_manifest(cfg, **changes): |
| path = resolve(cfg['paths']['results']) / 'run_manifest.json' |
| if path.exists(): |
| value = json.loads(path.read_text()) |
| value.update(changes) |
| save_json(path, value) |
|
|
|
|
| def seed_everything(seed, threads=8): |
| random.seed(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
| torch.set_num_threads(threads) |
| if torch.cuda.is_available(): |
| torch.cuda.manual_seed_all(seed) |
| torch.backends.cudnn.benchmark = False |
|
|
|
|
| def device_for(name): |
| if name in ('cuda', 'dcu'): |
| if not os.environ.get('SLURM_JOB_ID'): |
| raise RuntimeError('Accelerator execution must run inside a Slurm allocation.') |
| if not torch.cuda.is_available(): |
| raise RuntimeError('The allocated compute node has no accessible accelerator.') |
| if torch.cuda.device_count() != 1: |
| raise RuntimeError('Expected one Slurm-visible DCU; check allocation device visibility. ' |
| f'CUDA_VISIBLE_DEVICES={os.environ.get("CUDA_VISIBLE_DEVICES")} ' |
| f'HIP_VISIBLE_DEVICES={os.environ.get("HIP_VISIBLE_DEVICES")} ' |
| f'ROCR_VISIBLE_DEVICES={os.environ.get("ROCR_VISIBLE_DEVICES")} ' |
| f'SLURM_JOB_GPUS={os.environ.get("SLURM_JOB_GPUS")}') |
| print(f'Slurm job={os.environ["SLURM_JOB_ID"]} host={socket.gethostname()} ' |
| f'accelerator={torch.cuda.get_device_name()} count={torch.cuda.device_count()}', flush=True) |
| return torch.device('cuda') |
| return torch.device(name) |
|
|
|
|
| def load_splits(cfg, names): |
| """Read only the supplied .mat. No downloads, augmentation, or normalization. |
| |
| scipy decompresses the raw MATLAB array; only requested split tensors survive |
| this function. train.py requests train/validation and never constructs test. |
| """ |
| d = cfg['data'] |
| path = resolve(d['directory']) / d['filename'] |
| raw = loadmat(path, variable_names=['input', 'output', 'tspan', 'visc']) |
| if 'visc' not in raw or not np.isclose(float(raw['visc'].squeeze()), 0.01): |
| raise ValueError('Raw Burgers viscosity must equal 0.01') |
| initial, target = raw['input'], raw['output'] |
| times = raw['tspan'].reshape(-1) |
| if initial.shape != (1000, 1025) or target.shape != (1000, 101, 1025): |
| raise ValueError(f'Unexpected raw Burgers shapes: {initial.shape}, {target.shape}') |
| if times.shape != (101,) or not np.allclose(times, np.linspace(0, 1, 101), atol=1e-12): |
| raise ValueError('Raw tspan does not match the paper interval [0,1].') |
| counts = (d['train_samples'], d['validation_samples'], d['test_samples']) |
| if counts != (600, 200, 200): |
| raise ValueError('Formal Burgers protocol requires exactly 600/200/200 samples.') |
| if (d['space_stride'], d['time_stride']) != (32, 10): |
| raise ValueError('This experiment requires strides 32 in space and 10 in time.') |
| spans = {'train': (0, 600), 'validation': (600, 800), 'test': (800, 1000)} |
| splits = {} |
| for name in names: |
| begin, end = spans[name] |
| u = np.ascontiguousarray(initial[begin:end, ::32], dtype=np.float32) |
| y = np.ascontiguousarray(target[begin:end, ::10, ::32], dtype=np.float32) |
| if not np.isfinite(u).all() or not np.isfinite(y).all(): |
| raise ValueError(f'Nonfinite data in {name}') |
| mismatch = float(np.max(np.abs(u - y[:, 0]))) |
| if mismatch > 1e-5: |
| raise ValueError(f'input and output(t=0) disagree in {name}: {mismatch}') |
| splits[name] = (torch.from_numpy(u), torch.from_numpy(y), np.arange(begin, end)) |
| print(f'Data split={name} indices=[{begin},{end}) input={u.shape} ' |
| f'target={y.shape} initial_consistency_max={mismatch:.3e}', flush=True) |
| return splits, torch.tensor(times[::10], dtype=torch.float32), path |
|
|
|
|
| def digest(path): |
| sha = hashlib.sha256() |
| with open(path, 'rb') as stream: |
| for block in iter(lambda: stream.read(8 * 1024 * 1024), b''): |
| sha.update(block) |
| return sha.hexdigest() |
|
|
|
|
| def build_model(cfg, device): |
| model = RONOM(**cfg['model']).to(device) |
| if cfg['training'].get('jit_flow', False): |
| model.latent_flow = torch.jit.script(model.latent_flow) |
| count = sum(p.numel() for p in model.parameters()) |
| if count != 414217: |
| raise ValueError(f'Model parameter count {count} differs from paper Appendix B: 414217') |
| print(f'RONOM parameters={count}, latent=4x33=132', flush=True) |
| return model |
|
|
|
|
| def checkpoint_save(path, value): |
| path = resolve(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| tmp = path.with_suffix(path.suffix + '.tmp') |
| torch.save(value, tmp) |
| tmp.replace(path) |
|
|
|
|
| def load_best(cfg, device): |
| path = resolve(cfg['paths']['checkpoint']) |
| checkpoint = torch.load(path, map_location='cpu', weights_only=False) |
| |
| if checkpoint['model_config'] != cfg['model']: |
| raise ValueError('Model architecture in config differs from checkpoint.') |
| model = build_model(cfg, device) |
| model.load_state_dict(checkpoint['model_state_dict'], strict=True) |
| model.eval() |
| print(f'Loaded best checkpoint epoch={checkpoint["epoch"]} ' |
| f'validation_mse={checkpoint["validation_mse"]:.9e}', flush=True) |
| return model, checkpoint, path |
|
|
|
|
| def metrics(prediction, target): |
| err = prediction.double() - target.double() |
| mse_i = err.square().flatten(1).mean(1) |
| rmse_i = mse_i.sqrt() |
| relative = err.flatten(1).norm(dim=1) / target.double().flatten(1).norm(dim=1).clamp_min(1e-30) |
| return { |
| 'samples': len(rmse_i), |
| 'mse': mse_i.mean().item(), |
| 'global_rmse': mse_i.mean().sqrt().item(), |
| 'trajectory_rmse_mean': rmse_i.mean().item(), |
| 'trajectory_rmse_std_ddof1': rmse_i.std(correction=1).item() if len(rmse_i) > 1 else 0., |
| 'relative_l2_mean': relative.mean().item(), |
| 'rmse_by_time': err.square().mean(dim=(0, 2)).sqrt().cpu().tolist(), |
| }, rmse_i.cpu().numpy(), relative.cpu().numpy() |
|
|
|
|
| @torch.no_grad() |
| def evaluate(model, initial, target, times, batch_size, collect_latent=False): |
| predictions, latents = [], [] |
| model.eval() |
| for begin in range(0, len(initial), batch_size): |
| output = model(initial[begin:begin + batch_size], times, return_latent=collect_latent) |
| if collect_latent: |
| output, zt = output |
| latents.append(zt.cpu()) |
| predictions.append(output.cpu()) |
| predicted = torch.cat(predictions) |
| if not torch.isfinite(predicted).all(): |
| raise RuntimeError('Evaluation produced nonfinite predictions') |
| summary, per_sample, relative = metrics(predicted, target.cpu()) |
| return summary, predicted, per_sample, relative, torch.cat(latents) if latents else None |
|
|
