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

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

Covers: happy path (sensitivity ranking on linear function), boundary
(empty params / single trajectory / odd n_levels auto-corrected), anomaly
(unknown point_id report / missing objective metric), null inputs, registry
integration, and state() field contract.

All source is ASCII only. Run: python scripts/test_strategy_morris.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
    MorrisStrategy,
    get_strategy,
    is_registered,
    list_strategy_kinds,
)


def _linear_objective(params):
    """y = 2*a + 0.5*b -> parameter 'a' should rank higher than 'b'."""
    return 2.0 * params.get("a", 0.0) + 0.5 * params.get("b", 0.0)


class TestMorrisHappyPath(unittest.TestCase):
    def test_trajectory_point_count(self):
        s = MorrisStrategy(
            parameters=[{"name": "a", "min_value": 0, "max_value": 1},
                        {"name": "b", "min_value": 0, "max_value": 1}],
            n_trajectories=5, n_levels=4, objective_metric="y", rng_seed=1,
        )
        pts = s.select_next(1000)
        # n_trajectories * (n_params + 1) = 5 * 3 = 15
        self.assertEqual(len(pts), 15)

    def test_step_zero_has_no_changed_param(self):
        s = MorrisStrategy(
            parameters=[{"name": "a", "min_value": 0, "max_value": 1}],
            n_trajectories=2, n_levels=4, objective_metric="y", rng_seed=2,
        )
        pts = s.select_next(1000)
        step0 = [p for p in pts if p["step"] == 0]
        self.assertEqual(len(step0), 2)
        for p in step0:
            self.assertIsNone(p["changed_param"])

    def test_sensitivity_ranking_linear(self):
        s = MorrisStrategy(
            parameters=[{"name": "a", "min_value": 0, "max_value": 1},
                        {"name": "b", "min_value": 0, "max_value": 1}],
            n_trajectories=8, n_levels=4, objective_metric="y", rng_seed=3,
        )
        pts = s.select_next(1000)
        for p in pts:
            s.report(p["point_id"], {"y": _linear_objective(p["params"])}, "ok")
        st = s.state()
        ranking = st["sensitivity_ranking"]
        self.assertEqual(len(ranking), 2)
        # 'a' should rank first (higher mu_star)
        self.assertEqual(ranking[0]["parameter"], "a")
        self.assertGreater(ranking[0]["mu_star"], ranking[1]["mu_star"])
        # linear function -> sigma near zero
        self.assertAlmostEqual(ranking[0]["sigma"], 0.0, places=6)

    def test_converged_after_all_reported(self):
        s = MorrisStrategy(
            parameters=[{"name": "a", "min_value": 0, "max_value": 1}],
            n_trajectories=3, n_levels=4, objective_metric="y", rng_seed=4,
        )
        # take only one point first -> pending still non-empty -> not converged
        first = s.select_next(1)
        self.assertEqual(len(first), 1)
        self.assertFalse(s.is_converged())
        # take the rest and report all
        rest = s.select_next(1000)
        all_pts = first + rest
        for p in all_pts:
            s.report(p["point_id"], {"y": 1.0}, "ok")
        self.assertTrue(s.is_converged())


class TestMorrisBoundary(unittest.TestCase):
    def test_empty_parameters(self):
        s = MorrisStrategy(parameters=[], n_trajectories=5, objective_metric="y")
        pts = s.select_next(1000)
        self.assertEqual(pts, [])
        self.assertTrue(s.is_converged())
        st = s.state()
        self.assertEqual(st["n_parameters"], 0)
        self.assertEqual(st["sensitivity_ranking"], [])

    def test_none_parameters(self):
        s = MorrisStrategy(parameters=None, n_trajectories=3, objective_metric="y")
        self.assertEqual(s.select_next(1000), [])

    def test_single_trajectory(self):
        s = MorrisStrategy(
            parameters=[{"name": "a", "min_value": 0, "max_value": 1},
                        {"name": "b", "min_value": 0, "max_value": 1}],
            n_trajectories=1, n_levels=4, objective_metric="y", rng_seed=5,
        )
        pts = s.select_next(1000)
        self.assertEqual(len(pts), 3)  # 1 * (2+1)

    def test_odd_n_levels_auto_corrected(self):
        s = MorrisStrategy(
            parameters=[{"name": "a", "min_value": 0, "max_value": 1}],
            n_trajectories=1, n_levels=5, objective_metric="y", rng_seed=6,
        )
        self.assertEqual(s.n_levels, 6)  # odd -> next even
        self.assertAlmostEqual(s._delta, 6 / (2 * 5), places=6)

    def test_params_within_bounds(self):
        s = MorrisStrategy(
            parameters=[{"name": "a", "min_value": 0.5, "max_value": 2.0, "step": 0.1},
                        {"name": "b", "min_value": 10, "max_value": 20}],
            n_trajectories=4, n_levels=4, objective_metric="y", rng_seed=7,
        )
        pts = s.select_next(1000)
        for p in pts:
            self.assertGreaterEqual(p["params"]["a"], 0.5)
            self.assertLessEqual(p["params"]["a"], 2.0)
            self.assertGreaterEqual(p["params"]["b"], 10)
            self.assertLessEqual(p["params"]["b"], 20)


class TestMorrisAnomaly(unittest.TestCase):
    def test_report_unknown_point_id_no_crash(self):
        s = MorrisStrategy(
            parameters=[{"name": "a", "min_value": 0, "max_value": 1}],
            n_trajectories=2, n_levels=4, objective_metric="y", rng_seed=8,
        )
        s.select_next(1000)
        # should not raise
        s.report(999999, {"y": 1.0}, "ok")
        # state() must remain callable and well-formed
        st = s.state()
        self.assertIn("sensitivity_ranking", st)

    def test_missing_objective_metric_yields_zero_sensitivity(self):
        s = MorrisStrategy(
            parameters=[{"name": "a", "min_value": 0, "max_value": 1}],
            n_trajectories=3, n_levels=4, objective_metric="nonexistent", rng_seed=9,
        )
        pts = s.select_next(1000)
        for p in pts:
            s.report(p["point_id"], {"y": 1.0}, "ok")  # wrong metric key
        st = s.state()
        for r in st["sensitivity_ranking"]:
            self.assertEqual(r["mu_star"], 0.0)
            self.assertEqual(r["n_effects"], 0)

    def test_failed_status_points_excluded_from_effects(self):
        s = MorrisStrategy(
            parameters=[{"name": "a", "min_value": 0, "max_value": 1}],
            n_trajectories=2, n_levels=4, objective_metric="y", rng_seed=10,
        )
        pts = s.select_next(1000)
        # report first point as failed, rest as ok
        s.report(pts[0]["point_id"], {"y": 999.0}, "failed")
        for p in pts[1:]:
            s.report(p["point_id"], {"y": _linear_objective(p["params"])}, "ok")
        st = s.state()
        # failed point's trajectory may be partially excluded; no crash
        self.assertIn("sensitivity_ranking", st)


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

    def test_get_strategy_creates_instance(self):
        s = get_strategy("morris", parameters=[{"name": "x", "min_value": 0, "max_value": 1}],
                          n_trajectories=2, objective_metric="y", rng_seed=11)
        self.assertIsInstance(s, MorrisStrategy)
        self.assertEqual(s.kind, "morris")

    def test_state_field_contract(self):
        s = MorrisStrategy(
            parameters=[{"name": "a", "min_value": 0, "max_value": 1}],
            n_trajectories=2, n_levels=4, objective_metric="y", rng_seed=12,
        )
        st = s.state()
        for key in ("kind", "batch_size", "n_parameters", "n_trajectories",
                     "n_levels", "delta", "objective_metric", "total_points",
                     "pending", "reported", "sensitivity_ranking", "key_parameters"):
            self.assertIn(key, st, "missing state field: %s" % key)
        self.assertEqual(st["kind"], "morris")


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