| """Load validation-selected weights and evaluate the last 200 independent cases.""" |
| import csv |
| import os |
| import time |
|
|
| import numpy as np |
| from common import (parser, load_config, seed_everything, device_for, resolve, |
| load_splits, load_best, evaluate, save_json, digest, update_run_manifest) |
|
|
|
|
| def main(): |
| args = parser(__doc__).parse_args() |
| cfg = load_config(args.config) |
| seed_everything(cfg['seed'], cfg['training']['num_threads']) |
| device = device_for(cfg['inference']['device']) |
| model, checkpoint, checkpoint_path = load_best(cfg, device) |
| splits, times, data_path = load_splits(cfg, ['test']) |
| if checkpoint['data_sha256'] != digest(data_path): |
| raise ValueError('Raw test data checksum differs from the training source') |
| u, y, ids = splits['test'] |
| start = time.monotonic() |
| summary, pred, rmse_i, relative_i, latent = evaluate( |
| model, u.to(device), y.to(device), times, cfg['inference']['batch_size'], True) |
| summary.update(split='independent_test', indices=[int(ids[0]), int(ids[-1])], |
| checkpoint_epoch=checkpoint['epoch'], checkpoint_sha256=digest(checkpoint_path), |
| checkpoint_validation_mse=checkpoint['validation_mse'], |
| selection_split=checkpoint['selection_split'], |
| paper_rmse_mean=0.00134, paper_rmse_std=0.00162, |
| elapsed_seconds=time.monotonic() - start, |
| slurm_job_id=os.environ.get('SLURM_JOB_ID')) |
| result_dir = resolve(cfg['paths']['results']) |
| result_dir.mkdir(parents=True, exist_ok=True) |
| np.savez_compressed(result_dir / 'test_predictions.npz', prediction=pred.numpy(), |
| reference=y.numpy(), initial=u.numpy(), latent=latent.numpy(), |
| sample_ids=ids, x=np.linspace(0, 1, 33), times=times.numpy()) |
| with (result_dir / 'test_per_sample.csv').open('w', newline='') as file: |
| writer = csv.writer(file) |
| writer.writerow(['sample_id', 'rmse', 'relative_l2']) |
| writer.writerows(zip(ids, rmse_i, relative_i)) |
| save_json(result_dir / 'test_metrics.json', summary) |
| update_run_manifest(cfg, status='inference_completed', test_samples=len(ids), |
| test_indices=[int(ids[0]), int(ids[-1])], |
| test_metrics_file='test_metrics.json') |
| print(f'INDEPENDENT TEST n={len(ids)} RMSE mean={summary["trajectory_rmse_mean"]:.9e} ' |
| f'std(ddof1)={summary["trajectory_rmse_std_ddof1"]:.9e} ' |
| f'global_RMSE={summary["global_rmse"]:.9e} ' |
| f'relative_L2={summary["relative_l2_mean"]:.9e}', flush=True) |
| print(f'Table 1(a) units 1e-3: {summary["trajectory_rmse_mean"] * 1000:.4f} ' |
| f'({summary["trajectory_rmse_std_ddof1"] * 1000:.4f}); paper 1.34 (1.62)', flush=True) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
