├── ExactGaussianProcessV3.py ├── README.md ├── REMBO_demonstration.py ├── bayes_opt.py ├── bayes_opt_demonstration.py ├── dict_to_tensor_IO_demonstration.py ├── rembo.py ├── requirements.txt └── utils.py /ExactGaussianProcessV3.py: -------------------------------------------------------------------------------- 1 | from abc import ABC 2 | 3 | import torch 4 | import gpytorch 5 | 6 | from botorch.models import SingleTaskGP 7 | from botorch import fit_gpytorch_model 8 | 9 | # # use a GPU if available 10 | # device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 11 | # dtype = torch.float 12 | 13 | 14 | # GPyTorch 15 | class ExactGaussianProcess(SingleTaskGP): 16 | def __init__(self, train_x, train_y, 17 | likelihood=gpytorch.likelihoods.GaussianLikelihood(), 18 | covar_module=gpytorch.kernels.ScaleKernel( 19 | gpytorch.kernels.MaternKernel(nu=2.5)) 20 | ): 21 | """ 22 | 23 | :param train_x: (n, d) torch.Tensor 24 | :param train_y: (n, 1) torch.Tensor 25 | :param likelihood: 26 | :param covar_module: Default assumes that all dimensions of x are of the 27 | same scale. This assumption requires data preprocessing. 28 | """ 29 | train_X = train_x.float() 30 | train_Y = train_y.float() 31 | 32 | super().__init__(train_X=train_X, 33 | train_Y=train_Y, 34 | likelihood=likelihood, 35 | covar_module=covar_module) 36 | 37 | def fit(self, train_x_, train_y_): 38 | """ 39 | Fit the Gaussian Process to training data on 40 | the marginal log likelihood. (refits the model hyperparameters) 41 | 42 | Code based on the following GPyTorch tutorial: 43 | https://gpytorch.readthedocs.io/en/latest/examples/01_Exact_GPs/Simple_GP_Regression.html#Training-the-model 44 | 45 | :param train_x_: torch.Tensor (n, d) 46 | :param train_y_: torch.Tensor (n, 1) 47 | """ 48 | 49 | train_X = train_x_.float() 50 | train_Y = train_y_.float() 51 | 52 | # Update self.train_x and self.train_y 53 | self.set_train_data(inputs=train_X, targets=train_Y) 54 | 55 | # "Loss" for GPs - the marginal log likelihood 56 | mll = gpytorch.mlls.ExactMarginalLogLikelihood(self.likelihood, self) 57 | mll = mll.to(train_X) 58 | 59 | fit_gpytorch_model(mll) 60 | 61 | def set_train_data(self, inputs=None, targets=None, strict=True): 62 | """ 63 | ** Adapted from gpytorch.models.exactgp ** 64 | Set training data (does not re-fit model hyper-parameters). 65 | 66 | :param torch.Tensor inputs: The new training inputs. 67 | :param torch.Tensor targets: The new training targets. 68 | :param bool strict: (default True) If `True`, the new inputs and 69 | targets must have the same dtype/device as the current inputs and 70 | targets. Otherwise, any dtype/device are allowed. 71 | """ 72 | if inputs is not None: 73 | if torch.is_tensor(inputs): 74 | inputs = (inputs,) 75 | inputs = tuple(input_.unsqueeze(-1) if input_.ndimension() == 1 else input_ for input_ in inputs) 76 | if strict: 77 | for input_, t_input in zip(inputs, self.train_inputs or (None,)): 78 | for attr in {"dtype", "device"}: 79 | expected_attr = getattr(t_input, attr, None) 80 | found_attr = getattr(input_, attr, None) 81 | if expected_attr != found_attr: 82 | msg = "Cannot modify {attr} of inputs (expected {e_attr}, found {f_attr})." 83 | msg = msg.format(attr=attr, e_attr=expected_attr, f_attr=found_attr) 84 | raise RuntimeError(msg) 85 | self.train_inputs = inputs[0] 86 | if targets is not None: 87 | if strict: 88 | for attr in {"dtype", "device"}: 89 | expected_attr = getattr(self.train_targets, attr, None) 90 | found_attr = getattr(targets, attr, None) 91 | if expected_attr != found_attr: 92 | msg = "Cannot modify {attr} of targets (expected {e_attr}, found {f_attr})." 93 | msg = msg.format(attr=attr, e_attr=expected_attr, f_attr=found_attr) 94 | raise RuntimeError(msg) 95 | self.train_targets = targets 96 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # PyREMBO 2 | 3 | Python implementation of REMBO [1] built on BoTorch/GPyTorch/PyTorch, allowing 4 | for batched sampling. 5 | 6 | Some code adapted from 7 | https://github.com/jmetzen/bayesian_optimization/blob/master/bayesian_optimization/bayesian_optimization.py 8 | 9 | REMBO demonstration can be found in 'REMBO_demonstration.py'. 10 | Vanilla BO demonstration in 'bayes_opt.py'. 11 | utils.dict_to_tensor_IO demonstration in 'dict_to_tensor_IO_demonstration.py'. 12 | 13 | References \ 14 | [1] Ziyu Wang and Masrour Zoghi and Frank Hutter and David Matheson and 15 | Nando de Freitas Bayesian Optimization in High Dimensions via Random 16 | Embeddings. Proceedings of the 23rd international joint conference 17 | on Artificial Intelligence (IJCAI) 18 | 19 | Work in progress: 20 | - 21 | - Implement interleaved REMBO 22 | - Allow for initial data 23 | - Use 'gpytorch.likelihoods.FixedNoiseGaussianLikelihood' to add a fixed, 24 | known observation noise to the GP predictions 25 | -------------------------------------------------------------------------------- /REMBO_demonstration.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from rembo import REMBO 4 | 5 | 6 | def main(): 7 | n_dims = 20 8 | d_embedding = 4 9 | n_trials = 30 10 | 11 | ind = np.random.RandomState(seed=0).choice(n_dims, 2, replace=False) 12 | def f(X): # black-box objective function to minimize 13 | """ 14 | minimum value of 0 15 | """ 16 | x1 = X[0][ind[0]] 17 | x2 = X[0][ind[1]] 18 | return x1**2 + x2**2 19 | 20 | def ensure_not_1D(x): 21 | """ 22 | Ensure x is not 1D (i.e. make size (D,) data into size (1,D)) 23 | :param x: torch.Tensor 24 | :return: 25 | """ 26 | import torch 27 | 28 | if x.ndim == 1: 29 | if isinstance(x, np.ndarray): 30 | x = np.expand_dims(x, axis=0) 31 | elif isinstance(x, torch.Tensor): 32 | x = x.unsqueeze(0) 33 | return x 34 | 35 | original_boundaries = np.array([[-1, 1]] * n_dims) 36 | print("original_boundaries.shape: {}".format(original_boundaries.shape)) 37 | opt = REMBO(original_boundaries, d_embedding) 38 | 39 | # Perform optimization 40 | for i in range(n_trials): 41 | X_queries, X_queries_embedded = opt.select_query_point(batch_size=3) 42 | 43 | # Ensure not 1D (i.e. size (D,)) 44 | X_queries = ensure_not_1D(X_queries) 45 | 46 | # Evaluate the batch of query points 1-by-1 47 | for row_idx in range(len(X_queries)): 48 | X_query = X_queries[row_idx] 49 | X_query_embedded = X_queries_embedded[row_idx] 50 | 51 | # Ensure no 1D tensors (i.e. expand tensors of size (D,)) 52 | X_query = ensure_not_1D(X_query) 53 | X_query_embedded = ensure_not_1D(X_query_embedded) 54 | 55 | y_query = -f(X_query) 56 | opt.update(X_query, y_query, X_query_embedded) 57 | 58 | print("best y value: {}".format(opt.best_value())) 59 | print("best actual x: {}".format(opt.best_params()[0][ind[:2]])) 60 | print("best actual x values distance from 0: {}".format( 61 | np.linalg.norm(opt.best_params()[0][ind[:2]]))) 62 | print("---------------------") 63 | 64 | 65 | if __name__ == "__main__": 66 | main() 67 | -------------------------------------------------------------------------------- /bayes_opt.py: -------------------------------------------------------------------------------- 1 | import logging 2 | 3 | import torch 4 | import numpy as np 5 | from ExactGaussianProcessV3 import ExactGaussianProcess 6 | from sklearn.utils import check_random_state 7 | from botorch.acquisition.monte_carlo import qExpectedImprovement, \ 8 | qUpperConfidenceBound 9 | from botorch import fit_gpytorch_model 10 | from scipy.optimize import fmin_l_bfgs_b 11 | from botorch.acquisition.objective import ConstrainedMCObjective 12 | 13 | from utils import global_optimization, get_fitted_model 14 | 15 | class bayes_opt(): 16 | """ 17 | Maximize a black-box objective function. 18 | """ 19 | # TODO: implement functionality for discrete variables, 20 | # implement acquisition functions 21 | 22 | def __init__(self, boundaries, initial_x=None, 23 | initial_y=None, 24 | acquisition_func=qExpectedImprovement, 25 | maxf=1000, optimizer="random+GD", 26 | initial_random_samples=5, 27 | opt=None, fopt=None, random_embedding_seed=0, 28 | types=None, do_scaling=True): 29 | """ 30 | Vanilla Bayesian optimization. 31 | 32 | Parameters 33 | ---------- 34 | original_boundaries ((D, 2) np.array): Boundaries of the original search 35 | space (of dimension D). The first column is the minimum value for 36 | the corresponding dimension/row, and the second column is the 37 | maximum value. 38 | initial_x: np.array 39 | Initial data points (in original data space) TODO: Implement 40 | initial_y: np.array 41 | Initial function evaluations TODO: Implement 42 | acquisition_func (str): Acquisition function to use. # TODO: Implement 43 | maxf (int): Maximum number of acquisition function evaluations that the 44 | optimizer can make. 45 | optimizer (str): Method name to use for optimizing the acquisition 46 | function. 47 | opt: (N, D) numpy array 48 | The global optima of the objective function (if known). 49 | Allows to compute and plot the distance of the incumbent 50 | to the global optimum. 51 | fopt: (N, 1) numpy array 52 | Function value of the N global optima (if known). Useful 53 | to compute the immediate or cumulative regret. 54 | """ 55 | self.initial_random_samples = initial_random_samples 56 | self.acquisition_func = acquisition_func 57 | self.optimizer = optimizer 58 | self.maxf = maxf 59 | self.rng = check_random_state(random_embedding_seed) 60 | self.opt = opt # optimal point 61 | self.fopt = fopt # optimal function value 62 | self.boundaries = np.asarray(boundaries) 63 | 64 | # Dimensions of the original space 65 | self.d_orig = self.boundaries.shape[0] 66 | 67 | self.X = torch.Tensor() # running list of data 68 | self.y = torch.Tensor() # running list of function evaluations 69 | 70 | self.model = None 71 | self.boundaries_cache = {} 72 | 73 | def select_query_point(self, batch_size=1): 74 | """ 75 | 76 | :param 77 | batch_size (int): number of query points to return 78 | :return: 79 | (batch_size x d_orig) numpy array 80 | """ 81 | 82 | # TODO: Make the random initialization its own function so it can be done separately from the acquisition argmin 83 | # Initialize with random points 84 | if len(self.X) < self.initial_random_samples: 85 | 86 | # Select query point randomly from embedding_boundaries 87 | X_query = \ 88 | self.rng.uniform(size=self.boundaries.shape[0]) \ 89 | * (self.boundaries[:, 1] - self.boundaries[:, 0]) \ 90 | + self.boundaries[:, 0] 91 | X_query = torch.from_numpy(X_query).unsqueeze(0) 92 | 93 | # Query by maximizing the acquisition function 94 | else: 95 | print("---------------------") 96 | print('querying') 97 | 98 | print("self.X.shape: {}".format(self.X.shape)) 99 | print("self.y.shape: {}".format(self.y.shape)) 100 | # Initialize model 101 | if len(self.X) == self.initial_random_samples: 102 | self.model = ExactGaussianProcess( 103 | train_x=self.X.float(), 104 | train_y=self.y.float(), 105 | ) 106 | 107 | # Acquisition function 108 | qEI = qExpectedImprovement( 109 | model=self.model, 110 | best_f=torch.max(self.y).item(), 111 | ) 112 | # qUCB = qUpperConfidenceBound( 113 | # model=self.model, 114 | # beta=2.0, 115 | # ) 116 | 117 | print("batch_size: {}".format(batch_size)) 118 | 119 | # Optimize for a (batch_size x d_embedding) tensor query point 120 | X_query = global_optimization( 121 | objective_function=qEI, 122 | boundaries=torch.from_numpy(self.boundaries).float(), 123 | batch_size=batch_size, # number of query points to suggest 124 | ) 125 | 126 | print("batched X_query: {}".format(X_query)) 127 | print("batched X_query.shape: {}".format(X_query.shape)) 128 | 129 | print("X concatenated: {}".format(self.X.shape)) 130 | 131 | return X_query 132 | 133 | def update(self, X_query, y_query): 134 | """ Update internal model for observed (X, y) from true function. 135 | The function is meant to be used as follows. 136 | 1. Call 'select_query_point' to update self.X_embedded with a new 137 | embedded query point, and to return a query point X_query in the 138 | original (unscaled) search space 139 | 2. Evaluate X_query to get y_query 140 | 3. Call this function ('update') to update the surrogate model (e.g. 141 | Gaussian Process) 142 | 143 | Args: 144 | X_query ((1,d_orig) np.array): 145 | Point in original input space to query 146 | y_query (float): 147 | Value of black-box function evaluated at X_query 148 | X_query_embedded ((1, d_embedding) np.array): 149 | Point in embedding space which maps 1:1 with X_query 150 | """ 151 | print("X_query.shape: {}".format(X_query.shape)) 152 | print("y_query.shape: {}".format(y_query.shape)) 153 | 154 | # add new rows of data 155 | self.X = torch.cat([self.X.float(), 156 | X_query.float()], 157 | dim=0) 158 | self.y = torch.cat([self.y, torch.Tensor([[y_query]])], axis=0) 159 | 160 | print("self.X_embedded.shape: {}".format(self.X.shape)) 161 | print("self.y.shape: {}".format(self.y.shape)) 162 | self.model = get_fitted_model(self.X.float(), 163 | self.y.float()) 164 | 165 | def best_params(self): 166 | """ Returns the best parameters found so far.""" 167 | return self.X[np.argmax(self.y.numpy())] 168 | 169 | def best_value(self): 170 | """ Returns the optimal value found so far.""" 171 | return np.max(self.y.numpy()) 172 | 173 | def evaluate_f(self, x_query, black_box_function=None): 174 | """ 175 | Evaluates input point in embedded space by first projecting back to 176 | original space and then scaling it to its original boundaries. 177 | 178 | Args: 179 | :return: 180 | """ 181 | # BoTorch assumes a maximization problem 182 | if black_box_function is not None: 183 | return -black_box_function(x_query) 184 | -------------------------------------------------------------------------------- /bayes_opt_demonstration.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from bayes_opt import bayes_opt 4 | 5 | 6 | def main(): 7 | n_dims = 5 8 | n_trials = 30 9 | 10 | ind = np.random.RandomState(seed=0).choice(n_dims, 2, replace=False) 11 | def f(X): # black-box objective function to minimize 12 | """ 13 | minimum value of 0 14 | """ 15 | x1 = X[0][ind[0]] 16 | x2 = X[0][ind[1]] 17 | return x1**2 + x2**2 18 | 19 | def ensure_not_1D(x): 20 | """ 21 | Ensure x is not 1D (i.e. make size (D,) data into size (1,D)) 22 | :param x: torch.Tensor 23 | :return: 24 | """ 25 | import torch 26 | 27 | if x.ndim == 1: 28 | if isinstance(x, np.ndarray): 29 | x = np.expand_dims(x, axis=0) 30 | elif isinstance(x, torch.Tensor): 31 | x = x.unsqueeze(0) 32 | return x 33 | 34 | original_boundaries = np.array([[-1, 1]] * n_dims) 35 | print("original_boundaries.shape: {}".format(original_boundaries.shape)) 36 | opt = bayes_opt(original_boundaries) 37 | 38 | # Perform optimization 39 | for i in range(n_trials): 40 | X_queries = opt.select_query_point(batch_size=3) 41 | 42 | # Ensure not 1D (i.e. size (D,)) 43 | X_queries = ensure_not_1D(X_queries) 44 | 45 | # Evaluate the batch of query points 1-by-1 46 | for row_idx in range(len(X_queries)): 47 | X_query = X_queries[row_idx] 48 | 49 | # Ensure no 1D tensors (i.e. expand tensors of size (D,)) 50 | X_query = ensure_not_1D(X_query) 51 | 52 | y_query = -f(X_query) 53 | opt.update(X_query, y_query) 54 | 55 | print("best y value: {}".format(opt.best_value())) 56 | print("best actual x: {}".format(opt.best_params()[ind[:2]])) 57 | print("best actual x values distance from 0: {}".format( 58 | np.linalg.norm(opt.best_params()[ind[:2]]))) 59 | print("---------------------") 60 | 61 | 62 | if __name__ == "__main__": 63 | main() 64 | -------------------------------------------------------------------------------- /dict_to_tensor_IO_demonstration.py: -------------------------------------------------------------------------------- 1 | from utils import dict_to_tensor_IO 2 | 3 | def main(): 4 | dict_data = {'a': 10, 'b': 20, 'c': 30, 'd': 40} 5 | converter = dict_to_tensor_IO(dict_data=dict_data) 6 | test_tensor = converter.map_dict_state_point_to_tensor(dict_data=dict_data) 7 | print(test_tensor) 8 | 9 | test_dict = converter.map_tensor_state_point_to_dict(test_tensor) 10 | print(test_dict) 11 | 12 | if __name__ == "__main__": 13 | main() -------------------------------------------------------------------------------- /rembo.py: -------------------------------------------------------------------------------- 1 | import logging 2 | 3 | import torch 4 | import numpy as np 5 | from ExactGaussianProcessV3 import ExactGaussianProcess 6 | from sklearn.utils import check_random_state 7 | from botorch.acquisition.monte_carlo import qExpectedImprovement, \ 8 | qUpperConfidenceBound 9 | from botorch import fit_gpytorch_model 10 | from scipy.optimize import fmin_l_bfgs_b 11 | from botorch.acquisition.objective import ConstrainedMCObjective 12 | 13 | from utils import global_optimization, get_fitted_model 14 | 15 | class REMBO(): 16 | """ 17 | Maximize a black-box objective function. 18 | """ 19 | # TODO: implement functionality for discrete variables, 20 | # implement acquisition functions 21 | 22 | def __init__(self, original_boundaries, d_embedding, initial_x=None, 23 | initial_y=None, 24 | embedding_boundaries_setting="constant", 25 | acquisition_func=qExpectedImprovement, 26 | n_keep_dims=0, maxf=1000, optimizer="random+GD", 27 | initial_random_samples=5, 28 | opt=None, fopt=None, random_embedding_seed=0, 29 | types=None, do_scaling=True): 30 | """ 31 | Random EMbedding Bayesian Optimization [1] maps the original space 32 | in a lower dimensional subspace via a random embedding matrix and 33 | performs Bayesian Optimization only in this lower dimensional 34 | subspace. 35 | 36 | [1] Ziyu Wang and Masrour Zoghi and Frank Hutter and David Matheson 37 | and Nando de Freitas 38 | Bayesian Optimization in High Dimensions via Random Embeddings 39 | In: Proceedings of the 23rd international joint conference 40 | on Artificial Intelligence (IJCAI) 41 | 42 | Parameters 43 | ---------- 44 | original_boundaries ((D, 2) np.array): Boundaries of the original search 45 | space (of dimension D). This is used for rescaling. The first column 46 | is the minimum value for the corresponding dimension/row, and the 47 | second column is the maximum value. 48 | n_keep_dims (int): Number of dimensions in the original space that are 49 | preserved in the embedding. This is used if certain dimensions are 50 | expected to be independently relevant. Assume that these dimensions 51 | are the first parameters of the input space. 52 | d_embedding: int 53 | Number of dimensions for the lower dimensional subspace 54 | initial_x: np.array 55 | Initial data points (in original data space) TODO: Implement 56 | initial_y: np.array 57 | Initial function evaluations TODO: Implement 58 | embedding_boundaries_setting (str): "auto" causes embedding boundaries to be 59 | computed approximately. "constant" will set all dimensions of the 60 | embedding boundaries to be set to 61 | [-sqrt(d_embedding), sqrt(d_embedding)], as described in [1]. 62 | acquisition_func (str): Acquisition function to use. # TODO: Implement 63 | maxf (int): Maximum number of acquisition function evaluations that the 64 | optimizer can make. 65 | optimizer (str): Method name to use for optimizing the acquisition 66 | function. 67 | opt: (N, D) numpy array 68 | The global optima of the objective function (if known). 69 | Allows to compute and plot the distance of the incumbent 70 | to the global optimum. 71 | fopt: (N, 1) numpy array 72 | Function value of the N global optima (if known). Useful 73 | to compute the immediate or cumulative regret. 74 | do_scaling: boolean 75 | If set to true the input space is scaled to [-1, 1]. Useful 76 | to specify priors for the kernel lengthscale. 77 | """ 78 | self.initial_random_samples = initial_random_samples 79 | self.acquisition_func = acquisition_func 80 | self.optimizer = optimizer 81 | self.maxf = maxf 82 | self.rng = check_random_state(random_embedding_seed) 83 | self.opt = opt # optimal point 84 | self.fopt = fopt # optimal function value 85 | self.n_keep_dims = n_keep_dims 86 | self.original_boundaries = np.asarray(original_boundaries) 87 | self.embedding_boundaries_setting = embedding_boundaries_setting 88 | 89 | # Dimensions of the original space 90 | self.d_orig = self.original_boundaries.shape[0] 91 | 92 | # Dimension of the embedded space 93 | self.d_embedding = d_embedding 94 | 95 | # Draw random matrix from a standard normal distribution 96 | # (x = A @ x_embedding) 97 | self.A = self.rng.normal(loc=0.0, 98 | scale=1.0, 99 | size=(self.d_orig, 100 | self.d_embedding - self.n_keep_dims)) 101 | 102 | self.X = [] # running list of data 103 | self.X_embedded = torch.Tensor() # running list of embedded data 104 | self.y = torch.Tensor() # running list of function evaluations 105 | 106 | self.model = None 107 | self.boundaries_cache = {} 108 | 109 | def select_query_point(self, batch_size=1): 110 | """ 111 | 112 | :param 113 | batch_size (int): number of query points to return 114 | :return: 115 | (batch_size x d_orig) numpy array 116 | """ 117 | 118 | # Produces (d_embedding, 2) array 119 | if self.embedding_boundaries_setting == "auto": 120 | # Approximately compute boundaries on embedded space 121 | embedding_boundaries = self._compute_boundaries_embedding( 122 | self.original_boundaries) 123 | elif self.embedding_boundaries_setting == "constant": 124 | # As described in the original paper. This is default. 125 | embedding_boundaries = np.array( 126 | [[-np.sqrt(self.d_embedding), 127 | np.sqrt(self.d_embedding)]] * self.d_embedding) 128 | else: 129 | raise NotImplementedError("embedding_boundaries_setting must be " 130 | "'auto' or 'constant'.") 131 | 132 | # TODO: Make the random initialization its own function so it can be done separately from the acquisition argmin 133 | # Initialize with random points 134 | if len(self.X) < self.initial_random_samples: 135 | 136 | # Select query point randomly from embedding_boundaries 137 | X_query_embedded = \ 138 | self.rng.uniform(size=embedding_boundaries.shape[0]) \ 139 | * (embedding_boundaries[:, 1] - embedding_boundaries[:, 0]) \ 140 | + embedding_boundaries[:, 0] 141 | X_query_embedded = torch.from_numpy(X_query_embedded).unsqueeze(0) 142 | 143 | print("X_query_embedded.shape: {}".format(X_query_embedded.shape)) 144 | 145 | # Query by maximizing the acquisition function 146 | else: 147 | print("---------------------") 148 | print('querying') 149 | 150 | print("self.X_embedded.shape: {}".format(self.X_embedded.shape)) 151 | print("self.y.shape: {}".format(self.y.shape)) 152 | # Initialize model 153 | if len(self.X) == self.initial_random_samples: 154 | self.model = ExactGaussianProcess( 155 | train_x=self.X_embedded.float(), 156 | train_y=self.y.float(), 157 | ) 158 | 159 | # Acquisition function 160 | qEI = qExpectedImprovement( 161 | model=self.model, 162 | best_f=torch.max(self.y).item(), 163 | ) 164 | # qUCB = qUpperConfidenceBound( 165 | # model=self.model, 166 | # beta=2.0, 167 | # ) 168 | 169 | print("batch_size: {}".format(batch_size)) 170 | 171 | # Optimize for a (batch_size x d_embedding) tensor query point 172 | X_query_embedded = global_optimization( 173 | objective_function=qEI, 174 | boundaries=torch.from_numpy(embedding_boundaries).float(), 175 | batch_size=batch_size, # number of query points to suggest 176 | ) 177 | 178 | print("batched X_query_embedded: {}".format(X_query_embedded)) 179 | print("batched X_query_embedded.shape: {}".format(X_query_embedded.shape)) 180 | 181 | print("X_embedded concatenated: {}".format(self.X_embedded.shape)) 182 | 183 | # Map to higher dimensional space and clip to hard boundaries [-1, 1] 184 | X_query = np.clip(a=self._manifold_to_dataspace(X_query_embedded.numpy()), 185 | a_min=self.original_boundaries[:, 0], 186 | a_max=self.original_boundaries[:, 1]) 187 | 188 | return X_query, X_query_embedded 189 | 190 | def _manifold_to_dataspace(self, X_embedded): 191 | """ 192 | Map data from manifold to original data space. 193 | 194 | :param X_embedded: (1 x d_embedding) numpy.array 195 | :return: (1 x d_orig) numpy.array 196 | """ 197 | 198 | if self.A.shape[1] != X_embedded.shape[0]: 199 | X_embedded = X_embedded.T 200 | 201 | X_query_kd = self.A @ X_embedded[self.n_keep_dims:] 202 | X_query_kd = X_query_kd.T # make X_query_kd of shape (q x d) 203 | 204 | # concatenate column-wise 205 | if self.n_keep_dims > 0: 206 | X_query = np.hstack((X_embedded[:self.n_keep_dims], X_query_kd)) 207 | else: 208 | X_query = X_query_kd 209 | 210 | X_query = self._unscale(X_query) 211 | return X_query 212 | 213 | def update(self, X_query, y_query, X_query_embedded): 214 | """ Update internal model for observed (X, y) from true function. 215 | The function is meant to be used as follows. 216 | 1. Call 'select_query_point' to update self.X_embedded with a new 217 | embedded query point, and to return a query point X_query in the 218 | original (unscaled) search space 219 | 2. Evaluate X_query to get y_query 220 | 3. Call this function ('update') to update the surrogate model (e.g. 221 | Gaussian Process) 222 | 223 | Args: 224 | X_query ((1,d_orig) np.array): 225 | Point in original input space to query 226 | y_query (float): 227 | Value of black-box function evaluated at X_query 228 | X_query_embedded ((1, d_embedding) np.array): 229 | Point in embedding space which maps 1:1 with X_query 230 | """ 231 | print("X_query.shape: {}".format(X_query.shape)) 232 | print("y_query.shape: {}".format(y_query.shape)) 233 | 234 | # add new rows of data 235 | self.X.append(X_query) 236 | self.y = torch.cat([self.y, torch.Tensor([[y_query]])], axis=0) 237 | 238 | # Append to (n x d_embedding) Tensor 239 | self.X_embedded = torch.cat([self.X_embedded.float(), 240 | X_query_embedded.float()], 241 | dim=0) 242 | 243 | print("self.X_embedded.shape: {}".format(self.X_embedded.shape)) 244 | print("self.y.shape: {}".format(self.y.shape)) 245 | self.model = get_fitted_model(self.X_embedded.float(), 246 | self.y.float()) 247 | 248 | def best_params(self): 249 | """ Returns the best parameters found so far.""" 250 | return self.X[np.argmax(self.y.numpy())] 251 | 252 | def best_value(self): 253 | """ Returns the optimal value found so far.""" 254 | return np.max(self.y.numpy()) 255 | 256 | def _unscale(self, x_scaled, 257 | scaled_lower_bound=-1, scaled_upper_bound=1): 258 | """ 259 | Use this function to scale x_scaled (i.e. 'X_query' from 'select_query_point') 260 | to the original input space boundaries so that it can be evaluated by the 261 | function, 'evaluate_f'. 262 | 263 | :param x_scaled: (1, D) numpy.array 264 | :param scaled_lower_bound: int 265 | :param scaled_upper_bound: int 266 | :return: 267 | """ 268 | x_scaled_ = np.copy(x_scaled) 269 | if not x_scaled_.ndim == 2: 270 | # Add a dimension to be (1,D) if x_scaled is (D,) 271 | x_scaled_ = np.expand_dims(x_scaled_, axis=0) 272 | 273 | x_unscaled = np.empty(x_scaled_.shape) 274 | unscaled_ranges = self.original_boundaries[:, 1] \ 275 | - self.original_boundaries[:, 0] # (D,) max-min 276 | scaled_range = scaled_upper_bound - scaled_lower_bound 277 | 278 | # Iterate through each feature 279 | for dim in range(len(x_scaled_)): 280 | x_unscaled[:][dim] = (x_scaled_[:][dim] - scaled_lower_bound) \ 281 | * (unscaled_ranges[dim] / scaled_range) \ 282 | + self.original_boundaries[dim][0] 283 | return x_unscaled 284 | 285 | def evaluate_f(self, x_query, black_box_function=None): 286 | """ 287 | Evaluates input point in embedded space by first projecting back to 288 | original space and then scaling it to its original boundaries. 289 | 290 | Args: 291 | :return: 292 | """ 293 | # BoTorch assumes a maximization problem 294 | if black_box_function is not None: 295 | return -black_box_function(x_query) 296 | 297 | def _compute_boundaries_embedding(self, boundaries): 298 | """ Approximate box constraint boundaries on low-dimensional manifold 299 | 300 | Args: 301 | boundaries ((d_orig, 2) np.array): 302 | Returns: 303 | boundaries_embedded ((d_embedding, 2) np.narray): 304 | 305 | """ 306 | # Check if boundaries have been determined before 307 | boundaries_hash = hash(boundaries[self.n_keep_dims:].tostring()) 308 | if boundaries_hash in self.boundaries_cache: 309 | boundaries_embedded = \ 310 | np.array(self.boundaries_cache[boundaries_hash]) 311 | boundaries_embedded[:self.n_keep_dims] = \ 312 | boundaries[:self.n_keep_dims] # Overwrite keep-dim's boundaries 313 | return boundaries_embedded 314 | 315 | # Determine boundaries on embedded space 316 | boundaries_embedded = \ 317 | np.empty((self.n_keep_dims + self.d_embedding, 2)) 318 | boundaries_embedded[:self.n_keep_dims] = boundaries[:self.n_keep_dims] 319 | for dim in range(self.n_keep_dims, 320 | self.n_keep_dims + self.d_embedding): 321 | x_embedded = np.zeros(self.n_keep_dims + self.d_embedding) 322 | while True: 323 | x = self._manifold_to_dataspace(x_embedded) 324 | if np.sum(np.logical_or( 325 | x[self.n_keep_dims:] < boundaries[self.n_keep_dims:, 0], 326 | x[self.n_keep_dims:] > boundaries[self.n_keep_dims:, 327 | 1])) \ 328 | > (self.d_orig - self.n_keep_dims) / 2: 329 | break 330 | x_embedded[dim] -= 0.01 331 | boundaries_embedded[dim, 0] = x_embedded[dim] 332 | 333 | x_embedded = np.zeros(self.n_keep_dims + self.d_embedding) 334 | while True: 335 | x = self._manifold_to_dataspace(x_embedded) 336 | if np.sum(np.logical_or( 337 | x[self.n_keep_dims:] < boundaries[self.n_keep_dims:, 0], 338 | x[self.n_keep_dims:] > boundaries[self.n_keep_dims:, 1])) \ 339 | > (self.d_orig - self.n_keep_dims) / 2: 340 | break 341 | x_embedded[dim] += 0.01 342 | boundaries_embedded[dim, 1] = x_embedded[dim] 343 | 344 | self.boundaries_cache[boundaries_hash] = boundaries_embedded 345 | 346 | return boundaries_embedded 347 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | numpy 2 | torch 3 | gpytorch 4 | botorch -------------------------------------------------------------------------------- /utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import gpytorch 3 | import torch 4 | from sklearn.utils import check_random_state 5 | from botorch.models import SingleTaskGP 6 | from gpytorch.mlls.exact_marginal_log_likelihood import ExactMarginalLogLikelihood 7 | from botorch import fit_gpytorch_model 8 | 9 | 10 | def global_optimization(objective_function, boundaries, batch_size): 11 | """ 12 | 13 | :param objective_function: 14 | :param boundaries (torch.Tensor): A `2 x d` tensor of lower and upper bounds 15 | for each column of `X`. 16 | :param batch_size (int): Number of candidates to return. 17 | :param optimizer: 18 | :param maxf: 19 | :param random: 20 | Returns: 21 | `(num_restarts) x q x d`-dim tensor of generated candidates 22 | """ 23 | 24 | # # USING BOTORCH 25 | from botorch.optim import optimize_acqf 26 | 27 | # Boundaries must be (2 x d) for optimize_acqf to work 28 | if boundaries.shape[0] != 2: 29 | boundaries = boundaries.T 30 | 31 | # optimize 32 | candidates, _ = optimize_acqf( 33 | acq_function=objective_function, 34 | bounds=boundaries, 35 | q=batch_size, 36 | num_restarts=10, # number of initial points for optimization 37 | raw_samples=512, # used for initialization heuristic 38 | return_best_only=True # only returns the best of the n_restarts random restarts 39 | ) 40 | # what exactly is raw_samples? 41 | 42 | # Removes the 'candidates' variable from the computational graph 43 | new_x = candidates.detach() 44 | return new_x 45 | 46 | 47 | class ExactGPModel(gpytorch.models.ExactGP): 48 | def __init__(self, train_x, train_y, likelihood): 49 | super(ExactGPModel, self).__init__(train_x, train_y, likelihood) 50 | self.mean_module = gpytorch.means.ConstantMean() 51 | self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel()) 52 | 53 | def forward(self, x): 54 | mean_x = self.mean_module(x) 55 | covar_x = self.covar_module(x) 56 | return gpytorch.distributions.MultivariateNormal(mean_x, covar_x) 57 | 58 | 59 | def get_fitted_model(train_x, train_obj, state_dict=None): 60 | # initialize and fit model 61 | model = SingleTaskGP(train_X=train_x, train_Y=train_obj) 62 | 63 | # # initialize likelihood and model 64 | # likelihood = gpytorch.likelihoods.GaussianLikelihood() 65 | # model = ExactGPModel(train_x, train_obj, likelihood) 66 | # model.train() 67 | # likelihood.train() 68 | 69 | if state_dict is not None: 70 | model.load_state_dict(state_dict) 71 | mll = ExactMarginalLogLikelihood(model.likelihood, model) 72 | mll.to(train_x) 73 | fit_gpytorch_model(mll) 74 | 75 | return model 76 | 77 | 78 | class dict_to_tensor_IO(): 79 | """ 80 | Map back and forth between dictionary data and tensor data. 81 | 82 | Args: 83 | dict_data (dict) 84 | State point described with a key-value pairs. 85 | Returns: 86 | tensor_state_point ((1, d) torch.Tensor) 87 | A tensor version of dict_data, where 'd' is the number of 88 | keys/parameters/features/dimensions. 89 | 90 | """ 91 | 92 | def __init__(self, dict_data=None): 93 | self.column_indexes_to_keys = {} 94 | 95 | # Create the mapping as a dictionary of column_index:key pairs 96 | if dict_data is not None: 97 | for index, (key, value) in enumerate(dict_data.items()): 98 | self.column_indexes_to_keys[index] = key 99 | else: 100 | raise AttributeError('A state point dictionary was not passed.') 101 | 102 | def map_dict_state_point_to_tensor(self, dict_data=None): 103 | """ 104 | Convert state point 'dict_data' into a (1, d) tensor using the map, 105 | 'column_indexes_to_keys', and then return the tensor. 106 | """ 107 | if dict_data is not None: 108 | tensor_state_point = torch.Tensor() # will make this (1, d) 109 | for index, key in self.column_indexes_to_keys.items(): 110 | value = torch.Tensor([[dict_data[key]]]) 111 | tensor_state_point = torch.cat([tensor_state_point, value], 112 | dim=1) 113 | return tensor_state_point 114 | else: 115 | raise AttributeError('A state point dictionary was not passed.') 116 | 117 | def map_tensor_state_point_to_dict(self, tensor_state_point=None): 118 | """ 119 | Convert (1, d) state point 'tensor_data' into a dict using the map, 120 | 'column_indexes_to_keys', and then return the dict. 121 | """ 122 | dict_state_point = {} 123 | if (tensor_state_point is not None) or \ 124 | tensor_state_point.shape[0] != 1 or \ 125 | tensor_state_point.ndim != 2: 126 | 127 | for index, key in self.column_indexes_to_keys.items(): 128 | value = tensor_state_point[0][index] 129 | dict_state_point[key] = float(value) 130 | 131 | return dict_state_point 132 | else: 133 | raise AttributeError('A (1,d) state point tensor was not passed.') 134 | --------------------------------------------------------------------------------