new yaml schema, new TODO

.gitignore

@@ -1,5 +1,6 @@
 results/
 old/
+notebooks
 setup.txt
 
 *.mat

LICENSE

@@ -0,0 +1,19 @@
+Copyright (c) 2026 Sooyoung Cheong
+
+Permission is hereby granted, free of charge, to any person obtaining a
+copy of this software and associated documentation files (the "Software"),
+to deal in the Software without restriction, including without limitation the
+rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+sell copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.

TODO.md

@@ -1,13 +1,24 @@
-TODO:
-
-- Clean up BO states
-    - unify `BOEngine.__init__()`
-- BO train_x/y transfer between engines for single job
-- multi-GPU dispatcher
-- template job sequences / yamls
-    - mobo
-- metric() function(s) for each ptycho engine
-    - FRC score
-- fold_slice: load diffractions / hdf5 files; change only param per job
+# TODO:
+- fold_slice:
+    - load diffractions / hdf5 files; change only param per job
     - restructure `FoldSlicePtychoEngine.__init__()` to load data but not params
         - write new `prepare_data.m`
+    - start fold_slice from config.yaml instead of setup.txt
+- BO train_x/y transfer between engines for single job
+- metric() function(s) for each ptycho engine
+    - FRC score
+- separate `config` into `bo_config` and `ptycho_config`; let `BOEngine` have no knowledge of ptychography and vice versa.
+- add GPU version of ExamplePtychoEngine
+- change `PtychoEngine.metric()` to accept list of names and return dict
+- `GridSampler`
+    - `.__init__()` should create a grid
+    - `.ask()` should remove those items from the grid
+- change job.sub to job.sh
+    - dynamic jobname
+    - pre-check result directory
+
+# TODAY:
+- change yaml schema, use stages in config-new.yaml
+- separation of available GPUs and parallel sample batches
+- prepare next batch for efficient GPU use
+

bo/__init__.py

@@ -1,3 +0,0 @@
-from .base import BOEngine
-from .random import RandomBOEngine
-from .sobo import SingleObjectiveBOEngine

bo/base.py

@@ -1,15 +0,0 @@
-from abc import ABC, abstractmethod
-
-
-class BOEngine(ABC):
-    def __init__(self, config):
-        self.config = config
-        self.state = None
-
-    @abstractmethod
-    def ask(self):
-        pass
-
-    @abstractmethod
-    def tell(self, job_config, y_value):
-        pass

config-new.yaml

@@ -0,0 +1,71 @@
+io:
+  input_data_path: '/home/swim/shared/Si_project/data/Si2V1_2.mat'
+  result_dir: '/home/swim/bo-ptycho/results/test'
+  verbosity: 1
+
+ptycho:
+  engine: 'fake'
+  # path: '/home/swim/fold_slice-park'
+  params:
+    voltage: 200
+    alpha_max: 30
+    defocus: -250
+    rot_ang: 0.3
+    Nlayers: 20
+    thickness: 250
+    rbf: 37
+    tilt_x: 2
+    tilt_y: 0 
+    scan_step_size: 0.36
+
+    Niter: 10
+    Niter_save_results: 10
+
+    CBED_size: 192
+    ADU: 1
+    Nprobe: 1
+    N_scan_x: 64
+    N_scan_y: 64 
+    extra_print_info: 'FIB'
+    scan_number: 1
+    gpu_id: 1
+    roi_label: '0_Ndp64'
+    diff_pattern_blur: 1
+    probe_change_start: 1
+    object_change_start: 1
+    grouping: 64
+    probe_position_search: 1
+    regularize_layers: 0.2
+    variable_probe: false
+
+search:
+  metric: log_fourier
+  
+  params:
+    defocus:
+      radius: 100
+    Nlayers:
+      radius: 5
+      type: int
+    thickness:
+      radius: 150
+  
+  train_x:
+  train_y:
+
+  stages:
+    # - sampler: fixed
+    
+    # - sampler: grid
+    #   defocus: 7
+    #   Nlayers: 11
+    #   thickness: 7
+
+    - sampler: random
+      samples: 16
+
+    - sampler: sobo
+      samples: 256
+      batch: 4
+      acquisition: 'ucb'
+      beta: 0.1

config.yaml

@@ -1,35 +1,37 @@
 job:
-  type: "random+sobo"
-  random_iters: 3
-  sobo_iters: 30
+  type: 'random+sobo'
+  random_iters: 16
+  sobo_iters: 128
 
 io:
-  input_data_path: '/home/swim/Si_project/data/Si2V1_1.mat'
+  input_data_path: '/home/swim/shared/Si_project/data/Si2V1_2.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-park'
   params:
     voltage: 200
     alpha_max: 30
-    defocus: -200
+    defocus: -250
     rot_ang: 0.3
     Nlayers: 20
     thickness: 250
     rbf: 37
-    Nprobe: 1
-    N_scan_x: 64
-    N_scan_y: 64
     tilt_x: 2
     tilt_y: 0 
     scan_step_size: 0.36
+
     Niter: 100
     Niter_save_results: 100
+
     CBED_size: 192
     ADU: 1
-    extra_print_info: ''
+    Nprobe: 1
+    N_scan_x: 64
+    N_scan_y: 64 
+    extra_print_info: 'FIB'
     scan_number: 1
     gpu_id: 1
     roi_label: '0_Ndp64'
@@ -37,23 +39,23 @@ ptycho:
     probe_change_start: 1
     object_change_start: 1
     grouping: 64
-    probe_posiiton_search: 1
+    probe_position_search: 1
     regularize_layers: 0.2
-    variable_probe: 'false'
+    variable_probe: false
+
 
 bo:
-  mode: 'sobo'
+  batch: 4
   acquisition: 'ucb'
+  beta: 0.1
   metric: 'log_fourier'
   params:
-    alpha_max:
     defocus:
       radius: 100
-    rot_ang:
     Nlayers:
       radius: 5
       type: int
     thickness:
-      radius: 100
+      radius: 150
   train_x:
   train_y:

examples/260805-Si/Si_02V1_2.sub

@@ -1,25 +0,0 @@
-#!/bin/bash
-#SBATCH --job-name=Si2V
-#SBATCH --nodes=1
-#SBATCH --ntasks=1
-#SBATCH --cpus-per-task=1
-#SBATCH --gres=gpu:rtx-6000ada:1
-#SBATCH --time=100:00:00
-#SBATCH --output=/home/swim/slurm-logs/job_%j.log
-
-echo "2V"
-pwd
-hostname
-date
-
-export PATH=/home/shared/MATLAB/R2021a/bin:$PATH
-export PATH=/usr/local/cuda-11.4/bin:$PATH
-export LD_LIBRARY_PATH=/usr/local/cuda-11.4/lib64:$LD_LIBRARY_PATH
-
-source /home/swim/bo-ptycho/venv/bin/activate
-
-cat examples/260805-Si/Si_02V1_2.yaml
-python -u main.py examples/260805-Si/Si_02V1_2.yaml
-
-date
-echo "SLURM JOB FINISHED"

examples/260805-Si/Si_05V1_2.sub

@@ -1,25 +0,0 @@
-#!/bin/bash
-#SBATCH --job-name=Si5V
-#SBATCH --nodes=1
-#SBATCH --ntasks=1
-#SBATCH --cpus-per-task=1
-#SBATCH --gres=gpu:rtx-6000ada:1
-#SBATCH --time=100:00:00
-#SBATCH --output=/home/swim/slurm-logs/job_%j.log
-
-echo "5V"
-pwd
-hostname
-date
-
-export PATH=/home/shared/MATLAB/R2021a/bin:$PATH
-export PATH=/usr/local/cuda-11.4/bin:$PATH
-export LD_LIBRARY_PATH=/usr/local/cuda-11.4/lib64:$LD_LIBRARY_PATH
-
-source /home/swim/bo-ptycho/venv/bin/activate
-
-cat examples/260805-Si/Si_05V1_2.yaml
-python -u main.py examples/260805-Si/Si_05V1_2.yaml
-
-date
-echo "SLURM JOB FINISHED"

examples/260805-Si/Si_08V1_2.sub

@@ -1,25 +0,0 @@
-#!/bin/bash
-#SBATCH --job-name=Si8V
-#SBATCH --nodes=1
-#SBATCH --ntasks=1
-#SBATCH --cpus-per-task=1
-#SBATCH --gres=gpu:rtx-6000ada:1
-#SBATCH --time=100:00:00
-#SBATCH --output=/home/swim/slurm-logs/job_%j.log
-
-echo "8V"
-pwd
-hostname
-date
-
-export PATH=/home/shared/MATLAB/R2021a/bin:$PATH
-export PATH=/usr/local/cuda-11.4/bin:$PATH
-export LD_LIBRARY_PATH=/usr/local/cuda-11.4/lib64:$LD_LIBRARY_PATH
-
-source /home/swim/bo-ptycho/venv/bin/activate
-
-cat examples/260805-Si/Si_08V1_2.yaml
-python -u main.py examples/260805-Si/Si_08V1_2.yaml
-
-date
-echo "SLURM JOB FINISHED"

examples/260805-Si/Si_30V3_2.sub

@@ -1,25 +0,0 @@
-#!/bin/bash
-#SBATCH --job-name=Si30V
-#SBATCH --nodes=1
-#SBATCH --ntasks=1
-#SBATCH --cpus-per-task=1
-#SBATCH --gres=gpu:rtx-6000ada:1
-#SBATCH --time=100:00:00
-#SBATCH --output=/home/swim/slurm-logs/job_%j.log
-
-echo "2V"
-pwd
-hostname
-date
-
-export PATH=/home/shared/MATLAB/R2021a/bin:$PATH
-export PATH=/usr/local/cuda-11.4/bin:$PATH
-export LD_LIBRARY_PATH=/usr/local/cuda-11.4/lib64:$LD_LIBRARY_PATH
-
-source /home/swim/bo-ptycho/venv/bin/activate
-
-cat examples/260805-Si/Si_30V3_2.yaml
-python -u main.py examples/260805-Si/Si_30V3_2.yaml
-
-date
-echo "SLURM JOB FINISHED"

examples/silicon-fib/config-2V.yaml

@@ -1,16 +1,16 @@
 job:
-  type: "random+sobo"
-  random_iters: 30
-  sobo_iters: 300
+  type: 'random+sobo'
+  random_iters: 16
+  sobo_iters: 128
 
 io:
-  input_data_path: '/home/swim/Si_project/data/Si2V1_2.mat'
-  result_dir: '/home/swim/bo-ptycho/results/260805-si/Si2V1_2'
+  input_data_path: '/home/swim/shared/Si_project/data/Si2V1_2.mat'
+  result_dir: '/home/swim/bo-ptycho/results/Si2V1_2/260817'
   verbosity: 1
 
 ptycho:
   engine: 'fold_slice'
-  path: '/home/swim/fold_slice'
+  path: '/home/swim/fold_slice-park'
   params:
     voltage: 200
     alpha_max: 30
@@ -19,16 +19,18 @@ ptycho:
     Nlayers: 20
     thickness: 250
     rbf: 37
-    Nprobe: 1
-    N_scan_x: 64
-    N_scan_y: 64
     tilt_x: 2
     tilt_y: 0 
     scan_step_size: 0.36
+
     Niter: 100
     Niter_save_results: 100
+
     CBED_size: 192
     ADU: 1
+    Nprobe: 1
+    N_scan_x: 64
+    N_scan_y: 64 
     extra_print_info: 'FIB'
     scan_number: 1
     gpu_id: 1
@@ -37,23 +39,23 @@ ptycho:
     probe_change_start: 1
     object_change_start: 1
     grouping: 64
-    probe_posiiton_search: 1
+    probe_position_search: 1
     regularize_layers: 0.2
     variable_probe: false
 
+
 bo:
-  mode: 'sobo'
+  batch: 4
   acquisition: 'ucb'
+  beta: 0.1
   metric: 'log_fourier'
   params:
-    alpha_max:
     defocus:
       radius: 100
-    rot_ang:
     Nlayers:
       radius: 5
       type: int
     thickness:
-      radius: 100
+      radius: 150
   train_x:
   train_y:

examples/silicon-fib/config-30V.yaml

@@ -1,16 +1,16 @@
 job:
   type: "random+sobo"
-  random_iters: 30
-  sobo_iters: 300
+  random_iters: 16
+  sobo_iters: 512
 
 io:
-  input_data_path: '/home/swim/Si_project/data/Si30V3_2.mat'
-  result_dir: '/home/swim/bo-ptycho/results/260805-si/Si30V3_2'
+  input_data_path: '/home/swim/shared/Si_project/data/Si30V3_2.mat'
+  result_dir: '/home/swim/bo-ptycho/results/Si30V3_2/260816'
   verbosity: 1
 
 ptycho:
   engine: 'fold_slice'
-  path: '/home/swim/fold_slice'
+  path: '/home/swim/fold_slice-park'
   params:
     voltage: 200
     alpha_max: 30
@@ -18,17 +18,19 @@ ptycho:
     rot_ang: 0.1
     Nlayers: 30
     thickness: 650
-    rbf: 37
-    Nprobe: 1
-    N_scan_x: 64
-    N_scan_y: 64
+    rbf: 38
     tilt_x: 5
     tilt_y: 3 
     scan_step_size: 0.35
+
     Niter: 100
     Niter_save_results: 100
+
     CBED_size: 192
     ADU: 1
+    Nprobe: 1
+    N_scan_x: 64
+    N_scan_y: 64 
     extra_print_info: 'FIB'
     scan_number: 1
     gpu_id: 1
@@ -37,13 +39,14 @@ ptycho:
     probe_change_start: 1
     object_change_start: 1
     grouping: 64
-    probe_posiiton_search: 1
+    probe_position_search: 1
     regularize_layers: 0.2
     variable_probe: false
 
 bo:
-  mode: 'sobo'
+  batch: 4
   acquisition: 'ucb'
+  beta: 0.1
   metric: 'log_fourier'
   params:
     alpha_max:

examples/silicon-fib/config-5V.yaml

@@ -1,16 +1,16 @@
 job:
   type: "random+sobo"
-  random_iters: 30
-  sobo_iters: 300
+  random_iters: 16
+  sobo_iters: 128
 
 io:
-  input_data_path: '/home/swim/Si_project/data/Si5V1_2.mat'
-  result_dir: '/home/swim/bo-ptycho/results/260805-si/Si5V1_2'
+  input_data_path: '/home/swim/shared/Si_project/data/Si5V1_2.mat'
+  result_dir: '/home/swim/bo-ptycho/results/Si5V1_2/260815'
   verbosity: 1
 
 ptycho:
   engine: 'fold_slice'
-  path: '/home/swim/fold_slice'
+  path: '/home/swim/fold_slice-park'
   params:
     voltage: 200
     alpha_max: 30
@@ -19,16 +19,18 @@ ptycho:
     Nlayers: 20
     thickness: 250
     rbf: 37
-    Nprobe: 1
-    N_scan_x: 64
-    N_scan_y: 64
     tilt_x: 3
     tilt_y: -1 
     scan_step_size: 0.36
+
     Niter: 100
     Niter_save_results: 100
+
     CBED_size: 192
     ADU: 1
+    Nprobe: 1
+    N_scan_x: 64
+    N_scan_y: 64 
     extra_print_info: 'FIB'
     scan_number: 1
     gpu_id: 1
@@ -37,13 +39,14 @@ ptycho:
     probe_change_start: 1
     object_change_start: 1
     grouping: 64
-    probe_posiiton_search: 1
+    probe_position_search: 1
     regularize_layers: 0.2
     variable_probe: false
 
 bo:
-  mode: 'sobo'
+  batch: 4
   acquisition: 'ucb'
+  beta: 0.1
   metric: 'log_fourier'
   params:
     alpha_max:

examples/silicon-fib/config-8V.yaml

@@ -1,16 +1,16 @@
 job:
   type: "random+sobo"
-  random_iters: 30
-  sobo_iters: 300
+  random_iters: 16
+  sobo_iters: 128
 
 io:
-  input_data_path: '/home/swim/Si_project/data/Si8V1_2.mat'
-  result_dir: '/home/swim/bo-ptycho/results/260805-si/Si8V1_2'
+  input_data_path: '/home/swim/shared/Si_project/data/Si8V2_2.mat'
+  result_dir: '/home/swim/bo-ptycho/results/Si8V2_2/260815'
   verbosity: 1
 
 ptycho:
   engine: 'fold_slice'
-  path: '/home/swim/fold_slice'
+  path: '/home/swim/fold_slice-park'
   params:
     voltage: 200
     alpha_max: 30
@@ -19,16 +19,18 @@ ptycho:
     Nlayers: 25
     thickness: 350
     rbf: 37
-    Nprobe: 1
-    N_scan_x: 64
-    N_scan_y: 64
     tilt_x: 4
     tilt_y: 0 
     scan_step_size: 0.36
+
     Niter: 100
     Niter_save_results: 100
+
     CBED_size: 192
     ADU: 1
+    Nprobe: 1
+    N_scan_x: 64
+    N_scan_y: 64 
     extra_print_info: 'FIB'
     scan_number: 1
     gpu_id: 1
@@ -37,13 +39,14 @@ ptycho:
     probe_change_start: 1
     object_change_start: 1
     grouping: 64
-    probe_posiiton_search: 1
+    probe_position_search: 1
     regularize_layers: 0.2
     variable_probe: false
 
 bo:
-  mode: 'sobo'
+  batch: 4
   acquisition: 'ucb'
+  beta: 0.1
   metric: 'log_fourier'
   params:
     alpha_max:

job.sub

@@ -2,12 +2,11 @@
 #SBATCH --job-name=Si2V
 #SBATCH --nodes=1
 #SBATCH --ntasks=1
-#SBATCH --cpus-per-task=1
-#SBATCH --gres=gpu:rtx-6000ada:1
+#SBATCH --cpus-per-task=16
+#SBATCH --gres=gpu:rtx-6000ada:4
 #SBATCH --time=100:00:00
 #SBATCH --output=/home/swim/slurm-logs/job_%j.log
 
-echo "2V"
 pwd
 hostname
 date
@@ -18,8 +17,10 @@ export LD_LIBRARY_PATH=/usr/local/cuda-11.4/lib64:$LD_LIBRARY_PATH
 
 source /home/swim/bo-ptycho/venv/bin/activate
 
-cat config.yaml
-python -u main.py config.yaml
+YAML="config.yaml"
+
+cat $YAML
+python -u main.py $YAML
 
 date
 echo "SLURM JOB FINISHED"

main.py

@@ -1,6 +1,7 @@
 import os
 import sys
 import yaml
+import shutil
 
 import pipelines
 
@@ -8,11 +9,18 @@ def main(config_yaml):
     with open(config_yaml, 'r') as f:
         config = yaml.safe_load(f)
 
-    job_types = {
-        'random+sobo': pipelines.sobo_pipeline
-    }
+    result_dir = config["io"]["result_dir"]
+    if os.path.exists(result_dir):
+        if sys.stdin.isatty():
+            answer = input(f"Will delete {result_dir}: [y/N]\n> ").strip().lower() == 'y'
+            if not answer:
+                print("Aborted.")
+                return         
+        shutil.rmtree(result_dir)
     
-    job_types[config['job']['type']](config)
+    os.makedirs(result_dir, exist_ok=True)
+    shutil.copy(config_yaml, os.path.join(result_dir, os.path.basename(config_yaml)))
+    pipelines.job_types[config['job']['type']](config)
     
 
 if __name__ == "__main__":

notebooks/mat_preview.ipynb

@@ -1,127 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": 17,
-   "id": "5c8aa3fd",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from scipy.io import loadmat as scipy_loadmat\n",
-    "from mat73 import loadmat as mat73_loadmat\n",
-    "import numpy as np"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "id": "1d3c8601",
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "'/home/swim/bo-ptycho/notebooks'"
-      ]
-     },
-     "execution_count": 2,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "%pwd"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "id": "4a419040",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "f = \"/home/swim/Si_project/Si8V2_full/roi1_Ndp256/MLs_L1_p1_g64_Ndp192_pc1_noModel_Ns23_dz14.6087_reg0.2/Niter1000.mat\""
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "id": "405b77b4",
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "dict_keys(['__header__', '__version__', '__globals__', 'outputs', 'probe', 'object', 'p', '__function_workspace__'])"
-      ]
-     },
-     "execution_count": 6,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "contents = scipy_loadmat(f)\n",
-    "contents.keys()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 25,
-   "id": "8a7b58b2",
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "mappingproxy({'object_ROI': (dtype('O'), 0),\n",
-       "              'binning': (dtype('O'), 8),\n",
-       "              'detector': (dtype('O'), 16),\n",
-       "              'dx_spec': (dtype('O'), 24),\n",
-       "              'lambda': (dtype('O'), 32),\n",
-       "              'multi_slice_param': (dtype('O'), 40),\n",
-       "              'obj_init_param': (dtype('O'), 48),\n",
-       "              'init_probe_file': (dtype('O'), 56),\n",
-       "              'normalize_init_probe': (dtype('O'), 64)})"
-      ]
-     },
-     "execution_count": 25,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "contents['p'].dtype.fields"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "d8870f98",
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "lemon",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.12.12"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 5
-}

pipelines/__init__.py

@@ -1 +1,8 @@
-from .sobo import sobo_pipeline
+# from .sobo import sobo_pipeline # depreacated
+from .test import test_pipeline
+from .batched_sobo import sobo_pipeline
+
+job_types = {
+    'random+sobo': sobo_pipeline,
+    'test': test_pipeline,
+}

pipelines/batched_sobo.py

@@ -0,0 +1,94 @@
+import os
+import traceback
+import multiprocessing as mp
+
+import samplers
+import ptycho
+
+def run_ptycho_worker(worker_id, gpu_token, job_config, metric, run_id, result_queue):
+
+    # This process, and MATLAB launched from it, can see exactly one GPU.
+    os.environ["CUDA_VISIBLE_DEVICES"] = gpu_token
+
+    try:
+        ptycho_engine = ptycho.engines[job_config["ptycho"]["engine"]](job_config)
+        ptycho_engine.run(run_id=run_id)
+        y_value = ptycho_engine.metric(metric)
+        result_queue.put((worker_id, y_value, None))
+
+    except Exception:
+        result_queue.put((worker_id, None, traceback.format_exc()))
+
+
+def run_batch(ctx, gpu_tokens, job_configs, metric, iteration):
+
+    result_queue = ctx.Queue()
+    processes = []
+
+    for i, job_config in enumerate(job_configs):
+        # Important: unique run_id for simultaneous jobs.
+        run_id = f"bo-{iteration:03d}-{i:02d}"
+
+        p = ctx.Process(
+            target=run_ptycho_worker,
+            args=(i, gpu_tokens[i], job_config, metric, run_id, result_queue),
+        )
+        p.start()
+        processes.append(p)
+
+    results = [result_queue.get() for _ in processes]
+
+    # Synchronization barrier.
+    for p in processes:
+        p.join()
+
+    # Completion order is arbitrary.
+    results.sort(key=lambda x: x[0])
+
+    for worker_id, _, error in results:
+        if error is not None:
+            raise RuntimeError(f"Ptycho worker {worker_id} failed:\n{error}")
+
+    return [y_value for _, y_value, _ in results]
+
+
+def sobo_pipeline(config):
+
+    RANDOM_ITERS = config["job"].get("random_iters", 0)
+    SOBO_ITERS = config["job"].get("sobo_iters", 0)
+    METRIC = config["bo"]["metric"]
+    BO_BATCH = config["bo"]["batch"]
+
+    visible = os.environ.get("CUDA_VISIBLE_DEVICES")
+    if visible is None:
+        raise RuntimeError("CUDA_VISIBLE_DEVICES is not set")
+    gpu_tokens = [token.strip() for token in visible.split(",") if token.strip()]
+    if len(gpu_tokens) < BO_BATCH:
+        raise RuntimeError(f"Expected {BO_BATCH} allocated GPUs, got {len(gpu_tokens)}")
+
+    # Explicitly use spawn for CUDA / MATLAB isolation.
+    ctx = mp.get_context("spawn")
+
+
+    ############################ RANDOM SAMPLING ###############################
+    randombo = samplers.RandomSampler(config)
+
+    for j in range(RANDOM_ITERS):
+        print(f"RANDOM sampling; iteration {j}")
+        job_configs = randombo.ask(n=BO_BATCH)
+        y_values = run_batch(ctx=ctx, gpu_tokens=gpu_tokens, job_configs=job_configs, metric=METRIC, iteration=j)
+        for job_config, y_value in zip(job_configs, y_values):
+            randombo.tell(job_config, y_value)
+
+    ############################  SOBO SAMPLING  ###############################
+    sobo = samplers.SOBOSampler(config)
+    sobo.train_x = randombo.train_x
+    sobo.train_y = randombo.train_y
+
+    for j in range(RANDOM_ITERS, SOBO_ITERS+RANDOM_ITERS):
+        print(f"SOBO sampling iteration {j}")
+        job_configs = sobo.ask(n=BO_BATCH)
+        y_values = run_batch(ctx=ctx, gpu_tokens=gpu_tokens, job_configs=job_configs, metric=METRIC, iteration=j)
+
+        for job_config, y_value in zip(job_configs, y_values):
+            sobo.tell(job_config, y_value)

pipelines/sobo.py

@@ -1,47 +0,0 @@
-import os
-import bo
-import ptycho
-
-def sobo_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_config = randombo.ask()
-        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_config = sobo.ask()
-        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")

pipelines/test.py

@@ -0,0 +1,6 @@
+import samplers
+import ptycho
+
+
+def test_pipeline(config):
+    pass

ptycho/__init__.py

@@ -1,5 +1,4 @@
-from .base import PtychoEngine
-from .ptycho_example import ExamplePtychoEngine
+from .example import ExamplePtychoEngine
 from .fold_slice import FoldSlicePtychoEngine
 
 engines = {

ptycho/example.py

@@ -8,7 +8,6 @@ class ExamplePtychoEngine(PtychoEngine):
     def __init__(self, config):
         super().__init__(config)
 
-    # run single ptychography job based on `config`
     def run(self, run_id="") -> None:
         print(f"[{run_id}] [ExamplePtychoEngine] Sleeping for 0.1 second.")
         time.sleep(0.1)

ptycho/fold_slice.py

@@ -11,9 +11,7 @@ class FoldSlicePtychoEngine(PtychoEngine):
 
     def __init__(self, config):
         super().__init__(config)
-        self._output_dir = os.path.join(config['io']['result_dir'], 'fold_slice')
         self._fold_slice_path = self.config['ptycho']['path']
-        self._setup_txt_path = os.path.join(self._output_dir, 'setup.txt')
         self.metric_methods = {
             'log_fourier': self._log_fourier_metric,
         }
@@ -21,6 +19,9 @@ class FoldSlicePtychoEngine(PtychoEngine):
 
     def run(self, run_id="") -> None:
 
+        self._output_dir = os.path.join(self.config['io']['result_dir'], f'fold_slice-{run_id}')
+        self._setup_txt_path = os.path.join(self._output_dir, 'setup.txt')
+
         # generate setup.txt for fold_slice input
         fold_slice_dict = {}
         fold_slice_dict['raw_data'] = self.config['io']['input_data_path']
@@ -82,9 +83,11 @@ class FoldSlicePtychoEngine(PtychoEngine):
 
         log_fourier_error = self._log_fourier_metric()
         # os.makedirs(os.path.join(self.config['io']['result_dir'], "mat"), exist_ok=True)
-        os.makedirs(os.path.join(self.config['io']['result_dir'], "tiff"), exist_ok=True)
+        # os.makedirs(os.path.join(self.config['io']['result_dir'], "tiff"), exist_ok=True)
         # shutil.copy(mat_path, os.path.join(self.config['io']['result_dir'], "mat", f"{log_fourier_error:.4f}_{run_id}.mat")) # saving .mat files takes a lot of space (expect 20+ GB for 64*64 scan size, 300 iterations)
-        shutil.copy(image_path, os.path.join(self.config['io']['result_dir'], "tiff", f"{log_fourier_error:.4f}_{run_id}.tiff"))
+        # shutil.copy(image_path, os.path.join(self.config['io']['result_dir'], "tiff", f"{log_fourier_error:.4f}_{run_id}.tiff"))
+
+        shutil.rmtree(self._output_dir)
 
 
     def metric(self, names):

readme.md

@@ -0,0 +1,26 @@
+# bo-ptycho
+
+This code is a (re)implementation of Bayesian Optimized ptychography in python for use at LEMON lab, used for
+
+- Optimization-based Thickness Estimation and FIB-induced Damage Characterization of TEM Sample via Multislice Electron Ptychography (manuscript)
+
+
+
+## Code usage
+
+Prerequisites:
+1. Create a python virtual environment with [requirements.txt](./requirements.txt)
+2. Download [fold_slice](#)
+2. Download [experimental data](#)
+
+Reconstruction and parameter optimization:
+1. Edit [config.yaml](./config.yaml)
+    - Example configurations can be found under `examples/silicon-fib/`
+    - Set io.input_data_path and io.result_dir as needed (WARNING: the result directory will be completely removed)
+    - Set ptycho.path to your fold_slice path
+2. Edit [job.sub](./job.sub)
+    - Set the SBATCH configurations as needed (especially `gres` and `output`)
+    - The number of GPUs should match bo.batch in config.yaml
+    - Set the virtual environment path
+    - Set the config.yaml path
+3. Submit the job via SLURM

samplers/__init__.py

@@ -0,0 +1,8 @@
+from .base import Sampler
+from .random import RandomSampler
+from .sobo import SOBOSampler
+
+samplers = {
+    'sobo':   SOBOSampler,
+    'random': RandomSampler,
+}

samplers/base.py

@@ -1,14 +1,11 @@
+from abc import ABC, abstractmethod
 import os
-import copy
 import numpy as np
-from bo.base import BOEngine
 
 
-class RandomBOEngine(BOEngine):
-
+class Sampler(ABC):
     def __init__(self, config):
-        super().__init__(config)
-
+        self.config = 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']
@@ -37,23 +34,12 @@ class RandomBOEngine(BOEngine):
                 self.train_y = train_y
         
 
-    def ask(self):
-        config = self.config
-        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
-
-        return next_config
+    @abstractmethod
+    def ask(self, n: int = 1) -> list[dict]:
+        pass
 
 
-    def tell(self, job_config, y_value):
+    def tell(self, job_config, y_value) -> None:
         config = self.config
 
         x_value = []

samplers/random.py

@@ -0,0 +1,27 @@
+import copy
+import numpy as np
+from samplers.base import Sampler
+
+
+class RandomSampler(Sampler):
+
+    def __init__(self, config):
+        super().__init__(config)
+
+
+    def ask(self, n = 1):
+        next_configs = []
+
+        for _ in range(n):
+            next_config = copy.deepcopy(self.config)
+            for param in self.params:
+                radius = self.config['bo']['params'][param]['radius']
+                center = self.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_configs

samplers/sobo.py

@@ -15,44 +15,16 @@ from botorch.sampling.normal import SobolQMCNormalSampler
 from botorch.utils.rounding import approximate_round
 
 
-from bo.base import BOEngine
+from samplers.base import Sampler
 
 
-class SingleObjectiveBOEngine(BOEngine):
+class SOBOSampler(Sampler):
     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):
+    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 +70,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)
+        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)
         
-        for i, param in enumerate(self.params):
-            next_config['ptycho']['params'][param] = new_x[0,i].item()
-        
-        return next_config
+        return next_configs
 
 
     def _pr_post_processing(self, X):
@@ -149,34 +118,6 @@ class SingleObjectiveBOEngine(BOEngine):
         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('tain_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)
-