Source code for ax.modelbridge.transforms.relativize

#!/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 math import sqrt
from typing import Callable, Dict, List, Optional, Tuple, TYPE_CHECKING

import numpy as np
from ax.core.observation import Observation, ObservationData, ObservationFeatures
from ax.core.optimization_config import (
    MultiObjectiveOptimizationConfig,
    OptimizationConfig,
)
from ax.core.outcome_constraint import OutcomeConstraint
from ax.core.search_space import SearchSpace
from ax.modelbridge import ModelBridge
from ax.modelbridge.transforms.base import Transform
from ax.models.types import TConfig
from ax.utils.common.typeutils import not_none
from ax.utils.stats.statstools import relativize, unrelativize

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


[docs]class Relativize(Transform): """ Change the relative flag of the given relative optimization configuration to False. This is needed in order for the new opt config to pass ModelBridge that requires non-relativized opt config. Also transforms absolute data and opt configs to relative. Requires a modelbridge with a status quo set to work. """ MISSING_STATUS_QUO_ERROR = "Cannot relativize data without status quo data" def __init__( self, search_space: Optional[SearchSpace] = None, observations: Optional[List[Observation]] = None, modelbridge: Optional[modelbridge_module.base.ModelBridge] = None, config: Optional[TConfig] = None, ) -> None: assert observations is not None, "Relativize requires observations" super().__init__( search_space=search_space, observations=observations, modelbridge=modelbridge, config=config, ) # self.modelbridge should NOT be modified self.modelbridge: ModelBridge = not_none( modelbridge, "Relativize transform requires a modelbridge" )
[docs] def transform_optimization_config( self, optimization_config: OptimizationConfig, modelbridge: Optional[modelbridge_module.base.ModelBridge] = None, fixed_features: Optional[ObservationFeatures] = None, ) -> OptimizationConfig: r""" Change the relative flag of the given relative optimization configuration to False. This is needed in order for the new opt config to pass ModelBridge that requires non-relativized opt config. Args: opt_config: Optimization configuration relative to status quo. Returns: Optimization configuration relative to status quo with relative flag equal to false. """ # Getting constraints constraints = [ constraint.clone() for constraint in optimization_config.outcome_constraints ] if not all( constraint.relative for constraint in optimization_config.outcome_constraints ): raise ValueError( "All constraints must be relative to use the Relativize transform." ) for constraint in constraints: constraint.relative = False if isinstance(optimization_config, MultiObjectiveOptimizationConfig): # Getting objective thresholds obj_thresholds = [ obj_threshold.clone() for obj_threshold in optimization_config.objective_thresholds ] for obj_threshold in obj_thresholds: if not obj_threshold.relative: raise ValueError( "All objective thresholds must be relative to use " "the Relativize transform." ) obj_threshold.relative = False new_optimization_config = MultiObjectiveOptimizationConfig( objective=optimization_config.objective, outcome_constraints=constraints, objective_thresholds=obj_thresholds, ) else: new_optimization_config = OptimizationConfig( objective=optimization_config.objective, outcome_constraints=constraints, ) return new_optimization_config
[docs] def untransform_outcome_constraints( self, outcome_constraints: List[OutcomeConstraint], fixed_features: Optional[ObservationFeatures] = None, ) -> List[OutcomeConstraint]: for c in outcome_constraints: c.relative = True return outcome_constraints
[docs] def transform_observations( self, observations: List[Observation], ) -> List[Observation]: return self._rel_op_on_observations( observations=observations, rel_op=relativize )
[docs] def untransform_observations( self, observations: List[Observation] ) -> List[Observation]: """Unrelativize the data""" return self._rel_op_on_observations( observations=observations, rel_op=unrelativize )
def _rel_op_on_observations( self, observations: List[Observation], rel_op: Callable[..., Tuple[np.ndarray, np.ndarray]], ) -> List[Observation]: sq_data_by_trial: Dict[int, ObservationData] = not_none( self.modelbridge.status_quo_data_by_trial, self.MISSING_STATUS_QUO_ERROR ) missing_index = any(obs.features.trial_index is None for obs in observations) default_trial_idx: Optional[int] = None if missing_index: if len(sq_data_by_trial) == 1: default_trial_idx = next(iter(sq_data_by_trial)) else: raise ValueError( "Observations contain missing trial index that can't be inferred." ) def _get_relative_data_from_obs( obs: Observation, rel_op: Callable[..., Tuple[np.ndarray, np.ndarray]], ) -> ObservationData: idx = ( int(obs.features.trial_index) if obs.features.trial_index is not None else default_trial_idx ) if idx not in sq_data_by_trial: raise ValueError(self.MISSING_STATUS_QUO_ERROR) return self._get_relative_data( data=obs.data, status_quo_data=sq_data_by_trial[idx], rel_op=rel_op, ) return [ Observation( features=obs.features, data=_get_relative_data_from_obs(obs, rel_op), arm_name=obs.arm_name, ) for obs in observations ] @staticmethod def _get_relative_data( data: ObservationData, status_quo_data: ObservationData, rel_op: Callable[..., Tuple[np.ndarray, np.ndarray]], ) -> ObservationData: r""" Relativize or unrelativize `data` based on `status_quo_data` based on `rel_op` Args: data: ObservationData object to relativize status_quo_data: The status quo data (un)relativization is based upon rel_op: relativize or unrelativize operator. Returns: (un)relativized ObservationData """ L = len(data.metric_names) result = ObservationData( metric_names=data.metric_names, # zeros are just to create the shape so values can be set by index means=np.zeros(L), covariance=np.zeros((L, L)), ) for i, metric in enumerate(data.metric_names): try: j = next( k for k in range(L) if status_quo_data.metric_names[k] == metric ) except (IndexError, StopIteration): raise ValueError( "Relativization cannot be performed because " "ObservationData for status quo is missing metrics" ) means_t = data.means[i] sems_t = sqrt(data.covariance[i][i]) mean_c = status_quo_data.means[j] sem_c = sqrt(status_quo_data.covariance[j][j]) # if the is the status quo if means_t == mean_c and sems_t == sem_c: means_rel, sems_rel = 0, 0 else: means_rel, sems_rel = rel_op( means_t=means_t, sems_t=sems_t, mean_c=mean_c, sem_c=sem_c, as_percent=True, ) result.means[i] = means_rel result.covariance[i][i] = sems_rel**2 return result