mask-ddpm/example/run_all_full.py

#!/usr/bin/env python3
"""Run full pipeline + diagnostics in one command."""

import argparse
import json
import subprocess
import sys
from pathlib import Path

from platform_utils import safe_path, is_windows


def run(cmd):
    print("running:", " ".join(cmd))
    cmd = [safe_path(arg) for arg in cmd]
    if is_windows():
        subprocess.run(cmd, check=True, shell=False)
    else:
        subprocess.run(cmd, check=True)


def parse_args():
    parser = argparse.ArgumentParser(description="Run prepare -> train -> export -> eval -> diagnostics.")
    base_dir = Path(__file__).resolve().parent
    parser.add_argument("--config", default=str(base_dir / "config.json"))
    parser.add_argument("--device", default="auto", help="cpu, cuda, or auto")
    parser.add_argument("--skip-prepare", action="store_true")
    parser.add_argument("--skip-train", action="store_true")
    parser.add_argument("--skip-export", action="store_true")
    parser.add_argument("--skip-eval", action="store_true")
    parser.add_argument("--skip-diagnose", action="store_true")
    return parser.parse_args()


def resolve_config_path(base_dir: Path, cfg_arg: str) -> Path:
    p = Path(cfg_arg)
    if p.is_absolute():
        if p.exists():
            return p.resolve()
        raise SystemExit(f"config not found: {p}")

    repo_dir = base_dir.parent
    candidates = [p, base_dir / p, repo_dir / p]
    if p.parts and p.parts[0] == "example":
        trimmed = Path(*p.parts[1:]) if len(p.parts) > 1 else Path()
        if str(trimmed):
            candidates.extend([base_dir / trimmed, repo_dir / trimmed])

    for c in candidates:
        if c.exists():
            return c.resolve()

    tried = "\n".join(str(c) for c in candidates)
    raise SystemExit(f"config not found: {cfg_arg}\ntried:\n{tried}")


def main():
    args = parse_args()
    base_dir = Path(__file__).resolve().parent
    config_path = resolve_config_path(base_dir, args.config)
    with open(config_path, "r", encoding="utf-8") as f:
        cfg = json.load(f)

    timesteps = cfg.get("timesteps", 200)
    seq_len = cfg.get("sample_seq_len", cfg.get("seq_len", 64))
    batch_size = cfg.get("sample_batch_size", cfg.get("batch_size", 2))
    clip_k = cfg.get("clip_k", 5.0)

    if not args.skip_prepare:
        run([sys.executable, str(base_dir / "prepare_data.py")])
    if not args.skip_train:
        run([sys.executable, str(base_dir / "train.py"), "--config", str(config_path), "--device", args.device])
    if not args.skip_export:
        run(
            [
                sys.executable,
                str(base_dir / "export_samples.py"),
                "--include-time",
                "--device",
                args.device,
                "--config",
                str(config_path),
                "--timesteps",
                str(timesteps),
                "--seq-len",
                str(seq_len),
                "--batch-size",
                str(batch_size),
                "--clip-k",
                str(clip_k),
                "--use-ema",
            ]
        )
    if not args.skip_eval:
        ref = cfg.get("data_glob") or cfg.get("data_path") or ""
        if ref:
            run([sys.executable, str(base_dir / "evaluate_generated.py"), "--reference", str(ref)])
        else:
            run([sys.executable, str(base_dir / "evaluate_generated.py")])
        run([sys.executable, str(base_dir / "summary_metrics.py")])
        run([sys.executable, str(base_dir / "filtered_metrics.py")])
        run([sys.executable, str(base_dir / "ranked_ks.py")])
    if not args.skip_diagnose:
        run(
            [
                sys.executable,
                str(base_dir / "diagnose_ks.py"),
                "--generated",
                str(base_dir / "results" / "generated.csv"),
                "--reference",
                str(config_path),
                "--top-k",
                "10",
            ]
        )
        run(
            [
                sys.executable,
                str(base_dir / "program_stats.py"),
                "--generated",
                str(base_dir / "results" / "generated.csv"),
                "--reference",
                str(config_path),
                "--config",
                str(config_path),
            ]
        )
        run(
            [
                sys.executable,
                str(base_dir / "controller_stats.py"),
                "--generated",
                str(base_dir / "results" / "generated.csv"),
                "--reference",
                str(config_path),
                "--config",
                str(config_path),
            ]
        )
        run(
            [
                sys.executable,
                str(base_dir / "actuator_stats.py"),
                "--generated",
                str(base_dir / "results" / "generated.csv"),
                "--reference",
                str(config_path),
                "--config",
                str(config_path),
            ]
        )
        run(
            [
                sys.executable,
                str(base_dir / "pv_stats.py"),
                "--generated",
                str(base_dir / "results" / "generated.csv"),
                "--reference",
                str(config_path),
                "--config",
                str(config_path),
            ]
        )
        run(
            [
                sys.executable,
                str(base_dir / "aux_stats.py"),
                "--generated",
                str(base_dir / "results" / "generated.csv"),
                "--reference",
                str(config_path),
                "--config",
                str(config_path),
            ]
        )
        run(
            [
                sys.executable,
                str(base_dir / "postprocess_types.py"),
                "--generated",
                str(base_dir / "results" / "generated.csv"),
                "--reference",
                str(config_path),
                "--config",
                str(config_path),
                "--out",
                str(base_dir / "results" / "generated_post.csv"),
            ]
        )


if __name__ == "__main__":
    main()