Source code for ax.modelbridge.transforms.int_range_to_choice

#!/usr/bin/env python3
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

# pyre-strict

from numbers import Real
from typing import cast, Optional, TYPE_CHECKING

from ax.core.observation import Observation
from ax.core.parameter import ChoiceParameter, Parameter, ParameterType, RangeParameter
from ax.core.search_space import SearchSpace
from ax.modelbridge.transforms.base import Transform
from ax.modelbridge.transforms.utils import construct_new_search_space
from ax.models.types import TConfig

if TYPE_CHECKING:
    # import as module to make sphinx-autodoc-typehints happy
    from ax import modelbridge as modelbridge_module  # noqa F401


[docs] class IntRangeToChoice(Transform): """Convert a RangeParameter of type int to a ordered ChoiceParameter. Transform is done in-place. """ 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: assert search_space is not None, "IntRangeToChoice requires search space" config = config or {} self.max_choices: float = float( cast(Real, (config.get("max_choices", float("inf")))) ) # Identify parameters that should be transformed self.transform_parameters: set[str] = { p_name for p_name, p in search_space.parameters.items() if isinstance(p, RangeParameter) and p.parameter_type == ParameterType.INT and p.cardinality() <= self.max_choices } def _transform_search_space(self, search_space: SearchSpace) -> SearchSpace: transformed_parameters: dict[str, Parameter] = {} for p_name, p in search_space.parameters.items(): if ( p_name in self.transform_parameters and isinstance(p, RangeParameter) and p.parameter_type == ParameterType.INT and p.cardinality() <= self.max_choices ): values = list(range(int(p.lower), int(p.upper) + 1)) target_value = ( None if p.target_value is None else next(i for i, v in enumerate(values) if v == p.target_value) ) transformed_parameters[p_name] = ChoiceParameter( name=p_name, parameter_type=p.parameter_type, values=values, # pyre-fixme[6] is_ordered=True, is_fidelity=p.is_fidelity, target_value=target_value, ) else: transformed_parameters[p.name] = p return construct_new_search_space( search_space=search_space, parameters=list(transformed_parameters.values()), parameter_constraints=[ pc.clone_with_transformed_parameters( transformed_parameters=transformed_parameters ) for pc in search_space.parameter_constraints ], )