gpu/single-gpu/home_benjamin_hdbasin/code/hdbasin/experiments.py

from __future__ import annotations

import argparse
from pathlib import Path

import pandas as pd

from .core import (
    basin_slice_descent_search,
    HDBasinFlow,
    HDBasinFlowConfig,
    hybrid_hdbasin_slice_search,
    isobasinflow_search,
    optuna_tpe_search,
    random_search,
    sobol_quasi_random_search,
    successive_halving_search,
)
from .objectives import Objective, highdim_multiwell, rotated_active, synthetic_suite
from .real_ml import (
    run_ag_news,
    run_ag_news_distilbert,
    run_cifar10,
    run_fashion_mnist,
    run_imdb_distilbert,
    run_sst2_distilbert,
    run_tabular_credit,
)
from .real_data import probe_real_datasets
from .reporting import build_report, make_plots, summarize_metrics


ABLATIONS = [
    "full",
    "NoFlow",
    "NoActive",
    "NoShell",
    "NoCurvature",
    "NoLeakage",
    "NoRepair",
    "NoRandom",
    "MidpointOnly",
    "LossOnly",
    "NoIsoReentry",
    "IsoRandomDirectionsOnly",
]


def _optimizer_config(obj: Objective, seed: int, budget: int, ablation: str) -> HDBasinFlowConfig:
    return HDBasinFlowConfig(
        dim=obj.dim,
        budget=budget,
        n_initial=min(16, max(8, budget // 5)),
        k_neighbors=min(10, max(4, budget // 8)),
        seed=seed,
        ablation=ablation,
    )


def _pairwise_and_certificate_diagnostics(opt: HDBasinFlow, obj: Objective, seed: int) -> tuple[list[dict], list[dict]]:
    pairwise_rows = []
    for i, j, score, _ in opt.edges:
        xi = opt.records[i].x
        xj = opt.records[j].x
        midpoint = 0.5 * (xi + xj)
        mid_loss = float(obj.loss(midpoint))
        endpoint_best = min(opt.records[i].loss, opt.records[j].loss)
        pairwise_rows.append(
            {
                "experiment": obj.name,
                "seed": seed,
                "method": "HD-BasinFlow",
                "edge_i": i,
                "edge_j": j,
                "edge_score": score,
                "midpoint_loss": mid_loss,
                "endpoint_best_loss": endpoint_best,
                "segment_improved": bool(mid_loss < endpoint_best),
            }
        )
    cert_rows = []
    for basin in opt.basins:
        center_loss = float(obj.loss(basin.center))
        x = basin.center.copy()
        if obj.grad is not None:
            for _ in range(12):
                x = x - 0.05 * obj.grad(x)
                x = x.clip(0.0, 1.0)
        descent_loss = float(obj.loss(x))
        cert_rows.append(
            {
                "experiment": obj.name,
                "seed": seed,
                "method": "HD-BasinFlow",
                "basin_id": basin.basin_id,
                "accepted": basin.accepted,
                "inward_fraction": basin.inward_fraction,
                "coverage": basin.coverage,
                "curvature": basin.curvature,
                "leakage": basin.leakage,
                "center_loss": center_loss,
                "local_descent_loss": descent_loss,
                "local_improved": bool(descent_loss < center_loss),
                "saddle_false_acceptance": bool(obj.name == "saddle_rejection" and basin.accepted),
            }
        )
    return pairwise_rows, cert_rows


def run_one_objective(
    obj: Objective,
    seed: int,
    budget: int,
    ablations: list[str],
    baseline_methods: list[str],
    include_hdbasinflow: bool = True,
    include_basin_slice: bool = False,
    include_hybrid: bool = False,
    include_isobasinflow: bool = False,
) -> tuple[list[dict], list[dict], list[dict]]:
    rows: list[dict] = []
    pairwise_rows: list[dict] = []
    certificate_rows: list[dict] = []
    baseline_seed = seed * 1009 + len(obj.name)
    if "random" in baseline_methods:
        rows.extend(random_search(obj.dim, budget, obj.loss, obj.grad, baseline_seed, obj.optimum_loss, obj.name))
    if "sobol" in baseline_methods:
        rows.extend(sobol_quasi_random_search(obj.dim, budget, obj.loss, obj.grad, baseline_seed + 1, obj.optimum_loss, obj.name))
    if "optuna" in baseline_methods:
        rows.extend(optuna_tpe_search(obj.dim, budget, obj.loss, obj.grad, baseline_seed + 2, obj.optimum_loss, obj.name))
    if "asha" in baseline_methods:
        rows.extend(successive_halving_search(obj.dim, budget, obj.loss, obj.grad, baseline_seed + 3, obj.optimum_loss, obj.name))
    if not include_hdbasinflow:
        if include_basin_slice:
            rows.extend(
                basin_slice_descent_search(
                    obj.dim,
                    budget,
                    obj.loss,
                    obj.grad,
                    baseline_seed + 30,
                    obj.optimum_loss,
                    obj.name,
                    true_U=obj.metadata.get("true_U"),
                    certificate_loss_mode="surrogate",
                )
            )
        if include_hybrid:
            rows.extend(
                hybrid_hdbasin_slice_search(
                    obj.dim,
                    budget,
                    obj.loss,
                    obj.grad,
                    baseline_seed + 40,
                    obj.optimum_loss,
                    obj.name,
                    true_U=obj.metadata.get("true_U"),
                    certificate_loss_mode="surrogate",
                )
            )
        return rows, pairwise_rows, certificate_rows
    hdbasin_seed = baseline_seed + 20
    for ablation in ablations:
        cfg = _optimizer_config(obj, hdbasin_seed, budget, ablation)
        opt = HDBasinFlow(
            cfg,
            obj.loss,
            obj.grad,
            optimum_loss=obj.optimum_loss,
            experiment=obj.name,
            method="HD-BasinFlow",
            true_U=obj.metadata.get("true_U"),
        )
        rows.extend(opt.run())
        if ablation == "full":
            p_rows, c_rows = _pairwise_and_certificate_diagnostics(opt, obj, seed)
            pairwise_rows.extend(p_rows)
            certificate_rows.extend(c_rows)
    if include_basin_slice:
        rows.extend(
            basin_slice_descent_search(
                obj.dim,
                budget,
                obj.loss,
                obj.grad,
                baseline_seed + 30,
                obj.optimum_loss,
                obj.name,
                true_U=obj.metadata.get("true_U"),
                certificate_loss_mode="surrogate",
            )
        )
    if include_isobasinflow:
        for iso_ablation in ["full", "NoIsoReentry", "IsoRandomDirectionsOnly"]:
            rows.extend(
                isobasinflow_search(
                    obj.dim,
                    budget,
                    obj.loss,
                    obj.grad,
                    baseline_seed + 50 + len(iso_ablation),
                    obj.optimum_loss,
                    obj.name,
                    true_U=obj.metadata.get("true_U"),
                    ablation=iso_ablation,
                    certificate_loss_mode="surrogate",
                )
            )
    if include_hybrid:
        rows.extend(
            hybrid_hdbasin_slice_search(
                obj.dim,
                budget,
                obj.loss,
                obj.grad,
                baseline_seed + 40,
                obj.optimum_loss,
                obj.name,
                true_U=obj.metadata.get("true_U"),
                certificate_loss_mode="surrogate",
            )
        )
    return rows, pairwise_rows, certificate_rows


def run_synthetic(
    seeds: int,
    budget: int,
    ablations: list[str],
    baseline_methods: list[str] | None = None,
    include_hdbasinflow: bool = True,
    include_basin_slice: bool = False,
    include_isobasinflow: bool = False,
) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
    baseline_methods = baseline_methods or ["random", "sobol", "optuna", "asha"]
    all_rows: list[dict] = []
    all_pairwise: list[dict] = []
    all_certs: list[dict] = []
    for seed in range(seeds):
        for obj in synthetic_suite(seed):
            rows, pairwise_rows, certificate_rows = run_one_objective(
                obj,
                seed,
                budget,
                ablations,
                baseline_methods,
                include_hdbasinflow,
                include_basin_slice,
                include_isobasinflow=include_isobasinflow,
            )
            all_rows.extend(rows)
            all_pairwise.extend(pairwise_rows)
            all_certs.extend(certificate_rows)
    return pd.DataFrame(all_rows), pd.DataFrame(all_pairwise), pd.DataFrame(all_certs)


def run_multiwell_scaling(
    seeds: int,
    budget: int,
    dims: list[int],
    ablations: list[str],
    baseline_methods: list[str] | None = None,
    include_hdbasinflow: bool = True,
    include_basin_slice: bool = True,
    include_isobasinflow: bool = False,
) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
    baseline_methods = baseline_methods or ["random", "sobol", "optuna"]
    all_rows: list[dict] = []
    all_pairwise: list[dict] = []
    all_certs: list[dict] = []
    for seed in range(seeds):
        for dim in dims:
            obj = highdim_multiwell(dim, seed + 9000 + dim * 13)
            rows, pairwise_rows, certificate_rows = run_one_objective(
                obj,
                seed,
                budget,
                ablations,
                baseline_methods,
                include_hdbasinflow=include_hdbasinflow,
                include_basin_slice=include_basin_slice,
                include_isobasinflow=include_isobasinflow,
            )
            all_rows.extend(rows)
            all_pairwise.extend(pairwise_rows)
            all_certs.extend(certificate_rows)
    return pd.DataFrame(all_rows), pd.DataFrame(all_pairwise), pd.DataFrame(all_certs)


def run_dimension_scaling(
    seeds: int,
    budget: int,
    dims: list[int],
    active_dims: list[int],
    baseline_methods: list[str] | None = None,
    include_hdbasinflow: bool = True,
    include_basin_slice: bool = False,
    include_isobasinflow: bool = False,
) -> pd.DataFrame:
    baseline_methods = baseline_methods or ["random", "sobol"]
    rows: list[dict] = []
    for seed in range(seeds):
        for dim in dims:
            for active_dim in active_dims:
                if active_dim > dim:
                    continue
                obj = rotated_active(dim, active_dim, seed + 7000 + dim * 17 + active_dim)
                experiment = f"dimension_scaling_d{dim}_h{active_dim}"
                baseline_seed = seed * 1009 + dim * 37 + active_dim
                baseline_rows: list[dict] = []
                if "random" in baseline_methods:
                    baseline_rows.extend(random_search(dim, budget, obj.loss, obj.grad, baseline_seed, obj.optimum_loss, experiment))
                if "sobol" in baseline_methods:
                    baseline_rows.extend(sobol_quasi_random_search(dim, budget, obj.loss, obj.grad, baseline_seed + 1, obj.optimum_loss, experiment))
                if "optuna" in baseline_methods:
                    baseline_rows.extend(optuna_tpe_search(dim, budget, obj.loss, obj.grad, baseline_seed + 2, obj.optimum_loss, experiment))
                if "asha" in baseline_methods:
                    baseline_rows.extend(successive_halving_search(dim, budget, obj.loss, obj.grad, baseline_seed + 3, obj.optimum_loss, experiment))
                for row in baseline_rows:
                    row["config_json"] = row["config_json"][:-1] + f', "ambient_dim": {dim}, "true_active_dim": {active_dim}}}'
                    rows.append(row)
                if include_hdbasinflow:
                    cfg = HDBasinFlowConfig(
                        dim=dim,
                        budget=budget,
                        n_initial=min(16, max(8, budget // 5)),
                        k_neighbors=min(10, max(4, budget // 8)),
                        seed=baseline_seed + 10,
                        ablation="full",
                    )
                    opt = HDBasinFlow(cfg, obj.loss, obj.grad, obj.optimum_loss, experiment, "HD-BasinFlow", true_U=obj.metadata.get("true_U"))
                    for row in opt.run():
                        row["config_json"] = row["config_json"][:-1] + f', "ambient_dim": {dim}, "true_active_dim": {active_dim}}}'
                        rows.append(row)
                if include_basin_slice:
                    for row in basin_slice_descent_search(
                        dim,
                        budget,
                        obj.loss,
                        obj.grad,
                        baseline_seed + 30,
                        obj.optimum_loss,
                        experiment,
                        true_U=obj.metadata.get("true_U"),
                        certificate_loss_mode="surrogate",
                    ):
                        row["config_json"] = row["config_json"][:-1] + f', "ambient_dim": {dim}, "true_active_dim": {active_dim}}}'
                        rows.append(row)
                if include_isobasinflow:
                    for row in isobasinflow_search(
                        dim,
                        budget,
                        obj.loss,
                        obj.grad,
                        baseline_seed + 40,
                        obj.optimum_loss,
                        experiment,
                        true_U=obj.metadata.get("true_U"),
                        certificate_loss_mode="surrogate",
                    ):
                        row["config_json"] = row["config_json"][:-1] + f', "ambient_dim": {dim}, "true_active_dim": {active_dim}}}'
                        rows.append(row)
    return pd.DataFrame(rows)


def write_outputs(
    outdir: Path,
    evaluations: pd.DataFrame | None,
    real_status: pd.DataFrame | None,
    pairwise: pd.DataFrame | None = None,
    certificates: pd.DataFrame | None = None,
) -> None:
    raw = outdir / "raw"
    processed = outdir / "processed"
    figures = outdir / "figures"
    plot_data = outdir / "plot_data"
    reports = outdir / "reports"
    for path in [raw, processed, figures, plot_data, reports]:
        path.mkdir(parents=True, exist_ok=True)
    metrics = pd.DataFrame()
    if evaluations is not None and not evaluations.empty:
        evaluations.to_csv(raw / "evaluations.csv", index=False)
        metrics = summarize_metrics(evaluations)
        metrics.to_csv(processed / "metrics.csv", index=False)
        try:
            make_plots(evaluations, metrics, figures, plot_data, pairwise, certificates)
        except ModuleNotFoundError as exc:
            (reports / "plotting_error.txt").write_text(
                f"Plotting skipped because a plotting dependency is missing: {type(exc).__name__}: {exc}\n",
                encoding="utf-8",
            )
    if pairwise is not None and not pairwise.empty:
        pairwise.to_csv(processed / "pairwise_score_validation.csv", index=False)
    if certificates is not None and not certificates.empty:
        certificates.to_csv(processed / "certificate_validation.csv", index=False)
    if real_status is not None:
        real_status.to_csv(processed / "real_data_status.csv", index=False)
    build_report(reports / "experiment_report.md", evaluations, metrics, real_status, pairwise, certificates)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Run HD-BasinFlow validation experiments.")
    parser.add_argument("--synthetic", action="store_true", help="Run synthetic experiments.")
    parser.add_argument("--real-status", action="store_true", help="Probe Hugging Face dataset/GPU availability.")
    parser.add_argument("--ag-news", action="store_true", help="Run real AG News TF-IDF/logistic-regression HPO.")
    parser.add_argument("--ag-news-distilbert", action="store_true", help="Run real AG News DistilBERT fine-tuning HPO.")
    parser.add_argument("--fashion-mnist", action="store_true", help="Run real Fashion-MNIST tiny-CNN HPO.")
    parser.add_argument("--cifar10", action="store_true", help="Run real CIFAR-10 tiny-CNN HPO.")
    parser.add_argument("--tabular-credit", action="store_true", help="Run real tabular credit HistGradientBoosting HPO.")
    parser.add_argument("--sst2", action="store_true", help="Run real SST-2 DistilBERT fine-tuning HPO.")
    parser.add_argument("--imdb", action="store_true", help="Run real IMDB DistilBERT fine-tuning HPO.")
    parser.add_argument("--dimension-scaling", action="store_true", help="Run synthetic dimension scaling tests.")
    parser.add_argument("--multiwell-scaling", action="store_true", help="Run high-dimensional multi-well synthetic tests.")
    parser.add_argument("--basin-slice", action="store_true", help="Include Basin-Slice Descent as a competitor method.")
    parser.add_argument("--basin-slice-hybrid", action="store_true", help="Include hybrid HD-Basin + Basin-Slice as a competitor method.")
    parser.add_argument("--isobasinflow", action="store_true", help="Include IsoBasinFlow full and iso ablation methods.")
    parser.add_argument("--seeds", type=int, default=5)
    parser.add_argument("--seed-offset", type=int, default=0, help="Start real-workload seed loops at this offset.")
    parser.add_argument("--budget", type=int, default=80)
    parser.add_argument("--ag-train-size", type=int, default=1200)
    parser.add_argument("--ag-val-size", type=int, default=500)
    parser.add_argument("--ag-distilbert-train-size", type=int, default=256)
    parser.add_argument("--ag-distilbert-val-size", type=int, default=256)
    parser.add_argument("--fashion-train-size", type=int, default=800)
    parser.add_argument("--fashion-val-size", type=int, default=300)
    parser.add_argument("--fashion-epochs", type=int, default=1)
    parser.add_argument("--cifar-train-size", type=int, default=800)
    parser.add_argument("--cifar-val-size", type=int, default=300)
    parser.add_argument("--cifar-epochs", type=int, default=1)
    parser.add_argument("--tabular-train-size", type=int, default=1200)
    parser.add_argument("--tabular-val-size", type=int, default=400)
    parser.add_argument("--sst2-train-size", type=int, default=96)
    parser.add_argument("--sst2-val-size", type=int, default=96)
    parser.add_argument("--imdb-train-size", type=int, default=48)
    parser.add_argument("--imdb-val-size", type=int, default=48)
    parser.add_argument("--scaling-dims", default="2,5,10,20,50")
    parser.add_argument("--multiwell-dims", default="10,50,100")
    parser.add_argument("--scaling-active-dims", default="2,5,10")
    parser.add_argument("--outdir", type=Path, default=Path("runs/latest"))
    parser.add_argument("--ablations", default="full,NoFlow,NoActive,NoShell,NoCurvature,NoLeakage,NoRepair,NoRandom,MidpointOnly,LossOnly")
    parser.add_argument("--baselines", default="random,sobol,optuna,asha", help="Comma-separated baseline methods: random,sobol,optuna,asha.")
    parser.add_argument("--skip-hdbasinflow", action="store_true", help="Run only selected baselines for the chosen experiment groups.")
    args = parser.parse_args(argv)

    selected_ablations = [a.strip() for a in args.ablations.split(",") if a.strip()]
    bad = sorted(set(selected_ablations) - set(ABLATIONS))
    if bad:
        raise SystemExit(f"Unknown ablations: {bad}")
    selected_baselines = [a.strip() for a in args.baselines.split(",") if a.strip()]
    bad_baselines = sorted(set(selected_baselines) - {"random", "sobol", "optuna", "asha"})
    if bad_baselines:
        raise SystemExit(f"Unknown baselines: {bad_baselines}")
    frames = []
    status_frames = []
    pairwise_frames = []
    certificate_frames = []
    if args.synthetic:
        synth_eval, synth_pairwise, synth_certificates = run_synthetic(
            args.seeds,
            args.budget,
            selected_ablations,
            baseline_methods=selected_baselines,
            include_hdbasinflow=not args.skip_hdbasinflow,
            include_basin_slice=args.basin_slice,
            include_isobasinflow=args.isobasinflow,
        )
        frames.append(synth_eval)
        pairwise_frames.append(synth_pairwise)
        certificate_frames.append(synth_certificates)
    if args.ag_news:
        for seed in range(args.seed_offset, args.seed_offset + args.seeds):
            ag_eval, ag_status = run_ag_news(
                seed=seed,
                budget=args.budget,
                ablations=selected_ablations,
                train_size=args.ag_train_size,
                val_size=args.ag_val_size,
            )
            if not ag_eval.empty:
                frames.append(ag_eval)
            status_frames.append(ag_status)
    if args.ag_news_distilbert:
        for seed in range(args.seed_offset, args.seed_offset + args.seeds):
            ag_eval, ag_status = run_ag_news_distilbert(
                seed=seed,
                budget=args.budget,
                ablations=selected_ablations,
                train_size=args.ag_distilbert_train_size,
                val_size=args.ag_distilbert_val_size,
                include_basin_slice=args.basin_slice,
                include_hybrid=args.basin_slice_hybrid,
            )
            if not ag_eval.empty:
                frames.append(ag_eval)
            status_frames.append(ag_status)
    if args.fashion_mnist:
        for seed in range(args.seed_offset, args.seed_offset + args.seeds):
            fashion_eval, fashion_status = run_fashion_mnist(
                seed=seed,
                budget=args.budget,
                ablations=selected_ablations,
                train_size=args.fashion_train_size,
                val_size=args.fashion_val_size,
                epochs=args.fashion_epochs,
            )
            if not fashion_eval.empty:
                frames.append(fashion_eval)
            status_frames.append(fashion_status)
    if args.cifar10:
        for seed in range(args.seed_offset, args.seed_offset + args.seeds):
            cifar_eval, cifar_status = run_cifar10(
                seed=seed,
                budget=args.budget,
                ablations=selected_ablations,
                train_size=args.cifar_train_size,
                val_size=args.cifar_val_size,
                epochs=args.cifar_epochs,
                include_basin_slice=args.basin_slice,
                include_hybrid=args.basin_slice_hybrid,
            )
            if not cifar_eval.empty:
                frames.append(cifar_eval)
            status_frames.append(cifar_status)
    if args.tabular_credit:
        for seed in range(args.seed_offset, args.seed_offset + args.seeds):
            tabular_eval, tabular_status = run_tabular_credit(
                seed=seed,
                budget=args.budget,
                ablations=selected_ablations,
                train_size=args.tabular_train_size,
                val_size=args.tabular_val_size,
            )
            if not tabular_eval.empty:
                frames.append(tabular_eval)
            status_frames.append(tabular_status)
    if args.sst2:
        for seed in range(args.seed_offset, args.seed_offset + args.seeds):
            sst2_eval, sst2_status = run_sst2_distilbert(
                seed=seed,
                budget=args.budget,
                ablations=selected_ablations,
                train_size=args.sst2_train_size,
                val_size=args.sst2_val_size,
                include_basin_slice=args.basin_slice,
                include_hybrid=args.basin_slice_hybrid,
            )
            if not sst2_eval.empty:
                frames.append(sst2_eval)
            status_frames.append(sst2_status)
    if args.imdb:
        for seed in range(args.seed_offset, args.seed_offset + args.seeds):
            imdb_eval, imdb_status = run_imdb_distilbert(
                seed=seed,
                budget=args.budget,
                ablations=selected_ablations,
                train_size=args.imdb_train_size,
                val_size=args.imdb_val_size,
            )
            if not imdb_eval.empty:
                frames.append(imdb_eval)
            status_frames.append(imdb_status)
    if args.dimension_scaling:
        dims = [int(v) for v in args.scaling_dims.split(",") if v.strip()]
        active_dims = [int(v) for v in args.scaling_active_dims.split(",") if v.strip()]
        frames.append(
            run_dimension_scaling(
                args.seeds,
                args.budget,
                dims,
                active_dims,
                baseline_methods=selected_baselines,
                include_hdbasinflow=not args.skip_hdbasinflow,
                include_basin_slice=args.basin_slice,
                include_isobasinflow=args.isobasinflow,
            )
        )
    if args.multiwell_scaling:
        dims = [int(v) for v in args.multiwell_dims.split(",") if v.strip()]
        evals, pairwise, certs = run_multiwell_scaling(
            args.seeds,
            args.budget,
            dims,
            selected_ablations,
            baseline_methods=selected_baselines,
            include_hdbasinflow=not args.skip_hdbasinflow,
            include_basin_slice=args.basin_slice,
            include_isobasinflow=args.isobasinflow,
        )
        frames.append(evals)
        pairwise_frames.append(pairwise)
        certificate_frames.append(certs)
    evaluations = pd.concat(frames, ignore_index=True) if frames else None
    pairwise = pd.concat(pairwise_frames, ignore_index=True) if pairwise_frames else None
    certificates = pd.concat(certificate_frames, ignore_index=True) if certificate_frames else None
    real_status = pd.DataFrame(probe_real_datasets()) if args.real_status else None
    if status_frames:
        real_status = pd.concat(([real_status] if real_status is not None else []) + status_frames, ignore_index=True)
    if evaluations is None and real_status is None:
        raise SystemExit("Nothing selected. Use --synthetic, --dimension-scaling, --multiwell-scaling, --ag-news, --ag-news-distilbert, --fashion-mnist, --cifar10, --tabular-credit, --sst2, --imdb, and/or --real-status.")
    write_outputs(args.outdir, evaluations, real_status, pairwise, certificates)
    print(f"Wrote outputs to {args.outdir}")
    return 0


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