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