| """Load the selected checkpoint and query the three paper snapshot grids.""" |
| from pathlib import Path |
| import sys |
| import time |
|
|
| import numpy as np |
| import torch |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parents[1])) |
| from model import PINNShallowWater |
| from common import arguments, device_for, errors, load_config, output_path, runtime_info, save_json, sha256 |
| from fake_data import initial_values |
| from train import to_device, validation, weights |
|
|
|
|
| def main(): |
| args = arguments(__doc__).parse_args() |
| config, root = load_config(args.config) |
| device = device_for(config) |
| torch.set_num_threads(8) |
| checkpoint_path = output_path(root, config['paths']['best_checkpoint']) |
| checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=False) |
| if checkpoint.get('completed_iterations') != 5000 or not checkpoint.get('formal_run_complete'): |
| raise ValueError('Formal 5000-iteration training has not completed.') |
| if checkpoint['config'] != config or checkpoint['model_class'] != 'PINNShallowWater': |
| raise ValueError('Checkpoint/config/model identity mismatch.') |
| model = PINNShallowWater(**config['model']).to(device) |
| model.load_state_dict(checkpoint['model_state_dict'], strict=True) |
| model.eval() |
| with np.load(output_path(root, config['paths']['grid'])) as grid: |
| x, y, xy, times = [grid[k] for k in ['x', 'y', 'xy', 'times']] |
| batch_size = config['evaluation']['inference_batch_size'] |
| snapshots = [] |
| started = time.monotonic() |
| with torch.no_grad(): |
| for t in times: |
| coordinates = np.column_stack((xy, np.full(len(xy), t))).astype(np.float32) |
| batches = [] |
| for begin in range(0, len(coordinates), batch_size): |
| prediction = model(torch.as_tensor(coordinates[begin:begin + batch_size], device=device)) |
| batches.append(prediction.cpu().numpy()) |
| field = np.concatenate(batches).reshape(len(y), len(x), 3) |
| if not np.isfinite(field).all(): |
| raise FloatingPointError(f'Nonfinite prediction at t={t}') |
| snapshots.append(field) |
| print('INFERENCE_SNAPSHOT', {'t': float(t), 'shape': field.shape, |
| 'min_huv': field.reshape(-1, 3).min(axis=0).tolist(), |
| 'max_huv': field.reshape(-1, 3).max(axis=0).tolist()}, flush=True) |
| fields = np.stack(snapshots) |
| path = output_path(root, config['paths']['predictions']) |
| np.savez_compressed(path, x=x, y=y, times=times, fields=fields, |
| variables=np.array(['h', 'u', 'v']), units=np.array(['m', 'm/s', 'm/s'])) |
| initial_errors = errors(fields[0].reshape(-1, 3), initial_values(xy, config['physics'])) |
| samples_path = output_path(root, config['paths']['samples']) |
| if checkpoint['sampling_sha256'] != sha256(samples_path): |
| raise ValueError('Checkpoint sampling hash mismatch.') |
| with np.load(samples_path) as arrays: |
| val = to_device(arrays, 'val_', device) |
| check = validation(model, val, config, weights(config, device)) |
| metadata = {'status': 'completed', 'checkpoint': str(checkpoint_path), |
| 'checkpoint_sha256': sha256(checkpoint_path), 'selected_iteration': checkpoint['iteration'], |
| 'completed_training_iterations': 5000, 'state_dict_strict_load': True, |
| 'fields_shape': list(fields.shape), 'all_finite': True, 'initial_grid_errors': initial_errors, |
| 'best_checkpoint_validation': check, 'runtime': runtime_info(device), |
| 'elapsed_seconds': time.monotonic() - started, 'predictions_sha256': sha256(path)} |
| save_json(path.parent / 'inference_summary.json', metadata) |
| print('INFERENCE_COMPLETE', metadata, flush=True) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
