mirror of
https://github.com/c-sooyoung/bo-ptycho.git
synced 2026-09-17 23:29:07 +09:00
114 lines
3.9 KiB
Python
114 lines
3.9 KiB
Python
import os
|
|
import sys
|
|
import shutil
|
|
import subprocess
|
|
from typing import Any
|
|
from scipy.io import loadmat
|
|
import numpy as np
|
|
from ptycho.ptycho_base import PtychoEngine
|
|
|
|
|
|
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_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, header="") -> 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 = 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()
|
|
fourier_error = output['outputs']['fourier_error_out'][0][0].squeeze()
|
|
metric = float(fourier_error[-1])
|
|
self._metric = metric
|
|
return self._metric
|