initial commit

.gitignore

@@ -0,0 +1,358 @@
+*.mat
+*.tiff
+*.hdf5
+*.jpg
+*.dat
+*.asv
+
+
+# LINUX   #################################################################################
+
+*~
+
+# temporary files which can be created if a process still has a handle open of a deleted file
+.fuse_hidden*
+
+# Metadata left by Dolphin file manager, which comes with KDE Plasma
+.directory
+
+# Linux trash folder which might appear on any partition or disk
+.Trash-*
+
+# .nfs files are created when an open file is removed but is still being accessed
+.nfs*
+
+# Log files created by default by the nohup command
+nohup.out
+
+
+# WINDOWS #################################################################################
+
+# Windows thumbnail cache files
+Thumbs.db
+Thumbs.db:encryptable
+ehthumbs.db
+ehthumbs_vista.db
+
+# Dump file
+*.stackdump
+
+# Folder config file
+[Dd]esktop.ini
+
+# Recycle Bin used on file shares
+$RECYCLE.BIN/
+
+# Windows Installer files
+*.cab
+*.msi
+*.msix
+*.msm
+*.msp
+
+# Windows shortcuts
+*.lnk
+
+
+# MACOS   #################################################################################
+
+# General
+.DS_Store
+.localized
+__MACOSX/
+.AppleDouble
+.LSOverride
+Icon[
+]
+
+# Resource forks
+._*
+
+# Files and directories that might appear in the root of a volume
+.DocumentRevisions-V100
+.fseventsd
+.Spotlight-V100
+.TemporaryItems
+.Trashes
+.VolumeIcon.icns
+.com.apple.timemachine.donotpresent
+.com.apple.timemachine.supported
+.PKInstallSandboxManager
+.PKInstallSandboxManager-SystemSoftware
+.hotfiles.btree
+.vol
+.file
+.disk_label*
+lost+found
+.HFS+ Private Directory Data[
+]
+
+# Directories potentially created on remote AFP share
+.AppleDB
+.AppleDesktop
+Network Trash Folder
+Temporary Items
+.apdisk
+
+# Mac OS 6 to 9
+Desktop DB
+Desktop DF
+TheFindByContentFolder
+TheVolumeSettingsFolder
+.FBCIndex
+.FBCSemaphoreFile
+.FBCLockFolder
+
+# Quota system
+.quota.group
+.quota.user
+.quota.ops.group
+.quota.ops.user
+
+# TimeMachine
+Backups.backupdb
+.MobileBackups
+.MobileBackups.trash
+MobileBackups.trash
+tmbootpicker.efi
+
+
+# PYTHON  #################################################################################
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[codz]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#   Usually these files are written by a python script from a template
+#   before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py.cover
+*.lcov
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+# Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+# uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+# poetry.lock
+# poetry.toml
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#   pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
+#   https://pdm-project.org/en/latest/usage/project/#working-with-version-control
+# pdm.lock
+# pdm.toml
+.pdm-python
+.pdm-build/
+
+# pixi
+#   Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
+# pixi.lock
+#   Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
+#   in the .venv directory. It is recommended not to include this directory in version control.
+.pixi/*
+!.pixi/config.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule*
+celerybeat.pid
+
+# Redis
+*.rdb
+*.aof
+*.pid
+
+# RabbitMQ
+mnesia/
+rabbitmq/
+rabbitmq-data/
+
+# ActiveMQ
+activemq-data/
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.envrc
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# Ruff stuff:
+.ruff_cache/
+
+# PyPI configuration file
+.pypirc
+
+# Marimo
+marimo/_static/
+marimo/_lsp/
+__marimo__/
+
+# Streamlit
+.streamlit/secrets.toml
+
+
+# MATLAB  #################################################################################
+
+# Autosave files
+*.asv
+*.m~
+*.autosave
+*.slx.r*
+*.mdl.r*
+*.bak
+
+# Derived content-obscured files
+*.p
+
+# Compiled MEX files
+*.mex*
+
+# Packaged app and toolbox files
+*.mlappinstall
+*.mltbx
+
+# Deployable archives
+*.ctf
+
+# Generated helpsearch folders
+helpsearch*/
+
+# Code generation folders
+slprj/
+sccprj/
+codegen/
+
+# Cache files
+*.slxc
+
+# Cloud based storage dotfile
+.MATLABDriveTag
+

bo-ptycho.py

@@ -0,0 +1,90 @@
+import os
+import sys
+import yaml
+
+def ptycho_run(config):
+    ptycho_engine = config['ptycho']['engine']
+
+    if ptycho_engine == 'fold_slice':
+        from ptycho.fold_slice import run
+        run(config)
+
+    # Add more ptycho engines here as needed
+    # write their respective run functions and import them above
+
+
+def ptycho_error(config):
+    ptycho_engine = config['ptycho']['engine']
+
+    if ptycho_engine == 'fold_slice':
+        from ptycho.fold_slice import error
+        return error(config)
+
+
+
+
+
+
+def bo_initialize(config):
+    bo_algorithm = config['bo']['algorithm']
+
+    if bo_algorithm == 'ucb':
+        from bo.ucb import initialize
+        bo_state = initialize(config)
+    
+    else:
+        from bo.random import initialize
+        bo_state = initialize(config)
+    
+    return bo_state
+
+
+def bo_ask(config, bo_state):
+    bo_engine = config['bo']['algorithm']
+
+    if bo_engine == 'ucb':
+        from bo.ucb import ask
+        next_config = ask(config, bo_state)
+    
+    else:
+        from bo.random import ask
+        next_config = ask(config, bo_state)
+
+    # Add more BO engines here as needed
+    # write their respective ask functions and import them above
+
+    return next_config
+
+
+def bo_tell(config, job_config, bo_state, y_value):
+    bo_engine = config['bo']['algorithm']
+
+    if bo_engine == 'ucb':
+        from bo.ucb import tell
+        bo_config = tell(job_config, bo_state, y_value)
+    
+    else:
+        from bo.random import tell
+        bo_config = tell(job_config, bo_state, y_value)
+    
+    return bo_config
+
+
+
+
+
+
+if __name__ == "__main__":
+    if len(sys.argv) != 2:
+        print("Usage: python bo-ptycho.py <config_yaml>")
+        sys.exit(1)
+
+    config_yaml = sys.argv[1]
+    with open(config_yaml, 'r') as f:
+        config = yaml.safe_load(f)
+
+    # ptycho_run(config_yaml)
+    # results = ptycho_results(config_yaml)
+    # new_config = bo_loop(config)
+    # print(new_config) # type: ignore
+    

bo/random.py

@@ -0,0 +1,56 @@
+import numpy as np
+import copy
+
+
+def initialize(config):
+    bo_params = []
+
+    for key, value in config['bo']['params'].items():
+        if value is not None:
+            bo_params.append(key)
+
+    train_x = np.empty((0, len(bo_params)))
+    train_y = np.empty((0,))
+
+    bo_state = {
+        'params': bo_params,
+        'train_x': train_x,
+        'train_y': train_y,
+    }
+
+    return bo_state
+
+
+def ask(config, bo_state):
+    next_config = copy.deepcopy(config)
+
+    for param in bo_state['params']:
+        max_modulation = config['bo']['params'][param]
+        center_value = config['ptycho']['params'][param]
+
+        modulation = max_modulation * (np.random.rand() - 0.5) * 2
+        next_config['ptycho']['params'][param] = center_value + modulation
+
+    return next_config
+
+
+def tell(job_config, bo_state, y_value):
+    x_value = []
+
+    for param in bo_state['params']:
+        x_value.append(job_config['ptycho']['params'][param])
+
+    x_value = np.array(x_value).reshape(1, -1)
+    y_value = np.array([y_value])
+
+    bo_state['train_x'] = np.vstack([
+        bo_state['train_x'],
+        x_value,
+    ])
+
+    bo_state['train_y'] = np.concatenate([
+        bo_state['train_y'],
+        y_value,
+    ])
+
+    return bo_state

bo/ucb.py

@@ -0,0 +1,48 @@
+# import os
+# import sys
+# import shutil
+# import numpy as np
+# import subprocess
+# import h5py
+# from PIL import Image
+# import time
+# import random
+
+# import torch
+# from botorch.sampling.samplers import SobolQMCNormalSampler
+# from botorch.models import SingleTaskGP
+# from botorch.fit import fit_gpytorch_model
+# from gpytorch.mlls import ExactMarginalLogLikelihood
+# from botorch.optim import optimize_acqf
+# from botorch.acquisition import UpperConfidenceBound
+# from botorch.models.transforms.outcome import Standardize
+# from botorch.acquisition.monte_carlo import qUpperConfidenceBound
+# from botorch.utils.multi_objective.box_decompositions.non_dominated import NondominatedPartitioning
+# from botorch.acquisition.multi_objective.monte_carlo import qExpectedHypervolumeImprovement
+# from botorch.utils.transforms import unnormalize
+
+# from bo import bo_random_config
+
+
+def initialize(config):
+    pass
+    
+#     n_initial_jobs = config['bo']['parallel_jobs']
+#     n_var_params = len([p for p in config['bo']['params'].values() if p is not None])
+    
+#     train_x = np.zeros([n_initial_jobs, n_var_params])
+
+#     for i in range(n_initial_jobs):
+#         config_i = bo_random_config(config)
+#         train_x[i] = [p for p in config_i['bo']['params'].values() if p is not None]
+
+
+def ask(config, bo_state):
+
+    next_config = config.copy()
+
+    return next_config
+
+
+def tell(job_config, bo_state, y_value):
+    pass

config.yaml

@@ -0,0 +1,47 @@
+io:
+  input_data_path: '/home/swim/data/BTO_01_crop.mat'
+  result_dir: '/home/swim/Si_project/wrapper/results'
+  verbosity: 1
+
+ptycho:
+  engine: 'fold_slice'
+  path: '/home/swim/Si_project/lemon-ptychography/fold_slice'
+  params:
+    alpha_max: 30
+    defocus: -200
+    rot_ang: 1
+    Nlayers: 20
+    thickness: 250
+    voltage: 200
+    rbf: 37
+    Nprobe: 1
+    N_scan_x: 64
+    N_scan_y: 64
+    tilt_x: 2
+    tilt_y: 0 
+    scan_step_size: 0.36
+    Niter: 20
+    Niter_save_results: 10
+    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: 128
+    probe_posiiton_search: 1
+    regularize_layers: 0.1
+    variable_probe: 'false'
+
+bo:
+  parallel_jobs: 4
+  algorithm: ucb
+  params:
+    alpha_max:
+    defocus: 150
+    rot_ang: 1
+    Nlayers:
+    thickness: 150

job.sub

@@ -0,0 +1,23 @@
+#!/bin/bash
+#SBATCH --job-name=swim_test
+#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 "TEST JOB"
+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/Si_project/wrapper/venv/bin/activate
+
+python -u bo-ptycho.py config.yaml
+
+echo "SLURM JOB FINISHED"

ptycho/fold_slice.py

@@ -0,0 +1,73 @@
+import os
+import sys
+import shutil
+import subprocess
+import numpy as np
+from scipy.io import loadmat
+
+def fold_slice_translator(config):
+
+    result_dir = config['io']['result_dir']
+
+    fold_slice_dict = {}
+    fold_slice_dict['raw_data'] = config['io']['input_data_path']
+    fold_slice_dict['result_dir'] = os.path.join(result_dir, '')
+    fold_slice_dict.update(config['ptycho']['fold_slice'])
+
+    if os.path.exists(os.path.join(result_dir)):
+        shutil.rmtree(os.path.join(result_dir))
+    os.makedirs(os.path.join(result_dir))
+
+    setup_txt = os.path.join(result_dir, 'setup.txt')
+    with open(setup_txt, 'w') as f:
+        f.write('\n\n')
+        for key, value in fold_slice_dict.items():
+            f.write(f"{key} {value}\n")
+    
+    return setup_txt
+
+
+def run(config):
+    setup_txt = fold_slice_translator(config)
+    fold_slice_path = config['ptycho']['path']
+    verbosity = config['io'].get('verbosity', 0)
+
+    matlab_commands = [
+        f"cd('{fold_slice_path}');",
+        "cd('ptycho');",
+        f"prepare_data('{setup_txt}');"
+        f"run_multislice_new('{setup_txt}');"
+    ]
+    
+    p = subprocess.Popen(
+        ['matlab', '-batch', ' '.join(matlab_commands)],
+        stdout=subprocess.PIPE if verbosity > 0 else subprocess.DEVNULL,
+        stderr=subprocess.STDOUT if verbosity > 0 else subprocess.DEVNULL,
+        text=True
+    )
+
+    if not verbosity == 0:
+        header = '[fold slice]'
+        for line in p.stdout: # type: ignore
+            sys.stdout.write(f'{header} {line}')
+        p.stdout.close() # type: ignore
+    else:
+        print(f"fold_slice running. Set verbosity>0 for full fold_slice output.")
+
+    p.wait()
+
+
+def error(config):
+    result_dir = config['io']['result_dir']
+    roi_dir = os.path.join(
+        result_dir,
+        f"{config['ptycho']['fold_slice']['scan_number']}",
+        f"roi{config['ptycho']['fold_slice']['roi_label']}"
+    )
+    output_dir = os.path.join(roi_dir, next(os.walk(roi_dir))[1][0])
+    # image_path = os.path.join(output_dir, 'obj_phase_roi_sum', next(os.walk(os.path.join(output_dir, 'obj_phase_roi_sum')))[2][0])
+    result_mat = os.path.join(output_dir, f"Niter{config['ptycho']['fold_slice']['Niter']}.mat")
+    if not os.path.exists(result_mat):
+        raise FileNotFoundError(f"Result directory {result_mat} does not exist. Please check the fold_slice output.")
+    
+    return loadmat(result_mat)

ptycho/ptyrad.py

@@ -0,0 +1,2 @@
+def run(config, verbosity=0):
+    pass

requirements.txt

@@ -0,0 +1,10 @@
+botorch==0.6.1
+gpytorch==1.6.0
+multipledispatch>=0.6.0
+numpy>=1.22.2
+Pillow>=9.0.1
+scipy>=1.8.0
+six>=1.16.0
+typing_extensions>=4.1.1
+PyYAML
+h5py