changed ptycho/ to class/methods.

TODO: also change bo/, fix main.py

ptycho/fold_slice.py

@@ -2,39 +2,50 @@ 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')
+
+class FoldSlicePtychoEngine(PtychoEngine):
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.fold_slice_output_dir = os.path.join(config['io']['result_dir'], 'fold_slice')
+
+
+    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_result_dir, '')
+        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_result_dir)):
-        shutil.rmtree(os.path.join(fold_slice_result_dir))
-    os.makedirs(os.path.join(fold_slice_result_dir))
+        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))
 
-    setup_txt = os.path.join(fold_slice_result_dir, 'setup.txt')
-    with open(setup_txt, 'w') as f:
+        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()
 
-def run(config):
-    setup_txt = fold_slice_translator(config)
+        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"prepare_data('{setup_txt}');",
             f"run_multislice_new('{setup_txt}');"
         ]
         
@@ -51,23 +62,50 @@ def run(config):
                 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.")
+            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')
+    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_result_dir,
+            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])
-    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.")
+        output_mat_path = os.path.join(output_dir, f"Niter{config['ptycho']['params']['Niter']}.mat")
+        return output_mat_path
     
-    # return loadmat(result_mat)
-    return 1
+
+    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
+    
+
+    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

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

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