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

tomwer/core/process/reconstruction/nabu/nabuscores.py

@@ -1,36 +1,3 @@
-# coding: utf-8
-###########################################################################
-# Copyright (C) 2016 European Synchrotron Radiation Facility
-#
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-#
-# The above copyright notice and this permission notice shall be included in
-# all copies or substantial portions of the Software.
-#
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
-# THE SOFTWARE.
-#
-#############################################################################
-
-"""contain utils for score process
-"""
-
-__authors__ = [
-    "H.Payno",
-]
-__license__ = "MIT"
-__date__ = "28/10/2021"
-
 try:
     from nabu.pipeline.fullfield.reconstruction import (  # noqa F401
         FullFieldReconstructor,
@@ -49,13 +16,18 @@ else:
     has_nabu = True
 import logging
 import os
+import sys
 from copy import deepcopy
+import subprocess
 from typing import Iterable, Optional, Union
+import numpy
 
 from nabu.pipeline.fullfield.nabu_config import (
     nabu_config as nabu_fullfield_default_config,
 )
 from processview.core.manager.manager import ProcessManager
+from sluurp.job import SBatchScriptJob
+from sluurp.executor import submit as submit_to_slurm_cluster
 
 from tomwer.core.futureobject import FutureTomwerObject
 from tomwer.core.process.reconstruction.nabu.nabuslices import (
@@ -63,15 +35,21 @@ from tomwer.core.process.reconstruction.nabu.nabuslices import (
     _NabuBaseReconstructor,
     generate_nabu_configfile,
 )
+from tomwer.core.process.reconstruction.nabu.utils import (
+    slice_index_to_int,
+    get_nabu_multicor_file_prefix,
+)
 from tomwer.core.process.reconstruction.nabu.target import Target
 from tomwer.core.progress import Progress
 from tomwer.core.scan.scanbase import TomwerScanBase
-from tomwer.core.utils.slurm import is_slurm_available
+from tomwer.core.utils.slurm import get_slurm_script_name, is_slurm_available
 from tomwer.utils import docstring
+from tomwer.core.process.reconstruction.nabu.utils import nabu_std_err_has_error
+from tomwer.core.process.reconstruction.nabu.plane import NabuPlane
 
 from ..nabu import settings as nabu_settings
-from . import utils
-from .nabucommon import ResultsLocalRun, ResultSlurmRun, ResultsWithStd
+from . import utils, settings
+from .nabucommon import ResultsLocalRun, ResultSlurmRun, ResultsWithStd, ResultsRun
 
 _logger = logging.getLogger(__name__)
 
@@ -83,6 +61,7 @@ def run_nabu_one_slice_several_config(
     dry_run: bool,
     slice_index: Union[int, str],
     file_format: str,
+    axis: NabuPlane,
     advancement: Optional[Progress] = None,
     process_id: Optional[int] = None,
     instanciate_classes_only: bool = False,
@@ -132,6 +111,7 @@ def run_nabu_one_slice_several_config(
         cluster_config=cluster_config,
         process_name=process_name,
         output_file_prefix_pattern=output_file_prefix_pattern,
+        axis=axis,
     )
     if instanciate_classes_only:
         return (reconstructor,)
@@ -156,7 +136,7 @@ def run_nabu_one_slice_several_config(
                 std_outs.append(res.std_out)
                 std_errs.append(res.std_err)
             if isinstance(res, ResultsLocalRun):
-                recons_urls[var_value] = res.results_urls
+                recons_urls[var_value] = res.results_identifiers
             if isinstance(res, ResultSlurmRun):
                 future_tomo_obj = FutureTomwerObject(
                     tomo_obj=scan,
@@ -167,6 +147,87 @@ def run_nabu_one_slice_several_config(
         return success, recons_urls, std_outs, std_errs, future_tomo_objs
 
 
+def run_nabu_multicor(
+    scan: TomwerScanBase,
+    nabu_config: dict,
+    cors: tuple,
+    cluster_config: Optional[dict],
+    dry_run: bool,
+    slice_index: Union[int, str],
+    file_format: str,
+    process_id: Optional[int] = None,
+    instanciate_classes_only: bool = False,
+    output_file_prefix_pattern=None,
+):
+    if cluster_config in (None, {}):
+        target = Target.LOCAL
+    elif isinstance(cluster_config, dict):
+        if not is_slurm_available():
+            raise RuntimeError("Slurm computation requested but unvailable")
+        target = Target.SLURM
+    else:
+        raise TypeError(
+            f"cluster_config should be None or a dict not {type(cluster_config)}"
+        )
+
+    if process_id is not None:
+        try:
+            process_name = ProcessManager().get_process(process_id=process_id).name
+        except KeyError:
+            process_name = "unknow"
+    else:
+        process_name = ""
+
+    # TODO: FIXME small hack to make sure the configuration will be accepted when valie
+    # for now even if the cor values are given from a dedicated parameter nabu is still
+    # checking the value provided in the config file. If this value is invalid for
+    # half acquisition it will be refused
+    nabu_config["reconstruction"]["rotation_axis_position"] = numpy.mean(cors)
+
+    axis = nabu_config["reconstruction"].get("slice_plane", "XY")
+
+    reconstructor = _ReconstructorMultiCor(
+        scan=scan,
+        nabu_config=nabu_config,
+        cors=cors,
+        slice_index=slice_index,
+        target=target,
+        dry_run=dry_run,
+        file_format=file_format,
+        cluster_config=cluster_config,
+        process_name=process_name,
+        axis=axis,
+        output_file_prefix_pattern=output_file_prefix_pattern,
+    )
+    if instanciate_classes_only:
+        return reconstructor
+
+    try:
+        result = reconstructor.run()
+    except TimeoutError as e:
+        _logger.error(e)
+        return None
+    else:
+        recons_urls = {}
+        std_outs = []
+        std_errs = []
+        future_tomo_obj = None
+
+        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=process_id,
+                futures=result.future_slurm_jobs,
+            )
+        return success, recons_urls, (future_tomo_obj,), std_outs, std_errs
+
+
 class _Reconstructor(_NabuBaseReconstructor):
     def __init__(
         self,
@@ -174,6 +235,7 @@ class _Reconstructor(_NabuBaseReconstructor):
         nabu_configs: Iterable,
         advancement: Optional[Progress],
         slice_index: Union[int, str],
+        axis: Union[str, NabuPlane],
         target: Target,
         dry_run: bool,
         file_format: str,
@@ -182,6 +244,14 @@ class _Reconstructor(_NabuBaseReconstructor):
         output_file_prefix_pattern=None,
     ) -> None:
         """
+        :param scan: scan to reconstruct
+        :param nabu_configs: all the configuration to run
+        :param advancement: Progress object to notify advancement
+        :param slice_index: index of the slice to reconstruct.
+        :param axis: axis over which we want to do the reconstruction
+        :param target: is the reconstruction is to made locally or remotly
+        :param file_format: reconstructed volume file format
+        :param cluster_config: cluster configuration in the case of a remote execution
         :param str extra_output_file_pattern: possible extra file name pattern like for cor we want to add 'cor_' as prefix and cor value as suffix.
                                               To make the file name unique. For delta/beta it is already forseen to be unique. For now keywords are:
                                               * file_name: default file name according to db values and dataset name
@@ -193,12 +263,12 @@ class _Reconstructor(_NabuBaseReconstructor):
             target=target,
             cluster_config=cluster_config,
             process_name=process_name,
+            axis=axis,
         )
         if not isinstance(slice_index, (int, str)):
             raise TypeError(
                 f"slice_index should be an int or a string not {type(slice_index)}"
             )
-
         self.advancement = advancement
         self.slice_index = slice_index
         self.nabu_configs = nabu_configs
@@ -207,14 +277,9 @@ class _Reconstructor(_NabuBaseReconstructor):
 
     @docstring(_NabuBaseReconstructor)
     def run(self) -> Iterable:
-        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
+        self.slice_index = slice_index_to_int(
+            slice_index=self.slice_index, scan=self.scan
+        )
 
         results = {}
         if self.advancement:
@@ -223,8 +288,6 @@ class _Reconstructor(_NabuBaseReconstructor):
             if self._cancelled:
                 break
             config, conf_file = self.preprocess_config(deepcopy(config), var_value)
-            print("conf file is", conf_file)
-            print("confif output is", config.get("output"))
 
             # add some tomwer metadata and save the configuration
             # note: for now the section is ignored by nabu but shouldn't stay that way
@@ -240,8 +303,6 @@ class _Reconstructor(_NabuBaseReconstructor):
                 config_to_dump=config_to_dump,
                 config_file=conf_file,
                 file_format=self.file_format,
-                start_z=None,
-                end_z=None,
                 info="nabu slice reconstruction",
                 process_name=self.process_name,
             )
@@ -303,6 +364,7 @@ class _Reconstructor(_NabuBaseReconstructor):
                 pag=pag,
                 db=db,
                 ctf=ctf,
+                axis=self.axis,
             )
             file_prefix = self._format_file_prefix(file_prefix=file_prefix, value=value)
             _config["output"]["file_prefix"] = file_prefix
@@ -312,11 +374,22 @@ class _Reconstructor(_NabuBaseReconstructor):
 
         if "reconstruction" not in _config:
             _config["reconstruction"] = {}
-        _config["reconstruction"]["start_z"] = self.slice_index
-        _config["reconstruction"]["end_z"] = self.slice_index
+        if self.axis is NabuPlane.YZ:
+            _config["reconstruction"]["start_x"] = self.slice_index
+            _config["reconstruction"]["end_x"] = self.slice_index
+        elif self.axis is NabuPlane.XZ:
+            _config["reconstruction"]["start_y"] = self.slice_index
+            _config["reconstruction"]["end_y"] = self.slice_index
+        elif self.axis is NabuPlane.XY:
+            _config["reconstruction"]["start_z"] = self.slice_index
+            _config["reconstruction"]["end_z"] = self.slice_index
+        else:
+            raise ValueError(
+                f"self.axis has an invalid value: {self.axis} when expected to be in {NabuPlane.values()}"
+            )
         return _config, file_prefix
 
-    def preprocess_config(self, config, value: float):
+    def preprocess_config(self, config, value) -> tuple:
         dataset_params = self.scan.get_nabu_dataset_info()
         if "dataset" in config:
             dataset_params.update(config["dataset"])
@@ -344,3 +417,302 @@ class _Reconstructor(_NabuBaseReconstructor):
             cfg_folder, file_prefix + nabu_settings.NABU_CONFIG_FILE_EXTENSION
         )
         return config, conf_file
+
+
+class _ReconstructorMultiCor(_NabuBaseReconstructor):
+    def __init__(
+        self,
+        nabu_config: dict,
+        axis: Union[str, NabuPlane],
+        cors: tuple,
+        file_format,
+        slice_index: Union[int, str] = "middle",
+        output_file_prefix_pattern=None,
+        *args,
+        **kwargs,
+    ):
+        if not isinstance(cors, tuple):
+            raise TypeError(
+                f"cors are expected to be an instance of tuple. Get {type(cors)} instead"
+            )
+        self.__cors = cors
+        self.__slice_index = slice_index
+        self.__nabu_config = nabu_config
+        self.file_format = file_format
+        self._output_file_prefix_pattern = output_file_prefix_pattern
+
+        super().__init__(*args, **kwargs)
+
+    @property
+    def cors(self) -> tuple:
+        return self.__cors
+
+    @property
+    def slice_index(self) -> Union[int, str]:
+        return self.__slice_index
+
+    @property
+    def nabu_config(self) -> dict:
+        return self.__nabu_config
+
+    def _process_config(
+        self,
+        config_to_dump: dict,
+        config_file: str,
+        file_format: str,
+        info: Optional[str],
+        process_name: str,
+    ):
+        """
+        process provided configuration
+
+        :param str info:
+        """
+        if self.dry_run is True or self.only_create_config_file():
+            return ResultsRun(
+                success=True,
+                config=config_to_dump,
+            )
+        elif self.target is Target.LOCAL:
+            _logger.info(f"run {info} for {self.scan} with {config_to_dump}")
+            return self._run_nabu_multicor_locally(
+                conf_file=config_file,
+                file_format=file_format,
+                config_to_dump=config_to_dump,
+            )
+        elif self.target is Target.SLURM:
+            _logger.info(
+                f"run {info} on slurm for {self.scan.path} with {config_to_dump}"
+            )
+            return self._run_nabu_multicor_on_slurm(
+                conf_file=config_file,
+                config_to_dump=config_to_dump,
+                cluster_config=self.cluster_config.to_dict(),
+                process_name=process_name,
+                info=info,
+            )
+        else:
+            raise ValueError(f"{self.target} is not recognized as a valid target")
+
+    @staticmethod
+    def convert_cor_from_rel_to_abs(scan, cor):
+        if scan.dim_1 is not None:
+            return cor + scan.dim_1 / 2.0
+        else:
+            _logger.warning("enable to get image half width. Set it to 1024")
+            return cor + 1024
+
+    @staticmethod
+    def convert_cor_from_abs_to_rel(scan, cor):
+        if scan.dim_1 is not None:
+            return cor - scan.dim_1 / 2.0
+        else:
+            _logger.warning("enable to get image half width. Set it to 1024")
+            return cor - 1024
+
+    def _run_nabu_multicor_locally(
+        self,
+        conf_file: str,
+        file_format: str,
+        config_to_dump: dict,
+    ) -> ResultsLocalRun:
+        """
+        run locally nabu for a single configuration file.
+
+        :param str conf_file: path to the nabu .cfg file
+        :param str file_format: format of the generated file
+        :param dict config_to_dump: configuration saved in the .cfg as a dictionary
+        :return: results of the local run
+        :rtype: ResultsLocalRun
+        """
+        if not has_nabu:
+            raise ImportError("Fail to import nabu")
+        slice_index = slice_index_to_int(self.slice_index, scan=self.scan)
+
+        cor_in_nabu_ref = tuple(
+            [self.convert_cor_from_rel_to_abs(self.scan, cor) for cor in self.cors]
+        )
+        cor_in_nabu_ref = ",".join([str(cor) for cor in cor_in_nabu_ref])
+        command = " ".join(
+            (
+                sys.executable,
+                "-m",
+                settings.NABU_MULTICOR_PATH,
+                conf_file,  # input file
+                f"{slice_index}",  # slice
+                f"{cor_in_nabu_ref}",  # cor
+            )
+        )
+        _logger.info(f'call nabu from "{command}"')
+
+        self._process = subprocess.Popen(
+            command,
+            shell=True,
+            cwd=self.scan.path,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+        )
+        try:
+            outs, errs = self._process.communicate()
+        except (TimeoutError, KeyboardInterrupt):
+            self._process.kill()
+            outs, errs = self._process.communicate()
+
+        file_prefix = get_nabu_multicor_file_prefix(self.scan)
+
+        recons_vol_identifiers = utils.get_multi_cor_recons_volume_identifiers(
+            slice_index=slice_index,
+            location=config_to_dump["output"]["location"],
+            file_prefix=file_prefix,
+            scan=self.scan,
+            file_format=file_format,
+            cors=[
+                self.convert_cor_from_rel_to_abs(self.scan, cor) for cor in self.cors
+            ],
+        )
+        # convert back from abs ref to rel ref
+        recons_vol_identifiers = {
+            self.convert_cor_from_abs_to_rel(self.scan, cor): identifiers
+            for cor, identifiers in recons_vol_identifiers.items()
+        }
+        return ResultsLocalRun(
+            success=not nabu_std_err_has_error(errs),
+            results_identifiers=recons_vol_identifiers.values(),
+            std_out=outs,
+            std_err=errs,
+            config=config_to_dump,  # config_slices,
+        )
+
+    def _run_nabu_multicor_on_slurm(
+        self,
+        conf_file: str,
+        config_to_dump: dict,
+        cluster_config: dict,
+        process_name: str,
+        info: str,
+    ) -> ResultSlurmRun:
+        """
+        Run a nabu reconstruction on slurm of a single configuration
+
+        :return: results of the slurm run
+        :rtype: ResultSlurmRun
+        """
+        if not isinstance(conf_file, str):
+            raise TypeError(f"conf_file is expected to be a strg not {type(conf_file)}")
+        if not isinstance(config_to_dump, dict):
+            raise TypeError(
+                f"config_to_dump is expected to be a strg not {type(config_to_dump)}"
+            )
+        if not is_slurm_available():
+            raise RuntimeError("slurm not available")
+        if not isinstance(cluster_config, dict):
+            raise ValueError(
+                f"cluster config is expected to be a dict not {type(cluster_config)}"
+            )
+
+        # create slurm cluster
+        project_name = cluster_config.get(
+            "job_name", "tomwer_{scan}_-_{process}_-_{info}"
+        )
+        project_name = project_name.format(
+            scan=str(self.scan), process=process_name, info=info
+        )
+        # project name should not contain any spaces as it will be integrated in a script and interpreted.
+        project_name = project_name.replace(" ", "_")
+        cluster_config["job_name"] = project_name
+
+        slice_index = slice_index_to_int(self.slice_index, scan=self.scan)
+        cor_in_nabu_ref = tuple(
+            [self.convert_cor_from_rel_to_abs(self.scan, cor) for cor in self.cors]
+        )
+        cor_in_nabu_ref = ",".join([str(cor) for cor in cor_in_nabu_ref])
+
+        # submit job
+        script_name = get_slurm_script_name(prefix="nabu")
+        # for now force job name
+        cluster_config["job_name"] = f"tomwer-nabu {conf_file}"
+        job = SBatchScriptJob(
+            slurm_config=cluster_config,
+            script=(
+                f"python3 -m {settings.NABU_MULTICOR_PATH} {conf_file} {slice_index} {cor_in_nabu_ref}",
+            ),
+            script_path=os.path.join(self.scan.path, "slurm_scripts", script_name),
+            clean_script=False,
+            working_directory=self.scan.working_directory,
+        )
+        future_slurm_job = submit_to_slurm_cluster(job)
+
+        callbacks = self._get_futures_slurm_callback(config_to_dump)
+        assert isinstance(
+            callbacks, tuple
+        ), f"callbacks is expected to an instance of tuple and not {type(callbacks)}"
+        for callback in callbacks:
+            future_slurm_job.add_done_callback(callback.process)
+
+        return ResultSlurmRun(
+            success=True,
+            config=config_to_dump,
+            future_slurm_jobs=(future_slurm_job,),
+            std_out=None,
+            std_err=None,
+            job_id=job.job_id,
+        )
+
+    def preprocess_config(self, config):
+        dataset_params = self.scan.get_nabu_dataset_info()
+        if "dataset" in config:
+            dataset_params.update(config["dataset"])
+        config["dataset"] = dataset_params
+
+        config["resources"] = utils.get_nabu_resources_desc(
+            scan=self.scan, workers=1, method="local"
+        )
+        # force overwrite results
+        if "output" not in config:
+            config["output"] = {}
+        config["output"].update({"overwrite_results": 1})
+
+        cfg_folder = os.path.join(
+            config["output"]["location"],
+            nabu_settings.NABU_CFG_FILE_FOLDER,
+        )
+        os.makedirs(cfg_folder, exist_ok=True)
+
+        cfg_folder = os.path.join(
+            self.nabu_config["output"]["location"],
+            nabu_settings.NABU_CFG_FILE_FOLDER,
+        )
+        os.makedirs(cfg_folder, exist_ok=True)
+
+        conf_file = os.path.join(
+            cfg_folder,
+            f"{self.scan.get_dataset_basename()}_multi_cor"
+            + nabu_settings.NABU_CONFIG_FILE_EXTENSION,
+        )
+        return config, conf_file
+
+    def run(self) -> Iterable:
+        nabu_config, conf_file = self.preprocess_config(deepcopy(self.nabu_config))
+
+        # the policy is to save nabu .cfg file at the same location as the
+        # force overwrite results
+
+        # add some tomwer metadata and save the configuration
+        # note: for now the section is ignored by nabu but shouldn't stay that way
+        with utils.TomwerInfo(nabu_config) as config_to_dump:
+            generate_nabu_configfile(
+                conf_file,
+                nabu_fullfield_default_config,
+                config=config_to_dump,
+                options_level="advanced",
+            )
+
+        results = self._process_config(
+            config_to_dump=nabu_config,
+            config_file=conf_file,
+            file_format=self.file_format,
+            info="nabu sa-axis reconstruction",
+            process_name=self.process_name,
+        )
+
+        return results