Source code for ax.modelbridge.strategies.alebo

#!/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 Any, Dict, Optional

import numpy as np
import torch
from ax.core.data import Data
from ax.core.experiment import Experiment
from ax.core.search_space import SearchSpace
from ax.modelbridge.factory import DEFAULT_TORCH_DEVICE
from ax.modelbridge.generation_strategy import GenerationStep, GenerationStrategy
from ax.modelbridge.random import RandomModelBridge
from ax.modelbridge.registry import ALEBO_X_trans, ALEBO_Y_trans
from ax.modelbridge.torch import TorchModelBridge
from ax.models.random.alebo_initializer import ALEBOInitializer
from ax.models.torch.alebo import ALEBO


[docs]def get_ALEBOInitializer( search_space: SearchSpace, B: np.ndarray, seed: Optional[int] = None, **model_kwargs: Any, ) -> RandomModelBridge: return RandomModelBridge( search_space=search_space, model=ALEBOInitializer(B=B, seed=seed, **model_kwargs), transforms=ALEBO_X_trans, )
[docs]def get_ALEBO( experiment: Experiment, search_space: SearchSpace, data: Data, B: torch.Tensor, **model_kwargs: Any, ) -> TorchModelBridge: if search_space is None: search_space = experiment.search_space return TorchModelBridge( experiment=experiment, search_space=search_space, data=data, model=ALEBO(B=B, **model_kwargs), transforms=ALEBO_X_trans + ALEBO_Y_trans, torch_dtype=B.dtype, torch_device=B.device, )
[docs]class ALEBOStrategy(GenerationStrategy): """Generation strategy for Adaptive Linear Embedding BO. Both quasirandom initialization and BO are done with the same random projection. All function evaluations are done within that projection. Args: D: Dimensionality of high-dimensional space d: Dimensionality of low-dimensional space init_size: Size of random initialization name: Name of strategy dtype: torch dtype device: torch device random_kwargs: kwargs passed along to random model gp_kwargs: kwargs passed along to GP model gp_gen_kwargs: kwargs passed along to gen call on GP """ def __init__( self, D: int, d: int, init_size: int, name: str = "ALEBO", dtype: torch.dtype = torch.double, device: torch.device = DEFAULT_TORCH_DEVICE, random_kwargs: Optional[Dict[str, Any]] = None, gp_kwargs: Optional[Dict[str, Any]] = None, gp_gen_kwargs: Optional[Dict[str, Any]] = None, ) -> None: self.D = D self.d = d self.init_size = init_size self.dtype = dtype self.device = device self.random_kwargs = random_kwargs if random_kwargs is not None else {} self.gp_kwargs = gp_kwargs if gp_kwargs is not None else {} self.gp_gen_kwargs = gp_gen_kwargs B = self.gen_projection(d=d, D=D, device=device, dtype=dtype) self.gp_kwargs.update({"B": B}) self.random_kwargs.update({"B": B.cpu().numpy()}) steps = [ GenerationStep( model=get_ALEBOInitializer, num_trials=init_size, model_kwargs=self.random_kwargs, ), GenerationStep( model=get_ALEBO, num_trials=-1, model_kwargs=self.gp_kwargs, model_gen_kwargs=gp_gen_kwargs, ), ] super().__init__(steps=steps, name=name)
[docs] def clone_reset(self) -> "ALEBOStrategy": """Copy without state.""" return self.__class__( D=self.D, d=self.d, init_size=self.init_size, name=self.name, dtype=self.dtype, device=self.device, random_kwargs=self.random_kwargs, gp_kwargs=self.gp_kwargs, gp_gen_kwargs=self.gp_gen_kwargs, )
[docs] def gen_projection( self, d: int, D: int, dtype: torch.dtype, device: torch.device ) -> torch.Tensor: """Generate the projection matrix B as a (d x D) tensor""" B0 = torch.randn(d, D, dtype=dtype, device=device) B = B0 / torch.sqrt((B0 ** 2).sum(dim=0)) return B