gpu/single-gpu/home_benjamin_hdbasin/code/tests/test_target_savings.py

import pandas as pd

from hdbasin.target_savings import TargetSavingsConfig, build_target_savings_decisions


def _row(experiment, method, seed, iteration, best_loss, cost, walltime, ablation="baseline"):
    return {
        "experiment": experiment,
        "method": method,
        "ablation": ablation,
        "seed": seed,
        "iteration": iteration,
        "best_loss_so_far": best_loss,
        "cumulative_cost": cost,
        "walltime": walltime,
    }


def test_target_savings_passes_when_hd_reaches_baseline_quality_early():
    rows = []
    for seed in range(3):
        rows.extend(
            [
                _row("toy", "Random Search", seed, 0, 2.0, 1, 1.0),
                _row("toy", "Random Search", seed, 1, 1.0, 10, 1.0),
                _row("toy", "HD-BasinFlow", seed, 0, 2.0, 1, 1.0, "full"),
                _row("toy", "HD-BasinFlow", seed, 1, 1.005, 6, 1.0, "full"),
            ]
        )
    decisions = build_target_savings_decisions(
        pd.DataFrame(rows),
        TargetSavingsConfig(quality_tolerance=0.01, min_savings=0.15, target_savings=0.20, max_savings=0.30, min_runs=3),
    )
    row = decisions.iloc[0]
    assert row["decision"] == "pass"
    assert row["quality_pass"]
    assert row["operating_point_pass"]
    assert row["evaluation_savings_fraction"] == 0.4


def test_target_savings_fails_when_quality_misses_tolerance():
    rows = []
    for seed in range(3):
        rows.extend(
            [
                _row("toy", "Optuna TPE", seed, 0, 1.0, 10, 1.0),
                _row("toy", "HD-BasinFlow", seed, 0, 1.05, 5, 1.0, "full"),
            ]
        )
    decisions = build_target_savings_decisions(pd.DataFrame(rows), TargetSavingsConfig(min_runs=3))
    row = decisions.iloc[0]
    assert row["decision"] == "fail"
    assert not row["quality_pass"]