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
],
)