Si_2V1_2 batched bo job

TODO.md

@@ -1,13 +1,21 @@
-TODO:
+# TODO:
 
-- Clean up BO states
+- unify `BOEngine.__init__()`
-    - unify `BOEngine.__init__()`
 - 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
+
+
+# TODAY:
+
 - 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/base.py

@@ -7,7 +7,7 @@ class BOEngine(ABC):
         self.state = None
 
     @abstractmethod
-    def ask(self):
+    def ask(self, n: int = 1) -> list[dict]:
         pass
 
     @abstractmethod

bo/random.py

@@ -37,10 +37,12 @@ class RandomBOEngine(BOEngine):
                 self.train_y = train_y
         
 
-    def ask(self):
+    def ask(self, n = 1):
         config = self.config
-        next_config = copy.deepcopy(config)
+        next_configs = []
 
+        for _ in range(n):
+            next_config = copy.deepcopy(config)
             for param in self.params:
                 radius = config['bo']['params'][param]['radius']
                 center = config['ptycho']['params'][param]
@@ -49,8 +51,9 @@ class RandomBOEngine(BOEngine):
                 if self.param_types[param] == 'int':
                     next_value = round(next_value)
                 next_config['ptycho']['params'][param] = next_value
+            next_configs.append(next_config)
 
-        return next_config
+        return next_configs
 
 
     def tell(self, job_config, y_value):

bo/sobo.py

@@ -52,7 +52,7 @@ class SingleObjectiveBOEngine(BOEngine):
         self.acquisition = config['bo']['acquisition']
 
 
-    def ask(self):
+    def ask(self, n = 1):
 
         train_x = torch.from_numpy(self.train_x)
         train_y = torch.from_numpy(self.train_y).unsqueeze(-1) # shape: (BOiter, 1)
@@ -98,44 +98,41 @@ class SingleObjectiveBOEngine(BOEngine):
         sampler = SobolQMCNormalSampler(sample_shape=torch.Size([512]))
 
         if self.acquisition == 'ucb':
-            beta = 0.2
+            acqf = qUpperConfidenceBound(gp, beta=self.config['bo']['beta'], sampler=sampler)
-            print("Acquisition: UCB | Beta: {} (fixed)".format(beta))
-            acqf = qUpperConfidenceBound(gp, beta=beta, sampler=sampler)
         elif self.acquisition == 'ei':
-            best_f = train_y.max()
+            acqf = qLogExpectedImprovement(gp, best_f=train_y.max(), sampler=sampler)
-            print("Acquisition: LogEI  best_f: {:.6f}".format(best_f.item()))
-            acqf = qLogExpectedImprovement(gp, best_f=best_f, 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_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):
+        return next_configs
-            next_config['ptycho']['params'][param] = new_x[0,i].item()
-        
-        return next_config
 
 
     def _pr_post_processing(self, X):

config.yaml

@@ -1,16 +1,16 @@
 job:
   type: "random+sobo"
-  random_iters: 3
+  random_iters: 8
-  sobo_iters: 30
+  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'
+  result_dir: '/home/swim/bo-ptycho/results/260812-Si/Si2V1_2'
   verbosity: 1
 
 ptycho:
   engine: 'fold_slice'
-  path: '/home/swim/fold_slice'
+  path: '/home/swim/fold_slice-stable'
   params:
     voltage: 200
     alpha_max: 30
@@ -25,8 +25,10 @@ ptycho:
     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: ''
@@ -36,14 +38,15 @@ ptycho:
     diff_pattern_blur: 1
     probe_change_start: 1
     object_change_start: 1
-    grouping: 64
+    grouping: 512
-    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 # for ucb only
   metric: 'log_fourier'
   params:
     alpha_max:

examples/260809-Si/Si_02V1_2.sub

@@ -0,0 +1,26 @@
+#!/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_02V1_2.yaml

@@ -0,0 +1,59 @@
+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'
+  verbosity: 1
+
+ptycho:
+  engine: 'fold_slice'
+  path: '/home/swim/fold_slice-stable'
+  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/260809-Si/Si_05V1_2.sub

@@ -0,0 +1,26 @@
+#!/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_05V1_2.yaml

@@ -0,0 +1,59 @@
+job:
+  type: "random+sobo"
+  random_iters: 100
+  sobo_iters: 1000
+
+io:
+  input_data_path: '/home/swim/shared/Si_project/data/Si5V1_2.mat'
+  result_dir: '/home/swim/bo-ptycho/results/260809-Si/Si5V1_2'
+  verbosity: 1
+
+ptycho:
+  engine: 'fold_slice'
+  path: '/home/swim/fold_slice-stable'
+  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/260809-Si/Si_08V1_2.sub

@@ -0,0 +1,26 @@
+#!/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_08V1_2.yaml

@@ -0,0 +1,59 @@
+job:
+  type: "random+sobo"
+  random_iters: 100
+  sobo_iters: 1000
+
+io:
+  input_data_path: '/home/swim/shared/Si_project/data/Si8V1_2.mat'
+  result_dir: '/home/swim/bo-ptycho/results/260809-Si/Si8V1_2'
+  verbosity: 1
+
+ptycho:
+  engine: 'fold_slice'
+  path: '/home/swim/fold_slice-stable'
+  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/260809-Si/Si_30V3_2.sub

@@ -0,0 +1,26 @@
+#!/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/260809-Si/Si_30V3_2.yaml

@@ -0,0 +1,59 @@
+job:
+  type: "random+sobo"
+  random_iters: 100
+  sobo_iters: 1000
+
+io:
+  input_data_path: '/home/swim/shared/Si_project/data/Si30V3_2.mat'
+  result_dir: '/home/swim/bo-ptycho/results/260809-Si/Si30V3_2'
+  verbosity: 1
+
+ptycho:
+  engine: 'fold_slice'
+  path: '/home/swim/fold_slice-stable'
+  params:
+    voltage: 200
+    alpha_max: 30
+    defocus: 175
+    rot_ang: 0.1
+    Nlayers: 30
+    thickness: 680
+    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: 50
+    rot_ang:
+    Nlayers:
+      radius: 5
+      type: int
+    thickness:
+      radius: 100
+  train_x:
+  train_y:

examples/260812-Si/2V.yaml

@@ -0,0 +1,60 @@
+job:
+  type: "random+sobo"
+  random_iters: 8
+  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'
+  verbosity: 1
+
+ptycho:
+  engine: 'fold_slice'
+  path: '/home/swim/fold_slice-stable'
+  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: 512
+    probe_posiiton_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
+  train_x:
+  train_y:

examples/260812-Si/job.sub

@@ -0,0 +1,26 @@
+#!/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"

job.sub

@@ -3,7 +3,7 @@
 #SBATCH --nodes=1
 #SBATCH --ntasks=1
 #SBATCH --cpus-per-task=1
-#SBATCH --gres=gpu:rtx-6000ada:1
+#SBATCH --gres=gpu:rtx-6000ada:4
 #SBATCH --time=100:00:00
 #SBATCH --output=/home/swim/slurm-logs/job_%j.log
 

main.py

@@ -8,11 +8,7 @@ def main(config_yaml):
     with open(config_yaml, 'r') as f:
         config = yaml.safe_load(f)
 
-    job_types = {
+    pipelines.job_types[config['job']['type']](config)
-        'random+sobo': pipelines.sobo_pipeline
-    }
-
-    job_types[config['job']['type']](config)
     
 
 if __name__ == "__main__":

notebooks/convergence.ipynb

@@ -2,43 +2,44 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "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": 42,
+   "execution_count": null,
    "id": "90917d29",
    "metadata": {},
    "outputs": [],
    "source": [
-    "train_x = np.load(\"../results/260805-si/Si8V1_2/train_x.npy\").T\n",
+    "result_dir = \"../results/260812-Si/Si2V1_2\"\n",
-    "train_y = np.load(\"../results/260805-si/Si8V1_2/train_y.npy\").T"
+    "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": 45,
+   "execution_count": null,
-   "id": "47b5dd69",
+   "id": "05c0f656",
    "metadata": {},
-   "outputs": [
+   "outputs": [],
-    {
+   "source": [
-     "data": {
+    "train_x.shape"
-      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAmQAAAESCAYAAABTgsn9AAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjYsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvq6yFwwAAAAlwSFlzAAAPYQAAD2EBqD+naQAAPIhJREFUeJzt3XlcVPX+P/DXLMyMAkPisAoqLiiJW6Bs+a20KNLKunVdcqm0GyVdTfl+v5g+UnxodrPMrLDScikfN/u6lCU3obQkt1yYxCVEUVkECZAZUAGZOb8/+M3kCAMzCByYeT0fj/O4eeYs73MezLmv+XzO+RyJIAgCiIiIiEg0UrELICIiInJ2DGREREREImMgIyIiIhIZAxkRERGRyBjIiIiIiETGQEZEREQkMgYyIiIiIpHJxS6gPRmNRly+fBnu7u6QSCRil0NErUAQBFRWVsLf3x9Saef9jcnrE5Hjsef65FSB7PLlywgMDBS7DCJqA/n5+QgICBC7jBbj9YnIcdlyfXKqQObu7g6g/sSo1WqRqyGi1qDX6xEYGGj+fndWvD4ROR57rk9OFchM3QBqtZoXPCIH09m7+Xh9InJctlyfOu8NF0REREQOgoGMiIiISGQMZEREREQiYyAjIiIiEhkDGREREZHIGMiIiIiIROZUw14QdSaCINzR/7bGNlq6D6PR2ORkMBia/Ly2thZVVVV49NFHW3r6HFZVVRWOHTsGuVyOmJgYscsholbCQNYBXLx4EV9//TWKioqg1+tx7do1GAwGCILQKhOAVttWS/dt0t7BoDNum/5SV1cHmUwmdhkdyhdffIFXXnkFDz/8MH744QexyyGiVsJAJqLc3Fx89913WLhwIaqqqsQuh6jVyGQySKXSZidry8nlcqjVaty4cQNubm5iH06HMmLECADAkSNHIAhCpx8Ql4jqMZCJZOXKlUhMTDS3ioSGhmLs2LFQq9Vwc3ODTCaDRCK54wlAq2ynNWowuX1ea/0vty3+tqntDRkyBAqFAuXl5bhw4QL69OkjdklE1AoYyERw4sQJLFiwAIIgYOTIkXjkkUcwb948vi6FiJqlUCgwdOhQHDlyBEeOHGEgI3IQDGTtbPPmzXjhhRdQW1uL++67D3v37mXrAhHZZdCgQThy5Ahyc3PFLoWIWgmHvWhHtbW1mDt3LmprazF27Fhs2bKFYYyI7Obl5QUAKCsrE7kSImotbCFrJ6WlpUhOTkZJSQn8/PywY8cOuLi4iF0WEXVC3bt3B1B/XSEix8BA1g4EQcC4ceNw+PBhAEBiYiLDGBG1mEajAcBARuRI2GXZxqqqqhAbG2sOY++88w5ee+01kasios6MgYzI8TCQtaHffvsNwcHB+PHHHwEA//znPzFv3jzeN0bkoGpqajBs2DBIJBJotdo22w8DGZHjYSBrQ6tXr0ZRUREAYN26dVixYoXIFRFRW/qf//kf+Pv7t/l+eA8ZkeNhIGsjRqMRaWlpAIBdu3ZhxowZUCgUIldFRG3lP//5D9LS0vDOO++0+b5MLWQ6nQ43b95s8/0RUdtjIGsjy5Ytw59//glXV1c8+OCDYpdDRG3oypUrePHFF/HFF1+ga9euNq1TU1MDvV5vMdmqW7du5lsf0tPTW1QzEXUsDGRtYO/evVi0aBEAYPHixWwZI3JggiDgueeeQ3x8PMLDw21eb/ny5fDw8DBPgYGBNq8rk8nwyCOPmLdDRJ0fA1krO336NKZMmQJBEDBz5kwkJiaKXRIRtcDixYubfU/r0aNH8cEHH0Cv12P+/Pl2bX/+/PnQ6XTmKT8/3671582bBwAoLi62az0i6pg4DlkrMhgMeOSRR3D58mX069cP7733ntglEVELJSQkYOLEiU0u07t3byxduhSHDh2CUqm0+Cw8PBzPPvssNm7c2Oi6SqWywTr2MN1HVllZ2eJtEFHHwUDWivbv34/8/HwoFArs2bMHbm5uYpdERC2k0WjMoacpq1evxtKlS83/vnz5Mh5++GFs2bIFERERbVafu7s7AAYyIkfBQNZK6urqsGrVKgDAxIkT7bofhIg6r549e1r82/RDrG/fvggICGiz/ZoC2fXr12EwGCCTydpsX0TU9lp0D1lKSgqCgoKgUqkQFhaGjIwMm9bbv38/5HI5hg0b1uCziooKzJo1C35+flCpVAgJCUFqaqr588bu5/D19W1J+a2uvLwcI0eOxI4dOyCVShEfHy92SUTk4EyBDGArGZEjsLuFbMuWLZgzZw5SUlIQExODTz75BHFxcTh9+nSDX4q30ul0mDZtGsaMGYMrV65YfFZbW4uHHnoI3t7e2Lp1KwICApCfn29xwQGAQYMGmUe9B9AhfhH+9ttvePzxx3HlyhV07doVq1atQlRUlNhlEZFIevfuDUEQ2nw/SqUSLi4uuHnzJiorK3HXXXe1+T6JqO3YHchWrlyJGTNmYObMmQCAVatWYffu3VizZk2Tj1+/9NJLmDx5MmQyGb755huLzz7//HOUl5fjwIED5pdu9+rVq2GxcrldrWI1NTWoqakx/9uecX5skZWVhYcffhgVFRXo3bs3duzY0WjrHxFRa5NIJHB3d0d5eTlbyIgcgF1dlrW1tTh27BhiY2Mt5sfGxuLAgQNW11u/fj3Onz9vHpvrdjt37kRUVBRmzZoFHx8fhIaG4s0334TBYLBYLicnB/7+/ggKCsLEiRORm5vbZL13Ms6PLebPn4+KigpER0cjKyuLYYyI2hVv7CdyHHYFstLSUhgMBvj4+FjM9/HxsToWTk5ODpKSkrB582bI5Y03yOXm5mLr1q0wGAxITU3FwoUL8e6772LZsmXmZSIiIrBp0ybs3r0ba9euRXFxMaKjo1FWVma13jsd56cpRqMRv/76K4D6p6z4RCURtTcGMiLH0aKnLE2v7DARBKHBPKB+XK7JkycjOTkZwcHBVrdnNBrh7e2NTz/9FDKZDGFhYbh8+TJWrFiBN954AwAQFxdnXn7w4MGIiopC3759sXHjRsydO7fR7d7pOD9N+eOPP6DT6dClSxcMGTKkTfZBRNQUUyBr7dsxiKj92RXINBoNZDJZg9awkpKSBq1mQP2vtqNHjyIzMxMJCQkA6sOXIAiQy+VIS0vD6NGj4efnBxcXF4ub9ENCQlBcXIza2tpGXz3k6uqKwYMHIycnx55DaDXbt28HAIwcOdJ83xsRUXtSq9UA2EJG5Ajs6rJUKBQICwtr8DLb9PR0REdHN1herVYjKysLWq3WPMXHx2PAgAHQarXmQRNjYmJw7tw5GI1G87pnz56Fn5+f1fdA1tTU4MyZM/Dz87PnEFpFbm4u/vWvfwEAXnzxxXbfPxER8FcL2YULF0SuhIjulN3jkM2dOxfr1q3D559/jjNnzuC1115DXl6eeeyt+fPnY9q0afUbl0oRGhpqMXl7e0OlUiE0NBSurq4AgJdffhllZWWYPXs2zp49i127duHNN9/ErFmzzPtNTEzEL7/8ggsXLuDw4cN4+umnodfrMX369NY4D3aZN28eqqqqcO+992LSpEntvn8iIgDmoS6Sk5NRUVEhai1EdGfsvodswoQJKCsrw5IlS1BUVITQ0FCkpqaah6koKipCXl6eXdsMDAxEWloaXnvtNQwZMgQ9evTA7Nmz8b//+7/mZQoKCjBp0iSUlpbCy8sLkZGROHToUKPDY7Sl69ev44cffgAAfPjhh5BK+X52IhLHK6+8gnXr1gGo71UYOXKkyBURUUtJhPYYwbCD0Ov18PDwgE6nM997Ya/U1FSMHTsWgYGBuHTpUqMPMxBR+2mN73VH0NLjCAsLw/Hjx7Fr1y48+uijbVghEdnLnu81m3fslJWVBQAYNWoUwxgRic70AvTS0lKRKyGiO8FAZifTa5969OghciVERAxkRI6CgcxOpkDW2DAfRETtjYGMyDEwkNnJNAYbAxkRdQQMZESOgYHMTqYWMnteck5E1FYYyIgcAwOZndhlSUQdCQMZkWNgILPDzZs3zRc9BjIi6ggYyIgcAwOZHU6dOgUAkMvl6N69u8jVEBExkBE5CgYyO7z33nsAgPHjx1u8CJ2ISCymQFZWVmbxPmAi6lwYyGxkNBrx3XffAQBmz54tcjVERPVMrfVGo5HvsyTqxBjIbJSVlYWrV6/Czc0NERERYpdDRAQAUCgU5leysNuSqPNiILPR/v37AQAxMTFwcXERuRoior/wPjKizo+BzEYlJSUAgKCgIJErISKyxEBG1PkxkNno2rVrAABXV1eRKyEissRARtT5MZDZ6Pr16wAYyIio42EgI+r8GMhsZGoh69q1q8iVEBFZCggIAAAcP35c5EqIqKUYyGzELksi6qieeuopAMC3336L8vJykashopZgILMRuyyJqKO65557MHToUFRXVyMpKUnscoioBRjIbMQWMiLqqCQSCVasWAEA2LZtm8jVEFFLMJDZiPeQEVFHNmLECABAeXm5uUWfiDoPBjIbsYWMiDoyDw8P8/WpoKBA5GqIyF4MZDbiPWRE1JFJJBIEBgYCYCAj6owYyGzELksi6uhMw1/k5+eLXAkR2YuBzEbssiSijs4UyNhCRtT5tCiQpaSkICgoCCqVCmFhYcjIyLBpvf3790Mul2PYsGENPquoqMCsWbPg5+cHlUqFkJAQpKamtsp+75TRaMSNGzcAMJARUcdletfuiRMnRK6EiOxldyDbsmUL5syZgwULFiAzMxOjRo1CXFwc8vLymlxPp9Nh2rRpGDNmTIPPamtr8dBDD+HixYvYunUrsrOzsXbtWvTo0eOO99saTGEMYJclEXVcDz30EABg9+7duHnzpsjVEJE9JIIgCPasEBERgXvuuQdr1qwxzwsJCcH48eOxfPlyq+tNnDgR/fv3h0wmwzfffAOtVmv+7OOPP8aKFSvwxx9/wMXFpVX3eyu9Xg8PDw/odDqo1Wqb1gGAkpIS+Pj4AAAMBgOkUvb0EnUULf1edzSLine truncated
-      "text/plain": [
-       "<Figure size 700x300 with 2 Axes>"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "47b5dd69",
    "metadata": {},
-     "output_type": "display_data"
+   "outputs": [],
-    }
-   ],
    "source": [
     "xparams = [\n",
     "    'defocus [A]',\n",
@@ -51,40 +52,35 @@
     "x = train_x[:,sort_index]\n",
     "y = train_y[sort_index]\n",
     "\n",
-    "EPSILON = 1e-5 # choose best according to plot\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": 52,
+   "execution_count": null,
    "id": "4990770c",
    "metadata": {},
-   "outputs": [
+   "outputs": [],
-    {
-     "data": {
-      "image/png": "iVBORw0KGgoAAAANSUhEUgAAA6sAAAFUCAYAAADCuapgAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjYsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvq6yFwwAAAAlwSFlzAAAPYQAAD2EBqD+naQABAABJREFUeJzsnXV8FNfax3/nzKzFlZBAQnB3K1CKW42631L3W29veyvQvr2Xut26Uah7C5QCpRQrTnG3ACEhbuu7M+f9Y3Ynu1mNbuR8P5+F7MyZM8/s7px5nnMeIYwxBg6Hw+FwOBwOh8PhcJoRNNICcDgcDofD4XA4HA6HUxNurHI4HA6Hw+FwOBwOp9nBjVUOh8PhcDgcDofD4TQ7uLHK4XA4HA6Hw+FwOJxmBzdWORwOh8PhcDgcDofT7ODGKofD4XA4HA6Hw+Fwmh3cWOVwOBwOh8PhcDgcTrODG6scDofD4XA4HA6Hw2l2cGOVw+FwOBwOh8PhcDjNDm6scjgcDofD4XA4HA6n2cGNVQBr1qzBBRdcgIyMDBBC8PPPPzfq+ebMmQNCiNerffv2DdJ3SUkJOnbsCEIIysvLQ7bfsGEDJk6ciOjoaCQkJGD8+PGwWCwAgFWrVvnI6X5t2bJF7WPLli2YNGkSEhISkJiYiKlTp2LHjh0Ncj2BuO+++zB06FDodDoMGjSoUc/F4XA4HA6Hw+Fwmh5urAIwmUwYOHAg3nrrrSY7Z9++fZGfn6++du/e3SD93nzzzRgwYEBYbTds2IDp06dj6tSp2Lx5M7Zs2YJ77rkHlCo/i9GjR3vJmJ+fj1tuuQXZ2dkYNmwYAKCqqgrTpk1DVlYWNm3ahHXr1iEuLg7Tpk2Dw+FokGvyB2MMN910E6688spGOweHw+FwOBwOh8OJHGKkBWgOzJgxAzNmzAi4326348knn8QXX3yB8vJy9OvXDy+88ALGjx9f53OKothgq6lu3n33XZSXl+Ppp5/Gb7/9FrL9Aw88gHvvvRePPfaYuq179+7q31qt1ktGh8OBhQsX4p577gEhBABw8OBBlJWV4dlnn0VmZiYAYPbs2RgwYABOnjyJrl27AgDWr1+Pxx57DFu2bEFKSgouvvhizJ07F9HR0XW61jfffBMAUFRUhF27dtWpDw6Hw+FwOBwOh9N84SurYXDjjTfir7/+wtdff41du3bh8ssvx/Tp03H48OE693n48GFkZGSgc+fOuOqqq3Ds2LF6ybhv3z48++yzWLBggboyGozCwkJs2rQJ7dq1w+jRo5GWloZx48Zh3bp1AY9ZuHAhiouLccMNN6jbevbsiZSUFHz88cew2+2wWCz4+OOP0bdvX3Tq1AkAsHv3bkybNg2XXHIJdu3ahW+++Qbr1q3DPffcU69r5nA4HA6Hw+FwOK0XbqyG4OjRo/jqq6/w3XffYezYsejatSsefvhhnH322Zg3b16d+hw5ciQWLFiAZcuW4cMPP8SZM2cwevRolJSU1Kk/m82Gq6++Gi+99BKysrLCOsZtHM+ZMwe33norli5diiFDhmDSpEkBjfCPP/4Y06ZNU1dQASA2NharVq3C559/DoPBgJiYGCxbtgxLliyBKCoL9y+99BKuueYa3H///ejevTtGjx6NN998EwsWLIDVaq3TNXM4HA6Hw+FwOJzWDTdWQ/D333+DMYYePXogJiZGfa1evRpHjx4FAOTk5ARMROR+ea4izpgxA5deein69++PyZMn49dffwUAzJ8/X21zww03hOyzuLgYAPD444+jd+/euO6668K+LlmWAQC33347brzxRgwePBivvfYaevbsiU8++cSnfW5uLpYtW4abb77Za7vFYsFNN92EMWPGYOPGjfjrr7/Qt29fnHvuuWqipm3btuHTTz/1+vymTZsGWZZx/PjxWl8vh8PhcDgcDofDaf3wmNUQyLIMQRCwbds2CILgtS8mJgYA0KFDB+zfvz9oP4mJiQH3RUdHo3///l4rmnPnzvWKJQ3W58qVK7F79258//33AJTkQwCQkpKCJ554As8884zPsenp6QCAPn36eG3v3bs3Tp486dN+3rx5SE5OxoUXXui1/csvv0ROTg42bNiguh9/+eWXSExMxC+//IKrrroKsizj9ttvx7333uvTr3sluDbXy+FwOBwOh8PhcFo/3FgNweDBgyFJEgoLCzF27Fi/bTQaDXr16lXnc9hsNuzfv9+r//T0dNWgDMUPP/ygrmICSimZm266CWvXrlUTHNUkOzsbGRkZOHjwoNf2Q4cO+SSbYoxh3rx5uP7666HRaLz2mc1mUErVhEsA1Pfu1dshQ4Zg79696NatW8BrqM31cjgcDofD4XA4nNYPdwMGYDQasWPHDrU26PHjx7Fjxw6cPHkSPXr0wLXXXovrr78eP/74I44fP44tW7bghRdewJIlS+p0vocffhirV6/G8ePHsWnTJlx22WWorKzErFmz6tRf165d0a9fP/XVuXNnAMoqabt27QAAp0+fRq9evbB582YAACEEjzzyCN588018//33OHLkCJ566ikcOHDAx9V35cqVOH78uM92AJgyZQrKyspw9913Y//+/di7dy9uvPFGiKKICRMmAAD+9a9/YcOGDbj77ruxY8cOHD58GAsXLsQ///nPOl0vABw5cgQ7duzAmTNnYLFY1O/PbrfXuU8Oh8PhcDgcDofTfOArqwC2bt2qGlYA8OCDDwIAZs2ahU8//RTz5s3Dc889h4ceeginT59GcnIyRo0ahXPPPbdO58vNzcXVV1+N4uJipKam4qyzzsLGjRvV7LmNgcPhwMGDB2E2m9Vt999/P6xWKx544AGUlpZi4MCB+P33331WYz/++GOMHj0avXv39um3V69eWLRoEZ555hmMGjUKlFIMHjwYS5cuVVdKBwwYgNWrV+OJJ57A2LFjwRhD165d61Uj9ZZbbsHq1avV94MHDwagTDRkZ2fXuV8Oh8PhcDgcDofTPCDMHeDI4XA4HA6Hw+FwOBxOM4G7AXM4HA6Hw+FwOBwOp9nBjVUOh8PhcDgcDofD4TQ72mzMqizLyMvLQ2xsrFcmWw6Hw/EHYwxVVVXIyMhQyzQ1BXys4nA4tYGPVRwOpyUQ7ljVZo3VvLw8ZGZmRloMDofTwjh16hQ6duzYZOfjYxWHw6kLfKzicDgtgVBjVZs1VmNjYwEoH1BcXFyEpeFwOM2dyspKZGZmqmNHU8HHKg6HUxv4WMXhcFoC4Y5VbdZYdbuoxMXF8UGVw+GETVO7t/GxisPh1AU+VnE4nJZAqLGKJ1jicDgcDofD4XA4HE6zgxurHA6Hw+FwOBwOh8NpdnBjlcPhcDgcDofD4XA4zQ5urHI4HA6Hw+FwOBwOp9nRZhMscZoWy8k8nPjwGxSv+AtgDMnjR6LTbVcjqgtPc8/hNASOSiNyF/yE/B+WQqoyIXZAT3S67WoknjUo0qJxOBwOx4VksSLv68U4/dUiOErLEd2jM7JuuRLJE87i9Wk5HD9wY5XT6BQsXoltV94LJkmAJAMAKnfsx/E3PsWgBS8j4/JzIywhp63hNJlhPnoSVKtBdI/OIEGKUbcEjAePYePkf8BWUKJsYAxVew/h9Gc/o8tDN6PX3Ee4EsThtECcRhPMR0+BGnSI7p7N7+MWjjWvABunXA/ToRyAEkBmqNp7GPnf/YYO11yIgZ88DyIIkRaTw6k1jkojLMdzIUTpEdWtU4OOVdxY5TQq5uOnFEPV4QQYU7czSQIA7Lj+YcT27YHYPt0iJSKnDeGoNOLQ06/j1LzvIJmtAABDpw7o9q/bkXnLFS1SEZSdTmw+/1bYi8q87zGnco8de+VjxPbtgY7/uChCEnI4nNriKKvAgSdfRe6CnyBbbQCAqC5Z6PbvO5E565IIS8epK9uuvBfmY6eUN7IyXrvH6tNfLURMn27o9q/bIyUeh1Nr7MWlOPDEqzj9xS+QbXYAQHT3bHR/6h50uPqCBjlHy15O4DR7Trz/lbKa6qFE1yTnnc+bUCJOW8VpNGHjpOuQ894XqqEKKC7qu+96Godmvx454epB4ZJVsOTkqhNAPhCCoy9/BBbkHuRwOM0HR0UV1o+7Bqc+/k41VAFl8nfXLY/j8H/fiaB0nLpSvmUXyjfuUI1THxhw/I1PITscTSsYh1NH7KXl+GvsVcid/6NqqAKA6cgJ7Lj+YRx77ZMGOQ83VjmNStHytYGVaCgzikXL1jahRJy2Ss7bn6Ny10HVFV3FZcQdmfsejAeORkCy+lGyciOIJoiTDGMw7jsMR0lZ0wnF4XDqzPHX58F46Jjvs9M1Vh2a8ybMx09FQDJOfSj+cyMgBFe77UWlMO5vec8hTtvk6IsfwHLcz2S5a6za/9hLsOYV1Ps83FjlNCrMKYduE8SY5XAaihPvfQnIgX+PRBRw8pPvm1CihiHc+4fVNNI5HE6zgzFW7ZEUAEIpTn36QxNKxWkImCSBIHSoCR+rOS0BJkk4+dG3wXUQApya/2O9z8WNVU6jkjxuBIgYOFkAEQUkjxvRhBJx2iJMlmHNPRO8jVOC+djJJpKo4UgcNViJCQ8EITBkd4S2XXLTCcXhcOqEbLHCXlQatA1jDOajLW+sauskjR4ScnJRjI1GTK8uTSQRh1N3HOWVcFZUBW1DCGmQsYonWOI0KIwxVO3aD0tuPrQpici67SplljhQe6eE7Luua0IJWx+2wmLkfbsIlpOnoUmMR/ol5yKmZ9dIi9WsIJRCiDJAMlsCtxEFaOLjmlCqhqH9pdOhfei/sJdWALIMIhCAAExmgGuCvvO9s1pk8igOp61BdVoQjahOQBENUe5pBshOGZAAQgnE+NgIS9q2MB09AdPBoxCiDEg4awgEva7WfSSdMwIxvbrCdDjHv9FKKbJuvxqCQd8AEnM8sZzOR/63i2HNL4CuXTLSL78AUZ06RlqsFo0QHQVQGtRjDSDQJNRfr+LGKqfBKP1rC/Y/+hyMB46o2/QLine truncated
-      "text/plain": [
-       "<Figure size 950x350 with 3 Axes>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    }
-   ],
    "source": [
-    "cut = 45 # set to -1 to get all points\n",
+    "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[cut:], c=y_scaled[cut:], cmap='coolwarm')\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",
@@ -95,21 +91,10 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 53,
+   "execution_count": null,
    "id": "a2694f3c",
    "metadata": {},
-   "outputs": [
+   "outputs": [],
-    {
-     "data": {
-      "image/png": "iVBORw0KGgoAAAANSUhEUgAAA6kAAAOrCAYAAAC1F8C4AAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjYsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvq6yFwwAAAAlwSFlzAAAPYQAAD2EBqD+naQABAABJREFUeJzs3Xl8VPW9P/7X55zZJ8lkX8nCvqogqASqEBdAba32akVc4NZq+1XqvQbvvYV6KVCR26r0Vnur7c8FWum9KrebtfUiSrUWqoCgLEoECQlZSMg2WWc55/P7Y8jAkElIYDJnJnk9H495kDmfM2deE5KTec/ncz4fIaWUICIiIiIiIooBitEBiIiIiIiIiLqxSCUiIiIiIqKYwSKViIiIiIiIYgaLVCIiIiIiIooZLFKJiIiIiIgoZrBIJSIiIiIiopjBIpWIiIiIiIhihsnoAPFI13VUV1cjMTERQgij4xBRjJFSorW1Fbm5uVCUwfsskOciIjoXno+IKFYM5HzEIvU8VFdXIz8/3+gYRBTjKisrMWLEiEE7Ps9FRNRfPB8RUazoz/mIRep5SExMBBD4BiclJRmchohijdvtRn5+fvBcMVh4LiKic+H5iIhixUDORyxSz0P3MJakpCSeiImoV4M95I3nIiLqL56PiChW9Od8FDcTJ61duxazZs2Cw+FAcnJy2H0qKirwla98BU6nE+np6XjooYfg9XpD9tm3bx/mzJkDu92OvLw8rFmzBlLKKLwCIiIiIiIiOpe46Un1er247bbbUFxcjBdeeKFHu6ZpuPHGG5GRkYH3338fDQ0NWLx4MaSUeOaZZwAEupivu+46lJSUYOfOnSgrK8OSJUvgdDqxbNmyaL8kIiIiIiIiOkvcFKmrV68GAGzYsCFs+5YtW3Dw4EFUVlYiNzcXAPDUU09hyZIlWLt2LZKSkrBp0yZ0dXVhw4YNsFqtmDJlCsrKyrB+/XqUlpZyNjoiIiIiIiKDxc1w33PZsWMHpkyZEixQAWD+/PnweDzYvXt3cJ85c+bAarWG7FNdXY3y8vJej+3xeOB2u0NuREREREREFHlDpkitra1FVlZWyLaUlBRYLBbU1tb2uk/3/e59wlm3bh1cLlfwxinWiYiIiIiIBoehReqqVasghOjztmvXrn4fL9xwXSllyPaz9+meNKmvob7Lly9HS0tL8FZZWdnvTERERERERNR/hl6TunTpUixcuLDPfYqKivp1rOzsbHzwwQch25qamuDz+YK9pdnZ2T16TOvq6gCgRw/rmaxWa8gQYSIiokjQG2rg/+ITSK8HalYh1JGTIdS4mS5iSJBSQmtrgr+pFtB1qK50mJIzIcSQGWxGRL3QG2vh/3wvZFcHlKwCmMZcDGGyGB2LYHCRmp6ejvT09Igcq7i4GGvXrkVNTQ1ycnIABCZTslqtmD59enCfFStWwOv1wmKxBPfJzc3tdzFMREQUCb5978P30TvAqWJIO7wXYv922BYshrA5DE43PEgp4an4FL76CgCBEVW+k8ehJqTAPm4GhKIaG5CIBo3vwN/h3f7H4DkYZR/B9/F7sH3lPiiORGPDUfxck1pRUYG9e/eioqICmqZh79692Lt3L9ra2gAA8+bNw6RJk3D33Xdjz549ePvtt/HII4/gvvvuCy4qvWjRIlitVixZsgT79+/Hb3/7Wzz++OOc2ZeIiKJKb6oLFKgAIPXADYB0n4RvzzYDkw0vmvvkqQIVAOSpG6C1NcFbW25ULCIaZLq7Ed7tbwTunHkObm2G9+9/NjAZdYubInXlypWYNm0avv/976OtrQ3Tpk3DtGnTgtesqqqKN954AzabDbNnz8bXv/513HzzzXjyySeDx3C5XHjrrbdw/PhxzJgxAw888ABKS0tRWlpq1MsiIqJhyH90/+lP788kJfxffBL9QMOUr7EG3T2oPdoaqqIbhoiixv/FvvC/+lKH9sV+SF2LeiYKFTcXvmzYsKHXNVK7FRQU4I9//GOf+1x00UV47733IpiMiIhogPze3ts0f49J/2iQaBq6e0974JtUilEdB/bCveX38NZWwZJXANf8m2EfP8XoWPHF50WgSg3z+y/1wO8/h/sbKm56UomIiIYKJXtkcHhZCCGgZBWyQI0SNSm1lxYBNSktqlmI+sO97c+ofeJRdOzbDf+JanTs/RA1676L1u28TGAg1NxR4c/BEFDSczl5UgxgkUpERBRl6oixUDJGAGcWo0IAQsA8rcS4YMOMOS0PwupA6Lg/ASgKrDmjjYpFFJbe2YGG/37+1B095N+Gl5+D9PkMShZ/lNxRUPLGIOR3XwhAAJbL5xuWi05jkUpERBRlQlFgve4umCbPAmxOQDVDyRkF64IlUDPzjY43bAjVBMeEmTBn5gOqGVBUmFKy4JhYDMXmNDoeUYiusgOQXk/YNr2jHV1fHIpyovglhIBt3p0wXzoXwpEIqCYoOSNhu/FeqHn8gCoWxM01qUREREOJMFtgmX4NLNOvMTrKsKaYLbAVTIKtYJLRUYj6pvTdt8QlkwZGmMw8B8cw9qQSEREREcU42/gpEPZwaygLKEkuWEeNi3omosHCIpWIiIiIKMYpFisyliwNXDvZ3WuqqIAikPGPD0Go7EmloYPDfYmIiIiI4kDCFVfBnJWLlrffgP9ENcy5+XBd82VY8ouMjkYUUSxSiYiIiIjOU8cXX+Dkm3+G1tqKxEsuQerV10CxDN4SJtaiMci8958G7fhEsYBFKhERERHReah95X9w7D9/DJwaalv3u9+i+pcbMfHZn8Pschmcjih+8ZpUIiIiIqIB6qyowLGf/GfgjqYFbgA6jx1D5bP/ZVwwoiGARSoRERER0QA1vLUlMInR2XQdJ998E1LK6IeiuCalhL/yU3i2/xZd234Nz+4t0BtrjY5lCA73JSIiIiIaIK29HUIokNB7tEmPB9D14DBgov7wfbodetXnwfuysQbexmqYL7kaamaBgcmijz2pREREREQDlHTpdEjN37NBUZBw0UVcEoYGRG9rDilQAwK98f6yncOuZ55FKhERERHRACUXFyPh4osB5Yy306e+zv/2AwalonilN1b32iY7W4Gu9iimMR6LVCIiIiKiARKqigk//glyFt0JkysZwmRC4rRpmPhfP0PSpZcaHY/ijXKOqzCHWc88r0klIiIiIjoPqsOBggeXouDBpUZHoTinZubD/5kAegzrFRApmRAWuyG5jMKeVCIiIiIiIgMJix2micWn7ggAp2aONltgnlBsWC6jsCeViIiIiIjIYKa8cVCSMqBVl0F6OqEkpkHNGwthsRkdLepYpBIREREREcUAJTEFyvgrjI5hOA73JSIiIiIiopjBIpWIiIYkKfVht65cJEgpY/r7Fuv5iIjownG4LxERDSna8TL4D2yHdDcAZivUURfDNHkWhMo/eX3RfF601x6D190ASAmzMwmO7EKY7QlGRwMA+NwN6Kwth9bZBigqrKnZsGcX8f+ViGgIipue1LVr12LWrFlwOBxITk7u0f7xxx/jjjvuQH5+Pux2OyZOnIif/OQnIfuUl5dDCNHj9uabb0bpVRAR0WDSjn0K347XAwUqAPg80A7tgm/7H9j71gdd86Pli/3wtpwMLn/ga3ej5Yv98MfAAvI+dwPaju4PFKgAoGvwnKxC6xf7+P9KRDQExc3Hj16vF7fddhuKi4vxwgsv9GjfvXs3MjIy8PLLLyM/Px/bt2/H/fffD1VVsXRp6NpVW7duxeTJk4P3U1NTBz0/ERENLiklfPv/Gq4Feu1RyMZaiLScqOeKB56meug+T88GKdFRV4WkgnHRD3WGjpqjYbdrHW74WxthTkqLciIiIhpMcVOkrl69GgCwYcOGsO3f+MY3Qu6PGjUKO3bswG9+85seRWpaWhqys7MHJScRERmksw3oaO2lUUA/eRwKi9SwfB3u3tvaW6KYpCepa9B77c0V8Le3sEglIhpi4ma47/loaWkJ20t60003ITMzE7Nnz8bmzZvPeRyPxwO32x1yIyKiGGO29NEoAbM1alHiTV/XdRp+zadQTi1sH440Ph8REUXckC1Sd+zYgVdffRXf+ta3gtsSEhKwfv16bN68GX/6059wzTXX4Pbbb8fLL7/c57HWrVsHl8sVvOXn5w92fCIiGiBhtkLJHR2+oFFUqCOMHbIay2zJGb23pWRGMUlPQghYkrN6bbckG5uPiIgiz9AiddWqVWEnMjrztmvXrgEf98CBA/jqV7+KlStX4rrrrgtuT09Px8MPP4zLL78cM2bMwJo1a/DAAw/gRz/6UZ/HW758OVpaWoK3ysrKAWciIqLBZ770WgiHK3BHKAAEIBSYr7gRwmIzNFssMzuTYM/IO2NLoNA3JyTDHgNDpO25o6DYnKfunf4QwpE/Hgr/X4mIhhxDx8gsXboUCxcu7HOfoqKiAR3z4MGDuPrqq3Hffffh0UcfPef+M2fOxPPPP9/nPlarFVYrh4kREcU6YU+AZf5i6Mc/h95UC2FzQi2cCGFPNDpazHNmFcDqSoOnpQHQdZgTk2F2uiB6HWobPYrJjKRxl8LX0gB/ewuEyQxLShZUFqhEREOSoUVqeno60tPTI3a8AwcO4Oqrr8bixYuxdu3afj1mz549yMkx/lNiIiKKDKGaoBZOhFo40egoccdkc8IU7LGMLUIosCRnwNLH0GQiIhoa4ma2gYqKCjQ2NqKiogKapmHv3r0AgDFjxiAhIQEHDhxASUkJ5s2bh9LSUtTW1gIAVFVFRkbgD9rGjRthNpsxbdo0KIqC119/HU8//TR++MMfGvWyiIiIiIiI6AxxU6SuXLkSGzduDN6fNm0aAGDbtm2YO3cuXnvtNdTX12PTpk3YtGlTcL/CwkKUl5cH7z/22GM4duwYVFXFuHHj8OKLL+Kuu+6K2usgIiIiIiKi3gkppTQ6RLxxu91wuVxoaWlBUlKS0XGIKMZE6xzBcxERnQvPR0QUKwZLine truncated
-      "text/plain": [
-       "<Figure size 950x950 with 9 Axes>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    }
-   ],
    "source": [
     "\n",
     "fig, axs = plt.subplots(len(x), len(x), figsize=(3*len(x)+0.5, 3*len(x)+0.5))\n",
@@ -117,7 +102,7 @@
     "for i, axi in enumerate(axs):\n",
     "    for j, axij in enumerate(axi):\n",
     "        axij.scatter(\n",
-    "            x[i], x[j],\n",
+    "            x[j], x[i],\n",
     "            c = y_scaled, # color\n",
     "            cmap = 'coolwarm',\n",
     "            s = 20, # size\n",

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,190 @@
+import os
+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 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(
+            "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,
+            )

pipelines/sobo.py

@@ -1,47 +1,49 @@
-import os
+# depreacated, use pipelines.batched_sobo.sobo_pipeline()
-import bo
-import ptycho
 
-def sobo_pipeline(config):
+# import os
+# import bo
+# import ptycho
 
-    result_dir = config['io']['result_dir']
+# def sobo_pipeline(config):
-    os.makedirs(result_dir, exist_ok=True)
 
-    RANDOM_ITERS = config['job'].get('random_iters', 0)
+#     result_dir = config['io']['result_dir']
-    SOBO_ITERS   = config['job'].get('sobo_iters')
+#     os.makedirs(result_dir, exist_ok=True)
-    METRIC = config['bo']['metric']
+
-    PTYCHOENGINE = ptycho.engines[config['ptycho']['engine']]
+#     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)
+#     randombo = bo.RandomBOEngine(config)
 
-    bo_txt = os.path.join(result_dir, "bo.txt")
+#     bo_txt = os.path.join(result_dir, "bo.txt")
-    with open(bo_txt, "w") as f:
+#     with open(bo_txt, "w") as f:
-        f.write(f"    iter\tmetric\t{"\t".join([p[:7] for p in randombo.params])}\n")
+#         f.write(f"    iter\tmetric\t{"\t".join([p[:7] for p in randombo.params])}\n")
 
-    for j in range(RANDOM_ITERS):
+#     for j in range(RANDOM_ITERS):
-        print(f"RANDOM sampling; iteration {j}")
+#         print(f"RANDOM sampling; iteration {j}")
-        job_config = randombo.ask()
+#         job_config = randombo.ask()
-        ptycho_engine = PTYCHOENGINE(job_config)
+#         ptycho_engine = PTYCHOENGINE(job_config)
-        ptycho_engine.run(run_id=f"bo-{j:03d}")
+#         ptycho_engine.run(run_id=f"bo-{j:03d}")
-        y_value = ptycho_engine.metric(METRIC)
+#         y_value = ptycho_engine.metric(METRIC)
-        randombo.tell(job_config, y_value)
+#         randombo.tell(job_config, y_value)
-        with open(bo_txt, "a") as f:
+#         with open(bo_txt, "a") as f:
-            p = [f'{job_config['ptycho']['params'][key]:.2f}' for key in randombo.params]
+#             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")
+#             f.write(f"{j: 8d}\t{y_value:.4f}\t{"\t".join(p)}\n")
 
 
-    sobo = bo.SingleObjectiveBOEngine(config)
+#     sobo = bo.SingleObjectiveBOEngine(config)
-    sobo.train_x = randombo.train_x
+#     sobo.train_x = randombo.train_x
-    sobo.train_y = randombo.train_y
+#     sobo.train_y = randombo.train_y
 
-    for j in range(SOBO_ITERS):
+#     for j in range(SOBO_ITERS):
-        print(f"SOBO sampling; iteration {RANDOM_ITERS + j}")
+#         print(f"SOBO sampling; iteration {RANDOM_ITERS + j}")
-        job_config = sobo.ask()
+#         job_config = sobo.ask()
-        ptycho_engine = PTYCHOENGINE(job_config)
+#         ptycho_engine = PTYCHOENGINE(job_config)
-        ptycho_engine.run(run_id=f"bo-{RANDOM_ITERS + j:03d}")
+#         ptycho_engine.run(run_id=f"bo-{RANDOM_ITERS + j:03d}")
-        y_value = ptycho_engine.metric(METRIC)
+#         y_value = ptycho_engine.metric(METRIC)
-        sobo.tell(job_config, y_value)
+#         sobo.tell(job_config, y_value)
-        with open(bo_txt, "a") as f:
+#         with open(bo_txt, "a") as f:
-            p = [f'{job_config['ptycho']['params'][key]:.2f}' for key in sobo.params]
+#             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")
+#             f.write(f"{RANDOM_ITERS + j: 8d}\t{y_value:.4f}\t{"\t".join(p)}\n")

pipelines/test.py

@@ -0,0 +1,189 @@
+import os
+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):
+
+    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,
+            )

ptycho/__init__.py

@@ -1,5 +1,4 @@
-from .base import PtychoEngine
+from .example import ExamplePtychoEngine
-from .ptycho_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):