ablate.py

"""Run a registered post-hoc sparsification ablation on a trained run.

This is the registry-driven launcher for the post-hoc ablations
(``importance_prune``, ``topk_harden``, and any future ``@register_posthoc``
method). It loads a finished run, optionally estimates occupancy statistics,
sweeps the ablation's strength parameter, optionally fine-tunes after each
setting, and writes a CSV + plot - variant-agnostic and method-agnostic.

The *algorithms* live in ``sparsekan/core/posthoc.py``; this file only loads
the run, drives the sweep, and reports.

Sweep parameter by ablation:
    importance_prune : --keep-ratios   (fraction of active terms to keep)
    topk_harden      : --topk-values   (terms kept per active edge)

Example::

    python ablate.py --run-dir runs_unified/mnist_efficientkan/run0 \\
        --ablation importance_prune --keep-ratios 1.0 0.5 0.25 \\
        --finetune-epochs 2

    python ablate.py --run-dir runs_unified/mnist_efficientkan/run0 \\
        --ablation topk_harden --topk-values 8 4 2 1 --score-mode coeff
"""

import argparse
import copy
import sys
from pathlib import Path

import torch

REPO_ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(REPO_ROOT))

from sparsekan.core.data import build_loaders  # noqa: E402
from sparsekan.core.engine import (  # noqa: E402
    estimate_occupancies,
    evaluate,
    train_one_epoch,
)
from sparsekan.core.posthoc import POSTHOC_REGISTRY, apply_posthoc  # noqa: E402
from sparsekan.core.reporting import append_csv, global_cost_ratio  # noqa: E402
from sparsekan.core.utils import get_device, set_seed  # noqa: E402


def _run_args_from_config(config: dict):
    """Reconstruct the minimal arg namespace the engine helpers expect."""
    ns = argparse.Namespace()
    for k, v in config.items():
        setattr(ns, k, v)
    defaults = {
        "prune_threshold": 0.5, "base_cost": 1.0, "branch_cost": 0.0,
        "term_cost": 1.0, "grad_clip_norm": 1.0, "entropy_weight": 0.0,
        "sparsity_reg_type": "cost",
        "task": config.get("task", "classification"),
    }
    for k, v in defaults.items():
        if not hasattr(ns, k):
            setattr(ns, k, v)
    return ns


def _sweep_values(args):
    """Return (label_prefix, [values], hparam_name) for the chosen ablation."""
    if args.ablation == "importance_prune":
        return "keep", args.keep_ratios, "keep_ratio"
    if args.ablation == "topk_harden":
        return "k", [int(v) for v in args.topk_values], "k"
    raise SystemExit(f"No sweep parameter wired for ablation {args.ablation!r}.")


def _cost_ratio(model, run_args) -> float:
    """Cost ratio via sparsity_report (same call the legacy ablation used)."""
    report = model.sparsity_report(
        getattr(run_args, "prune_threshold", 0.5),
        getattr(run_args, "base_cost", 1.0),
        getattr(run_args, "branch_cost", 0.0),
        getattr(run_args, "term_cost", 1.0),
    )
    return global_cost_ratio(report)


def main() -> None:
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--run-dir", type=str, required=True)
    p.add_argument("--ablation", type=str, required=True,
                   choices=sorted(POSTHOC_REGISTRY),
                   help="Registered post-hoc ablation to run.")
    p.add_argument("--keep-ratios", type=float, nargs="+",
                   default=[1.0, 0.75, 0.5, 0.25, 0.1],
                   help="importance_prune: active-term keep fractions to sweep.")
    p.add_argument("--topk-values", type=int, nargs="+",
                   default=[8, 4, 2, 1],
                   help="topk_harden: terms kept per active edge to sweep.")
    p.add_argument("--score-mode", type=str, default="coeff",
                   choices=["gate", "coeff", "occupancy_coeff", "quant_coeff"],
                   help="topk_harden scoring mode.")
    p.add_argument("--score-quant-bits", type=int, default=0,
                   help="QAS Level 1: >0 quantizes coefficients inside the "
                        "term score at this bit-width (0 = FP32 scoring, "
                        "previous behavior). quant_coeff mode defaults to 4 "
                        "bits unless overridden.")
    p.add_argument("--finetune-epochs", type=int, default=2)
    p.add_argument("--finetune-lr", type=float, default=1e-4)
    p.add_argument("--batch-size", type=int, default=None)
    p.add_argument("--data-dir", type=str, default=None)
    p.add_argument("--num-workers", type=int, default=2)
    p.add_argument("--max-occupancy-batches", type=int, default=8)
    p.add_argument("--seed", type=int, default=42)
    args = p.parse_args()

    set_seed(args.seed)
    device = get_device()
    run_dir = Path(args.run_dir)

    # Reuse the same run loader as benchmark_cost / the legacy ablation script.
    from benchmark_cost import rebuild_from_run

    base_model, config, spec = rebuild_from_run(run_dir, device)
    run_args = _run_args_from_config(config)
    threshold = float(getattr(run_args, "prune_threshold", 0.5))

    train_loader, _, test_loader, _ = build_loaders(
        dataset=config["dataset"],
        data_dir=args.data_dir or config.get("data_dir", "./data"),
        batch_size=args.batch_size or config.get("batch_size", 128),
        num_workers=args.num_workers,
        val_size=0,
        seed=args.seed,
    )

    # Occupancy needed only for occupancy-based scoring.
    needs_occ = (
        args.ablation == "importance_prune"
        or (args.ablation == "topk_harden"
            and args.score_mode in ("occupancy_coeff", "quant_coeff"))
    )
    occupancies = (
        estimate_occupancies(base_model, train_loader, device,
                             args.max_occupancy_batches)
        if needs_occ else {}
    )

    label_prefix, values, hparam = _sweep_values(args)
    out_csv = run_dir / f"{args.ablation}_ablation.csv"

    # Baseline (no ablation).
    base_eval = evaluate(base_model, test_loader, device, run_args)
    append_csv(out_csv, {
        "ablation": args.ablation, "setting": "baseline",
        "value": 1.0 if args.ablation == "importance_prune" else 0,
        "terms_closed": 0, "test_metric": base_eval["metric"],
        "cost_ratio": _cost_ratio(base_model, run_args),
    })
    print(f"[ablate] baseline: metric={base_eval['metric']:.4f}")

    for val in values:
        model = copy.deepcopy(base_model)
        kw = {hparam: val}
        if args.ablation == "topk_harden":
            kw["score_mode"] = args.score_mode
        sqb = args.score_quant_bits or (
            4 if args.score_mode == "quant_coeff"
            and args.ablation == "topk_harden" else 0)
        if sqb > 0:
            kw["score_quant_bits"] = sqb
        closed = apply_posthoc(model, args.ablation, occupancies,
                               threshold=threshold, **kw)

        if args.finetune_epochs > 0:
            opt = torch.optim.AdamW(model.parameters(), lr=args.finetune_lr)
            zero = {"base": 0.0, "branch": 0.0, "term": 0.0}
            for _ in range(args.finetune_epochs):
                train_one_epoch(model, train_loader, opt, device, zero, run_args)

        ev = evaluate(model, test_loader, device, run_args)
        row = {
            "ablation": args.ablation, "setting": f"{label_prefix}_{val}",
            "value": val, "terms_closed": closed,
            "test_metric": ev["metric"],
            "cost_ratio": _cost_ratio(model, run_args),
        }
        append_csv(out_csv, row)
        print(f"[ablate] {label_prefix}={val}: closed={closed} "
              f"metric={ev['metric']:.4f} cost_ratio={row['cost_ratio']:.4f}")

    print(f"[ablate] CSV: {out_csv}")


if __name__ == "__main__":
    main()