added hotfix for batched bo

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

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):

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):

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:

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__":

pipelines/__init__.py

@@ -1 +1,7 @@
-from .sobo import sobo_pipeline
+from .sobo import sobo_pipeline
+from .test import test_pipeline
+
+job_types = {
+    'random+sobo': sobo_pipeline,
+    'test': test_pipeline,
+}

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")