scripts/06_fusion_ablation.py

#!/usr/bin/env python
"""Stage 06 — fusion ablation.

Trains and evaluates all three fusion strategies (concat, gated,
cross_attention) at matched hyperparameters, using the SAME cached
embeddings from stage 02 (no re-encoding, no new GPU-heavy work). Produces
a comparison report so cross-attention's benefit over simpler fusion can
be quantified rather than just asserted.

Each variant gets its own checkpoint_dir (checkpoints/ablation_<type>) and
its own generated config (configs/ablation/<type>.yaml) so runs never
overwrite each other and remain individually reproducible.

Usage:
    python scripts/06_fusion_ablation.py --config configs/base.yaml
"""

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

import yaml

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from src.utils.config import load_config
from src.utils.exceptions import PricePredictorError
from src.utils.logging import get_logger

logger = get_logger(__name__)

FUSION_TYPES = ["concat", "gated", "cross_attention"]


def _run_subprocess(cmd: list) -> None:
    result = subprocess.run(cmd)
    if result.returncode != 0:
        raise PricePredictorError(f"Command failed (exit {result.returncode}): {' '.join(cmd)}")


def run(base_config_path: str, output_report: str) -> dict:
    base_config = load_config(base_config_path)  # validates the base config once up front

    repo_root = Path(__file__).resolve().parents[1]
    ablation_configs_dir = repo_root / "configs" / "ablation"
    ablation_configs_dir.mkdir(parents=True, exist_ok=True)

    results = {}
    for fusion_type in FUSION_TYPES:
        logger.info("=== Ablation run: fusion.type=%s ===", fusion_type)

        config = copy.deepcopy(base_config)
        config["fusion"]["type"] = fusion_type
        config["checkpoint_dir"] = f"checkpoints/ablation_{fusion_type}"

        config_path = ablation_configs_dir / f"{fusion_type}.yaml"
        with config_path.open("w") as f:
            yaml.safe_dump(config, f)

        _run_subprocess([sys.executable, "scripts/03_train.py", "--config", str(config_path)])
        _run_subprocess([sys.executable, "scripts/04_evaluate.py", "--config", str(config_path)])

        eval_report_path = repo_root / "reports" / f"ablation_{fusion_type}" / "eval_report.json"
        if not eval_report_path.exists():
            raise PricePredictorError(
                f"Expected eval report not found at {eval_report_path} — "
                "stage 04 may have failed silently for this fusion type."
            )
        with eval_report_path.open() as f:
            metrics = json.load(f)

        results[fusion_type] = metrics
        logger.info("Result for %s: SMAPE=%.4f MAE=%.4f", fusion_type, metrics["smape"], metrics["mae"])

    output_path = Path(output_report)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with output_path.open("w") as f:
        json.dump(results, f, indent=2)

    best_fusion = min(results, key=lambda k: results[k]["smape"])
    logger.info("=== Fusion ablation comparison (lower SMAPE is better) ===")
    for fusion_type, metrics in results.items():
        marker = "  <-- best" if fusion_type == best_fusion else ""
        logger.info("%-16s SMAPE=%.4f  MAE=%.4f%s", fusion_type, metrics["smape"], metrics["mae"], marker)
    logger.info("Wrote comparison report to %s", output_path)

    return results


def main() -> None:
    parser = argparse.ArgumentParser(description="Stage 06: train+evaluate all fusion strategies for comparison")
    parser.add_argument("--config", default="configs/base.yaml")
    parser.add_argument("--output", default="reports/fusion_ablation_comparison.json")
    args = parser.parse_args()

    try:
        run(args.config, args.output)
    except PricePredictorError as e:
        logger.error("Fusion ablation failed: %s", e)
        sys.exit(1)
    except Exception as e:
        logger.exception("Unexpected error during fusion ablation: %s", e)
        sys.exit(1)


if __name__ == "__main__":
    main()