Source code for ax.modelbridge.transforms.int_range_to_choice

#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.

from typing import List, Optional, Set

from ax.core.observation import ObservationData, ObservationFeatures
from ax.core.parameter import ChoiceParameter, Parameter, ParameterType, RangeParameter
from ax.core.search_space import SearchSpace
from ax.core.types import TConfig
from ax.modelbridge.transforms.base import Transform


[docs]class IntRangeToChoice(Transform):
    """Convert a RangeParameter of type int to a ChoiceParameter.

    Transform is done in-place.
    """

    def __init__(
        self,
        search_space: SearchSpace,
        observation_features: List[ObservationFeatures],
        observation_data: List[ObservationData],
        config: Optional[TConfig] = None,
    ) -> None:
        # 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
        }

[docs]    def transform_search_space(self, search_space: SearchSpace) -> SearchSpace:
        transformed_parameters: List[Parameter] = []
        for p in search_space.parameters.values():
            if p.name in self.transform_parameters:
                # pyre: p_cast is declared to have type `RangeParameter` but
                # pyre-fixme[9]: is used as type `Parameter`.
                p_cast: RangeParameter = p
                transformed_parameters.append(
                    ChoiceParameter(
                        name=p_cast.name,
                        parameter_type=p_cast.parameter_type,
                        # Expected `List[Optional[typing.Union[bool, float, str]]]` for
                        # 4th parameter `values` to call
                        # `ax.core.parameter.ChoiceParameter.__init__` but got
                        # `List[int]`.
                        # pyre-fixme[6]:
                        values=list(range(p_cast.lower, p_cast.upper + 1)),
                    )
                )
            else:
                transformed_parameters.append(p)
        return SearchSpace(
            parameters=transformed_parameters,
            parameter_constraints=[
                pc.clone() for pc in search_space.parameter_constraints
            ],
        )