airfRANS-model-exploration/src/airfrans_frontier/cli.py

from __future__ import annotations

import argparse
import sys

from airfrans_frontier.paths import DEFAULT_RAW_DATA_DIR, DEFAULT_RAW_MANIFEST_PATH, resolve_path
from airfrans_frontier.raw.inspect import format_raw_inspection, inspect_raw_subset


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="airfrans-frontier")
    subparsers = parser.add_subparsers(required=True)

    inspect_raw = subparsers.add_parser("inspect-raw", help="inspect the local raw AirfRANS subset")
    inspect_raw.add_argument("--data-dir", default=str(DEFAULT_RAW_DATA_DIR))
    inspect_raw.add_argument("--manifest", default=str(DEFAULT_RAW_MANIFEST_PATH))
    inspect_raw.add_argument("--sample-limit", type=int, default=5)
    inspect_raw.set_defaults(command="inspect-raw")

    process_raw = subparsers.add_parser("process-raw", help="convert raw OpenFOAM cases into training tensors")
    process_raw.add_argument("--raw-dir", default=str(DEFAULT_RAW_DATA_DIR))
    process_raw.add_argument("--output-dir", default="data/processed/full")
    process_raw.add_argument("--limit", type=int)
    process_raw.add_argument("--force", action="store_true")
    process_raw.set_defaults(command="process-raw")

    train = subparsers.add_parser("train", help="train a configured baseline model")
    train.add_argument("config", help="path to a training config TOML file")
    train.add_argument("--resume", help="path to checkpoint_latest.pt to resume from")
    train.set_defaults(command="train")

    return parser


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)

    if args.command == "inspect-raw":
        if args.sample_limit < 0:
            print("error: --sample-limit must be non-negative", file=sys.stderr)
            return 1

        data_dir = resolve_path(args.data_dir)
        manifest_path = resolve_path(args.manifest)
        try:
            report = inspect_raw_subset(data_dir, manifest_path)
        except (FileNotFoundError, NotADirectoryError, ValueError) as exc:
            print(f"error: {exc}", file=sys.stderr)
            return 1

        print(format_raw_inspection(report, sample_limit=args.sample_limit))
        return 0 if report.matches_manifest else 1

    if args.command == "process-raw":
        if args.limit is not None and args.limit <= 0:
            print("error: --limit must be positive", file=sys.stderr)
            return 1
        from airfrans_frontier.raw.process import process_raw_dataset

        try:
            result = process_raw_dataset(
                resolve_path(args.raw_dir),
                resolve_path(args.output_dir),
                limit=args.limit,
                force=args.force,
            )
        except (FileNotFoundError, NotADirectoryError, ValueError) as exc:
            print(f"error: {exc}", file=sys.stderr)
            return 1
        print(f"output_dir: {result.output_dir}")
        print(f"case_count: {result.case_count}")
        print(f"total_points: {result.total_points}")
        print(f"manifest: {result.manifest_path}")
        return 0

    if args.command == "train":
        from airfrans_frontier.runtime import remove_pythonpath_entries

        remove_pythonpath_entries()
        from airfrans_frontier.training.loop import train_from_config_path

        try:
            result = train_from_config_path(resolve_path(args.config), resume_path=resolve_path(args.resume) if args.resume else None)
        except (FileNotFoundError, NotADirectoryError, ValueError, RuntimeError) as exc:
            print(f"error: {exc}", file=sys.stderr)
            return 1

        print(f"run_dir: {result.run_dir}")
        print(f"final_metrics: {result.run_dir / 'final_metrics.json'}")
        return 0

    parser.error(f"unknown command: {args.command}")
    return 2


if __name__ == "__main__":
    raise SystemExit(main())