Add LogIntToFloat transform

Summary:
This is a simple subclass of `IntToFloat` that only transforms log-scale parameters.

Replacing `IntToFloat` with `LogIntToFloat` will avoid unnecessary use of continuous relaxation across the board, and allow us to utilize the various optimizers available in `Acquisition.optimize`.

Differential Revision: D66244582

Author: saitcakmak

ax/core/search_space.py

@@ -67,7 +67,7 @@ class SearchSpace(Base):
 
     def __init__(
         self,
-        parameters: list[Parameter],
+        parameters: Sequence[Parameter],
         parameter_constraints: list[ParameterConstraint] | None = None,
     ) -> None:
         """Initialize SearchSpace

ax/modelbridge/transforms/int_to_float.py

@@ -12,6 +12,7 @@
 from ax.core.observation import Observation, ObservationFeatures
 from ax.core.parameter import Parameter, ParameterType, RangeParameter
 from ax.core.search_space import SearchSpace
+from ax.exceptions.core import UserInputError
 from ax.modelbridge.transforms.base import Transform
 from ax.modelbridge.transforms.rounding import (
     contains_constrained_integer,
@@ -65,18 +66,22 @@ def __init__(
         self.min_choices: int = checked_cast(int, config.get("min_choices", 0))
 
         # Identify parameters that should be transformed
-        self.transform_parameters: set[str] = {
+        self.transform_parameters: set[str] = self._get_transform_parameters()
+        if contains_constrained_integer(self.search_space, self.transform_parameters):
+            self.rounding = "randomized"
+            self.contains_constrained_integer: bool = True
+        else:
+            self.contains_constrained_integer: bool = False
+
+    def _get_transform_parameters(self) -> set[str]:
+        """Identify parameters that should be transformed."""
+        return {
             p_name
             for p_name, p in self.search_space.parameters.items()
             if isinstance(p, RangeParameter)
             and p.parameter_type == ParameterType.INT
             and ((p.cardinality() >= self.min_choices) or p.log_scale)
         }
-        if contains_constrained_integer(self.search_space, self.transform_parameters):
-            self.rounding = "randomized"
-            self.contains_constrained_integer: bool = True
-        else:
-            self.contains_constrained_integer: bool = False
 
     def transform_observation_features(
         self, observation_features: list[ObservationFeatures]
@@ -183,3 +188,41 @@ def untransform_observation_features(
                         obsf.parameters[p_name] = rounded_parameters[p_name]
 
         return observation_features
+
+
+class LogIntToFloat(IntToFloat):
+    """Convert a log-scale RangeParameter of type int to type float.
+
+    The behavior of this transform mirrors ``IntToFloat`` with the key difference
+    being that it only operates on log-scale parameters.
+    """
+
+    def __init__(
+        self,
+        search_space: SearchSpace | None = None,
+        observations: list[Observation] | None = None,
+        modelbridge: Optional["modelbridge_module.base.ModelBridge"] = None,
+        config: TConfig | None = None,
+    ) -> None:
+        if config is not None and "min_choices" in config:
+            raise UserInputError(
+                "`min_choices` cannot be specified for `LogIntToFloat` transform. "
+            )
+        super().__init__(
+            search_space=search_space,
+            observations=observations,
+            modelbridge=modelbridge,
+            config=config,
+        )
+        # Delete the attribute to avoid it presenting a misleading value.
+        del self.min_choices
+
+    def _get_transform_parameters(self) -> set[str]:
+        """Identify parameters that should be transformed."""
+        return {
+            p_name
+            for p_name, p in self.search_space.parameters.items()
+            if isinstance(p, RangeParameter)
+            and p.parameter_type == ParameterType.INT
+            and p.log_scale
+        }

ax/modelbridge/transforms/tests/test_int_to_float_transform.py

@@ -13,8 +13,8 @@
 from ax.core.parameter import ChoiceParameter, Parameter, ParameterType, RangeParameter
 from ax.core.parameter_constraint import OrderConstraint, SumConstraint
 from ax.core.search_space import RobustSearchSpace, SearchSpace
-from ax.exceptions.core import UnsupportedError
-from ax.modelbridge.transforms.int_to_float import IntToFloat
+from ax.exceptions.core import UnsupportedError, UserInputError
+from ax.modelbridge.transforms.int_to_float import IntToFloat, LogIntToFloat
 from ax.utils.common.testutils import TestCase
 from ax.utils.common.typeutils import checked_cast
 from ax.utils.testing.core_stubs import get_robust_search_space
@@ -324,3 +324,33 @@ def test_w_parameter_distributions(self) -> None:
         )
         with self.assertRaisesRegex(UnsupportedError, "transform is not supported"):
             t.transform_search_space(rss)
+
+
+class LogIntToFloatTransformTest(TestCase):
+    def test_log_int_to_float(self) -> None:
+        parameters = [
+            RangeParameter("x", lower=1, upper=3, parameter_type=ParameterType.INT),
+            RangeParameter("y", lower=1, upper=50, parameter_type=ParameterType.INT),
+            RangeParameter(
+                "z", lower=1, upper=50, parameter_type=ParameterType.INT, log_scale=True
+            ),
+        ]
+        search_space = SearchSpace(parameters=parameters)
+        with self.assertRaisesRegex(UserInputError, "min_choices"):
+            LogIntToFloat(search_space=search_space, config={"min_choices": 5})
+        t = LogIntToFloat(search_space=search_space)
+        self.assertFalse(hasattr(t, "min_choices"))
+        self.assertEqual(t.transform_parameters, {"z"})
+        t_ss = t.transform_search_space(search_space)
+        self.assertEqual(t_ss.parameters["x"], parameters[0])
+        self.assertEqual(t_ss.parameters["y"], parameters[1])
+        self.assertEqual(
+            t_ss.parameters["z"],
+            RangeParameter(
+                name="z",
+                lower=0.50001,
+                upper=50.49999,
+                parameter_type=ParameterType.FLOAT,
+                log_scale=True,
+            ),
+        )