pcb-afm-simulation-system/scripts/test_strategy_surrogate.py

"""P5-M4: unit tests for SurrogateGuidedStrategy.

Covers: happy path (initial LHS + surrogate-guided convergence on bowl
function), boundary (empty params / budget exhaustion / n_initial > budget),
anomaly (unknown point_id / missing objective), null inputs, budget-adaptive
batch sizing, both maximize and minimize directions, registry integration,
and state() field contract.

All source is ASCII only. Run: python scripts/test_strategy_surrogate.py
exit 0 = PASS.
"""
import os
import sys
import unittest

_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, os.path.join(_ROOT, "src"))

from afmcore.strategies import (  # noqa: E402
    SurrogateGuidedStrategy,
    get_strategy,
    is_registered,
    list_strategy_kinds,
)


def _bowl(params):
    """Convex bowl: y = (a-0.5)^2 + (b-0.5)^2, minimum at (0.5, 0.5)."""
    return (params.get("a", 0.0) - 0.5) ** 2 + (params.get("b", 0.0) - 0.5) ** 2


def _hill(params):
    """Inverse bowl: y = 1 - bowl, maximum at (0.5, 0.5)."""
    return 1.0 - _bowl(params)


class TestSurrogateHappyPath(unittest.TestCase):
    def test_initial_lhs_batch(self):
        s = SurrogateGuidedStrategy(
            parameters=[{"name": "a", "min_value": 0, "max_value": 1},
                        {"name": "b", "min_value": 0, "max_value": 1}],
            objective_metric="y", objective_direction="minimize",
            n_initial=8, batch_size=4, budget=20, rng_seed=1,
        )
        self.assertEqual(s.state()["phase"], "initial")
        batch = s.select_next(1000)
        self.assertEqual(len(batch), 8)
        for p in batch:
            self.assertIn("point_id", p)
            self.assertIn("params", p)
            self.assertIn("a", p["params"])
            self.assertIn("b", p["params"])

    def test_convergence_on_bowl_minimize(self):
        s = SurrogateGuidedStrategy(
            parameters=[{"name": "a", "min_value": 0, "max_value": 1},
                        {"name": "b", "min_value": 0, "max_value": 1}],
            objective_metric="y", objective_direction="minimize",
            n_initial=10, batch_size=3, max_batch_size=5, budget=40,
            n_candidates=60, rng_seed=2,
        )
        # run full budget
        used = 0
        while used < 40:
            batch = s.select_next()
            if not batch:
                break
            for p in batch:
                s.report(p["point_id"], {"y": _bowl(p["params"])}, "ok")
                used += 1
        best = min(r["metrics"]["y"] for r in s._done.values() if r["status"] == "ok")
        # should find a point reasonably close to the minimum (0.0)
        self.assertLess(best, 0.05)
        self.assertEqual(s.state()["phase"], "surrogate_guided")

    def test_maximize_direction(self):
        s = SurrogateGuidedStrategy(
            parameters=[{"name": "a", "min_value": 0, "max_value": 1},
                        {"name": "b", "min_value": 0, "max_value": 1}],
            objective_metric="y", objective_direction="maximize",
            n_initial=8, batch_size=3, budget=25, n_candidates=50, rng_seed=3,
        )
        used = 0
        while used < 25:
            batch = s.select_next()
            if not batch:
                break
            for p in batch:
                s.report(p["point_id"], {"y": _hill(p["params"])}, "ok")
                used += 1
        best = max(r["metrics"]["y"] for r in s._done.values() if r["status"] == "ok")
        # hill maximum is 1.0
        self.assertGreater(best, 0.95)


class TestSurrogateBoundary(unittest.TestCase):
    def test_empty_parameters(self):
        s = SurrogateGuidedStrategy(parameters=[], n_initial=5, budget=10, objective_metric="y")
        batch = s.select_next(1000)
        self.assertEqual(batch, [])
        self.assertTrue(s.is_converged())

    def test_none_parameters(self):
        s = SurrogateGuidedStrategy(parameters=None, n_initial=5, budget=10, objective_metric="y")
        self.assertEqual(s.select_next(1000), [])

    def test_budget_exhaustion(self):
        s = SurrogateGuidedStrategy(
            parameters=[{"name": "a", "min_value": 0, "max_value": 1}],
            objective_metric="y", objective_direction="minimize",
            n_initial=3, batch_size=2, budget=5, rng_seed=4,
        )
        used = 0
        while True:
            batch = s.select_next()
            if not batch:
                break
            for p in batch:
                s.report(p["point_id"], {"y": p["params"]["a"] ** 2}, "ok")
                used += 1
        self.assertLessEqual(used, 5)
        self.assertTrue(s.is_converged())
        self.assertEqual(s.state()["phase"], "exhausted")

    def test_n_initial_greater_than_budget(self):
        s = SurrogateGuidedStrategy(
            parameters=[{"name": "a", "min_value": 0, "max_value": 1}],
            objective_metric="y", n_initial=20, batch_size=4, budget=5, rng_seed=5,
        )
        batch = s.select_next(1000)
        # initial pending is 20 but budget is 5; select_next serves pending
        # regardless (pending points were already generated)
        self.assertEqual(len(batch), 20)

    def test_adaptive_batch_size_within_bounds(self):
        s = SurrogateGuidedStrategy(
            parameters=[{"name": "a", "min_value": 0, "max_value": 1},
                        {"name": "b", "min_value": 0, "max_value": 1}],
            objective_metric="y", objective_direction="minimize",
            n_initial=6, batch_size=2, max_batch_size=6, budget=30,
            n_candidates=40, rng_seed=6,
        )
        # initial
        batch = s.select_next(1000)
        for p in batch:
            s.report(p["point_id"], {"y": _bowl(p["params"])}, "ok")
        # surrogate batches
        for _ in range(4):
            b = s.select_next()
            if not b:
                break
            self.assertGreaterEqual(len(b), 1)
            self.assertLessEqual(len(b), 6)
            for p in b:
                s.report(p["point_id"], {"y": _bowl(p["params"])}, "ok")


class TestSurrogateAnomaly(unittest.TestCase):
    def test_report_unknown_point_id_no_crash(self):
        s = SurrogateGuidedStrategy(
            parameters=[{"name": "a", "min_value": 0, "max_value": 1}],
            objective_metric="y", n_initial=3, budget=10, rng_seed=7,
        )
        s.select_next(1000)
        s.report(999999, {"y": 1.0}, "ok")  # should not raise
        # state() must remain callable and well-formed
        st = s.state()
        self.assertIn("surrogate", st)

    def test_missing_objective_metric_excluded_from_training(self):
        s = SurrogateGuidedStrategy(
            parameters=[{"name": "a", "min_value": 0, "max_value": 1}],
            objective_metric="nonexistent", n_initial=3, budget=10, rng_seed=8,
        )
        batch = s.select_next(1000)
        for p in batch:
            s.report(p["point_id"], {"y": 0.5}, "ok")  # wrong metric
        # surrogate should have 0 training points (no objective values)
        train = s._training_data()
        self.assertEqual(len(train), 0)

    def test_failed_status_excluded_from_training(self):
        s = SurrogateGuidedStrategy(
            parameters=[{"name": "a", "min_value": 0, "max_value": 1}],
            objective_metric="y", n_initial=4, budget=10, rng_seed=9,
        )
        batch = s.select_next(1000)
        s.report(batch[0]["point_id"], {"y": 999.0}, "failed")
        for p in batch[1:]:
            s.report(p["point_id"], {"y": 0.5}, "ok")
        train = s._training_data()
        self.assertEqual(len(train), 3)  # failed excluded


class TestSurrogateRegistry(unittest.TestCase):
    def test_registered(self):
        self.assertTrue(is_registered("surrogate_guided"))
        self.assertIn("surrogate_guided", list_strategy_kinds())

    def test_get_strategy_creates_instance(self):
        s = get_strategy("surrogate_guided",
                          parameters=[{"name": "x", "min_value": 0, "max_value": 1}],
                          objective_metric="y", budget=10, rng_seed=10)
        self.assertIsInstance(s, SurrogateGuidedStrategy)
        self.assertEqual(s.kind, "surrogate_guided")

    def test_state_field_contract(self):
        s = SurrogateGuidedStrategy(
            parameters=[{"name": "a", "min_value": 0, "max_value": 1}],
            objective_metric="y", n_initial=3, budget=10, rng_seed=11,
        )
        st = s.state()
        for key in ("kind", "batch_size", "max_batch_size", "budget", "used_budget",
                     "remaining_budget", "n_initial", "n_parameters", "objective_metric",
                     "objective_direction", "phase", "pending", "reported",
                     "last_batch_size", "surrogate", "idw_power", "kappa"):
            self.assertIn(key, st, "missing state field: %s" % key)
        self.assertEqual(st["kind"], "surrogate_guided")
        self.assertIn("n_train", st["surrogate"])


class TestSurrogateIDWInternals(unittest.TestCase):
    def test_idw_prediction_interpolation(self):
        s = SurrogateGuidedStrategy(
            parameters=[{"name": "a", "min_value": 0, "max_value": 1}],
            objective_metric="y", n_initial=0, budget=10, rng_seed=12,
        )
        train = [((0.0,), 10.0), ((1.0,), 20.0)]
        pred, unc = s._idw_predict((0.5,), train)
        # midpoint should be between 10 and 20
        self.assertGreater(pred, 10.0)
        self.assertLess(pred, 20.0)
        self.assertGreater(unc, 0.0)

    def test_distance_calculation(self):
        d = SurrogateGuidedStrategy._distance((0.0, 0.0), (3.0, 4.0))
        self.assertAlmostEqual(d, 5.0, places=6)

    def test_normalize_denormalize_roundtrip(self):
        s = SurrogateGuidedStrategy(
            parameters=[{"name": "a", "min_value": 0, "max_value": 10},
                        {"name": "b", "min_value": -5, "max_value": 5}],
            objective_metric="y", n_initial=0, budget=10, rng_seed=13,
        )
        phys = {"a": 5.0, "b": 0.0}
        norm = s._normalize(phys)
        self.assertAlmostEqual(norm[0], 0.5, places=6)
        self.assertAlmostEqual(norm[1], 0.5, places=6)
        back = s._denormalize(norm)
        self.assertAlmostEqual(back["a"], 5.0, places=6)
        self.assertAlmostEqual(back["b"], 0.0, places=6)


if __name__ == "__main__":
    unittest.main(verbosity=2)