Use array_to_params to translate categoricals back for max()

Author: NathanDeMaria

bayes_opt/parameter.py

@@ -195,7 +195,7 @@ def to_param(self, value: float | NDArray[Float]) -> float:
         Any
             The canonical representation of the parameter.
         """
-        return value.flatten()[0]
+        return value.flatten()[0].item()
 
     def to_string(self, value: float, str_len: int) -> str:
         """Represent a parameter value as a string.

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]

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)