Source code for ax.models.torch.botorch_modular.multi_fidelity

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

# pyre-strict

from typing import Any, Dict, Mapping, Optional

from ax.core.search_space import SearchSpaceDigest
from ax.exceptions.core import UnsupportedError
from ax.models.torch.botorch_modular.acquisition import Acquisition
from ax.models.torch.botorch_modular.surrogate import Surrogate
from ax.models.torch_base import TorchOptConfig
from ax.utils.common.constants import Keys
from botorch.acquisition.cost_aware import InverseCostWeightedUtility
from botorch.acquisition.utils import (
    expand_trace_observations,
    project_to_target_fidelity,
)
from botorch.models.cost import AffineFidelityCostModel
from torch import Tensor


[docs]class MultiFidelityAcquisition(Acquisition): # NOTE: Here, we do not consider using `IIDNormalSampler` and always # use the `SobolQMCNormalSampler`.
[docs] def compute_model_dependencies( self, surrogates: Mapping[str, Surrogate], search_space_digest: SearchSpaceDigest, torch_opt_config: TorchOptConfig, options: Optional[Dict[str, Any]] = None, ) -> Dict[str, Any]: if torch_opt_config.risk_measure is not None: raise UnsupportedError( f"{self.__class__.__name__} does not support risk measures." ) target_fidelities = { k: v for k, v in search_space_digest.target_values.items() if k in search_space_digest.fidelity_features } if not target_fidelities: raise ValueError( "Target fidelities are required for {self.__class__.__name__}." ) dependencies = super().compute_model_dependencies( surrogates=surrogates, search_space_digest=search_space_digest, torch_opt_config=torch_opt_config, options=options, ) options = options or {} fidelity_weights = options.get(Keys.FIDELITY_WEIGHTS, None) if fidelity_weights is None: fidelity_weights = {f: 1.0 for f in target_fidelities} if not set(target_fidelities) == set(fidelity_weights): raise RuntimeError( "Must provide the same indices for target_fidelities " f"({set(target_fidelities)}) and fidelity_weights " f" ({set(fidelity_weights)})." ) cost_intercept = options.get(Keys.COST_INTERCEPT, 1.0) cost_model = AffineFidelityCostModel( fidelity_weights=fidelity_weights, fixed_cost=cost_intercept ) cost_aware_utility = InverseCostWeightedUtility(cost_model=cost_model) # pyre-fixme[53]: Captured variable `target_fidelities` is not annotated. def project(X: Tensor) -> Tensor: return project_to_target_fidelity(X=X, target_fidelities=target_fidelities) # pyre-fixme[53]: Captured variable `target_fidelities` is not annotated. def expand(X: Tensor) -> Tensor: return expand_trace_observations( X=X, fidelity_dims=sorted(target_fidelities), # pyre-fixme[16]: `Optional` has no attribute `get`. num_trace_obs=options.get(Keys.NUM_TRACE_OBSERVATIONS, 0), ) dependencies.update( # pyre-fixme[6]: For 1st param expected `SupportsKeysAndGetItem[str, # typing.Any]` but got `Dict[Keys, typing.Callable[[Named(X, Tensor)], # typing.Any]]`. { Keys.COST_AWARE_UTILITY: cost_aware_utility, Keys.PROJECT: project, Keys.EXPAND: expand, } ) return dependencies