| |
| """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() |
|
|
