Source code for ax.modelbridge.map_torch

# 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 Dict, List, Optional, Tuple, Type

import numpy as np

import torch
from ax.core.data import Data
from ax.core.experiment import Experiment
from ax.core.map_data import MapData
from ax.core.observation import (
    Observation,
    ObservationData,
    ObservationFeatures,
    observations_from_map_data,
    separate_observations,
)
from ax.core.optimization_config import OptimizationConfig
from ax.core.search_space import SearchSpace
from ax.core.types import TCandidateMetadata
from ax.modelbridge.base import GenResults
from ax.modelbridge.modelbridge_utils import (
    array_to_observation_data,
    observation_features_to_array,
    parse_observation_features,
)
from ax.modelbridge.torch import FIT_MODEL_ERROR, TorchModelBridge
from ax.modelbridge.transforms.base import Transform
from ax.models.torch_base import TorchModel
from ax.models.types import TConfig
from ax.utils.common.constants import Keys
from ax.utils.common.typeutils import not_none


# A mapping from map_key to its target (or final) value; by default,
# we assume normalization to [0, 1], making 1.0 the target value.
# Used in both generation and prediction.
DEFAULT_TARGET_MAP_VALUES = {"steps": 1.0}


[docs]class MapTorchModelBridge(TorchModelBridge): """A model bridge for using torch-based models that fit on MapData. Most of the `TorchModelBridge` functionality is retained, except that this class should be used in the case where `model` makes use of map_key values. For example, the use case of fitting a joint surrogate model on `(parameters, map_key)`, while candidate generation is only for `parameters`. """ def __init__( self, experiment: Experiment, search_space: SearchSpace, data: Data, model: TorchModel, transforms: List[Type[Transform]], transform_configs: Optional[Dict[str, TConfig]] = None, torch_dtype: Optional[torch.dtype] = None, torch_device: Optional[torch.device] = None, status_quo_name: Optional[str] = None, status_quo_features: Optional[ObservationFeatures] = None, optimization_config: Optional[OptimizationConfig] = None, fit_out_of_design: bool = False, default_model_gen_options: Optional[TConfig] = None, map_data_limit_total_rows: Optional[int] = None, map_data_limit_rows_per_group: Optional[int] = None, ) -> None: """ Applies transforms and fits model. Args: experiment: Is used to get arm parameters. Is not mutated. search_space: Search space for fitting the model. Constraints need not be the same ones used in gen. data: Ax Data. model: Interface will be specified in subclass. If model requires initialization, that should be done prior to its use here. transforms: List of uninitialized transform classes. Forward transforms will be applied in this order, and untransforms in the reverse order. transform_configs: A dictionary from transform name to the transform config dictionary. torch_dtype: Torch data type. torch_device: Torch device. status_quo_name: Name of the status quo arm. Can only be used if Data has a single set of ObservationFeatures corresponding to that arm. status_quo_features: ObservationFeatures to use as status quo. Either this or status_quo_name should be specified, not both. optimization_config: Optimization config defining how to optimize the model. fit_out_of_design: If specified, all training data is returned. Otherwise, only in design points are returned. default_model_gen_options: Options passed down to `model.gen(...)`. map_data_limit_total_rows: Subsample the map data so that the total number of rows is limited by this value. map_data_limit_rows_per_group: Subsample the map data so that the number of rows in the `map_key` column for each (arm, metric) is limited by this value. """ if not isinstance(data, MapData): raise ValueError( # pragma: no cover "`MapTorchModelBridge expects `MapData` instead of `Data`." ) # pyre-fixme[4]: Attribute must be annotated. self._map_key_features = data.map_keys self._map_data_limit_total_rows = map_data_limit_total_rows self._map_data_limit_rows_per_group = map_data_limit_rows_per_group super().__init__( experiment=experiment, search_space=search_space, data=data, model=model, transforms=transforms, transform_configs=transform_configs, torch_dtype=torch_dtype, torch_device=torch_device, status_quo_name=status_quo_name, status_quo_features=status_quo_features, optimization_config=optimization_config, fit_out_of_design=fit_out_of_design, default_model_gen_options=default_model_gen_options, ) @property def parameters_with_map_keys(self) -> List[str]: return self.parameters + self._map_key_features def _predict( self, observation_features: List[ObservationFeatures] ) -> List[ObservationData]: """This method is updated from `TorchModelBridge._predict(...) in that it will accept observation features with or without map_keys. If observation features do not contain map_keys, it will insert them based on `target_map_values`. """ if not self.model: # pragma: no cover raise ValueError(FIT_MODEL_ERROR.format(action="_model_predict")) # The fitted model expects map_keys. If they do not exist, we use the # target values. target_map_values = self._default_model_gen_options.get( "target_map_values", DEFAULT_TARGET_MAP_VALUES ) for p in self._map_key_features: for obsf in observation_features: if p not in obsf.parameters: obsf.parameters[p] = target_map_values[p] # pyre-ignore[16] # Convert observation features to array X = observation_features_to_array( self.parameters_with_map_keys, observation_features ) f, cov = not_none(self.model).predict(X=self._array_to_tensor(X)) f = f.detach().cpu().clone().numpy() cov = cov.detach().cpu().clone().numpy() # Convert resulting arrays to observations return array_to_observation_data(f=f, cov=cov, outcomes=self.outcomes) def _fit( self, model: TorchModel, search_space: SearchSpace, observations: List[Observation], parameters: Optional[List[str]] = None, ) -> None: """The difference from `TorchModelBridge._fit(...)` is that we use `self.parameters_with_map_keys` instead of `self.parameters`. """ self.parameters = list(search_space.parameters.keys()) if parameters is None: parameters = self.parameters_with_map_keys super()._fit( model=model, search_space=search_space, observations=observations, parameters=parameters, ) def _gen( self, n: int, search_space: SearchSpace, pending_observations: Dict[str, List[ObservationFeatures]], fixed_features: ObservationFeatures, model_gen_options: Optional[TConfig] = None, optimization_config: Optional[OptimizationConfig] = None, ) -> GenResults: """An updated version of `TorchModelBridge._gen(...) that first injects `map_dim_to_target` (e.g., `{-1: 1.0}`) into `model_gen_options` so that the target values of the map_keys are known during candidate generation. """ model_gen_options = self._add_map_dim_to_target(options=model_gen_options or {}) return super()._gen( n=n, search_space=search_space, pending_observations=pending_observations, fixed_features=fixed_features, model_gen_options=model_gen_options, optimization_config=optimization_config, ) def _array_to_observation_features( self, X: np.ndarray, candidate_metadata: Optional[List[TCandidateMetadata]] ) -> List[ObservationFeatures]: """The difference b/t this method and TorchModelBridge._update(...) is that this one makes use of `self.parameters_with_map_keys`. """ return parse_observation_features( X=X, param_names=self.parameters_with_map_keys, candidate_metadata=candidate_metadata, ) def _update( self, search_space: SearchSpace, observations: List[Observation], parameters: Optional[List[str]] = None, ) -> None: """The difference b/t this method and TorchModelBridge._update(...) is that this one makes use of `self.parameters_with_map_keys`. """ return super()._update( # pragma: no cover search_space=search_space, observations=observations, parameters=self.parameters_with_map_keys, ) def _prepare_observations( self, experiment: Optional[Experiment], data: Optional[Data] ) -> List[Observation]: """The difference b/t this method and ModelBridge._prepare_observations(...) is that this one uses `observations_from_map_data`. """ if experiment is None or data is None: return [] # pragma: no cover return observations_from_map_data( experiment=experiment, map_data=data, # pyre-ignore[6]: Checked in __init__. map_keys_as_parameters=True, include_abandoned=self._fit_abandoned, limit_total_rows=self._map_data_limit_total_rows, limit_rows_per_group=self._map_data_limit_rows_per_group, ) def _compute_in_design( self, search_space: SearchSpace, observations: List[Observation] ) -> List[bool]: """The difference b/t this method and ModelBridge._compute_in_design(...) is that this one correctly excludes map_keys when checking membership in search space (as map_keys are not explicitly in the search space). """ return [ # pragma: no cover search_space.check_membership( # Exclude map key features when checking { p: v for p, v in obs.features.parameters.items() if p not in self._map_key_features } ) for obs in observations ] def _cross_validate( self, search_space: SearchSpace, cv_training_data: List[Observation], cv_test_points: List[ObservationFeatures], parameters: Optional[List[str]] = None, ) -> List[ObservationData]: """Make predictions at cv_test_points using only the data in obs_feats and obs_data. The difference from `TorchModelBridge._cross_validate` is that here we do cross validation on the parameters + map_keys. There is some extra logic to filter out out-of-design points in the map_key dimension. """ if parameters is None: parameters = self.parameters_with_map_keys cv_test_data = super()._cross_validate( search_space=search_space, cv_training_data=cv_training_data, cv_test_points=cv_test_points, parameters=parameters, # we pass the map_keys too by default ) observation_features, observation_data = separate_observations(cv_training_data) # Since map_keys are used as features, there can be the possibility that # models for different outcomes were fit on different ranges of map_key # values; for example, this is the case if we (1) mix learning curve data with # standard data (taking default map value), or (2) are in a situation where # some learning curves start later than others. These prediction results are # "out-of-design" in the map_key dimension, so we should filter them out. map_key_ranges = self._get_map_key_ranges( observation_features=observation_features, observation_data=observation_data, ) cv_test_data = self._filter_outcomes_out_of_map_range( observation_features=cv_test_points, observation_data=cv_test_data, map_key_ranges=map_key_ranges, ) return cv_test_data def _filter_outcomes_out_of_map_range( self, observation_features: List[ObservationFeatures], observation_data: List[ObservationData], # pyre-fixme[24]: Generic type `tuple` expects at least 1 type parameter. map_key_ranges: Dict[str, Dict[str, Optional[Tuple]]], ) -> List[ObservationData]: """Uses `map_key_ranges` to detect which `observation_features` have out-of-range map_keys and filters out the corresponding outcomes in `observation_data`. """ filtered_observation_data = [] for obsf, obsd in zip(observation_features, observation_data): metric_names = obsd.metric_names means = obsd.means covariance = obsd.covariance for o in self.outcomes: if o in metric_names: for p in self._map_key_features: map_key_value = obsf.parameters[p] map_key_range = map_key_ranges[o][p] if map_key_range is not None: range_min, range_max = map_key_range if map_key_value < range_min or map_key_value > range_max: p_idx = metric_names.index(o) metric_names.pop(p_idx) means = np.delete(means, p_idx, axis=0) covariance = np.delete(covariance, p_idx, axis=0) covariance = np.delete(covariance, p_idx, axis=1) break new_obsd = ObservationData( metric_names=metric_names, means=means, covariance=covariance ) filtered_observation_data.append(new_obsd) return filtered_observation_data def _get_map_key_ranges( self, observation_features: List[ObservationFeatures], observation_data: List[ObservationData], # pyre-fixme[24]: Generic type `tuple` expects at least 1 type parameter. ) -> Dict[str, Dict[str, Optional[Tuple]]]: """Get ranges of map_key values in observation features. Returns a dict of the form: {"outcome": {"map_key": (min_val, max_val)}}. """ map_values = {o: {p: [] for p in self._map_key_features} for o in self.outcomes} for obsd, obsf in zip(observation_data, observation_features): for p in self._map_key_features: param_value = obsf.parameters[p] for o in obsd.metric_names: map_values[o][p].append(param_value) # pyre-fixme[3]: Return type must be annotated. # pyre-fixme[24]: Generic type `list` expects 1 type parameter, use # `typing.List` to avoid runtime subscripting errors. def get_range(values: List): return (min(values), max(values)) if len(values) > 0 else None return { o: {p: get_range(map_values[o][p]) for p in self._map_key_features} for o in self.outcomes } def _add_map_dim_to_target(self, options: TConfig) -> TConfig: """Convert `target_map_values` to `map_dim_to_target`, a form useable by the acquisition function and insert into options dict. """ target_map_values = self._default_model_gen_options.get("target_map_values") if target_map_values is None: target_map_values = DEFAULT_TARGET_MAP_VALUES # pragma: no cover param_and_map = self.parameters_with_map_keys map_dim_to_target = { param_and_map.index(p): target_map_values[p] # pyre-ignore[16] for p in self._map_key_features } options[Keys.ACQF_KWARGS] = { # pyre-ignore[32] **options.get(Keys.ACQF_KWARGS, {}), "map_dim_to_target": map_dim_to_target, } return options