Fix .max() for categorical (#601)

Author: NathanDeMaria

bayes_opt/target_space.py

@@ -643,7 +643,7 @@ def max(self) -> dict[str, Any] | None:
         params = self.params[self.mask]
         target_max_idx = np.argmax(target)
 
-        res = {"target": target_max, "params": dict(zip(self.keys, params[target_max_idx]))}
+        res = {"target": target_max, "params": self.array_to_params(params[target_max_idx])}
 
         if self._constraint is not None:
             constraint_values = self.constraint_values[self.mask]
@@ -672,7 +672,7 @@ def res(self) -> list[dict[str, Any]]:
 
             return [{"target": target, "params": param} for target, param in zip(self.target, params)]
 
-        params = [dict(zip(self.keys, p)) for p in self.params]
+        params = [self.array_to_params(p) for p in self.params]
 
         return [
             {"target": target, "constraint": constraint_value, "params": param, "allowed": allowed}

tests/test_target_space.py

@@ -253,6 +253,26 @@ def test_max_with_constraint_identical_target_value():
     assert space.max() == {"params": {"p1": 2, "p2": 3}, "target": 5, "constraint": -1}
 
 
+def test_max_categorical() -> None:
+    PBOUNDS = {
+        "first_float": (0.0, 1.0),
+        "categorical_value": ("a", "b", "c", "d"),
+        "second_float": (0.0, 1.0),
+    }
+
+    def _f(first_float: float, categorical_value: str, second_float: float) -> float:
+        return second_float if categorical_value == "c" else first_float
+
+    space = TargetSpace(_f, PBOUNDS)
+    space.probe(params={"first_float": 0.1, "categorical_value": "a", "second_float": 0.1})
+    space.probe(params={"first_float": 0.1, "categorical_value": "b", "second_float": 0.9})
+    space.probe(params={"first_float": 0.1, "categorical_value": "c", "second_float": 0.8})
+    space.probe(params={"first_float": 0.1, "categorical_value": "d", "second_float": 0.9})
+
+    expected = {"first_float": 0.1, "categorical_value": "c", "second_float": 0.8}
+    assert space.max()["params"] == expected
+
+
 def test_res():
     PBOUNDS = {"p1": (0, 10), "p2": (1, 100)}
     space = TargetSpace(target_func, PBOUNDS)
@@ -273,6 +293,54 @@ def test_res():
     assert space.res() == expected_res
 
 
+def test_res_categorical() -> None:
+    PBOUNDS = {"p1": (0, 10), "p2": ["a", "b", "c"]}
+
+    def _f(p1: float, p2: str) -> float:
+        return p1 + len(p2)
+
+    space = TargetSpace(_f, PBOUNDS)
+
+    assert space.res() == []
+    space.probe(params={"p1": 1, "p2": "a"})
+    space.probe(params={"p1": 5, "p2": "b"})
+    space.probe(params={"p1": 2, "p2": "c"})
+    space.probe(params={"p1": 2, "p2": "a"})
+
+    expected_res = [
+        {"params": {"p1": 1, "p2": "a"}, "target": 2},
+        {"params": {"p1": 5, "p2": "b"}, "target": 6},
+        {"params": {"p1": 2, "p2": "c"}, "target": 3},
+        {"params": {"p1": 2, "p2": "a"}, "target": 3},
+    ]
+    assert len(space.res()) == 4
+    assert space.res() == expected_res
+
+
+def test_res_categorical_with_constraints() -> None:
+    PBOUNDS = {"p1": (0, 10), "p2": ["a", "b", "c"]}
+
+    def _f(p1: float, p2: str) -> float:
+        return p1 + len(p2)
+
+    space = TargetSpace(_f, PBOUNDS, constraint=NonlinearConstraint(lambda p1, p2: p1 - 2, 0, 5))
+
+    assert space.res() == []
+    space.probe(params={"p1": 1, "p2": "a"})
+    space.probe(params={"p1": 5, "p2": "b"})
+    space.probe(params={"p1": 2, "p2": "c"})
+    space.probe(params={"p1": 2, "p2": "a"})
+
+    expected_res = [
+        {"params": {"p1": 1, "p2": "a"}, "target": 2, "allowed": False, "constraint": -1},
+        {"params": {"p1": 5, "p2": "b"}, "target": 6, "allowed": True, "constraint": 3},
+        {"params": {"p1": 2, "p2": "c"}, "target": 3, "allowed": True, "constraint": 0},
+        {"params": {"p1": 2, "p2": "a"}, "target": 3, "allowed": True, "constraint": 0},
+    ]
+    assert len(space.res()) == 4
+    assert space.res() == expected_res
+
+
 def test_set_bounds():
     pbounds = {"p1": (0, 1), "p3": (0, 3), "p2": (0, 2), "p4": (0, 4)}
     space = TargetSpace(target_func, pbounds)