tomwer 1.2.9__py3-none-any.whl → 1.3.0a0__py3-none-any.whl

tomwer/core/process/reconstruction/saaxis/saaxis.py

@@ -32,27 +32,26 @@ __license__ = "MIT"
 __date__ = "10/02/2021"
 
 
-import copy
 import logging
 import os
 from typing import Optional
 
 import h5py
 import numpy
+from multiprocessing import Pool
 from processview.core.manager import DatasetState, ProcessManager
 from processview.core.superviseprocess import SuperviseProcess
 from silx.io.url import DataUrl
-from silx.utils.deprecation import deprecated_warning
+from tomwer.core.utils.deprecation import deprecated_warning
 
 from tomoscan.io import HDF5File
 
 import tomwer.version
 from tomwer.core.process.reconstruction.axis import AxisRP
 from tomwer.core.process.reconstruction.nabu.nabuscores import (
-    run_nabu_one_slice_several_config,
+    run_nabu_multicor,
 )
 from tomwer.core.process.reconstruction.nabu.nabuslices import (
-    SingleSliceRunner,
     interpret_tomwer_configuration,
 )
 from tomwer.core.process.reconstruction.scores import (
@@ -63,9 +62,7 @@ from tomwer.core.process.reconstruction.scores import (
 )
 from tomwer.core.process.reconstruction.scores.params import ScoreMethod
 from tomwer.core.process.task import Task
-from tomwer.core.progress import Progress
-from tomwer.core.scan.edfscan import EDFTomoScan
-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
@@ -82,6 +79,11 @@ from tomwer.core.process.reconstruction.nabu.nabucommon import (
 )
 from tomwer.core.futureobject import FutureTomwerObject
 from tomwer.core.process.reconstruction.saaxis.params import ReconstructionMode
+from tomwer.core.process.reconstruction.nabu.utils import (
+    get_multi_cor_recons_volume_identifiers,
+    get_nabu_multicor_file_prefix,
+)
+from tomwer.core.process.reconstruction.nabu.nabuscores import _ReconstructorMultiCor
 
 from ..nabu import utils
 from .params import SAAxisParams
@@ -139,6 +141,8 @@ class SAAxisTask(
         "dump_roi",
         "dump_process",
         "serialize_output_data",
+        "compute_scores",  # for GUI we want to post pone the score calculation
+        "pool_size",
     ),
 ):
     """
@@ -148,6 +152,8 @@ class SAAxisTask(
     As the saaxis is integrating the score calculation we will never get a future_tomo_scan as output
     """
 
+    DEFAULT_POOL_SIZE = 10
+
     def __init__(
         self, process_id=None, inputs=None, varinfo=None, node_attrs=None, execinfo=None
     ):
@@ -214,7 +220,7 @@ class SAAxisTask(
         return best_cor
 
     def _config_preprocessing(
-        self, scan, config, cor_positions, file_format, output_dir, cluster_config
+        self, scan, config, file_format, output_dir, cluster_config
     ):
         """convert general configuration to nabu - single reconstruction - configuration"""
         nabu_configurations = interpret_tomwer_configuration(config, scan=None)
@@ -232,17 +238,8 @@ class SAAxisTask(
         # work on file name...
         if output_dir is None:
             output_dir = os.path.join(scan.path, DEFAULT_RECONS_FOLDER)
-        if scan.process_file is not None:
-            steps_file_basename, _ = os.path.splitext(scan.process_file)
-            steps_file_basename = "_".join(
-                ("steps_file_basename", "nabu", "sinogram", "save", "step")
-            )
-            steps_file_basename = steps_file_basename + ".hdf5"
-            steps_file = os.path.join(output_dir, steps_file_basename)
-        else:
-            steps_file = ""
 
-        base_config = nabu_configurations[0][0]
+        nabu_configuration = nabu_configurations[0][0]
         if cluster_config == {}:
             cluster_config = None
         is_cluster_job = cluster_config is not None
@@ -250,98 +247,73 @@ class SAAxisTask(
             raise ValueError(
                 "job on cluster requested but no access to slurm cluster found"
             )
-        configs = {}
-
-        for i_cor, cor in enumerate(cor_positions):
-            nabu_configuration = copy.deepcopy(base_config)
-            nabu_configuration["pipeline"] = {
-                "save_steps": "sinogram" if i_cor == 0 else "",
-                "resume_from_step": "sinogram",
-                "steps_file": steps_file,
-            }
-            # convert cor from tomwer ref to nabu ref
-            if scan.dim_1 is not None:
-                cor_nabu_ref = cor + scan.dim_1 / 2.0
-            else:
-                _logger.warning("enable to get image half width. Set it to 1024")
-                cor_nabu_ref = cor + 1024
-            # handle reconstruction section
-            if "reconstruction" not in nabu_configuration:
-                nabu_configuration["reconstruction"] = {}
-            nabu_configuration["reconstruction"]["rotation_axis_position"] = str(
-                cor_nabu_ref
-            )
-            # handle output section
-            if "output" not in nabu_configuration:
-                nabu_configuration["output"] = {}
-            nabu_configuration["output"]["location"] = output_dir
-            nabu_configuration["output"]["file_format"] = file_format
-            # handle resources section
-            nabu_configuration["resources"] = utils.get_nabu_resources_desc(
-                scan=scan, workers=1, method="local"
-            )
-            configs[cor] = nabu_configuration
-        return configs
 
-    def _run_slice_recons_per_cor(
+        # handle reconstruction section
+        if "reconstruction" not in nabu_configuration:
+            nabu_configuration["reconstruction"] = {}
+        nabu_configuration["reconstruction"]["rotation_axis_position"] = ""
+        # handle output section
+        if "output" not in nabu_configuration:
+            nabu_configuration["output"] = {}
+        nabu_configuration["output"]["location"] = output_dir
+        nabu_configuration["output"]["file_format"] = file_format
+        # handle resources section
+        nabu_configuration["resources"] = utils.get_nabu_resources_desc(
+            scan=scan, workers=1, method="local"
+        )
+        return nabu_configuration
+
+    def _run_nabu_multicor(
         self,
         scan,
-        configs,
+        nabu_config,
+        cors,
         slice_index,
         file_format,
-        advancement,
-        cluster_config,
+        cluster_config: Optional[dict],
         dry_run=False,
     ):
-        runners = run_nabu_one_slice_several_config(
-            nabu_configs=configs,
+        if not (cluster_config is None or isinstance(cluster_config, dict)):
+            raise TypeError(
+                f"cluster_config is expected to be a dict. Get {type(cluster_config)} instead."
+            )
+        runner = run_nabu_multicor(
+            nabu_config=nabu_config,
             scan=scan,
+            cors=cors,
             slice_index=slice_index,
             dry_run=dry_run,
             file_format=file_format,
-            advancement=advancement,
-            cluster_config=cluster_config.to_dict()
-            if cluster_config is not None
-            else None,
+            cluster_config=cluster_config if cluster_config is not None else None,
             process_id=self.process_id,
             instanciate_classes_only=True,
             output_file_prefix_pattern="cor_{file_name}_{value}",  # as the cor is evolving, create different files to make sure the name will be unique
         )
 
-        future_tomo_objs = {}
-        success = True
-        recons_urls = {}
+        future_tomo_obj = None
+        recons_urls = dict()
         std_outs = []
         std_errs = []
 
-        for runner in runners:
-            if self._cancelled:
-                break
-            self._current_processing = runner
-            try:
-                results = runner.run()
-            except TimeoutError as e:
-                _logger.error(e)
-            else:
-                assert isinstance(
-                    results, dict
-                ), "results should be a dictionary with cor as key and urls as value"
-
-                for cor, res in results.items():
-                    success = success and res.success
-                    if isinstance(res, ResultsWithStd):
-                        std_outs.append(res.std_out)
-                        std_errs.append(res.std_err)
-                    if isinstance(res, ResultsLocalRun):
-                        recons_urls[cor] = res.results_urls
-                    if isinstance(res, ResultSlurmRun):
-                        future_tomo_obj = FutureTomwerObject(
-                            tomo_obj=scan,
-                            process_requester_id=self.process_id,
-                            futures=res.future_slurm_jobs,
-                        )
-                        future_tomo_objs[cor] = future_tomo_obj
-        return success, recons_urls, future_tomo_objs, std_outs, std_errs
+        self._current_processing = runner
+        try:
+            result = runner.run()
+        except TimeoutError as e:
+            _logger.error(e)
+        else:
+            success = result.success
+            if isinstance(result, ResultsWithStd):
+                std_outs.append(result.std_out)
+                std_errs.append(result.std_err)
+            if isinstance(result, ResultsLocalRun):
+                recons_urls = result.results_identifiers
+            if isinstance(result, ResultSlurmRun):
+                future_tomo_obj = FutureTomwerObject(
+                    tomo_obj=scan,
+                    process_requester_id=self.process_id,
+                    futures=result.future_slurm_jobs,
+                )
+        return success, recons_urls, (future_tomo_obj,), std_outs, std_errs
 
     def _resolve_futures(
         self,
@@ -349,6 +321,7 @@ class SAAxisTask(
         nabu_config,
         slice_index,
         file_format,
+        cors,
         cor_reconstructions,
         future_tomo_objs: dict,
         output_dir,
@@ -359,54 +332,50 @@ class SAAxisTask(
         if output_dir is None:
             output_dir = os.path.join(scan.path, DEFAULT_RECONS_FOLDER)
 
-        db = None
-        pag = False
-        ctf = False
-        if "phase" in nabu_config:
-            phase_method = nabu_config["phase"].get("method", "").lower()
-            if phase_method in ("pag", "paganin"):
-                pag = True
-            elif phase_method in ("ctf",):
-                ctf = True
-            if "delta_beta" in nabu_config["phase"]:
-                db = round(float(nabu_config["phase"]["delta_beta"]))
-
-        for cor, future_tomo_obj in future_tomo_objs.items():
+        file_prefix = get_nabu_multicor_file_prefix(scan)
+
+        for future_tomo_obj in future_tomo_objs:
             if self._cancelled:
                 break
+
+            if future_tomo_obj is None:
+                continue
+
             future_tomo_obj.results()
-            # for saaxis we need to retrieve reconstruction url
             if future_tomo_obj.cancelled() or future_tomo_obj.exceptions():
                 continue
-            else:
-                _file_name = SingleSliceRunner.get_file_basename_reconstruction(
+
+            for cor in cors:
+                cor_nabu_ref = _ReconstructorMultiCor.convert_cor_from_rel_to_abs(
                     scan=scan,
-                    slice_index=slice_index,
-                    pag=pag,
-                    db=db,
-                    ctf=ctf,
+                    cor=cor,
                 )
-                file_prefix = f"cor_{_file_name}_{cor}"
-
-                recons_vol_id = utils.get_recons_volume_identifier(
+                volume_identifiers = get_multi_cor_recons_volume_identifiers(
                     scan=scan,
-                    file_format=file_format,
+                    slice_index=slice_index,
+                    location=nabu_config["output"]["location"],
                     file_prefix=file_prefix,
-                    location=output_dir,
-                    slice_index=None,
-                    start_z=None,
-                    end_z=None,
-                    expects_single_slice=True,
+                    file_format=file_format,
+                    cors=(cor_nabu_ref,),
                 )
-                assert len(recons_vol_id) == 1, "only one volume reconstructed expected"
-                cor_reconstructions[cor] = recons_vol_id
+                volume_identifier = volume_identifiers.get(cor_nabu_ref, None)
+                if volume_identifier is None:
+                    _logger.warning(
+                        f"failed to load volume for {cor}. Something went wrong on slurm submission job"
+                    )
+                cor_reconstructions[cor] = volume_identifier
 
-    def _post_processing(self, scan, slice_index, cor_reconstructions):
+    def _post_processing(self, scan, slice_index, cor_reconstructions: dict):
         """
         compute score along the different slices
+
+        :param dict cor_reconstructions: key is expected to be a float with the cor value and the value is expected to be a volume identifier (volume with a single frame)
         """
         post_processing = _PostProcessing(
-            slice_index=slice_index, scan=scan, cor_reconstructions=cor_reconstructions
+            slice_index=slice_index,
+            scan=scan,
+            cor_reconstructions=cor_reconstructions,
+            pool_size=self.get_input_value("pool_size", self.DEFAULT_POOL_SIZE),
         )
         post_processing._cancelled = self._cancelled
         self._current_processing = post_processing
@@ -435,6 +404,16 @@ class SAAxisTask(
         else:
             return list(slice_index.values())[0]
 
+    def get_output_dir(self, params: SAAxisParams, scan: TomwerScanBase):
+        output_dir = params.output_dir or None
+        if 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:
@@ -451,10 +430,10 @@ class SAAxisTask(
         configuration = self.inputs.sa_axis_params
         params = SAAxisParams.from_dict(configuration)
         # insure output dir is created
-        if params.output_dir in (None, ""):
-            params.output_dir = os.path.join(scan.path, "saaxis_results")
-            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)
+
         # try to find an estimated cor
         #  from a previously computed cor
         if params.estimated_cor is None and scan.axis_params is not None:
@@ -480,9 +459,7 @@ class SAAxisTask(
         if mode is not ReconstructionMode.VERTICAL:
             raise ValueError(f"{mode} is not handled for now")
 
-        output_dir = params.output_dir
-        if output_dir is None:
-            output_dir = os.path.join(scan.path, DEFAULT_RECONS_FOLDER)
+        nabu_config = configuration.get("nabu_params", {})
         nabu_output_config = configuration.get("output", {})
         file_format = nabu_output_config.get("file_format", "hdf5")
         slice_index = self._preprocess_slice_index(
@@ -493,18 +470,14 @@ class SAAxisTask(
         dry_run = self._dry_run
 
         # step one: complete nabu configuration(s)
-        configs = self._config_preprocessing(
+        nabu_config = self._config_preprocessing(
             scan=scan,
-            config=configuration,
-            cor_positions=params.cors,
+            config=nabu_config,
             file_format=file_format,
-            output_dir=output_dir,
+            output_dir=params.output_dir,
             cluster_config=cluster_config,
         )
         # step 2: run reconstructions
-        advancement = Progress(
-            f"sa-axis - slice {slice_index} of {scan.get_identifier().short_description()}"
-        )
         cors_res = {}
         rois = {}
 
@@ -515,12 +488,12 @@ class SAAxisTask(
                 future_tomo_objs,
                 self._std_outs,
                 self._std_errs,
-            ) = self._run_slice_recons_per_cor(
+            ) = self._run_nabu_multicor(
                 scan=scan,
-                configs=configs,
+                nabu_config=nabu_config,
+                cors=tuple(params.cors),
                 slice_index=slice_index,
                 file_format=file_format,
-                advancement=advancement,
                 cluster_config=cluster_config,
                 dry_run=dry_run,
             )
@@ -532,29 +505,35 @@ class SAAxisTask(
             # step 3: wait for future if any
             self._resolve_futures(
                 scan=scan,
-                nabu_config=configuration,
+                nabu_config=nabu_config,
                 slice_index=slice_index,
                 file_format=file_format,
                 cor_reconstructions=cor_reconstructions,
+                cors=tuple(params.cors),
                 future_tomo_objs=future_tomo_objs,
-                output_dir=output_dir,
+                output_dir=params.output_dir,
             )
 
             # step 4: run post processing (compute score for each slice)
-            try:
-                cors_res, rois = self._post_processing(
-                    scan=scan,
-                    slice_index=slice_index,
-                    cor_reconstructions=cor_reconstructions,
-                )
-            except Exception as e:
-                _logger.error(e)
-                mess = f"sa-axis -post-processing- computation for {str(scan)} failed."
-                state = DatasetState.FAILED
-                cors_res = {}
+            if self.get_input_value("compute_scores", True):
+                try:
+                    cors_res, rois = self._post_processing(
+                        scan=scan,
+                        slice_index=slice_index,
+                        cor_reconstructions=cor_reconstructions,
+                    )
+                except Exception as e:
+                    _logger.error(e)
+                    mess = (
+                        f"sa-axis -post-processing- computation for {str(scan)} failed."
+                    )
+                    state = DatasetState.FAILED
+                    cors_res = {}
+                else:
+                    state = DatasetState.WAIT_USER_VALIDATION
+                    mess = "sa-axis computation succeeded"
             else:
-                state = DatasetState.WAIT_USER_VALIDATION
-                mess = "sa-axis computation succeeded"
+                cors_res = {}
 
         if self._cancelled:
             state = DatasetState.CANCELLED
@@ -671,7 +650,7 @@ class SAAxisTask(
     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
 
             cor = None
@@ -749,128 +728,119 @@ class SAAxisTask(
 class _PostProcessing:
     """class used to run SA-axis post-processing on reconstructed slices"""
 
-    def __init__(self, cor_reconstructions, slice_index, scan) -> None:
+    def __init__(self, cor_reconstructions, slice_index, scan, pool_size) -> None:
         self._cor_reconstructions = cor_reconstructions
         self._slice_index = slice_index
         self._scan = scan
         self._cancelled = False
+        self.pool_size = pool_size
+
+    @staticmethod
+    def compute_score(item: tuple):
+        cor, (url, data), mask_disk_radius, cancelled = item
+        if cancelled:
+            return (None, None), None
+
+        if data is None:
+            score = None
+            data_roi = None
+        else:
+            if not isinstance(data, numpy.ndarray):
+                raise TypeError(
+                    f"data should be a numpy array. Get {type(data)} instead"
+                )
+            assert data.ndim == 2, f"data should be 2D. Get {data.ndim} instead"
+            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),
+                tomo_consistency=None,
+            )
+        return {cor: (url, score)}, {cor: 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 cor, (url, data) in datasets.items():
-            if self._cancelled:
-                break
-
-            if data is None:
-                score = None
-            else:
-                assert data.ndim == 2
-                data_roi = apply_roi(data=data, radius=mask_disk_radius, url=url)
-                rois[cor] = 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)
-
-                if isinstance(self._scan, EDFTomoScan):
-                    _logger.info("tomo consistency is not handled for EDF scan")
-                    tomo_consistency_score = None
-                else:
-                    try:
-                        projections_with_angle = self._scan.projections_with_angle()
-                        angles_ = [
-                            frame_angle
-                            for frame_angle, frame in projections_with_angle.items()
-                        ]
-                        angles = []
-                        for angle in angles_:
-                            if not isinstance(angle, str):
-                                angles.append(angle)
-                        if self._slice_index == "middle":
-                            if self._scan.dim_2 is not None:
-                                self._slice_index = self._scan.dim_2 // 2
-                            else:
-                                _logger.warning(
-                                    "scan.dim_2 returns None, unable to deduce middle "
-                                    "pick 1024"
-                                )
-                                self._slice_index = 1024
-                        tomo_consistency_score = compute_score(
-                            data=data,
-                            method=ScoreMethod.TOMO_CONSISTENCY,
-                            angles=angles,
-                            original_sinogram=self._scan.get_sinogram(
-                                self._slice_index
-                            ),
-                            detector_width=self._scan.dim_1,
-                            original_axis_position=cor + self._scan.dim_1 / 2.0,
-                        )
-                    except Exception as e:
-                        _logger.error(e)
-                        tomo_consistency_score = None
-                score = ComputedScore(
-                    tv=compute_score(data=data_roi_int, method=ScoreMethod.TV),
-                    std=compute_score(data=data_roi_int, method=ScoreMethod.STD),
-                    tomo_consistency=tomo_consistency_score,
-                )
-            scores[cor] = (url, score)
+        for mydict in res:
+            myscores, myrois = mydict
+            scores.update(myscores)
+            rois.update(myrois)
         return scores, rois
 
-    def load_datasets(self):
-        datasets_ = {}
-        for cor, volume_identifiers in self._cor_reconstructions.items():
-            if self._cancelled:
-                break
-
-            if len(volume_identifiers) == 0:
-                # in the case failed to load the url
-                continue
-            elif len(volume_identifiers) > 1:
-                raise ValueError("only one slice reconstructed expected per cor")
-            volume = VolumeFactory.create_tomo_object_from_identifier(
-                volume_identifiers[0]
+    @staticmethod
+    def _load_dataset(item: tuple):
+        cor, volume_identifier = item
+        if volume_identifier is None:
+            return {cor: (None, None)}
+
+        volume = VolumeFactory.create_tomo_object_from_identifier(volume_identifier)
+        urls = tuple(volume.browse_data_urls())
+        if len(urls) == 0:
+            _logger.error(
+                f"volume {volume.get_identifier().to_str()} has no url / slices. Unable to load any data."
             )
-            urls = tuple(volume.browse_data_urls())
-            if len(urls) != 1:
-                raise ValueError(
-                    f"volume is expected to have at most one url (single slice volume). get {len(urls)} - most likely nabu reconstruction failed. Do you have GPU ? Are the requested COR values valid ? - Especially for Half-acquisition"
-                )
-            url = urls[0]
-            if not isinstance(url, (DataUrl, str)):
-                raise TypeError(
-                    f"url is expected to be a str or DataUrl not {type(url)}"
-                )
+            return {cor: (None, None)}
+        if len(urls) != 1:
+            _logger.error(
+                f"volume is expected to have at most one url (single slice volume). get {len(urls)} - most likely nabu reconstruction failed. Do you have GPU ? Are the requested COR values valid ? - Especially for Half-acquisition"
+            )
+        url = urls[0]
+        if not isinstance(url, (DataUrl, str)):
+            raise TypeError(f"url is expected to be a str or DataUrl not {type(url)}")
 
-            try:
-                data = get_slice_data(url=url)
-            except Exception as e:
-                _logger.error(
-                    f"Fail to compute a score for {url.path()}. Reason is {e}"
-                )
-                datasets_[cor] = (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
+        try:
+            data = get_slice_data(url=url)
+        except Exception as e:
+            _logger.error(f"Fail to compute a score for {url.path()}. Reason is {e}")
+            return {cor: (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("Data is expected to be 2D. Not {data.ndim}D")
+                    raise ValueError(f"Data is expected to be 2D. Not {data.ndim}D")
+            elif data.ndim == 2:
+                pass
+            else:
+                raise ValueError("Data is expected to be 2D. Not {data.ndim}D")
+            return {cor: (url, data)}
 
-                datasets_[cor] = (url, data)
+    def load_datasets(self):
+        with Pool(self.pool_size) as pool:
+            res = pool.map(
+                self._load_dataset,
+                self._cor_reconstructions.items(),
+            )
+        datasets_ = {}
+        for mydict in res:
+            datasets_.update(mydict)
         return datasets_
 
     def cancel(self):