Port is_unordered_choice and can_map_to_binary to OSS

ax/adapter/adapter_utils.py

@@ -35,7 +35,7 @@
     OutcomeConstraint,
     ScalarizedOutcomeConstraint,
 )
-from ax.core.parameter import ChoiceParameter, ParameterType, RangeParameter
+from ax.core.parameter import ChoiceParameter, Parameter, ParameterType, RangeParameter
 from ax.core.parameter_constraint import ParameterConstraint
 from ax.core.search_space import SearchSpace, SearchSpaceDigest
 from ax.core.types import TBounds, TCandidateMetadata
@@ -1321,3 +1321,56 @@ def _consolidate_comparisons(X: Tensor, Y: Tensor) -> tuple[Tensor, Tensor]:
 
     X, Y, _ = consolidate_duplicates(X, Y)
     return X, Y
+
+
+def is_unordered_choice(
+    p: Parameter, min_choices: int | None = None, max_choices: int | None = None
+) -> bool:
+    """Returns whether a parameter is an unordered choice (categorical) parameter.
+
+    You can also specify `min_choices` and `max_choices` to restrict how many
+    possible values the parameter can take on.
+
+    Args:
+        p: Parameter.
+        min_choices: The minimum number of possible values for the parameter.
+        max_choices: The maximum number of possible values for the parameter.
+
+    Returns:
+        A boolean indicating whether p is an unordered choice parameter or not.
+    """
+    if min_choices is not None and min_choices < 0:
+        raise UserInputError("`min_choices` must be a non-negative integer.")
+    if max_choices is not None and max_choices < 0:
+        raise UserInputError("`max_choices` must be a non-negative integer.")
+    if (
+        min_choices is not None
+        and max_choices is not None
+        and min_choices > max_choices
+    ):
+        raise UserInputError("`min_choices` cannot be larger than `max_choices`.")
+    return (
+        isinstance(p, ChoiceParameter)
+        and not p.is_ordered
+        and (min_choices is None or min_choices <= len(p.values))
+        and (max_choices is None or max_choices >= len(p.values))
+    )
+
+
+def can_map_to_binary(p: Parameter) -> bool:
+    """Returns whether a parameter can be transformed to a binary parameter.
+
+    Any choice/range parameters with exactly two values can be transformed to a
+    binary parameter.
+
+    Args:
+        p: Parameter.
+
+    Returns
+        A boolean indicating whether p can be transformed to a binary parameter.
+    """
+    return (isinstance(p, ChoiceParameter) and len(p.values) == 2) or (
+        isinstance(p, RangeParameter)
+        and p.parameter_type == ParameterType.INT
+        and p.lower == p.upper - 1
+    )

ax/adapter/tests/test_adapter_utils.py

@@ -12,8 +12,10 @@
 from ax.adapter.adapter_utils import (
     _get_adapter_training_data,
     arm_to_np_array,
+    can_map_to_binary,
     extract_search_space_digest,
     feasible_hypervolume,
+    is_unordered_choice,
     process_contextual_datasets,
     transform_search_space,
     validate_and_apply_final_transform,
@@ -414,3 +416,112 @@ def test_validate_and_apply_final_transform_none_target_point(self) -> None:
 
         # Assert: confirm target point remains None
         self.assertIsNone(target_p)
+
+    def test_is_unordered_choice(self) -> None:
+        # Test cases where is_unordered_choice should return True
+        # (with min_choices=3, max_choices=5)
+        for p in [
+            # Unordered choice (INT), with 3 choices
+            ChoiceParameter("p", ParameterType.INT, values=[0, 1, 2], is_ordered=False),
+            # Unordered choice (STRING), with 5 choices
+            ChoiceParameter(
+                "p",
+                ParameterType.STRING,
+                values=["0", "1", "2", "4", "5"],
+                is_ordered=False,
+            ),
+            # Unordered choice (STRING), with 4 choices
+            ChoiceParameter(
+                "p", ParameterType.STRING, values=["a", "b", "c", "d"], is_ordered=False
+            ),
+        ]:
+            with self.subTest(p=p):
+                self.assertTrue(is_unordered_choice(p, min_choices=3, max_choices=5))
+
+        # Test cases where is_unordered_choice should return False
+        # (with min_choices=3, max_choices=5)
+        for p in [
+            # Too few choices
+            ChoiceParameter("p", ParameterType.INT, values=[0, 1], is_ordered=False),
+            # Ordered choice (INT)
+            ChoiceParameter(
+                "p", ParameterType.INT, values=[0, 1, 2, 4], is_ordered=True
+            ),
+            # Range parameter (non-choice)
+            RangeParameter("p", parameter_type=ParameterType.INT, lower=0, upper=3),
+            # Ordered choice (STRING)
+            ChoiceParameter(
+                "p", ParameterType.STRING, values=["0", "1", "2"], is_ordered=True
+            ),
+        ]:
+            with self.subTest(p=p):
+                self.assertFalse(is_unordered_choice(p, min_choices=3, max_choices=5))
+
+        # Test error cases
+        p = ChoiceParameter("p", ParameterType.INT, values=[0, 1, 2], is_ordered=False)
+        with self.assertRaisesRegex(
+            UserInputError, "`min_choices` must be a non-negative integer."
+        ):
+            is_unordered_choice(p, min_choices=-3)
+        with self.assertRaisesRegex(
+            UserInputError, "`max_choices` must be a non-negative integer."
+        ):
+            is_unordered_choice(p, max_choices=-1)
+        with self.assertRaisesRegex(
+            UserInputError, "`min_choices` cannot be larger than `max_choices`."
+        ):
+            is_unordered_choice(p, min_choices=3, max_choices=2)
+
+    def test_can_map_to_binary(self) -> None:
+        # Test cases where can_map_to_binary should return True
+        for p in [
+            # Int range with exactly 2 values
+            RangeParameter(
+                name="p", parameter_type=ParameterType.INT, lower=0, upper=1
+            ),
+            RangeParameter(
+                name="p", parameter_type=ParameterType.INT, lower=3, upper=4
+            ),
+            # Choice with exactly 2 values
+            ChoiceParameter(
+                name="p",
+                parameter_type=ParameterType.INT,
+                values=[0, 1],
+                is_ordered=False,
+            ),
+            ChoiceParameter(
+                name="p",
+                parameter_type=ParameterType.STRING,
+                values=["a", "b"],
+                is_ordered=False,
+            ),
+        ]:
+            with self.subTest(p=p):
+                self.assertTrue(can_map_to_binary(p))
+
+        # Test cases where can_map_to_binary should return False
+        for p in [
+            # Float range (continuous, not binary)
+            RangeParameter(
+                name="p", parameter_type=ParameterType.FLOAT, lower=0, upper=1
+            ),
+            # Int range with more than 2 values
+            RangeParameter(
+                name="p", parameter_type=ParameterType.INT, lower=0, upper=3
+            ),
+            # Choice with more than 2 values
+            ChoiceParameter(
+                name="p",
+                parameter_type=ParameterType.INT,
+                values=[0, 1, 2],
+                is_ordered=False,
+            ),
+            ChoiceParameter(
+                name="p",
+                parameter_type=ParameterType.STRING,
+                values=["a", "b", "c"],
+                is_ordered=False,
+            ),
+        ]:
+            with self.subTest(p=p):
+                self.assertFalse(can_map_to_binary(p))