| |
| """Evaluate a trained CNO on the paper's ID and OOD Navier--Stokes sets.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import sys |
| from pathlib import Path |
| from typing import Any |
|
|
| import numpy as np |
| import torch |
| from torch.utils.data import DataLoader |
|
|
|
|
| PROJECT_ROOT = Path(__file__).resolve().parents[1] |
| if str(PROJECT_ROOT) not in sys.path: |
| sys.path.insert(0, str(PROJECT_ROOT)) |
|
|
| from models.FNO import build_model |
| from scripts.common import ( |
| MinMaxNormalizer, |
| NavierStokesH5Dataset, |
| atomic_json_dump, |
| atomic_npz_save, |
| data_file, |
| load_config, |
| numeric_sample_ids, |
| project_path, |
| relative_l1_per_sample, |
| select_device, |
| set_reproducibility, |
| ) |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument( |
| "--config", default=str(PROJECT_ROOT / "config" / "config.yaml") |
| ) |
| parser.add_argument("--checkpoint", default=None) |
| parser.add_argument("--device", default=None) |
| parser.add_argument("--batch-size", type=int, default=None) |
| parser.add_argument("--output-dir", default=None) |
| parser.add_argument( |
| "--splits", |
| nargs="+", |
| choices=("id", "ood"), |
| default=("id", "ood"), |
| ) |
| return parser.parse_args() |
|
|
|
|
| def load_trained_model( |
| config: dict[str, Any], checkpoint_path: Path, device: torch.device |
| ) -> tuple[torch.nn.Module, MinMaxNormalizer, dict[str, Any]]: |
| if not checkpoint_path.is_file(): |
| raise FileNotFoundError(f"checkpoint not found: {checkpoint_path}") |
| checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False) |
| required = { |
| "model_state_dict", |
| "normalization", |
| "epoch", |
| "best_val_relative_l1", |
| } |
| missing = sorted(required.difference(checkpoint)) |
| if missing: |
| raise KeyError(f"checkpoint is missing required keys: {missing}") |
| checkpoint_config = checkpoint.get("config", {}) |
| if checkpoint_config and checkpoint_config.get("model") != config["model"]: |
| raise ValueError("checkpoint model configuration differs from config.yaml") |
| model = build_model(config["model"]).to(device) |
| model.load_state_dict(checkpoint["model_state_dict"], strict=True) |
| model.eval() |
| normalizer = MinMaxNormalizer.from_state(checkpoint["normalization"]) |
| return model, normalizer, checkpoint |
|
|
|
|
| @torch.inference_mode() |
| def evaluate_split( |
| split_name: str, |
| model: torch.nn.Module, |
| loader: DataLoader, |
| normalizer: MinMaxNormalizer, |
| device: torch.device, |
| epsilon: float, |
| ) -> tuple[dict[str, np.ndarray], dict[str, Any]]: |
| all_ids: list[np.ndarray] = [] |
| all_inputs: list[np.ndarray] = [] |
| all_targets: list[np.ndarray] = [] |
| all_predictions: list[np.ndarray] = [] |
| all_ratios: list[np.ndarray] = [] |
| completed = 0 |
| total = len(loader.dataset) |
|
|
| for batch_index, (inputs, targets, sample_ids) in enumerate(loader, start=1): |
| inputs = inputs.to(device, non_blocking=True) |
| targets = targets.to(device, non_blocking=True) |
| predictions = model(inputs) |
| physical_inputs = normalizer.denormalize_input(inputs) |
| physical_targets = normalizer.denormalize_output(targets) |
| physical_predictions = normalizer.denormalize_output(predictions) |
| ratios = relative_l1_per_sample( |
| physical_predictions, physical_targets, epsilon |
| ) |
|
|
| all_ids.append(np.asarray(sample_ids, dtype=np.int64)) |
| all_inputs.append(physical_inputs.cpu().numpy().astype(np.float32)) |
| all_targets.append(physical_targets.cpu().numpy().astype(np.float32)) |
| all_predictions.append(physical_predictions.cpu().numpy().astype(np.float32)) |
| all_ratios.append(ratios.cpu().numpy().astype(np.float64)) |
| completed += inputs.shape[0] |
| running = np.concatenate(all_ratios) * 100.0 |
| print( |
| f"inference split={split_name} batch={batch_index}/{len(loader)} " |
| f"samples={completed}/{total} running_rel_l1_median={np.median(running):.6f}%", |
| flush=True, |
| ) |
|
|
| arrays = { |
| "sample_ids": np.concatenate(all_ids), |
| "inputs": np.concatenate(all_inputs), |
| "targets": np.concatenate(all_targets), |
| "predictions": np.concatenate(all_predictions), |
| "relative_l1": np.concatenate(all_ratios), |
| } |
| percentages = arrays["relative_l1"] * 100.0 |
| metrics = { |
| "sample_count": int(percentages.size), |
| "resolution": [int(arrays["inputs"].shape[-2]), int(arrays["inputs"].shape[-1])], |
| "relative_l1_median_percent": float(np.median(percentages)), |
| "relative_l1_mean_percent": float(np.mean(percentages)), |
| "relative_l1_std_percent": float(np.std(percentages)), |
| "relative_l1_min_percent": float(np.min(percentages)), |
| "relative_l1_max_percent": float(np.max(percentages)), |
| } |
| return arrays, metrics |
|
|
|
|
| def main() -> None: |
| args = parse_args() |
| config = load_config(args.config) |
| inference_config = config["inference"] |
| device = select_device(args.device or str(inference_config["device"])) |
| batch_size = int( |
| args.batch_size if args.batch_size is not None else inference_config["batch_size"] |
| ) |
| if batch_size < 1: |
| raise ValueError("batch_size must be positive") |
| seed = int(config["experiment"]["seed"]) |
| set_reproducibility(seed, bool(config["experiment"].get("deterministic", True))) |
|
|
| checkpoint_path = ( |
| Path(args.checkpoint).expanduser().resolve() |
| if args.checkpoint |
| else project_path(config["paths"]["checkpoint"]) |
| ) |
| output_dir = ( |
| Path(args.output_dir).expanduser().resolve() |
| if args.output_dir |
| else project_path(config["paths"]["results_dir"]) |
| ) |
| output_dir.mkdir(parents=True, exist_ok=True) |
| model, normalizer, checkpoint = load_trained_model(config, checkpoint_path, device) |
|
|
| split_specs = { |
| "id": ("id_test_file", "test_id"), |
| "ood": ("ood_test_file", "test_ood"), |
| } |
| all_metrics: dict[str, Any] = { |
| "schema_version": "cno-navier-stokes-metrics-v1", |
| "checkpoint": str(checkpoint_path), |
| "checkpoint_epoch": int(checkpoint["epoch"]), |
| "best_validation_relative_l1_percent": float( |
| checkpoint["best_val_relative_l1"] |
| ), |
| "normalization": normalizer.state_dict(), |
| "paper_reference": config.get("paper_reference", {}), |
| "splits": {}, |
| } |
| print( |
| f"inference device={device} checkpoint_epoch={checkpoint['epoch']} " |
| f"best_val={float(checkpoint['best_val_relative_l1']):.6f}%", |
| flush=True, |
| ) |
|
|
| for split_name in args.splits: |
| filename_key, split_key = split_specs[split_name] |
| dataset = NavierStokesH5Dataset( |
| data_file(config, filename_key), |
| numeric_sample_ids(config["data"][split_key]), |
| normalizer, |
| str(config["data"]["input_key"]), |
| str(config["data"]["output_key"]), |
| ) |
| loader = DataLoader( |
| dataset, |
| batch_size=batch_size, |
| shuffle=False, |
| num_workers=int(inference_config["num_workers"]), |
| pin_memory=device.type == "cuda", |
| persistent_workers=int(inference_config["num_workers"]) > 0, |
| ) |
| arrays, metrics = evaluate_split( |
| split_name, |
| model, |
| loader, |
| normalizer, |
| device, |
| float(inference_config["metric_epsilon"]), |
| ) |
| artifact_path = output_dir / f"{split_name}_predictions.npz" |
| atomic_npz_save(artifact_path, **arrays) |
| metrics["predictions_file"] = str(artifact_path) |
| all_metrics["splits"][split_name] = metrics |
| print( |
| f"evaluation split={split_name} n={metrics['sample_count']} " |
| f"rel_l1_median={metrics['relative_l1_median_percent']:.6f}% " |
| f"mean={metrics['relative_l1_mean_percent']:.6f}% " |
| f"std={metrics['relative_l1_std_percent']:.6f}% " |
| f"saved={artifact_path}", |
| flush=True, |
| ) |
|
|
| metrics_path = output_dir / "metrics.json" |
| atomic_json_dump(all_metrics, metrics_path) |
| print(f"metrics saved path={metrics_path}", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
