Source code for ax.models.random.rembo_initializer

#!/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 Callable, Dict, List, Optional, Tuple

import numpy as np
from ax.models.random.uniform import UniformGenerator
from ax.models.types import TConfig
from ax.utils.common.docutils import copy_doc


[docs]class REMBOInitializer(UniformGenerator): """Sample in a low-dimensional linear embedding. Generates points in [-1, 1]^D by generating points in a d-dimensional embedding, with box bounds as specified. When points are projected up, if they fall outside [-1, 1]^D they are clamped to those bounds. Args: A: A (Dxd) linear embedding bounds_d: Box bounds in the low-d space seed: seed for UniformGenerator """ def __init__( self, A: np.ndarray, bounds_d: List[Tuple[float, float]], seed: Optional[int] = None, ) -> None: self.bounds_d = bounds_d self.A = A # pyre-fixme[4]: Attribute must be annotated. self.X_d_gen = [] # Store points in low-d space generated here super().__init__(seed=seed, deduplicate=False)
[docs] def project_up(self, X: np.ndarray) -> np.ndarray: """Project to high-dimensional space.""" Z = np.transpose(self.A @ np.transpose(X)) return np.clip(Z, a_min=-1, a_max=1)
[docs] @copy_doc(UniformGenerator.gen) def gen( self, n: int, bounds: List[Tuple[float, float]], linear_constraints: Optional[Tuple[np.ndarray, np.ndarray]] = None, fixed_features: Optional[Dict[int, float]] = None, model_gen_options: Optional[TConfig] = None, rounding_func: Optional[Callable[[np.ndarray], np.ndarray]] = None, ) -> Tuple[np.ndarray, np.ndarray]: # The projection is from [-1, 1]^D. for b in bounds: assert b == (-1, 1) # The following can be easily handled in the future when needed assert linear_constraints is None assert fixed_features is None # Do gen in the low-dimensional space. First on [0, 1]^d, X_01, w = super().gen(n=n, bounds=[(0.0, 1.0)] * len(self.bounds_d)) # Then map to bounds_d lw, up = zip(*self.bounds_d) lw = np.array(lw) up = np.array(up) X_d = X_01 * (up - lw) + lw # Store self.X_d_gen.extend(list(X_d)) # And finally project up return self.project_up(X_d), w