diff --git a/bo/base.py b/bo/base.py index 5d1c4a4..7cd3efe 100644 --- a/bo/base.py +++ b/bo/base.py @@ -7,7 +7,7 @@ class BOEngine(ABC): self.state = None @abstractmethod - def ask(self): + def ask(self, n: int = 1) -> list[dict]: pass @abstractmethod diff --git a/bo/random.py b/bo/random.py index 871ef54..28a1aac 100644 --- a/bo/random.py +++ b/bo/random.py @@ -37,20 +37,23 @@ class RandomBOEngine(BOEngine): self.train_y = train_y - def ask(self): + def ask(self, n = 1): config = self.config - next_config = copy.deepcopy(config) + next_configs = [] - for param in self.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 self.param_types[param] == 'int': - next_value = round(next_value) - next_config['ptycho']['params'][param] = next_value + for _ in range(n): + next_config = copy.deepcopy(config) + for param in self.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 self.param_types[param] == 'int': + next_value = round(next_value) + next_config['ptycho']['params'][param] = next_value + next_configs.append(next_config) - return next_config + return next_configs def tell(self, job_config, y_value): diff --git a/bo/sobo.py b/bo/sobo.py index 4783783..e639cb2 100644 --- a/bo/sobo.py +++ b/bo/sobo.py @@ -52,7 +52,7 @@ class SingleObjectiveBOEngine(BOEngine): self.acquisition = config['bo']['acquisition'] - def ask(self): + def ask(self, n = 1): train_x = torch.from_numpy(self.train_x) train_y = torch.from_numpy(self.train_y).unsqueeze(-1) # shape: (BOiter, 1) @@ -98,44 +98,41 @@ class SingleObjectiveBOEngine(BOEngine): 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) + acqf = qUpperConfidenceBound(gp, beta=self.config['bo']['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) + acqf = qLogExpectedImprovement(gp, best_f=train_y.max(), 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( + candidates, _ = optimize_acqf( acq_function=acqf, bounds=acqf_bounds, - q=1, + q=n, 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] + new_xs = candidates.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]) + new_xs[:, i] = torch.round(new_xs[:, i]) - next_config = copy.deepcopy(self.config) - - for i, param in enumerate(self.params): - next_config['ptycho']['params'][param] = new_x[0,i].item() + next_configs = [] + for i in range(n): + next_config = copy.deepcopy(self.config) + for j, param in enumerate(self.params): + next_config['ptycho']['params'][param] = new_xs[i,j].item() + next_configs.append(next_config) - return next_config + return next_configs def _pr_post_processing(self, X): diff --git a/config.yaml b/config.yaml index 1d8e453..8967680 100644 --- a/config.yaml +++ b/config.yaml @@ -1,22 +1,22 @@ job: - type: "random+sobo" - random_iters: 3 - sobo_iters: 30 + type: "test" + random_iters: 5 + sobo_iters: 10 io: - input_data_path: '/home/swim/Si_project/data/Si2V1_1.mat' + input_data_path: '/home/swim/shared/Si_project/data/Si2V1_1.mat' result_dir: '/home/swim/bo-ptycho/results/test' verbosity: 1 ptycho: engine: 'fold_slice' - path: '/home/swim/fold_slice' + path: '/home/swim/fold_slice-stable' params: voltage: 200 alpha_max: 30 defocus: -200 rot_ang: 0.3 - Nlayers: 20 + Nlayers: 5 thickness: 250 rbf: 37 Nprobe: 1 @@ -25,8 +25,8 @@ ptycho: tilt_x: 2 tilt_y: 0 scan_step_size: 0.36 - Niter: 100 - Niter_save_results: 100 + Niter: 5 + Niter_save_results: 5 CBED_size: 192 ADU: 1 extra_print_info: '' @@ -36,24 +36,26 @@ ptycho: diff_pattern_blur: 1 probe_change_start: 1 object_change_start: 1 - grouping: 64 - probe_posiiton_search: 1 + grouping: inf + probe_position_search: 1 regularize_layers: 0.2 variable_probe: 'false' + verbosity bo: - mode: 'sobo' + batch: 4 acquisition: 'ucb' + beta: 0.1 # for ucb only metric: 'log_fourier' params: alpha_max: defocus: - radius: 100 + # radius: 100 rot_ang: Nlayers: - radius: 5 + radius: 2 type: int thickness: - radius: 100 + # radius: 100 train_x: train_y: diff --git a/main.py b/main.py index 0156487..0b6079f 100644 --- a/main.py +++ b/main.py @@ -8,11 +8,7 @@ def main(config_yaml): with open(config_yaml, 'r') as f: config = yaml.safe_load(f) - job_types = { - 'random+sobo': pipelines.sobo_pipeline - } - - job_types[config['job']['type']](config) + pipelines.job_types[config['job']['type']](config) if __name__ == "__main__": diff --git a/pipelines/__init__.py b/pipelines/__init__.py index cfde87d..0e95477 100644 --- a/pipelines/__init__.py +++ b/pipelines/__init__.py @@ -1 +1,7 @@ -from .sobo import sobo_pipeline \ No newline at end of file +from .sobo import sobo_pipeline +from .test import test_pipeline + +job_types = { + 'random+sobo': sobo_pipeline, + 'test': test_pipeline, +} diff --git a/pipelines/test.py b/pipelines/test.py new file mode 100644 index 0000000..a944d78 --- /dev/null +++ b/pipelines/test.py @@ -0,0 +1,49 @@ +import os +import bo +import ptycho + + +def test_pipeline(config): + + result_dir = config['io']['result_dir'] + os.makedirs(result_dir, exist_ok=True) + + RANDOM_ITERS = config['job'].get('random_iters', 0) + SOBO_ITERS = config['job'].get('sobo_iters') + METRIC = config['bo']['metric'] + PTYCHOENGINE = ptycho.engines[config['ptycho']['engine']] + + randombo = bo.RandomBOEngine(config) + + # bo_txt = os.path.join(result_dir, "bo.txt") + # with open(bo_txt, "w") as f: + # f.write(f" iter\tmetric\t{"\t".join([p[:7] for p in randombo.params])}\n") + + for j in range(RANDOM_ITERS): + print(f"RANDOM sampling; iteration {j}") + job_configs = randombo.ask(n=4) + for job_config in job_configs: + ptycho_engine = PTYCHOENGINE(job_config) + ptycho_engine.run(run_id=f"bo-{j:03d}") + y_value = ptycho_engine.metric(METRIC) + randombo.tell(job_config, y_value) + # with open(bo_txt, "a") as f: + # p = [f'{job_config['ptycho']['params'][key]:.2f}' for key in randombo.params] + # f.write(f"{j: 8d}\t{y_value:.4f}\t{"\t".join(p)}\n") + + + sobo = bo.SingleObjectiveBOEngine(config) + sobo.train_x = randombo.train_x + sobo.train_y = randombo.train_y + + for j in range(SOBO_ITERS): + print(f"SOBO sampling; iteration {RANDOM_ITERS + j}") + job_configs = sobo.ask(n=4) + for job_config in job_configs: + ptycho_engine = PTYCHOENGINE(job_config) + ptycho_engine.run(run_id=f"bo-{RANDOM_ITERS + j:03d}") + y_value = ptycho_engine.metric(METRIC) + sobo.tell(job_config, y_value) + # with open(bo_txt, "a") as f: + # p = [f'{job_config['ptycho']['params'][key]:.2f}' for key in sobo.params] + # f.write(f"{RANDOM_ITERS + j: 8d}\t{y_value:.4f}\t{"\t".join(p)}\n")