changed ptycho/ to class/methods.

TODO: also change bo/, fix main.py

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
-def fold_slice_translator(config):
+from ptycho.ptycho_base import PtychoEngine
-    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'])
-
-    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))
-
-    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")
-    
-    return setup_txt
 
 
-def run(config):
+class FoldSlicePtychoEngine(PtychoEngine):
-    setup_txt = fold_slice_translator(config)
-    fold_slice_path = config['ptycho']['path']
-    verbosity = config['io'].get('verbosity', 0)
 
-    matlab_commands = [
+    def __init__(self, config):
-        f"cd('{fold_slice_path}');",
+        super().__init__(config)
-        "cd('ptycho');",
+        self.fold_slice_output_dir = os.path.join(config['io']['result_dir'], 'fold_slice')
-        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):
+    def fold_slice_prepare(self) -> None:
-    fold_slice_result_dir = os.path.join(config['io']['result_dir'], 'fold_slice')
+        config = self.config
-    roi_dir = os.path.join(
+        fold_slice_output_dir = self.fold_slice_output_dir
-        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.")
 
-    # return loadmat(result_mat)
+        fold_slice_dict = {}
-    return 1
+        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")
+    
+
+    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 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
+    
+
+    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