completed abstract BO loop

.gitignore

@@ -1,3 +1,6 @@
+results/
+old/
+
 *.mat
 *.tiff
 *.hdf5

bo-ptycho.py

@@ -21,10 +21,6 @@ def ptycho_error(config):
         return error(config)
 
 
-
-
-
-
 def bo_initialize(config):
     bo_algorithm = config['bo']['algorithm']
 
@@ -61,17 +57,33 @@ def bo_tell(config, job_config, bo_state, y_value):
 
     if bo_engine == 'ucb':
         from bo.ucb import tell
-        bo_config = tell(job_config, bo_state, y_value)
+        bo_state = tell(job_config, bo_state, y_value)
     
     else:
         from bo.random import tell
-        bo_config = tell(job_config, bo_state, y_value)
+        bo_state = tell(config, job_config, bo_state, y_value)
     
-    return bo_config
+    return bo_state
 
 
+def main(config):
 
+    bo_state = bo_initialize(config)
 
+    max_iterations = config['bo']['max_iterations']
+
+    for _ in range(max_iterations):
+        job_config = bo_ask(config, bo_state)
+
+        ptycho_run(job_config)
+        y_value = ptycho_error(job_config)
+
+        bo_state = bo_tell(config, job_config, bo_state, y_value)
+
+        print(bo_state['train_x'])
+        print(bo_state['train_y'])
+
+    return bo_state
 
 
 if __name__ == "__main__":
@@ -83,8 +95,5 @@ if __name__ == "__main__":
     with open(config_yaml, 'r') as f:
         config = yaml.safe_load(f)
     
-    # ptycho_run(config_yaml)
-    # results = ptycho_results(config_yaml)
-    # new_config = bo_loop(config)
-    # print(new_config) # type: ignore
+    main(config)
     

bo/random.py

@@ -1,23 +1,40 @@
-import numpy as np
+import os
 import copy
+import numpy as np
 
 
 def initialize(config):
-    bo_params = []
 
-    for key, value in config['bo']['params'].items():
-        if value is not None:
-            bo_params.append(key)
-
-    train_x = np.empty((0, len(bo_params)))
-    train_y = np.empty((0,))
+    bo_params = [
+        key
+        for key, value in config['bo']['params'].items()
+        if value is not None
+    ]
 
     bo_state = {
+        'algorithm': 'random',
         'params': bo_params,
-        'train_x': train_x,
-        'train_y': train_y,
+        'train_x': np.empty((0, len(bo_params))),
+        'train_y': np.empty((0,)),
     }
 
+    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)
+
+            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
+
     return bo_state
 
 
@@ -34,7 +51,7 @@ def ask(config, bo_state):
     return next_config
 
 
-def tell(job_config, bo_state, y_value):
+def tell(config, job_config, bo_state, y_value):
     x_value = []
 
     for param in bo_state['params']:
@@ -53,4 +70,9 @@ def tell(job_config, bo_state, y_value):
         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'])
+
+    
     return bo_state

config.yaml

@@ -20,11 +20,11 @@ ptycho:
     tilt_x: 2
     tilt_y: 0 
     scan_step_size: 0.36
-    Niter: 20
-    Niter_save_results: 10
+    Niter: 100
+    Niter_save_results: 100
     CBED_size: 192
     ADU: 1
-    extra_print_info: 'FIB'
+    extra_print_info: 'test'
     scan_number: 1
     gpu_id: 1
     roi_label: '0_Ndp64'
@@ -38,10 +38,14 @@ ptycho:
 
 bo:
   parallel_jobs: 4
-  algorithm: ucb
+  max_iterations: 100
+  algorithm: 'random'
   params:
     alpha_max:
     defocus: 150
     rot_ang: 1
     Nlayers:
     thickness: 150
+  train_x: '/home/swim/Si_project/wrapper/results/train_x.npy'
+  train_y: '/home/swim/Si_project/wrapper/results/train_y.npy'
+  

job.sub

@@ -18,6 +18,7 @@ export LD_LIBRARY_PATH=/usr/local/cuda-11.4/lib64:$LD_LIBRARY_PATH
 
 source /home/swim/Si_project/wrapper/venv/bin/activate
 
+cat config.yaml
 python -u bo-ptycho.py config.yaml
 
 echo "SLURM JOB FINISHED"

ptycho/fold_slice.py

@@ -6,19 +6,18 @@ import numpy as np
 from scipy.io import loadmat
 
 def fold_slice_translator(config):
-
-    result_dir = config['io']['result_dir']
+    fold_slice_result_dir = os.path.join(config['io']['result_dir'], 'fold_slice')
 
     fold_slice_dict = {}
     fold_slice_dict['raw_data'] = config['io']['input_data_path']
-    fold_slice_dict['result_dir'] = os.path.join(result_dir, '')
-    fold_slice_dict.update(config['ptycho']['fold_slice'])
+    fold_slice_dict['result_dir'] = os.path.join(fold_slice_result_dir, '')
+    fold_slice_dict.update(config['ptycho']['params'])
 
-    if os.path.exists(os.path.join(result_dir)):
-        shutil.rmtree(os.path.join(result_dir))
-    os.makedirs(os.path.join(result_dir))
+    if os.path.exists(os.path.join(fold_slice_result_dir)):
+        shutil.rmtree(os.path.join(fold_slice_result_dir))
+    os.makedirs(os.path.join(fold_slice_result_dir))
 
-    setup_txt = os.path.join(result_dir, 'setup.txt')
+    setup_txt = os.path.join(fold_slice_result_dir, 'setup.txt')
     with open(setup_txt, 'w') as f:
         f.write('\n\n')
         for key, value in fold_slice_dict.items():
@@ -58,16 +57,17 @@ def run(config):
 
 
 def error(config):
-    result_dir = config['io']['result_dir']
+    fold_slice_result_dir = os.path.join(config['io']['result_dir'], 'fold_slice')
     roi_dir = os.path.join(
-        result_dir,
-        f"{config['ptycho']['fold_slice']['scan_number']}",
-        f"roi{config['ptycho']['fold_slice']['roi_label']}"
+        fold_slice_result_dir,
+        f"{config['ptycho']['params']['scan_number']}",
+        f"roi{config['ptycho']['params']['roi_label']}"
     )
     output_dir = os.path.join(roi_dir, next(os.walk(roi_dir))[1][0])
     # image_path = os.path.join(output_dir, 'obj_phase_roi_sum', next(os.walk(os.path.join(output_dir, 'obj_phase_roi_sum')))[2][0])
-    result_mat = os.path.join(output_dir, f"Niter{config['ptycho']['fold_slice']['Niter']}.mat")
+    result_mat = os.path.join(output_dir, f"Niter{config['ptycho']['params']['Niter']}.mat")
     if not os.path.exists(result_mat):
         raise FileNotFoundError(f"Result directory {result_mat} does not exist. Please check the fold_slice output.")
     
-    return loadmat(result_mat)
+    # return loadmat(result_mat)
+    return 1