File size: 2,894 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
"""Reload the selected checkpoint and evaluate the held-out trajectories."""
import hashlib
import json
import sys
import time
from pathlib import Path
sys.path.insert(0,str(Path(__file__).resolve().parents[1]))
import numpy as np
import torch
from model.vrba import VRBADeepONet
from common import arguments,config,path,accelerator,load_data,metrics,write_json


def run(cfg,checkpoint=None):
    torch.set_num_threads(1)
    device=accelerator(cfg['training']['device'])
    checkpoint=path(checkpoint or cfg['paths']['best_checkpoint'])
    ck=torch.load(checkpoint,map_location='cpu',weights_only=False)
    model=VRBADeepONet(**ck['model_config']).to(device)
    model.load_state_dict(ck['model_state_dict'],strict=True); model.eval()
    if cfg['data']['sensors']!=ck['model_config']['sensors']:raise ValueError('Pressure sensor count mismatch.')
    data=load_data(cfg,stats=ck['normalization'])
    s=data['splits']['test']; x=torch.as_tensor(data['x'][s],device=device)
    t=torch.as_tensor(data['time'],device=device)
    begin=time.perf_counter(); torch.cuda.synchronize()
    with torch.inference_mode():
        pred=torch.cat([model(x[i:i+50],t) for i in range(0,len(x),50)]).cpu().numpy()
    torch.cuda.synchronize(); duration=time.perf_counter()-begin
    truth=data['radius'][s]; pred_physical=pred.astype(np.float64)*ck['normalization']['radius_std']+ck['normalization']['radius_mean']
    per_case=np.linalg.norm(pred_physical-truth,axis=1)/np.linalg.norm(truth,axis=1)
    normalized=metrics(pred,data['y'][s]); physical=metrics(pred_physical,truth)
    result=path(cfg['paths']['results']);result.mkdir(parents=True,exist_ok=True)
    np.savez_compressed(result/'predictions.npz',prediction=pred_physical,target=truth,
        prediction_normalized=pred,target_normalized=data['y'][s],time=data['output_t'],
        pressure=data['pressure'][s],pressure_time=data['pressure_t'],
        original_indices=np.arange(s.start,s.stop),per_trajectory_relative_l2=per_case)
    report=dict(checkpoint_name=checkpoint.name,checkpoint_sha256=hashlib.sha256(checkpoint.read_bytes()).hexdigest(),
        checkpoint_epoch=ck['epoch'],best_validation=ck['best_validation'],test_count=len(x),
        inference_seconds=duration,physical=physical,normalized=normalized,
        paper_reference=dict(relative_l2=0.00516,residual_variance=1.94e-5,residual_linf=0.0822),
        comparability='Different reproducible split; authors original indices unknown. Report both physical and train-normalized metrics; table aggregation is not fully specified.',
        worst_original_index=int(s.start+np.argmax(per_case)))
    write_json(result/'metrics.json',report)
    print('TEST_METRICS '+json.dumps(report),flush=True)
    return report


if __name__=='__main__':
    p=arguments(__doc__);p.add_argument('--checkpoint',default=None)
    a=p.parse_args();run(config(a.config),a.checkpoint)