tomwer 1.2.1__py3-none-any.whl → 1.3.12__py3-none-any.whl

tomwer/core/process/reconstruction/sadeltabeta/sadeltabeta.py

@@ -39,6 +39,7 @@ from typing import Optional, Union
 
 import h5py
 import numpy
+from multiprocessing import Pool
 from nabu.pipeline.config import get_default_nabu_config
 from nabu.pipeline.fullfield.nabu_config import (
     nabu_config as nabu_fullfield_default_config,
@@ -65,7 +66,7 @@ from tomwer.core.process.reconstruction.scores import (
 from tomwer.core.utils.locker import FileLockerManager
 from tomwer.core.process.task import Task
 from tomwer.core.progress import Progress
-from tomwer.core.scan.hdf5scan import HDF5TomoScan
+from tomwer.core.scan.nxtomoscan import NXtomoScan
 from tomwer.core.scan.scanbase import TomwerScanBase
 from tomwer.core.scan.scanfactory import ScanFactory
 from tomwer.core.utils import logconfig
@@ -86,7 +87,7 @@ from .params import SADeltaBetaParams
 _logger = logging.getLogger(__name__)
 
 
-DEFAULT_RECONS_FOLDER = "sadeltabeta_results"
+DEFAULT_RECONS_FOLDER = "multi_delta_beta_results"
 
 
 def one_slice_several_db(
@@ -135,6 +136,8 @@ class SADeltaBetaTask(
         "dump_roi",
         "dump_process",
         "serialize_output_data",
+        "compute_scores",  # for GUI we want to post pone the score calculation
+        "pool_size",
     ),
 ):
     """
@@ -142,6 +145,8 @@ class SADeltaBetaTask(
     several Center Of Rotation (cor) values
     """
 
+    DEFAULT_POOL_SIZE = 10
+
     def __init__(
         self,
         process_id=None,
@@ -221,6 +226,17 @@ class SADeltaBetaTask(
         scan.sa_delta_beta_params.value = best_db
         return best_db
 
+    def get_output_dir(self, params: SADeltaBetaParams, scan: TomwerScanBase):
+        output_dir = params.output_dir or None
+
+        if params.output_dir is None:
+            output_dir = (
+                params.nabu_params.get("output", {}).get("location", None) or None
+            )
+            if output_dir is None:
+                output_dir = os.path.join(scan.path, DEFAULT_RECONS_FOLDER)
+        return output_dir
+
     def run(self):
         scan = data_identifier_to_scan(self.inputs.data)
         if scan is None:
@@ -234,6 +250,7 @@ class SADeltaBetaTask(
             raise ValueError(f"input type of {scan}: {type(scan)} is not managed")
 
         config = copy(self.inputs.sa_delta_beta_params)
+        axis = config.get("reconstruction", {}).get("slice_plane", "XY")
         params = SADeltaBetaParams.from_dict(config)
 
         # insure scan contains some parameter regarding sa delta / beta
@@ -247,14 +264,12 @@ class SADeltaBetaTask(
             scan.axis_params = AxisRP()
 
         # create dir if does not exists
-        if params.output_dir in (None, ""):
-            params.output_dir = os.path.join(scan.path, DEFAULT_RECONS_FOLDER)
-            if not os.path.exists(params.output_dir):
-                os.makedirs(params.output_dir)
+        params.output_dir = self.get_output_dir(params=params, scan=scan)
+        if not os.path.exists(params.output_dir):
+            os.makedirs(params.output_dir)
 
         slice_index = self._preprocess_slice_index(params.slice_indexes)
         delta_beta_s = params.delta_beta_values
-        output_dir = params.output_dir
         # TODO: check: dry run should only be settable at one location
         dry_run = self._dry_run or params.dry_run
         cluster_config = params.cluster_config
@@ -265,7 +280,7 @@ class SADeltaBetaTask(
             scan=scan,
             config=nabu_config,
             delta_beta_s=delta_beta_s,
-            output_dir=output_dir,
+            output_dir=params.output_dir,
         )
         # step 2: run reconstructions
         advancement = Progress(
@@ -289,6 +304,7 @@ class SADeltaBetaTask(
                 slice_index=slice_index,
                 dry_run=dry_run,
                 cluster_config=cluster_config,
+                axis=axis,
             )
         except Exception as e:
             _logger.error(e)
@@ -304,23 +320,29 @@ class SADeltaBetaTask(
                 slice_index=slice_index,
                 db_reconstructions=dbs_res,
                 future_tomo_objs=future_tomo_objs,
+                axis=axis,
             )
 
             # step 4: run post processing (compute score for each slice)
-            try:
-                dbs_res, rois = self._post_processing(
-                    scan=scan,
-                    db_reconstructions=dbs_res,
-                )
-            except Exception as e:
-                _logger.error(e)
-                mess = f"sa-delta-beta -post-processing- computation for {str(scan)} failed."
-                state = DatasetState.FAILED
-                dbs_res = {}
+            if self.get_input_value("compute_scores", True):
+                try:
+                    dbs_res, rois = self._post_processing(
+                        scan=scan,
+                        db_reconstructions=dbs_res,
+                    )
+                except Exception as e:
+                    _logger.error(e)
+                    mess = f"sa-delta-beta -post-processing- computation for {str(scan)} failed."
+                    state = DatasetState.FAILED
+                    dbs_res = {}
+                else:
+                    state = DatasetState.WAIT_USER_VALIDATION
+                    self.delta_beta_s = scan.sa_delta_beta_params.autofocus
+                    mess = "sa-delta-beta computation succeeded"
             else:
-                state = DatasetState.WAIT_USER_VALIDATION
-                self.delta_beta_s = scan.sa_delta_beta_params.autofocus
-                mess = "sa-delta-beta computation succeeded"
+                dbs_res = {}
+                state = DatasetState.FAILED
+                mess = "couldn't find 'compute_scores'"
 
         if self._cancelled:
             state = DatasetState.CANCELLED
@@ -407,7 +429,8 @@ class SADeltaBetaTask(
         slice_index,
         advancement,
         dry_run,
-        cluster_config,
+        axis,
+        cluster_config: Optional[dict],
     ):
         future_tomo_objs = {}
         success = True
@@ -415,6 +438,11 @@ class SADeltaBetaTask(
         std_outs = []
         std_errs = []
 
+        if not isinstance(cluster_config, (dict, type(None))):
+            raise TypeError(
+                f"'cluster_config' is expected to be a dict or None. Get {type(cluster_config)} instead."
+            )
+
         runners = run_nabu_one_slice_several_config(
             nabu_configs=configs,
             scan=scan,
@@ -422,12 +450,11 @@ class SADeltaBetaTask(
             dry_run=dry_run,
             file_format="hdf5",
             advancement=advancement,
-            cluster_config=cluster_config.to_dict()
-            if cluster_config is not None
-            else None,
+            cluster_config=cluster_config,
             process_id=self.process_id,
             instanciate_classes_only=True,
             output_file_prefix_pattern=None,
+            axis=axis,
         )
 
         for runner in runners:
@@ -449,9 +476,14 @@ class SADeltaBetaTask(
                     if isinstance(res, ResultsWithStd):
                         std_outs.append(res.std_out)
                         std_errs.append(res.std_err)
-                    if isinstance(res, ResultsLocalRun) and len(res.results_urls) > 0:
-                        assert len(res.results_urls) == 1, "only one slice expected"
-                        recons_urls[db] = res.results_urls[0]
+                    if (
+                        isinstance(res, ResultsLocalRun)
+                        and len(res.results_identifiers) > 0
+                    ):
+                        assert (
+                            len(res.results_identifiers) == 1
+                        ), "only one slice expected"
+                        recons_urls[db] = res.results_identifiers[0]
                     if isinstance(res, ResultSlurmRun):
                         future_tomo_obj = FutureTomwerObject(
                             tomo_obj=scan,
@@ -473,6 +505,7 @@ class SADeltaBetaTask(
         post_processing = _PostProcessing(
             scan=scan,
             db_reconstructions=db_reconstructions,
+            pool_size=self.get_input_value("pool_size", self.DEFAULT_POOL_SIZE),
         )
         post_processing._cancelled = self._cancelled
         self._current_processing = post_processing
@@ -484,6 +517,7 @@ class SADeltaBetaTask(
         nabu_config: dict,
         db_reconstructions,
         slice_index,
+        axis,
         future_tomo_objs: dict,
     ):
         assert isinstance(nabu_config, dict)
@@ -510,6 +544,7 @@ class SADeltaBetaTask(
                 pag=pag,
                 db=int(db) if db is not None else None,
                 ctf=ctf,
+                axis=axis,
             )
             # retrieve url
             volume_identifier = nabu_utils.get_recons_volume_identifier(
@@ -518,9 +553,7 @@ class SADeltaBetaTask(
                 file_format=nabu_config.get("file_format", "hdf5"),
                 scan=scan,
                 slice_index=None,
-                start_z=None,
-                end_z=None,
-                expects_single_slice=True,
+                axis=axis,
             )
 
             assert len(volume_identifier) <= 1, "only one slice expected"
@@ -635,7 +668,7 @@ class SADeltaBetaTask(
     def process_to_tomwer_processes(scan):
         if scan.process_file is not None:
             entry = "entry"
-            if isinstance(scan, HDF5TomoScan):
+            if isinstance(scan, NXtomoScan):
                 entry = scan.entry
 
             db = None
@@ -712,76 +745,107 @@ class SADeltaBetaTask(
 class _PostProcessing:
     """class used to run SA-delta-beta post-processing on reconstructed slices"""
 
-    def __init__(self, db_reconstructions, scan) -> None:
+    DEFAULT_POOL_SIZE = 10
+
+    def __init__(self, db_reconstructions, scan, pool_size) -> None:
         self._db_reconstructions = db_reconstructions
         self._scan = scan
         self._cancelled = False
+        self.pool_size = pool_size
+
+    @staticmethod
+    def compute_score(item: tuple):
+        db, (url, data), mask_disk_radius, cancelled = item
+        if data is None:
+            score = None
+            data_roi = None
+        else:
+            assert data.ndim == 2
+            data_roi = apply_roi(data=data, radius=mask_disk_radius, url=url)
+
+            # move data_roi to [0-1] range
+            #  preprocessing: get percentile 0 and 99 from image and
+            #  "clean" highest and lowest pixels from it
+            min_p, max_p = numpy.percentile(data_roi, (1, 99))
+            data_roi_int = data_roi[...]
+            data_roi_int[data_roi_int < min_p] = min_p
+            data_roi_int[data_roi_int > max_p] = max_p
+            data_roi_int = (data_roi_int - min_p) / (max_p - min_p)
+
+            score = ComputedScore(
+                tv=compute_score(data=data_roi_int, method=ScoreMethod.TV),
+                std=compute_score(data=data_roi_int, method=ScoreMethod.STD),
+            )
+        return {db: (url, score)}, {db: data_roi}
 
     def run(self):
         datasets = self.load_datasets()
-
+        assert isinstance(datasets, dict)
         mask_disk_radius = get_disk_mask_radius(datasets)
+        with Pool(self.pool_size) as pool:
+            res = pool.map(
+                self.compute_score,
+                [
+                    (
+                        *item,
+                        mask_disk_radius,
+                        self._cancelled,
+                    )
+                    for item in datasets.items()
+                ],
+            )
         scores = {}
         rois = {}
-        for db, (url, data) in datasets.items():
-            if data is None:
-                score = None
-            else:
-                assert data.ndim == 2
-                data_roi = apply_roi(data=data, radius=mask_disk_radius, url=url)
-                rois[db] = data_roi
-
-                # move data_roi to [0-1] range
-                #  preprocessing: get percentile 0 and 99 from image and
-                #  "clean" highest and lowest pixels from it
-                min_p, max_p = numpy.percentile(data_roi, (1, 99))
-                data_roi_int = data_roi[...]
-                data_roi_int[data_roi_int < min_p] = min_p
-                data_roi_int[data_roi_int > max_p] = max_p
-                data_roi_int = (data_roi_int - min_p) / (max_p - min_p)
-
-                score = ComputedScore(
-                    tv=compute_score(data=data_roi_int, method=ScoreMethod.TV),
-                    std=compute_score(data=data_roi_int, method=ScoreMethod.STD),
-                )
-            scores[db] = (url, score)
+        for mydict in res:
+            myscores, myrois = mydict
+            scores.update(myscores)
+            rois.update(myrois)
         return scores, rois
 
-    def load_datasets(self):
-        datasets_ = {}
-        for db, volume_identifier in self._db_reconstructions.items():
-            slice_url = None
-            # in case the try processing fails
-            try:
-                volume = VolumeFactory.create_tomo_object_from_identifier(
-                    volume_identifier
-                )
-                volumes_urls = tuple(volume.browse_data_urls())
-                if len(volumes_urls) > 1:
-                    _logger.warning(
-                        f"found a volume with mode that one url ({volumes_urls})"
-                    )
-                slice_url = volumes_urls[0]
-                data = get_data(slice_url)
-            except Exception as e:
-                _logger.error(
-                    f"Fail to compute a score for {volume_identifier}. Reason is {e}"
+    @staticmethod
+    def _load_dataset(item: tuple):
+        db, volume_identifier = item
+        slice_url = None
+        # in case the try processing fails
+        try:
+            volume = VolumeFactory.create_tomo_object_from_identifier(volume_identifier)
+            volumes_urls = tuple(volume.browse_data_urls())
+            if len(volumes_urls) > 1:
+                _logger.warning(
+                    f"found a volume with mode that one url ({volumes_urls})"
                 )
-                datasets_[db] = (slice_url, None)
-            else:
-                if data.ndim == 3:
-                    if data.shape[0] == 1:
-                        data = data.reshape(data.shape[1], data.shape[2])
-                    elif data.shape[2] == 1:
-                        data = data.reshape(data.shape[0], data.shape[1])
-                    else:
-                        raise ValueError(f"Data is expected to be 2D. Not {data.ndim}D")
-                elif data.ndim == 2:
-                    pass
+            slice_url = volumes_urls[0]
+            data = get_data(slice_url)
+        except Exception as e:
+            _logger.error(
+                f"Fail to compute a score for {volume_identifier}. Reason is {e}"
+            )
+            return {db: (slice_url, None)}
+        else:
+            if data.ndim == 3:
+                if data.shape[0] == 1:
+                    data = data.reshape(data.shape[1], data.shape[2])
+                elif data.shape[2] == 1:
+                    data = data.reshape(data.shape[0], data.shape[1])
                 else:
                     raise ValueError(f"Data is expected to be 2D. Not {data.ndim}D")
+            elif data.ndim == 2:
+                pass
+            else:
+                raise ValueError(f"Data is expected to be 2D. Not {data.ndim}D")
+
+            return {db: (slice_url, data)}
+
+    def load_datasets(self):
+        with Pool(self.pool_size) as pool:
+            res = pool.map(
+                self._load_dataset,
+                self._db_reconstructions.items(),
+            )
+        datasets_ = {}
+        for mydict in res:
+            datasets_.update(mydict)
 
-                datasets_[db] = (slice_url, data)
         return datasets_
 
     def cancel(self):

tomwer/core/process/reconstruction/scores/params.py

@@ -117,6 +117,7 @@ class SABaseParams:
 
     @property
     def output_dir(self) -> Union[str, None]:
+        """nabu cfg_files output dir. If not provided will use nabu slice output with saaxis/cfg_files as postfix"""
         return self._output_dir
 
     @output_dir.setter
@@ -163,7 +164,9 @@ class SABaseParams:
 
     @cluster_config.setter
     def cluster_config(self, config: Optional[dict]):
-        if not isinstance(config, (dict, type(None), SlurmClusterConfiguration)):
+        if isinstance(config, SlurmClusterConfiguration):
+            config = config.to_dict()
+        if not isinstance(config, (dict, type(None))):
             raise TypeError(
                 "config is expected to be None, a dict or SlurmClusterConfiguration"
             )
@@ -176,9 +179,9 @@ class SABaseParams:
             "dry_run": self.dry_run,
             "output_dir": self.output_dir or "",
             "score_method": self.score_method.value,
-            "cluster_config": self.cluster_config
-            if self.cluster_config is not None
-            else "",
+            "cluster_config": (
+                self.cluster_config if self.cluster_config is not None else ""
+            ),
         }
 
     def load_from_dict(self, dict_: dict):

tomwer/core/process/reconstruction/scores/scores.py

@@ -79,6 +79,19 @@ class ComputedScore:
         else:
             raise ValueError(f"{method} is an unrecognized method")
 
+    def __eq__(self, __value: object) -> bool:
+        if not isinstance(__value, ComputedScore):
+            return False
+        else:
+            return (
+                self.total_variation == __value.total_variation
+                and self.std == __value.std
+                and self.tomo_consistency == __value.tomo_consistency
+            )
+
+    def __str__(self) -> str:
+        return f"std: {self.std} - tv: {self.total_variation} - tomo-consistency: {self.tomo_consistency}"
+
 
 def compute_tomo_consistency(
     data: numpy.ndarray,

tomwer/core/process/reconstruction/test/test_axis_params.py

@@ -7,7 +7,7 @@ from tomwer.core.process.reconstruction.axis.params import (
     AxisCalculationInput,
     AxisResource,
 )
-from tomwer.core.utils.scanutils import MockHDF5
+from tomwer.core.utils.scanutils import MockNXtomo
 
 
 def test_axis_resource(tmp_path):
@@ -15,7 +15,7 @@ def test_axis_resource(tmp_path):
     Test AxisResource class
     """
 
-    mock = MockHDF5(
+    mock = MockNXtomo(
         scan_path=os.path.join(tmp_path, "scan1"),
         n_proj=10,
         n_ini_proj=10,

tomwer/core/process/reconstruction/test/test_darkref.py

@@ -31,7 +31,7 @@ __date__ = "24/11/2021"
 from tomwer.core.process.reconstruction.darkref.darkrefs import (
     requires_reduced_dark_and_flat,
 )
-from tomwer.core.utils.scanutils import MockEDF, MockHDF5
+from tomwer.core.utils.scanutils import MockEDF, MockNXtomo
 
 
 def test_quick_run_necessary_edf(tmpdir):
@@ -45,8 +45,8 @@ def test_quick_run_necessary_edf(tmpdir):
 
 
 def test_quick_run_necessary_hdf5(tmpdir):
-    """test the `quick_run_necessary` function for HDF5TomoScan"""
-    scan = MockHDF5(
+    """test the `quick_run_necessary` function for NXtomoScan"""
+    scan = MockNXtomo(
         scan_path=tmpdir,
         n_proj=20,
         n_ini_proj=20,

tomwer/core/process/reconstruction/test/test_darkref_copy.py

@@ -5,7 +5,7 @@ import pytest
 from silx.io.dictdump import dicttoh5
 
 from tomwer.core.process.reconstruction.darkref.darkrefscopy import DarkRefsCopy
-from tomwer.core.utils.scanutils import MockHDF5
+from tomwer.core.utils.scanutils import MockNXtomo
 
 
 @pytest.mark.parametrize(
@@ -32,20 +32,20 @@ def test_register_and_copy_darks_and_flats(
     for my_dir in (save_dir, scan_folder_with_raw, scan_folder_without_raw):
         os.makedirs(my_dir)
 
-    scan_with_raw = MockHDF5(
+    scan_with_raw = MockNXtomo(
         scan_path=scan_folder_with_raw,
         create_ini_dark=True,
-        create_ini_ref=True,
-        create_final_ref=False,
+        create_ini_flat=True,
+        create_final_flat=False,
         n_proj=10,
         n_ini_proj=10,
         dim=12,
     ).scan
-    scan_without_raw = MockHDF5(
+    scan_without_raw = MockNXtomo(
         scan_path=scan_folder_without_raw,
         create_ini_dark=False,
-        create_ini_ref=False,
-        create_final_ref=False,
+        create_ini_flat=False,
+        create_final_flat=False,
         n_proj=10,
         n_ini_proj=10,
         dim=12,

tomwer/core/process/reconstruction/test/test_saaxis.py

@@ -41,7 +41,7 @@ from tomwer.core.process.reconstruction.scores.scores import (
     _METHOD_TO_FCT,
     compute_score_contrast_std,
 )
-from tomwer.core.utils.scanutils import MockHDF5
+from tomwer.core.utils.scanutils import MockNXtomo
 
 
 class TestScoreFunctions(unittest.TestCase):
@@ -69,7 +69,7 @@ class TestSAAxisProcess(unittest.TestCase):
     def setUp(self) -> None:
         self.tempdir = tempfile.mkdtemp()
         dim = 10
-        mock = MockHDF5(
+        mock = MockNXtomo(
             scan_path=self.tempdir, n_proj=10, n_ini_proj=10, scan_range=180, dim=dim
         )
         mock.add_alignment_radio(index=10, angle=90)
@@ -96,8 +96,7 @@ class TestSAAxisProcess(unittest.TestCase):
                 "serialize_output_data": False,
             }
         )
-        with self.assertRaises(ValueError):
-            process.run()
+
         self._default_saaxis_params.estimated_cor = 11
         self._default_saaxis_params.research_width = 2
         process = SAAxisTask(

tomwer/core/process/reconstruction/test/test_sadeltabeta.py

@@ -37,7 +37,7 @@ from tomwer.core.process.reconstruction.sadeltabeta.sadeltabeta import (
     SADeltaBetaParams,
     SADeltaBetaTask,
 )
-from tomwer.core.utils.scanutils import MockHDF5
+from tomwer.core.utils.scanutils import MockNXtomo
 
 
 class TestSADeltaBetaProcess(unittest.TestCase):
@@ -47,7 +47,7 @@ class TestSADeltaBetaProcess(unittest.TestCase):
         super().setUp()
         self.tempdir = tempfile.mkdtemp()
         dim = 10
-        mock = MockHDF5(
+        mock = MockNXtomo(
             scan_path=self.tempdir, n_proj=10, n_ini_proj=10, scan_range=180, dim=dim
         )
         self.scan = mock.scan

tomwer/core/process/stitching/nabustitcher.py

@@ -21,7 +21,7 @@ from nabu.stitching.slurm_utils import split_stitching_configuration_to_slurm_jo
 from processview.core.superviseprocess import SuperviseProcess
 
 from tomwer.core.process.task import Task
-from tomwer.core.scan.hdf5scan import HDF5TomoScan
+from tomwer.core.scan.nxtomoscan import NXtomoScan
 from tomwer.core.scan.scanfactory import ScanFactory
 from tomwer.core.futureobject import FutureTomwerObject
 from tomwer.core.volume.volumefactory import VolumeFactory
@@ -112,7 +112,7 @@ class StitcherTask(
             if config.stitching_type is StitchingType.Z_PREPROC:
                 output_nx_tomo_file = config.output_file_path
                 output_nx_tomo_entry = config.output_data_path
-                tomwer_scan = HDF5TomoScan(
+                tomwer_scan = NXtomoScan(
                     scan=output_nx_tomo_file, entry=output_nx_tomo_entry
                 )
                 self.outputs.future_tomo_obj = FutureTomwerObject(
@@ -139,32 +139,33 @@ class StitchingPostProcAggregation(_StitchingPostProcAggregation):
     once stitching is done remotly we need to 'aggregate' the different part in the correct order
     """
 
-    def __init__(self, futures, stitching_config) -> None:
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
         self._n_finished = 0
-        self._futures = futures
-        if isinstance(stitching_config, dict):
+
+        if isinstance(self.stitching_config, dict):
             stitching_type = StitchingType.from_value(
-                stitching_config["stitching"]["type"]
+                self.stitching_config["stitching"]["type"]
             )
             if stitching_type is StitchingType.Z_PREPROC:
                 self._stitching_config = PreProcessedZStitchingConfiguration.from_dict(
-                    stitching_config
+                    self.stitching_config
                 )
             elif stitching_type is StitchingType.Z_POSTPROC:
                 self._stitching_config = PostProcessedZStitchingConfiguration.from_dict(
-                    stitching_config
+                    self.stitching_config
                 )
             else:
                 raise NotImplementedError("stitching type not handled")
-        elif not isinstance(stitching_config, StitchingConfiguration):
+        elif not isinstance(self.stitching_config, StitchingConfiguration):
             raise TypeError(
-                f"stitching_config is expected to be an instance of {StitchingConfiguration}. {type(stitching_config)} provided instead"
+                f"stitching_config is expected to be an instance of {StitchingConfiguration}. {type(self.stitching_config)} provided instead"
             )
-        else:
-            self._stitching_config = stitching_config
 
     def onePartFinished(self, *args, **kwargs):
         self._n_finished += 1
+        # note: for now we only consider the user of the futures.
+        # if users want to only run post-processing agregation then they will use the nabu CLI for it.
         if self._n_finished == len(self._futures):
             _logger.info("All slurm job finished. Will aggregate results")
             self.process()