diff --git a/.gitignore b/.gitignore index 15350ae..cf1fd97 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,6 @@ +results/ +old/ + *.mat *.tiff *.hdf5 diff --git a/bo-ptycho.py b/bo-ptycho.py index 8d4b423..879caf4 100644 --- a/bo-ptycho.py +++ b/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__": @@ -82,9 +94,6 @@ if __name__ == "__main__": config_yaml = sys.argv[1] 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) diff --git a/bo/random.py b/bo/random.py index da144ca..3cf9748 100644 --- a/bo/random.py +++ b/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 \ No newline at end of file diff --git a/config.yaml b/config.yaml index 04262b2..73d7fc2 100644 --- a/config.yaml +++ b/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' + diff --git a/job.sub b/job.sub index 0d593b9..d3afa68 100644 --- a/job.sub +++ b/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" diff --git a/ptycho/fold_slice.py b/ptycho/fold_slice.py index 30c882c..fb1e671 100644 --- a/ptycho/fold_slice.py +++ b/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) \ No newline at end of file + # return loadmat(result_mat) + return 1 \ No newline at end of file