File size: 5,250 Bytes
3c8d50e | 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 | """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')
|