| """Reload the selected checkpoint and evaluate all 100 original test cases.""" |
| from __future__ import annotations |
|
|
| import csv |
| import time |
|
|
| import torch |
|
|
| from common import (DarcyEnergy, arguments, build_model, configuration, error_statistics, |
| evaluate, initialize, output_path, read_data, write_json) |
|
|
|
|
| def main(): |
| parser = arguments(__doc__) |
| parser.add_argument("--checkpoint", type=str) |
| args = parser.parse_args() |
| cfg = configuration(args) |
| device = initialize(cfg) |
| fields, metadata = read_data(cfg) |
| checkpoint_path = output_path(args.checkpoint or cfg["output"]["checkpoint"]) |
| checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False) |
| if checkpoint["data_sha256"] != metadata["sha256"]: |
| raise ValueError("The checkpoint was not trained using this exact dataset") |
| if checkpoint["config"]["model"] != cfg["model"]: |
| raise ValueError("Model configuration differs from checkpoint") |
| model = build_model(cfg, device) |
| model.load_state_dict(checkpoint["model_state_dict"], strict=True) |
| energy = DarcyEnergy(cfg["data"]["resolution"], cfg["physics"]["source"]).to(device) |
| if device.type == "cuda": |
| torch.cuda.synchronize() |
| started = time.perf_counter() |
| prediction, errors, energies = evaluate(model, fields["P_test"].to(device), fields["U_test"].to(device), energy, cfg["training"]["batch_size"]) |
| if device.type == "cuda": |
| torch.cuda.synchronize() |
| inference_seconds = time.perf_counter() - started |
| boundary = max(prediction[:, 0].abs().max().item(), prediction[:, -1].abs().max().item(), |
| prediction[:, :, 0].abs().max().item(), prediction[:, :, -1].abs().max().item()) |
| if boundary != 0 or not torch.isfinite(prediction).all(): |
| raise ValueError("Prediction failed zero-boundary or finiteness check") |
| metrics = {**error_statistics(errors), "checkpoint": str(checkpoint_path), "checkpoint_epoch": checkpoint["epoch"], |
| "test_energy_mean": energies.double().mean().item(), "boundary_max_abs": boundary, |
| "inference_seconds": inference_seconds, "data_sha256": metadata["sha256"], |
| "checkpoint_selection": "minimum test-set mean relative L2 across 500 epochs (user requested)", |
| "paper_table1_mean_percent": cfg["reference"]["table1_mean_percent"], |
| "paper_table1_std_percent": cfg["reference"]["table1_std_percent"]} |
| results = output_path(cfg["output"]["results_dir"]) |
| results.mkdir(parents=True, exist_ok=True) |
| torch.save({"raw_permeability": fields["F_test"], "nodal_permeability": fields["P_test"], |
| "reference": fields["U_test"], "prediction": prediction, "relative_l2": errors, |
| "energy": energies, "checkpoint_epoch": checkpoint["epoch"], "data_sha256": metadata["sha256"]}, results / "test_predictions.pt") |
| with (results / "per_sample_metrics.csv").open("w", newline="") as stream: |
| writer = csv.writer(stream) |
| writer.writerow(["test_index", "relative_l2", "relative_l2_percent", "energy"]) |
| for index, (error, value) in enumerate(zip(errors.tolist(), energies.tolist())): |
| writer.writerow([index, error, 100 * error, value]) |
| print(f"Test {index:03d}: relative_L2={error:.8f} ({100*error:.4f}%)", flush=True) |
| write_json(results / "metrics.json", metrics) |
| print(f"Best epoch {checkpoint['epoch']}: test relative L2 = {metrics['relative_l2_mean_percent']:.6f}% " |
| f"+/- {metrics['relative_l2_std_percent']:.6f}% (sample std, n=100); boundary={boundary}", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
