Source code for ax.modelbridge.transforms.base
#!/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 __future__ import annotations
from typing import List, Optional, TYPE_CHECKING
import numpy as np
from ax.core.observation import ObservationData, ObservationFeatures
from ax.core.optimization_config import OptimizationConfig
from ax.core.outcome_constraint import ObjectiveThreshold
from ax.core.search_space import RobustSearchSpace, SearchSpace
from ax.exceptions.core import UnsupportedError
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 Transform:
"""Defines the API for a transform that is applied to search_space,
observation_features, observation_data, and optimization_config.
Transforms are used to adapt the search space and data into the types
and structures expected by the model. When Transforms are used (for
instance, in ModelBridge), it is always assumed that they may potentially
mutate the transformed object in-place.
Forward transforms are defined for all four of those quantities. Reverse
transforms are defined for observation_data and observation.
The forward transform for observation features must accept a partial
observation with not all features recorded.
Forward and reverse transforms for observation data accept a list of
observation features as an input, but they will not be mutated.
The forward transform for optimization config accepts the modelbridge and
fixed features as inputs, but they will not be mutated.
This class provides an identify transform.
"""
config: TConfig
modelbridge: Optional[modelbridge_module.base.ModelBridge]
def __init__(
self,
search_space: Optional[SearchSpace],
observation_features: List[ObservationFeatures],
observation_data: List[ObservationData],
modelbridge: Optional[modelbridge_module.base.ModelBridge] = None,
config: Optional[TConfig] = None,
) -> None:
"""Do any initial computations for preparing the transform.
This takes in search space and observations, but they are not modified.
Args:
search_space: The search space
observation_features: Observation features
observation_data: Observation data
modelbridge: ModelBridge for referencing experiment, status quo, etc...
config: A dictionary of options specific to each transform
"""
if config is None:
config = {}
self.config = config
self.modelbridge = modelbridge
[docs] def transform_search_space(self, search_space: SearchSpace) -> SearchSpace:
"""Transform search space.
The transforms are typically done in-place. This calls two private methods,
`_transform_search_space`, which transforms the core search space attributes,
and `_transform_parameter_distributions`, which transforms the distributions
when using a `RobustSearchSpace`.
Args:
search_space: The search space
Returns: transformed search space.
"""
self._transform_parameter_distributions(search_space=search_space)
return self._transform_search_space(search_space=search_space)
def _transform_search_space(self, search_space: SearchSpace) -> SearchSpace:
"""Transform search space.
This is typically done in-place. This class implements the identity
transform (does nothing).
Args:
search_space: The search space
Returns: transformed search space.
"""
return search_space
[docs] def transform_optimization_config(
self,
optimization_config: OptimizationConfig,
modelbridge: Optional[modelbridge_module.base.ModelBridge],
fixed_features: ObservationFeatures,
) -> OptimizationConfig:
"""Transform optimization config.
This is typically done in-place. This class implements the identity
transform (does nothing).
Args:
optimization_config: The optimization config
Returns: transformed optimization config.
"""
return optimization_config
[docs] def transform_observation_features(
self, observation_features: List[ObservationFeatures]
) -> List[ObservationFeatures]:
"""Transform observation features.
This is typically done in-place. This class implements the identity
transform (does nothing).
Args:
observation_features: Observation features
Returns: transformed observation features
"""
return observation_features
[docs] def transform_observation_data(
self,
observation_data: List[ObservationData],
observation_features: List[ObservationFeatures],
) -> List[ObservationData]:
"""Transform observation features.
This is typically done in-place. This class implements the identity
transform (does nothing).
This takes in observation_features, so that data transforms can be
conditional on features, but observation_features are notmutated.
Args:
observation_data: Observation data
observation_features: Corresponding observation features
Returns: transformed observation data
"""
return observation_data
[docs] def untransform_observation_features(
self, observation_features: List[ObservationFeatures]
) -> List[ObservationFeatures]:
"""Untransform observation features.
This is typically done in-place. This class implements the identity
transform (does nothing).
Args:
observation_features: Observation features in the transformed space
Returns: observation features in the original space
"""
return observation_features
[docs] def untransform_observation_data(
self,
observation_data: List[ObservationData],
observation_features: List[ObservationFeatures],
) -> List[ObservationData]:
"""Untransform observation data.
This is typically done in-place. This class implements the identity
transform (does nothing).
Args:
observation_data: Observation data, in transformed space
observation_features: Corresponding observation features, in same
space.
Returns: observation data in original space.
"""
return observation_data
[docs] def untransform_objective_thresholds(
self,
objective_thresholds: List[ObjectiveThreshold],
observation_features: List[ObservationFeatures],
) -> List[ObjectiveThreshold]:
"""Untransforms objective thresholds.
By default, we untransform objective thresholds in the same way as the
observation data.
Args:
objective_thresholds: Objective thresholds in transformed space.
observation_features: Observation features in transformed space. Required
to correctly untransform thresholds for stratified observation data.
"""
# Create dummy ObservationData from objective_thresholds so we can easily
# untransform them using the existing Transform methods.
means = np.array([t.bound for t in objective_thresholds])
metric_names = [t.metric.name for t in objective_thresholds]
observation_data = [
ObservationData(
means=means,
metric_names=metric_names,
covariance=np.zeros((len(metric_names), len(metric_names))),
)
]
observation_data = self.untransform_observation_data(
observation_data, observation_features
)[0]
untransformed_thresholds = []
for threshold, bound in zip(objective_thresholds, observation_data.means):
if not np.isnan(bound):
untransformed_thresholds.append(
ObjectiveThreshold(
metric=threshold.metric,
bound=bound,
relative=False,
op=threshold.op,
)
)
return untransformed_thresholds
def _transform_parameter_distributions(self, search_space: SearchSpace) -> None:
"""Transform the parameter distributions of the given search space, in-place.
This method should be called in transform_search_space before parameters
are transformed.
The base implementation is a no-op for most transforms, except for those
that have a `transform_parameters` attribute, in which case this will
raise an `UnsupportedError` if a parameter with an associated distribution
is being transformed.
"""
if isinstance(search_space, RobustSearchSpace) and hasattr(
self, "transform_parameters"
):
# pyre-ignore Undefined attribute [16]
for p_name in self.transform_parameters:
if p_name in search_space._distributional_parameters:
raise UnsupportedError(
f"{self.__class__.__name__} transform is not supported for "
"parameters with an associated distribution. Consider updating "
"the transform config."
)