diff --git a/ptycho/__init__.py b/ptycho/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/ptycho/fold_slice.py b/ptycho/fold_slice.py index fb1e671..d79a51f 100644 --- a/ptycho/fold_slice.py +++ b/ptycho/fold_slice.py @@ -2,72 +2,110 @@ import os import sys import shutil import subprocess -import numpy as np +from typing import Any from scipy.io import loadmat +import numpy as np +from ptycho.ptycho_base import PtychoEngine -def fold_slice_translator(config): - fold_slice_result_dir = os.path.join(config['io']['result_dir'], 'fold_slice') - fold_slice_dict = {} - fold_slice_dict['raw_data'] = config['io']['input_data_path'] - fold_slice_dict['result_dir'] = os.path.join(fold_slice_result_dir, '') - fold_slice_dict.update(config['ptycho']['params']) +class FoldSlicePtychoEngine(PtychoEngine): - if os.path.exists(os.path.join(fold_slice_result_dir)): - shutil.rmtree(os.path.join(fold_slice_result_dir)) - os.makedirs(os.path.join(fold_slice_result_dir)) + def __init__(self, config): + super().__init__(config) + self.fold_slice_output_dir = os.path.join(config['io']['result_dir'], 'fold_slice') - setup_txt = os.path.join(fold_slice_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") + + def fold_slice_prepare(self) -> None: + config = self.config + fold_slice_output_dir = self.fold_slice_output_dir + + fold_slice_dict = {} + fold_slice_dict['raw_data'] = config['io']['input_data_path'] + fold_slice_dict['result_dir'] = os.path.join(fold_slice_output_dir, '') + fold_slice_dict.update(config['ptycho']['params']) + + if os.path.exists(os.path.join(fold_slice_output_dir)): + shutil.rmtree(os.path.join(fold_slice_output_dir)) + os.makedirs(os.path.join(fold_slice_output_dir)) + + with open(os.path.join(fold_slice_output_dir, '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(self) -> None: + self.fold_slice_prepare() + + config = self.config + fold_slice_path = config['ptycho']['path'] + setup_txt = os.path.join(self.fold_slice_output_dir, 'setup.txt') + 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 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}');" - ] + def get_output_mat_path(self) -> str: + config = self.config + fold_slice_output_dir = self.fold_slice_output_dir + roi_dir = os.path.join( + fold_slice_output_dir, + f"{config['ptycho']['params']['scan_number']}", + f"roi{config['ptycho']['params']['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]) + output_mat_path = os.path.join(output_dir, f"Niter{config['ptycho']['params']['Niter']}.mat") + return output_mat_path - 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): - fold_slice_result_dir = os.path.join(config['io']['result_dir'], 'fold_slice') - roi_dir = os.path.join( - fold_slice_result_dir, - f"{config['ptycho']['params']['scan_number']}", - f"roi{config['ptycho']['params']['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']['params']['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.") + def output(self) -> Any: + if self._output is None: + output_mat_path = self.get_output_mat_path() + if not os.path.exists(output_mat_path): + raise FileNotFoundError(f"Result directory {output_mat_path} does not exist. Please check the fold_slice output.") + self._output = loadmat(output_mat_path) + return self._output - # return loadmat(result_mat) - return 1 \ No newline at end of file + + def recon_object(self) -> np.ndarray: + if self._recon_object is None: + output = self.output() + self._recon_object = output['object'] # type: ignore + return self._recon_object + + + def recon_probe(self) -> np.ndarray: + if self._recon_probe is None: + output = self.output() + self._recon_probe = output['probe'] # type: ignore + return self._recon_probe + + + def metric(self) -> float: + if self._metric is None: + output = self.output() + self._metric = float(np.asarray(output['outputs']['fourier_error_out'][0]).squeeze()) # type: ignore + return self._metric diff --git a/ptycho/ptycho_base.py b/ptycho/ptycho_base.py new file mode 100644 index 0000000..5da347a --- /dev/null +++ b/ptycho/ptycho_base.py @@ -0,0 +1,34 @@ +from abc import ABC, abstractmethod +from typing import Any + + +class PtychoEngine(ABC): + name = None + + def __init__(self, config): + self.config = config + self._output = None + self._recon_object = None + self._recon_probe = None + self._metric = None + + + @abstractmethod + def run(self) -> None: + pass + + @abstractmethod + def output(self) -> Any: + pass + + @abstractmethod + def recon_object(self) -> Any: + pass + + @abstractmethod + def recon_probe(self) -> Any: + pass + + @abstractmethod + def metric(self) -> float: + pass diff --git a/ptycho/ptycho_example.py b/ptycho/ptycho_example.py new file mode 100644 index 0000000..8f906cd --- /dev/null +++ b/ptycho/ptycho_example.py @@ -0,0 +1,41 @@ +from typing import Any +from ptycho.ptycho_base import PtychoEngine + +class ExamplePtychoEngine(PtychoEngine): + + # initialize job/engine-specific variables here + # e.g. self.output_dir + + def __init__(self, config): + super().__init__(config) + + # run single ptychography job based on `config` + def run(self) -> None: + pass + + # save outputs to self._output + def output(self) -> Any: + if self._output is None: + pass + return self._output + + # return reconstructed object (numpy array) + def recon_object(self) -> Any: + if self._recon_object is None: + output = self.output() + pass + return self._recon_object + + # return reconstructed probe (numpy array) + def recon_probe(self) -> Any: + if self._recon_probe is None: + output = self.output() + pass + return self._recon_probe + + # return metric value for Bayesian optimization + def metric(self) -> float: + if self._metric is None: + output = self.output() + self._metric = 0.0 + return self._metric