File size: 1,677 Bytes
a3441f1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 | from __future__ import annotations
import argparse
import json
from pathlib import Path
import numpy as np
from common import DEFAULT_CONFIG, load_config, prepare_config
def main() -> int:
parser = argparse.ArgumentParser(description="Summarize PDENNEval FNO inference outputs.")
parser.add_argument("--config", default=str(DEFAULT_CONFIG), help="Path to conf/config.yaml")
parser.add_argument("--prediction-dir", default=None)
parser.add_argument("--output-dir", default=None)
args = parser.parse_args()
cfg = prepare_config(load_config(args.config))
prediction_dir = Path(args.prediction_dir or cfg.result.prediction_dir).expanduser().resolve()
output_dir = Path(args.output_dir or cfg.result.output_dir).expanduser().resolve()
files = sorted(prediction_dir.glob("*.npz"))
if not files:
raise FileNotFoundError(f"no prediction files found in {prediction_dir}")
mse_values = []
mae_values = []
for path in files:
data = np.load(path)
pred = data["prediction"]
target = data["target"]
diff = pred - target
mse_values.append(float(np.mean(diff ** 2)))
mae_values.append(float(np.mean(np.abs(diff))))
metrics = {
"num_files": len(files),
"mse": float(np.mean(mse_values)),
"mae": float(np.mean(mae_values)),
}
output_dir.mkdir(parents=True, exist_ok=True)
metrics_path = output_dir / "metrics.json"
metrics_path.write_text(json.dumps(metrics, indent=2), encoding="utf-8")
print(json.dumps(metrics, indent=2))
print(f"wrote {metrics_path}")
return 0
if __name__ == "__main__":
raise SystemExit(main())
|