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,21 +1,24 @@
 # TODO:
-
+- fold_slice:
-- unify `BOEngine.__init__()`
+    - 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
-- multi-GPU dispatcher
-    - synchronous batched BO
-    - asynchronous BO
-- template job sequences / yamls
-    - mobo
 - 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
 
-- 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`

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, n: int = 1) -> list[dict]:
-        pass
-
-    @abstractmethod
-    def tell(self, job_config, y_value):
-        pass

config-new.yaml

@@ -1,16 +1,11 @@
-job:
-  type: "random+sobo"
-  random_iters: 100
-  sobo_iters: 1000
-
 io:
   input_data_path: '/home/swim/shared/Si_project/data/Si2V1_2.mat'
-  result_dir: '/home/swim/bo-ptycho/results/260809-Si/Si2V1_2'
+  result_dir: '/home/swim/bo-ptycho/results/test'
   verbosity: 1
 
 ptycho:
-  engine: 'fold_slice'
+  engine: 'fake'
-  path: '/home/swim/fold_slice-stable'
+  # path: '/home/swim/fold_slice-park'
   params:
     voltage: 200
     alpha_max: 30
@@ -19,16 +14,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
+    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
@@ -37,23 +34,38 @@ 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:
+search:
-  mode: 'sobo'
+  metric: log_fourier
-  acquisition: 'ucb'
+  
-  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:
+
+  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,27 +1,24 @@
 job:
-  type: "random+sobo"
+  type: 'random+sobo'
-  random_iters: 8
+  random_iters: 16
   sobo_iters: 128
 
 io:
   input_data_path: '/home/swim/shared/Si_project/data/Si2V1_2.mat'
-  result_dir: '/home/swim/bo-ptycho/results/260812-Si/Si2V1_2'
+  result_dir: '/home/swim/bo-ptycho/results/test'
   verbosity: 1
 
 ptycho:
   engine: 'fold_slice'
-  path: '/home/swim/fold_slice-stable'
+  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
@@ -31,32 +28,34 @@ ptycho:
 
     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'
     diff_pattern_blur: 1
     probe_change_start: 1
     object_change_start: 1
-    grouping: 512
+    grouping: 64
     probe_position_search: 1
     regularize_layers: 0.2
-    variable_probe: 'false'
+    variable_probe: false
+
 
 bo:
   batch: 4
   acquisition: 'ucb'
-  beta: 0.1 # for ucb only
+  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_02V1_2.yaml

@@ -1,59 +0,0 @@
-job:
-  type: "random+sobo"
-  random_iters: 30
-  sobo_iters: 300
-
-io:
-  input_data_path: '/home/swim/Si_project/data/Si2V1_2.mat'
-  result_dir: '/home/swim/bo-ptycho/results/260805-si/Si2V1_2'
-  verbosity: 1
-
-ptycho:
-  engine: 'fold_slice'
-  path: '/home/swim/fold_slice'
-  params:
-    voltage: 200
-    alpha_max: 30
-    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: '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_posiiton_search: 1
-    regularize_layers: 0.2
-    variable_probe: false
-
-bo:
-  mode: 'sobo'
-  acquisition: 'ucb'
-  metric: 'log_fourier'
-  params:
-    alpha_max:
-    defocus:
-      radius: 100
-    rot_ang:
-    Nlayers:
-      radius: 5
-      type: int
-    thickness:
-      radius: 100
-  train_x:
-  train_y:

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_05V1_2.yaml

@@ -1,59 +0,0 @@
-job:
-  type: "random+sobo"
-  random_iters: 30
-  sobo_iters: 300
-
-io:
-  input_data_path: '/home/swim/Si_project/data/Si5V1_2.mat'
-  result_dir: '/home/swim/bo-ptycho/results/260805-si/Si5V1_2'
-  verbosity: 1
-
-ptycho:
-  engine: 'fold_slice'
-  path: '/home/swim/fold_slice'
-  params:
-    voltage: 200
-    alpha_max: 30
-    defocus: -200
-    rot_ang: 1
-    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
-    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_posiiton_search: 1
-    regularize_layers: 0.2
-    variable_probe: false
-
-bo:
-  mode: 'sobo'
-  acquisition: 'ucb'
-  metric: 'log_fourier'
-  params:
-    alpha_max:
-    defocus:
-      radius: 100
-    rot_ang:
-    Nlayers:
-      radius: 5
-      type: int
-    thickness:
-      radius: 150
-  train_x:
-  train_y:

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_08V1_2.yaml

@@ -1,59 +0,0 @@
-job:
-  type: "random+sobo"
-  random_iters: 30
-  sobo_iters: 300
-
-io:
-  input_data_path: '/home/swim/Si_project/data/Si8V1_2.mat'
-  result_dir: '/home/swim/bo-ptycho/results/260805-si/Si8V1_2'
-  verbosity: 1
-
-ptycho:
-  engine: 'fold_slice'
-  path: '/home/swim/fold_slice'
-  params:
-    voltage: 200
-    alpha_max: 30
-    defocus: -200
-    rot_ang: 1.3
-    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
-    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_posiiton_search: 1
-    regularize_layers: 0.2
-    variable_probe: false
-
-bo:
-  mode: 'sobo'
-  acquisition: 'ucb'
-  metric: 'log_fourier'
-  params:
-    alpha_max:
-    defocus:
-      radius: 100
-    rot_ang:
-    Nlayers:
-      radius: 5
-      type: int
-    thickness:
-      radius: 150
-  train_x:
-  train_y:

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/260805-Si/Si_30V3_2.yaml

@@ -1,59 +0,0 @@
-job:
-  type: "random+sobo"
-  random_iters: 30
-  sobo_iters: 300
-
-io:
-  input_data_path: '/home/swim/Si_project/data/Si30V3_2.mat'
-  result_dir: '/home/swim/bo-ptycho/results/260805-si/Si30V3_2'
-  verbosity: 1
-
-ptycho:
-  engine: 'fold_slice'
-  path: '/home/swim/fold_slice'
-  params:
-    voltage: 200
-    alpha_max: 30
-    defocus: 200
-    rot_ang: 0.1
-    Nlayers: 30
-    thickness: 650
-    rbf: 37
-    Nprobe: 1
-    N_scan_x: 64
-    N_scan_y: 64
-    tilt_x: 5
-    tilt_y: 3 
-    scan_step_size: 0.35
-    Niter: 100
-    Niter_save_results: 100
-    CBED_size: 192
-    ADU: 1
-    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_posiiton_search: 1
-    regularize_layers: 0.2
-    variable_probe: false
-
-bo:
-  mode: 'sobo'
-  acquisition: 'ucb'
-  metric: 'log_fourier'
-  params:
-    alpha_max:
-    defocus:
-      radius: 100
-    rot_ang:
-    Nlayers:
-      radius: 5
-      type: int
-    thickness:
-      radius: 150
-  train_x:
-  train_y:

examples/260809-Si/Si_02V1_2.sub

@@ -1,26 +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
-
-YAML=examples/260809-Si/Si_02V1_2.yaml
-cat $YAML
-python -u main.py $YAML
-
-date
-echo "SLURM JOB FINISHED"

examples/260809-Si/Si_05V1_2.sub

@@ -1,26 +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
-
-YAML=examples/260809-Si/Si_05V1_2.yaml
-cat $YAML
-python -u main.py $YAML
-
-date
-echo "SLURM JOB FINISHED"

examples/260809-Si/Si_08V1_2.sub

@@ -1,26 +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
-
-YAML=examples/260809-Si/Si_08V1_2.yaml
-cat $YAML
-python -u main.py $YAML
-
-date
-echo "SLURM JOB FINISHED"

examples/260809-Si/Si_30V3_2.sub

@@ -1,26 +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 "30V"
-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
-
-YAML=examples/260809-Si/Si_30V3_2.yaml
-cat $YAML
-python -u main.py $YAML
-
-date
-echo "SLURM JOB FINISHED"

examples/260812-Si/job.sub

@@ -1,26 +0,0 @@
-#!/bin/bash
-#SBATCH --job-name=Si2V
-#SBATCH --nodes=1
-#SBATCH --ntasks=1
-#SBATCH --cpus-per-task=1
-#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
-
-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
-
-YAML=examples/260812-Si/2V.yaml
-cat $YAML
-python -u main.py $YAML
-
-date
-echo "SLURM JOB FINISHED"

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

@@ -1,16 +1,16 @@
 job:
-  type: "random+sobo"
+  type: 'random+sobo'
-  random_iters: 8
+  random_iters: 16
   sobo_iters: 128
 
 io:
   input_data_path: '/home/swim/shared/Si_project/data/Si2V1_2.mat'
-  result_dir: '/home/swim/bo-ptycho/results/260812-Si/Si2V1_2'
+  result_dir: '/home/swim/bo-ptycho/results/Si2V1_2/260817'
   verbosity: 1
 
 ptycho:
   engine: 'fold_slice'
-  path: '/home/swim/fold_slice-stable'
+  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
@@ -36,25 +38,24 @@ ptycho:
     diff_pattern_blur: 1
     probe_change_start: 1
     object_change_start: 1
-    grouping: 512
+    grouping: 64
-    probe_posiiton_search: 1
+    probe_position_search: 1
     regularize_layers: 0.2
     variable_probe: false
 
+
 bo:
   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,34 +1,36 @@
 job:
   type: "random+sobo"
-  random_iters: 100
+  random_iters: 16
-  sobo_iters: 1000
+  sobo_iters: 512
 
 io:
   input_data_path: '/home/swim/shared/Si_project/data/Si30V3_2.mat'
-  result_dir: '/home/swim/bo-ptycho/results/260809-Si/Si30V3_2'
+  result_dir: '/home/swim/bo-ptycho/results/Si30V3_2/260816'
   verbosity: 1
 
 ptycho:
   engine: 'fold_slice'
-  path: '/home/swim/fold_slice-stable'
+  path: '/home/swim/fold_slice-park'
   params:
     voltage: 200
     alpha_max: 30
-    defocus: 175
+    defocus: 200
     rot_ang: 0.1
     Nlayers: 30
-    thickness: 680
+    thickness: 650
-    rbf: 37
+    rbf: 38
-    Nprobe: 1
-    N_scan_x: 64
-    N_scan_y: 64
     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,23 +39,24 @@ 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: 50
+      radius: 100
     rot_ang:
     Nlayers:
       radius: 5
       type: int
     thickness:
-      radius: 100
+      radius: 150
   train_x:
   train_y:

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

@@ -1,16 +1,16 @@
 job:
   type: "random+sobo"
-  random_iters: 100
+  random_iters: 16
-  sobo_iters: 1000
+  sobo_iters: 128
 
 io:
   input_data_path: '/home/swim/shared/Si_project/data/Si5V1_2.mat'
-  result_dir: '/home/swim/bo-ptycho/results/260809-Si/Si5V1_2'
+  result_dir: '/home/swim/bo-ptycho/results/Si5V1_2/260815'
   verbosity: 1
 
 ptycho:
   engine: 'fold_slice'
-  path: '/home/swim/fold_slice-stable'
+  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: 100
+  random_iters: 16
-  sobo_iters: 1000
+  sobo_iters: 128
 
 io:
-  input_data_path: '/home/swim/shared/Si_project/data/Si8V1_2.mat'
+  input_data_path: '/home/swim/shared/Si_project/data/Si8V2_2.mat'
-  result_dir: '/home/swim/bo-ptycho/results/260809-Si/Si8V1_2'
+  result_dir: '/home/swim/bo-ptycho/results/Si8V2_2/260815'
   verbosity: 1
 
 ptycho:
   engine: 'fold_slice'
-  path: '/home/swim/fold_slice-stable'
+  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 --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
+YAML="config.yaml"
-python -u main.py 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,6 +9,17 @@ def main(config_yaml):
     with open(config_yaml, 'r') as f:
         config = yaml.safe_load(f)
 
+    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)
+    
+    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)
     
 

notebooks/convergence.ipynb

@@ -1,154 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "65e78297",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import os\n",
-    "import numpy as np\n",
-    "import matplotlib.pyplot as plt"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "90917d29",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "result_dir = \"../results/260812-Si/Si2V1_2\"\n",
-    "train_x = np.load(os.path.join(result_dir, 'train_x.npy')).T\n",
-    "train_y = np.load(os.path.join(result_dir, 'train_y.npy')).T\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "05c0f656",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "train_x.shape"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "47b5dd69",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "xparams = [\n",
-    "    'defocus [A]',\n",
-    "    'layers',\n",
-    "    'thickness [A]'\n",
-    "]\n",
-    "\n",
-    "sort_index = np.argsort(train_y)\n",
-    "\n",
-    "x = train_x[:,sort_index]\n",
-    "y = train_y[sort_index]\n",
-    "\n",
-    "EPSILON = 1e-3 # choose best according to plot\n",
-    "\n",
-    "train_y_scaled = np.log(-train_y + y[-1] + EPSILON)\n",
-    "y_scaled = np.log(-y + y[-1] + EPSILON)\n",
-    "\n",
-    "fig, ax = plt.subplots(1, 2, figsize=(7,3))\n",
-    "ax[0].plot(y, 'k')\n",
-    "ax[1].plot(y_scaled, 'k')\n",
-    "plt.show()\n",
-    "\n",
-    "fig, ax = plt.subplots(1, 2, figsize=(7,3))\n",
-    "ax[0].plot(train_y, 'k')\n",
-    "ax[1].plot(train_y_scaled, 'k')\n",
-    "plt.show()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "4990770c",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "cut = 0 # set to -1 to get all points\n",
-    "\n",
-    "fig, axs = plt.subplots(1, len(x), figsize=(3*len(x)+0.5, 3.5), sharey=True)\n",
-    "\n",
-    "for i, axi in enumerate(axs):\n",
-    "    axi.scatter(x[i][cut:], y_scaled[cut:], c=y_scaled[cut:], cmap='coolwarm')\n",
-    "    axi.set_xlabel(xparams[i])\n",
-    "\n",
-    "axs[0].set_ylabel('$- \\\\log (\\ \\\\mathtt{fourier\\_error}\\ )$')\n",
-    "\n",
-    "plt.tight_layout()\n",
-    "plt.show()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "a2694f3c",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "\n",
-    "fig, axs = plt.subplots(len(x), len(x), figsize=(3*len(x)+0.5, 3*len(x)+0.5))\n",
-    "\n",
-    "for i, axi in enumerate(axs):\n",
-    "    for j, axij in enumerate(axi):\n",
-    "        axij.scatter(\n",
-    "            x[j], x[i],\n",
-    "            c = y_scaled, # color\n",
-    "            cmap = 'coolwarm',\n",
-    "            s = 20, # size\n",
-    "            alpha = 0 if i == j else 1 # make diagonal transparent\n",
-    "            )\n",
-    "        if i == len(x)-1:\n",
-    "            axij.set_xlabel(xparams[j])\n",
-    "        else:\n",
-    "            axij.set_xticks([])\n",
-    "        if j == 0:\n",
-    "            axij.set_ylabel(xparams[i])\n",
-    "        else:\n",
-    "            axij.set_yticks([])\n",
-    "\n",
-    "plt.tight_layout()\n",
-    "plt.show()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "295ca65c",
-   "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
-}

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/batched_sobo.py

@@ -1,47 +1,27 @@
 import os
 import traceback
 import multiprocessing as mp
-import shutil
 
-import bo
+import samplers
 import ptycho
 
-def run_ptycho_worker(
+def run_ptycho_worker(worker_id, gpu_token, job_config, metric, run_id, result_queue):
-    worker_id,
+
-    gpu_token,
+    # This process, and MATLAB launched from it, can see exactly one GPU.
-    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:
-        PTYCHOENGINE = ptycho.engines[job_config["ptycho"]["engine"]]
+        ptycho_engine = ptycho.engines[job_config["ptycho"]["engine"]](job_config)
-
-        ptycho_engine = PTYCHOENGINE(job_config)
         ptycho_engine.run(run_id=run_id)
         y_value = ptycho_engine.metric(metric)
-
+        result_queue.put((worker_id, y_value, None))
-        result_queue.put(
-            (worker_id, y_value, None)
-        )
 
     except Exception:
-        result_queue.put(
+        result_queue.put((worker_id, None, traceback.format_exc()))
-            (worker_id, None, traceback.format_exc())
-        )
 
 
-def run_batch(
+def run_batch(ctx, gpu_tokens, job_configs, metric, iteration):
-    ctx,
+
-    gpu_tokens,
-    job_configs,
-    metric,
-    iteration,
-):
     result_queue = ctx.Queue()
     processes = []
 
@@ -51,25 +31,12 @@ def run_batch(
 
         p = ctx.Process(
             target=run_ptycho_worker,
-            args=(
+            args=(i, gpu_tokens[i], job_config, metric, run_id, result_queue),
-                i,
-                gpu_tokens[i],
-                job_config,
-                metric,
-                run_id,
-                result_queue,
-            ),
         )
-
         p.start()
         processes.append(p)
 
-    # Four jobs are now running concurrently.
+    results = [result_queue.get() for _ in processes]
-
-    results = [
-        result_queue.get()
-        for _ in processes
-    ]
 
     # Synchronization barrier.
     for p in processes:
@@ -80,111 +47,48 @@ def run_batch(
 
     for worker_id, _, error in results:
         if error is not None:
-            raise RuntimeError(
+            raise RuntimeError(f"Ptycho worker {worker_id} failed:\n{error}")
-                f"Ptycho worker {worker_id} failed:\n{error}"
-            )
 
     return [y_value for _, y_value, _ in results]
 
 
 def sobo_pipeline(config):
 
-    result_dir = config["io"]["result_dir"]
-    if os.path.exists(result_dir):
-        shutil.rmtree(result_dir)
-    os.makedirs(result_dir, exist_ok=True)
-
     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"]
 
-    # SLURM should expose the four GPUs allocated to this job.
     visible = os.environ.get("CUDA_VISIBLE_DEVICES")
-
     if visible is None:
-        raise RuntimeError(
+        raise RuntimeError("CUDA_VISIBLE_DEVICES is not set")
-            "CUDA_VISIBLE_DEVICES is not set"
+    gpu_tokens = [token.strip() for token in visible.split(",") if token.strip()]
-        )
-
-    gpu_tokens = [
-        token.strip()
-        for token in visible.split(",")
-        if token.strip()
-    ]
-
     if len(gpu_tokens) < BO_BATCH:
-        raise RuntimeError(
+        raise RuntimeError(f"Expected {BO_BATCH} allocated GPUs, got {len(gpu_tokens)}")
-            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 = bo.RandomBOEngine(config)
+    ############################ RANDOM SAMPLING ###############################
+    randombo = samplers.RandomSampler(config)
 
     for j in range(RANDOM_ITERS):
-        print(
+        print(f"RANDOM sampling; iteration {j}")
-            f"RANDOM sampling; iteration {j}",
-            flush=True,
-        )
-
         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)
 
-        y_values = run_batch(
+    ############################  SOBO SAMPLING  ###############################
-            ctx=ctx,
+    sobo = samplers.SOBOSampler(config)
-            gpu_tokens=gpu_tokens,
-            job_configs=job_configs,
-            metric=METRIC,
-            iteration=j,
-        )
-
-        # Only the parent touches BO state / train_x / train_y.
-        for job_config, y_value in zip(
-            job_configs,
-            y_values,
-        ):
-            randombo.tell(
-                job_config,
-                y_value,
-            )
-
-    # ---------------------------------------------------------
-    # SOBO
-    # ---------------------------------------------------------
-
-    sobo = bo.SingleObjectiveBOEngine(config)
-
     sobo.train_x = randombo.train_x
     sobo.train_y = randombo.train_y
 
-    for j in range(SOBO_ITERS):
+    for j in range(RANDOM_ITERS, SOBO_ITERS+RANDOM_ITERS):
-        iteration = RANDOM_ITERS + j
+        print(f"SOBO sampling iteration {j}")
-
-        print(
-            f"SOBO sampling; iteration {iteration}",
-            flush=True,
-        )
-
         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)
 
-        y_values = run_batch(
+        for job_config, y_value in zip(job_configs, y_values):
-            ctx=ctx,
+            sobo.tell(job_config, y_value)
-            gpu_tokens=gpu_tokens,
-            job_configs=job_configs,
-            metric=METRIC,
-            iteration=iteration,
-        )
-
-        for job_config, y_value in zip(
-            job_configs,
-            y_values,
-        ):
-            sobo.tell(
-                job_config,
-                y_value,
-            )

pipelines/sobo.py

@@ -1,49 +0,0 @@
-# depreacated, use pipelines.batched_sobo.sobo_pipeline()
-
-# 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

@@ -1,189 +1,6 @@
-import os
+import samplers
-import traceback
-import multiprocessing as mp
-import shutil
-
-import bo
 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:
-        PTYCHOENGINE = ptycho.engines[job_config["ptycho"]["engine"]]
-
-        ptycho_engine = PTYCHOENGINE(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)
-
-    # Four jobs are now running concurrently.
-
-    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 test_pipeline(config):
-
+    pass
-    result_dir = config["io"]["result_dir"]
-    shutil.rmtree(result_dir)
-    os.makedirs(result_dir, exist_ok=True)
-
-    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"]
-
-    # SLURM should expose the four GPUs allocated to this job.
-    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 = bo.RandomBOEngine(config)
-
-    for j in range(RANDOM_ITERS):
-        print(
-            f"RANDOM sampling; iteration {j}",
-            flush=True,
-        )
-
-        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,
-        )
-
-        # Only the parent touches BO state / train_x / train_y.
-        for job_config, y_value in zip(
-            job_configs,
-            y_values,
-        ):
-            randombo.tell(
-                job_config,
-                y_value,
-            )
-
-    # ---------------------------------------------------------
-    # SOBO
-    # ---------------------------------------------------------
-
-    sobo = bo.SingleObjectiveBOEngine(config)
-
-    sobo.train_x = randombo.train_x
-    sobo.train_y = randombo.train_y
-
-    for j in range(SOBO_ITERS):
-        iteration = RANDOM_ITERS + j
-
-        print(
-            f"SOBO sampling; iteration {iteration}",
-            flush=True,
-        )
-
-        job_configs = sobo.ask(n=BO_BATCH)
-
-        y_values = run_batch(
-            ctx=ctx,
-            gpu_tokens=gpu_tokens,
-            job_configs=job_configs,
-            metric=METRIC,
-            iteration=iteration,
-        )
-
-        for job_config, y_value in zip(
-            job_configs,
-            y_values,
-        ):
-            sobo.tell(
-                job_config,
-                y_value,
-            )

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,26 +34,12 @@ class RandomBOEngine(BOEngine):
                 self.train_y = train_y
         
 
-    def ask(self, n = 1):
+    @abstractmethod
-        config = self.config
+    def ask(self, n: int = 1) -> list[dict]:
-        next_configs = []
+        pass
-
-        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_configs
 
 
-    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,40 +15,12 @@ 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']
 
 
@@ -146,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)
-