Source code for ax.modelbridge.transforms.remove_fixed

#!/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.

from typing import TYPE_CHECKING, Dict, List, Optional, Union

from ax.core.observation import ObservationData, ObservationFeatures
from ax.core.parameter import ChoiceParameter, FixedParameter, RangeParameter
from ax.core.search_space import SearchSpace
from ax.modelbridge.transforms.base import Transform
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  # pragma: no cover


[docs]class RemoveFixed(Transform): """Remove fixed parameters. Fixed parameters should not be included in the SearchSpace. This transform removes these parameters, leaving only tunable parameters. Transform is done in-place. """ def __init__( self, search_space: SearchSpace, observation_features: List[ObservationFeatures], observation_data: List[ObservationData], modelbridge: Optional["modelbridge_module.base.ModelBridge"] = None, config: Optional[TConfig] = None, ) -> None: # Identify parameters that should be transformed self.fixed_parameters: Dict[str, FixedParameter] = { p_name: p for p_name, p in search_space.parameters.items() if isinstance(p, FixedParameter) }
[docs] def transform_observation_features( self, observation_features: List[ObservationFeatures] ) -> List[ObservationFeatures]: for obsf in observation_features: for p_name, fixed_p in self.fixed_parameters.items(): if p_name in obsf.parameters: if obsf.parameters[p_name] != fixed_p.value: raise ValueError( f"Fixed parameter {p_name} with out of design value: " f"{obsf.parameters[p_name]} passed to `RemoveFixed`." ) obsf.parameters.pop(p_name) return observation_features
[docs] def transform_search_space(self, search_space: SearchSpace) -> SearchSpace: tunable_parameters: List[Union[ChoiceParameter, RangeParameter]] = [] for p in search_space.parameters.values(): if p.name not in self.fixed_parameters: # If it's not in fixed_parameters, it must be a tunable param. # pyre: p_ is declared to have type `Union[ChoiceParameter, # pyre: RangeParameter]` but is used as type `ax.core. # pyre-fixme[9]: parameter.Parameter`. p_: Union[ChoiceParameter, RangeParameter] = p tunable_parameters.append(p_) return SearchSpace( # Expected `List[ax.core.parameter.Parameter]` for 2nd parameter # `parameters` to call `ax.core.search_space.SearchSpace.__init__` # but got `List[Union[ChoiceParameter, RangeParameter]]`. # pyre-fixme[6]: parameters=tunable_parameters, parameter_constraints=[ pc.clone() for pc in search_space.parameter_constraints ], )
[docs] def untransform_observation_features( self, observation_features: List[ObservationFeatures] ) -> List[ObservationFeatures]: for obsf in observation_features: for p_name, p in self.fixed_parameters.items(): obsf.parameters[p_name] = p.value return observation_features