File size: 3,256 Bytes
eae86f2 | 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 | #!/usr/bin/env python3
"""Evaluate best and final trained checkpoints at every original SU2 point."""
import argparse
import json
import os
from pathlib import Path
import sys
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
import numpy as np
import torch
from scripts.common import load_config, build_model, read_reference, predict, errors, atomic_json, require_compute_node
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--config', default='config/config.yaml')
parser.add_argument('--checkpoint', type=Path)
args = parser.parse_args()
cfg, root = load_config(args.config)
require_compute_node(cfg)
device = torch.device(cfg['training']['device'])
results = root / cfg['paths']['results']
results.mkdir(parents=True, exist_ok=True)
xy, ref, meta, mesh = read_reference(cfg['evaluation']['reference_vtu'], cfg['evaluation']['expected_points'])
paths = [('best', root / cfg['paths']['best_checkpoint']), ('final', root / cfg['paths']['last_checkpoint'])]
if args.checkpoint:
paths = [('custom', args.checkpoint)]
metrics = {'reference': meta, 'definition': 'norm(predicted Velocity.x - SU2 Velocity.x) / norm(SU2 Velocity.x)',
'all_original_nodes': True, 'volume_weighting': False, 'checkpoints': {}}
for label, path in paths:
ck = torch.load(path, map_location=device, weights_only=False)
if ck['step'] < cfg['training']['steps'] and label == 'final':
raise RuntimeError(f'Formal final checkpoint contains only {ck["step"]} steps')
model = build_model(ck['config'], device)
model.load_state_dict(ck['model_state_dict'], strict=True)
model.eval()
pred = predict(model, xy, device, cfg['evaluation']['batch_size'])
if not np.isfinite(pred).all():
raise FloatingPointError('Nonfinite inference prediction')
result = dict(checkpoint=str(path), checkpoint_step=ck['step'], **errors(pred, ref))
metrics['checkpoints'][label] = result
np.savez_compressed(results / f'predictions_{label}.npz', xy=xy, prediction_uvp=pred,
reference_uvp=ref, point_indices=np.arange(len(xy)), checkpoint_step=ck['step'])
from vtk import vtkXMLUnstructuredGridWriter
from vtk.util.numpy_support import numpy_to_vtk
for name, array in [('PINN_u', pred[:, 0]), ('PINN_v', pred[:, 1]),
('PINN_p', pred[:, 2]), ('absolute_error_u', abs(pred[:, 0] - ref[:, 0]))]:
a = numpy_to_vtk(np.ascontiguousarray(array), deep=True)
a.SetName(name)
mesh.GetPointData().AddArray(a)
writer = vtkXMLUnstructuredGridWriter()
writer.SetFileName(str(results / f'prediction_{label}.vtu'))
writer.SetInputData(mesh)
if writer.Write() != 1:
raise IOError('VTU prediction export failed')
print(json.dumps({'event': 'inference', 'selection': label, **result}), flush=True)
metrics['paper_relative_l2_u'] = cfg['evaluation']['paper_relative_l2_u']
metrics['slurm_job_id'] = os.environ.get('SLURM_JOB_ID')
atomic_json(results / 'metrics.json', metrics)
if __name__ == '__main__':
main()
|