| """Configuration, data and checkpoint contracts shared by all entrypoints.""" |
| import argparse |
| import hashlib |
| import json |
| import os |
| from pathlib import Path |
| import numpy as np |
| import torch |
| import yaml |
| from scipy.interpolate import interp1d |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
|
|
|
|
| def path(value): |
| p = Path(value).expanduser() |
| return p if p.is_absolute() else ROOT / p |
|
|
|
|
| def arguments(description): |
| p = argparse.ArgumentParser(description=description) |
| p.add_argument('--config', default='config/config.yaml') |
| return p |
|
|
|
|
| def config(filename): |
| with path(filename).open() as f: |
| cfg = yaml.safe_load(f) |
| assert cfg['data']['sensors'] == cfg['model']['sensors'] |
| assert sum(cfg['data'][s] for s in ['train','validation','test']) == cfg['data']['selected_trajectories'] |
| return cfg |
|
|
|
|
| def accelerator(name): |
| if name.startswith('cuda'): |
| if not os.environ.get('SLURM_JOB_ID'): |
| raise RuntimeError('Accelerator work must run inside a Slurm allocation.') |
| if not torch.cuda.is_available(): |
| raise RuntimeError('Configured accelerator is unavailable; CPU fallback forbidden.') |
| return torch.device(name) |
|
|
|
|
| def load_data(cfg, stats=None): |
| d = cfg['data'] |
| source = path(d['directory']) / d['file'] |
| with np.load(source, allow_pickle=False) as f: |
| p, radius, t = f['del_p'], f['R'], f['t'] |
| if p.shape != (1500,1000) or radius.shape != p.shape or t.shape != (1000,): |
| raise ValueError(f'Unexpected public data schema: {p.shape}, {radius.shape}, {t.shape}') |
| if p.dtype != np.float64 or radius.dtype != np.float64: |
| raise ValueError('Expected original float64 public trajectories.') |
| if not all(np.isfinite(a).all() for a in [p,radius,t]) or not np.all(np.diff(t)>0): |
| raise ValueError('Nonfinite or unordered input data.') |
| if not np.allclose([t[0],t[-1]],[0,0.0005],rtol=1e-12,atol=1e-14): |
| raise ValueError('Unexpected time range.') |
| n = d['selected_trajectories'] |
| pressure_t = np.linspace(t[0],t[-1],d['sensors']) |
| output_t = np.linspace(t[0],t[-1],d['output_points']) |
| x = interp1d(t,p[:n],kind=d['interpolation'],axis=-1)(pressure_t) |
| y = interp1d(t,radius[:n],kind=d['interpolation'],axis=-1)(output_t) |
| nt = d['train'] |
| if stats is None: |
| stats = dict(pressure_mean=float(x[:nt].mean()),pressure_std=float(x[:nt].std()), |
| radius_mean=float(y[:nt].mean()),radius_std=float(y[:nt].std()), |
| time_min=float(t[0]),time_span=float(t[-1]-t[0])) |
| for key in ['pressure_std','radius_std','time_span']: |
| if stats[key] <= 0: raise ValueError(f'Invalid normalization: {key}') |
| xn = ((x-stats['pressure_mean'])/stats['pressure_std']).astype(np.float32) |
| yn = ((y-stats['radius_mean'])/stats['radius_std']).astype(np.float32) |
| tn = ((output_t-stats['time_min'])/stats['time_span']).astype(np.float32)[:,None] |
| val_end = nt+d['validation'] |
| splits = {'train':slice(0,nt),'validation':slice(nt,val_end),'test':slice(val_end,n)} |
| return dict(x=xn,y=yn,time=tn,radius=y,pressure=x,output_t=output_t, |
| pressure_t=pressure_t,stats=stats,splits=splits,source=source) |
|
|
|
|
| def metrics(prediction, target): |
| e = np.asarray(prediction,dtype=np.float64)-np.asarray(target,dtype=np.float64) |
| rel = np.linalg.norm(e,axis=1)/np.linalg.norm(target,axis=1) |
| return dict(relative_l2_global=float(np.linalg.norm(e)/np.linalg.norm(target)), |
| relative_l2_mean=float(rel.mean()),relative_l2_std=float(rel.std()), |
| relative_l2_max=float(rel.max()),mse=float(np.mean(e*e)), |
| rmse=float(np.sqrt(np.mean(e*e))),mae=float(np.mean(np.abs(e))), |
| residual_variance_signed=float(np.var(e)), |
| residual_variance_absolute=float(np.var(np.abs(e))), |
| residual_linf=float(np.abs(e).max())) |
|
|
|
|
| def write_json(filename, value): |
| filename = Path(filename) |
| filename.parent.mkdir(parents=True,exist_ok=True) |
| temporary = filename.with_suffix(filename.suffix+'.tmp') |
| temporary.write_text(json.dumps(value,indent=2,ensure_ascii=False,allow_nan=False)+'\n') |
| temporary.replace(filename) |
|
|
|
|
| def save_checkpoint(filename, value): |
| filename = Path(filename) |
| filename.parent.mkdir(parents=True,exist_ok=True) |
| temp = filename.with_suffix('.tmp') |
| torch.save(value,temp) |
| temp.replace(filename) |
|
|
|
|
| def dataset_manifest(data,cfg): |
| h=hashlib.sha256() |
| with data['source'].open('rb') as f: |
| for block in iter(lambda:f.read(1024*1024),b''):h.update(block) |
| return dict(source=str(data['source']),sha256=h.hexdigest(),original_shape=[1500,1000], |
| selection='first 1000 trajectories in original file order; remaining 500 excluded', |
| train_indices=list(range(800)),validation_indices=list(range(800,900)),test_indices=list(range(900,1000)), |
| split_provenance='User-specified reproducible convention; authors original indices are unknown.', |
| pressure_sensor_times=data['pressure_t'].tolist(),radius_output_times=data['output_t'].tolist(), |
| interpolation=cfg['data']['interpolation'],normalization=data['stats'],statistics_fit_on='training split only') |
|
|
