new TODO

.gitignore

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

TODO.md

@@ -1,21 +1,28 @@
 # TODO:
 
-- unify `BOEngine.__init__()`
+- fold_slice:
+    - load diffractions / hdf5 files; change only param per job
+    - restructure `FoldSlicePtychoEngine.__init__()` to load data but not params
+        - write new `prepare_data.m`
+    - start fold_slice from config.yaml instead of setup.txt
 - BO train_x/y transfer between engines for single job
 - 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
 
+- rename `BOEngine` to `Sampler`
+- organize results/
+    - possibly: `[dataset name]/[salient config]-[date]/`
+
 
 # TODAY:
-
+- separation of available GPUs and parallel BO batches
-- fold_slice: load diffractions / hdf5 files; change only param per job
+- prepare next batch for efficient GPU use
-    - restructure `FoldSlicePtychoEngine.__init__()` to load data but not params
+- `GridSampler`
-        - write new `prepare_data.m`
+    - `.__init__()` should create a grid
+    - `.ask()` should remove those items from the grid

bo/base.py

@@ -1,15 +1,68 @@
 from abc import ABC, abstractmethod
+import os
+import numpy as np
 
 
 class BOEngine(ABC):
     def __init__(self, config):
         self.config = config
-        self.state = None
+        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
+        
 
     @abstractmethod
     def ask(self, n: int = 1) -> list[dict]:
         pass
 
-    @abstractmethod
+
-    def tell(self, job_config, y_value):
+    def tell(self, job_config, y_value) -> None:
-        pass
+        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('train_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)

bo/random.py

@@ -1,4 +1,3 @@
-import os
 import copy
 import numpy as np
 from bo.base import BOEngine
@@ -9,43 +8,15 @@ class RandomBOEngine(BOEngine):
     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
-        
 
     def ask(self, n = 1):
-        config = self.config
         next_configs = []
 
         for _ in range(n):
-            next_config = copy.deepcopy(config)
+            next_config = copy.deepcopy(self.config)
             for param in self.params:
-                radius = config['bo']['params'][param]['radius']
+                radius = self.config['bo']['params'][param]['radius']
-                center = config['ptycho']['params'][param]
+                center = self.config['ptycho']['params'][param]
                 modulation = radius * (np.random.rand() - 0.5) * 2
                 next_value = center + modulation
                 if self.param_types[param] == 'int':
@@ -54,32 +25,3 @@ class RandomBOEngine(BOEngine):
             next_configs.append(next_config)
 
         return next_configs
-
-
-    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('train_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)

bo/sobo.py

@@ -21,34 +21,6 @@ from bo.base import BOEngine
 class SingleObjectiveBOEngine(BOEngine):
     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)
-
     
 
 

config.yaml

@@ -1,52 +1,52 @@
 job:
   type: "random+sobo"
-  random_iters: 8
+  random_iters: 16
-  sobo_iters: 128
+  sobo_iters: 512
 
 io:
-  input_data_path: '/home/swim/shared/Si_project/data/Si2V1_2.mat'
+  input_data_path: '/home/swim/shared/Si_project/data/Si30V3_2.mat'
-  result_dir: '/home/swim/bo-ptycho/results/260812-Si/Si2V1_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: -200
+    defocus: 200
-    rot_ang: 0.3
+    rot_ang: 0.1
-    Nlayers: 20
+    Nlayers: 30
-    thickness: 250
+    thickness: 650
-    rbf: 37
+    rbf: 38
-    Nprobe: 1
+    tilt_x: 5
-    N_scan_x: 64
+    tilt_y: 3 
-    N_scan_y: 64
+    scan_step_size: 0.35
-    tilt_x: 2
-    tilt_y: 0 
-    scan_step_size: 0.36
 
     Niter: 100
     Niter_save_results: 100
 
     CBED_size: 192
     ADU: 1
-    extra_print_info: ''
+    Nprobe: 1
+    N_scan_x: 64
+    N_scan_y: 64 
+    extra_print_info: 'FIB'
     scan_number: 1
     gpu_id: 1
     roi_label: '0_Ndp64'
     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:
@@ -57,6 +57,6 @@ bo:
       radius: 5
       type: int
     thickness:
-      radius: 100
+      radius: 150
   train_x:
   train_y:

job.sub

@@ -1,13 +1,12 @@
 #!/bin/bash
-#SBATCH --job-name=Si2V
+#SBATCH --job-name=Si30V
 #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,12 @@ 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):
+        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 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
-    # ---------------------------------------------------------
 
+    ############################ RANDOM SAMPLING ###############################
     randombo = bo.RandomBOEngine(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,
-            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,9 @@
-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.
+    pass
-    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,
-            )