From b61d3f8f19290f82121726bbd3a8708be02ab3e4 Mon Sep 17 00:00:00 2001 From: Sooyoung Cheong <64125280+c-sooyoung@users.noreply.github.com> Date: Sun, 19 Jul 2026 18:16:16 +0900 Subject: [PATCH] added SOBO --- bo/__init__.py | 1 + bo/bo_base.py | 6 -- bo/random.py | 116 ++++++++++++++++++------------ bo/sobo.py | 186 +++++++++++++++++++++++++++++++++++++++++++++++++ bo/ucb.py | 48 ------------- 5 files changed, 257 insertions(+), 100 deletions(-) create mode 100644 bo/sobo.py delete mode 100644 bo/ucb.py diff --git a/bo/__init__.py b/bo/__init__.py index f2d3920..d7e45fc 100644 --- a/bo/__init__.py +++ b/bo/__init__.py @@ -1,2 +1,3 @@ from .bo_base import BOEngine from .random import RandomBOEngine +from .sobo import SingleObjectiveBOEngine diff --git a/bo/bo_base.py b/bo/bo_base.py index 0e9ca4c..5d1c4a4 100644 --- a/bo/bo_base.py +++ b/bo/bo_base.py @@ -2,16 +2,10 @@ from abc import ABC, abstractmethod class BOEngine(ABC): - name = None - def __init__(self, config): self.config = config self.state = None - @abstractmethod - def initialize(self): - pass - @abstractmethod def ask(self): pass diff --git a/bo/random.py b/bo/random.py index e5db64f..dfa020f 100644 --- a/bo/random.py +++ b/bo/random.py @@ -8,81 +8,105 @@ class RandomBOEngine(BOEngine): def __init__(self, config): super().__init__(config) - - def initialize(self): - config = self.config - bo_params = [ - key - for key, value in config['bo']['params'].items() - if value is not None + key for key, spec in config["bo"]["params"].items() if spec is not None ] - bo_state = { - 'algorithm': 'random', - 'params': bo_params, - 'train_x': np.empty((0, len(bo_params))), - 'train_y': np.empty((0,)), + bo_param_types = { + key: config["bo"]["params"][key].get("type", "float") for key in bo_params } - train_x_path = config['bo'].get('train_x') - train_y_path = config['bo'].get('train_y') + integer_params = [ + key for key in bo_params if bo_param_types[key] == "int" + ] + + integer_indices = [ + bo_params.index(key) for key in integer_params + ] + + bounds = np.empty((2, len(bo_params))) + + for i, param in enumerate(bo_params): + center = config["ptycho"]["params"][param] + radius = config["bo"]["params"][param]["radius"] + bounds[0, i] = center - radius + bounds[1, i] = center + radius + + state = { + "method": "random", + "acquisition": "", + "params": bo_params, + "param_types": bo_param_types, + "integer_params": integer_params, + "integer_indices": integer_indices, + "bounds": bounds, # shape: (2, BOparam) + "train_x": np.empty((0, len(bo_params))), # shape: (BOiter, BOparam) + "train_y": np.empty((0,)), # shape: (BOiter,) + "train_info": [] # shape: (BOiter,) + } + + train_x_path = config["bo"].get("train_x") + train_y_path = config["bo"].get("train_y") if train_x_path is not None and train_y_path is not None: if os.path.exists(train_x_path) and os.path.exists(train_y_path): train_x = np.load(train_x_path) train_y = np.load(train_y_path) + assert train_x.ndim == 2, "loaded train_x must be 2D" + assert train_x.shape[1] == len(bo_params), "loaded train_x shape(1) does not match number of variable parameters" + assert train_y.ndim == 1, "loaded train_y must be 1D" + assert train_y.shape[0] == train_x.shape[0], "loaded train_x and train_y shape(0) have unequal iterations" + state["train_x"] = train_x + state["train_y"] = train_y - if ( - train_x.ndim == 2 - and train_x.shape[1] == len(bo_params) - and train_y.ndim == 1 - and train_y.shape[0] == train_x.shape[0] - ): - bo_state['train_x'] = train_x - bo_state['train_y'] = train_y - - self.state = bo_state + self.state = state def ask(self): config = self.config - bo_state = self.state + state = self.state next_config = copy.deepcopy(config) - for param in bo_state['params']: - max_modulation = config['bo']['params'][param] - center_value = config['ptycho']['params'][param] - - modulation = max_modulation * (np.random.rand() - 0.5) * 2 - next_config['ptycho']['params'][param] = center_value + modulation + for param in state['params']: + radius = config['bo']['params'][param]['radius'] + center = config['ptycho']['params'][param] + modulation = radius * (np.random.rand() - 0.5) * 2 + next_value = center + modulation + if state['param_types'][param] == 'int': + next_value = round(next_value) + next_config['ptycho']['params'][param] = next_value return next_config def tell(self, job_config, y_value): config = self.config - bo_state = self.state + state = self.state x_value = [] - - for param in bo_state['params']: + for param in state['params']: x_value.append(job_config['ptycho']['params'][param]) - x_value = np.array(x_value).reshape(1, -1) - y_value = np.array([y_value]) - - bo_state['train_x'] = np.vstack([ - bo_state['train_x'], - x_value, + state['train_x'] = np.vstack([ + state['train_x'], + np.array(x_value).reshape(1, -1) ]) - bo_state['train_y'] = np.concatenate([ - bo_state['train_y'], - y_value, + state['train_y'] = np.concatenate([ + state['train_y'], + np.array([y_value]) ]) - result_dir = config['io']['result_dir'] - np.save(os.path.join(result_dir, 'train_x.npy'), bo_state['train_x']) - np.save(os.path.join(result_dir, 'train_y.npy'), bo_state['train_y']) + state['train_info'].append(state['method']) + + + train_x_path = config['bo'].get('train_x') + train_y_path = config['bo'].get('train_y') + if train_x_path is not None and train_y_path is not None: + np.save(train_x_path, state['train_x']) + np.save(train_y_path, state['train_y']) + else: + result_dir = config['io']['result_dir'] + np.save(os.path.join(result_dir, 'train_x.npy'), state['train_x']) + np.save(os.path.join(result_dir, 'train_y.npy'), state['train_y']) diff --git a/bo/sobo.py b/bo/sobo.py new file mode 100644 index 0000000..42adb9b --- /dev/null +++ b/bo/sobo.py @@ -0,0 +1,186 @@ +import os +import numpy as np +import copy + +import torch +from botorch.models import SingleTaskGP +from botorch.fit import fit_gpytorch_mll +from gpytorch.mlls import ExactMarginalLogLikelihood +from botorch.optim import optimize_acqf +from botorch.models.transforms.outcome import Standardize +from botorch.models.transforms.input import Normalize, Round, ChainedInputTransform +from botorch.acquisition.monte_carlo import qUpperConfidenceBound +from botorch.acquisition.logei import qLogExpectedImprovement +from botorch.sampling.normal import SobolQMCNormalSampler +from botorch.utils.rounding import approximate_round + + +from bo.bo_base import BOEngine + + +class SingleObjectiveBOEngine(BOEngine): + def __init__(self, config): + super().__init__(config) + + self.params = [key for key, spec in config["bo"]["params"].items() if spec is not None] + self.param_types = {key: config["bo"]["params"][key].get("type", "float") for key in self.params} + self.integer_indices = [i for i, param in enumerate(self.params) if self.param_types[param] == 'int'] + self.bounds = np.empty((2, len(self.params))) + + for i, param in enumerate(self.params): + center = config["ptycho"]["params"][param] + radius = config["bo"]["params"][param]["radius"] + self.bounds[0, i] = center - radius + self.bounds[1, i] = center + radius + + self.train_x = np.empty((0, len(self.params))) # shape: (BOiter, BOparam) + self.train_y = np.empty((0,)) # shape: (BOiter,) + + train_x_path = config["bo"].get("train_x") + train_y_path = config["bo"].get("train_y") + + if train_x_path is not None and train_y_path is not None: + if os.path.exists(train_x_path) and os.path.exists(train_y_path): + train_x = np.load(train_x_path) + train_y = np.load(train_y_path) + assert train_x.ndim == 2, "loaded train_x must be 2D" + assert train_x.shape[1] == len(self.params), "loaded train_x shape(1) does not match number of variable parameters" + assert train_y.ndim == 1, "loaded train_y must be 1D" + assert train_y.shape[0] == train_x.shape[0], "loaded train_x and train_y shape(0) have unequal iterations" + self.train_x = train_x + self.train_y = train_y + + self.acquisition = config['bo']['acquisition'] + + + def ask(self): + + train_x = torch.from_numpy(self.train_x) + train_y = torch.from_numpy(self.train_y).unsqueeze(-1) # shape: (BOiter, 1) + bounds = torch.from_numpy(self.bounds) + + assert self.train_x.shape[0] > 0 + + # Optimizing in [0, 1) unit cube is standard for BO; also numerically more stable. + # See also acqf_bounds + train_x_normalized = (train_x - bounds[0]) / (bounds[1] - bounds[0]) + + input_transform = ChainedInputTransform( + unnormalize = Normalize( + d=train_x.shape[1], + bounds=bounds, + transform_on_train=True, transform_on_eval=True, + reverse=True + ), + round = Round( + integer_indices=self.integer_indices, + transform_on_train=True, transform_on_eval=True, + approximate=True, tau=1e-3, + ), + normalize = Normalize( + d=train_x.shape[1], + bounds=bounds, + transform_on_train=True, transform_on_eval=True + ) + ) + + outcome_transform = Standardize(m=1, min_stdv=1e-8) + + gp = SingleTaskGP( + train_x_normalized, + train_y, + input_transform=input_transform, + outcome_transform=outcome_transform + ) + mll = ExactMarginalLogLikelihood(gp.likelihood, gp) + fit_gpytorch_mll(mll) + + + sampler = SobolQMCNormalSampler(sample_shape=torch.Size([512])) + + if self.acquisition == 'ucb': + beta = 0.2 + print("Acquisition: UCB | Beta: {} (fixed)".format(beta)) + acqf = qUpperConfidenceBound(gp, beta=beta, sampler=sampler) + elif self.acquisition == 'ei': + best_f = train_y.max() + print("Acquisition: LogEI best_f: {:.6f}".format(best_f.item())) + acqf = qLogExpectedImprovement(gp, best_f=best_f, sampler=sampler) + else: + raise NotImplementedError(f"Acquisition function {self.acquisition} is not implemented. Current options: 'ucb', 'ei'") + + # Full [0,1]^d search (trust region disabled) + acqf_bounds = torch.stack([ + torch.zeros(train_x.shape[1], dtype=torch.double), + torch.ones(train_x.shape[1], dtype=torch.double), + ]) + + candidate, _ = optimize_acqf( + acq_function=acqf, + bounds=acqf_bounds, + q=1, + num_restarts=20, + raw_samples=1024, + post_processing_func=self._pr_post_processing, # PR applied here + sequential=True, + ) + + new_x = candidate.detach() * (bounds[1] - bounds[0]) + bounds[0] + + # Hard-round integer dims (final guarantee) + for i in self.integer_indices: + new_x[:, i] = torch.round(new_x[:, i]) + + next_config = copy.deepcopy(self.config) + + for i, param in enumerate(self.params): + next_config['ptycho']['params'][param] = new_x[0,i].item() + + return next_config + + + def _pr_post_processing(self, X): + """Apply differentiable rounding to integer dims (PR forward pass).""" + X_out = X.clone() + for idx in self.integer_indices: + # Unnormalize -> approximate_round -> renormalize + raw = X_out[..., idx] * (self.bounds[1][idx] - self.bounds[0][idx]) + self.bounds[0][idx] + rounded = approximate_round(raw) + X_out[..., idx] = (rounded - self.bounds[0][idx]) / (self.bounds[1][idx] - self.bounds[0][idx]) + return X_out + + + def tell(self, job_config, y_value): + config = self.config + + x_value = [] + for param in self.params: + x_value.append(job_config['ptycho']['params'][param]) + + self.train_x = np.vstack([ + self.train_x, + np.array(x_value).reshape(1, -1) + ]) + + self.train_y = np.concatenate([ + self.train_y, + np.array([y_value]) + ]) + + + train_x_path = config['bo'].get('train_x') + train_y_path = config['bo'].get('train_y') + if train_x_path is not None and train_y_path is not None: + np.save(train_x_path, self.train_x) + np.save(train_y_path, self.train_y) + else: + result_dir = config['io']['result_dir'] + np.save(os.path.join(result_dir, 'train_x.npy'), self.train_x) + np.save(os.path.join(result_dir, 'train_y.npy'), self.train_y) + + + + + + + diff --git a/bo/ucb.py b/bo/ucb.py deleted file mode 100644 index 80c21e7..0000000 --- a/bo/ucb.py +++ /dev/null @@ -1,48 +0,0 @@ -# import os -# import sys -# import shutil -# import numpy as np -# import subprocess -# import h5py -# from PIL import Image -# import time -# import random - -# import torch -# from botorch.sampling.samplers import SobolQMCNormalSampler -# from botorch.models import SingleTaskGP -# from botorch.fit import fit_gpytorch_model -# from gpytorch.mlls import ExactMarginalLogLikelihood -# from botorch.optim import optimize_acqf -# from botorch.acquisition import UpperConfidenceBound -# from botorch.models.transforms.outcome import Standardize -# from botorch.acquisition.monte_carlo import qUpperConfidenceBound -# from botorch.utils.multi_objective.box_decompositions.non_dominated import NondominatedPartitioning -# from botorch.acquisition.multi_objective.monte_carlo import qExpectedHypervolumeImprovement -# from botorch.utils.transforms import unnormalize - -# from bo import bo_random_config - - -def initialize(config): - pass - -# n_initial_jobs = config['bo']['parallel_jobs'] -# n_var_params = len([p for p in config['bo']['params'].values() if p is not None]) - -# train_x = np.zeros([n_initial_jobs, n_var_params]) - -# for i in range(n_initial_jobs): -# config_i = bo_random_config(config) -# train_x[i] = [p for p in config_i['bo']['params'].values() if p is not None] - - -def ask(config, bo_state): - - next_config = config.copy() - - return next_config - - -def tell(job_config, bo_state, y_value): - pass \ No newline at end of file