nabu 2024.1.0rc4__py3-none-any.whl → 2024.1.1__py3-none-any.whl

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/nabu/stitching/slurm_utils.py

@@ -0,0 +1,204 @@
+import os
+import copy
+from typing import Optional, Union
+import numpy
+from silx.io.url import DataUrl
+from tomoscan.tomoobject import TomoObject
+from tomoscan.esrf.scan.nxtomoscan import NXtomoScan
+from tomoscan.esrf import EDFTomoScan
+from tomoscan.esrf.volume import HDF5Volume, MultiTIFFVolume
+from tomoscan.esrf.volume.singleframebase import VolumeSingleFrameBase
+from ..app.bootstrap_stitching import _SECTIONS_COMMENTS
+from ..pipeline.config import generate_nabu_configfile
+from .config import (
+    StitchingConfiguration,
+    get_default_stitching_config,
+    PreProcessedZStitchingConfiguration,
+    PostProcessedZStitchingConfiguration,
+    SLURM_SECTION,
+)
+
+try:
+    from sluurp.job import SBatchScriptJob
+except ImportError:
+    has_sluurp = False
+else:
+    has_sluurp = True
+
+
+def split_stitching_configuration_to_slurm_job(
+    configuration: StitchingConfiguration, yield_configuration: bool = False
+):
+    """
+    generator to split a StitchingConfiguration into several SBatchScriptJob.
+
+    This will handle:
+       * division into several jobs according to `slices` and `n_job`
+       * creation of SBatchScriptJob handling slurm configuration and command to be launched
+
+    :param StitchingConfiguration configuration: configuration of the stitching to launch (into several jobs)
+    :param bool yield_configuration: if True then yield (SBatchScriptJob, StitchingConfiguration) else yield only SBatchScriptJob
+    """
+    if not isinstance(configuration, StitchingConfiguration):
+        raise TypeError(
+            f"configuration is expected to be an instance of {StitchingConfiguration}. {type(configuration)} provided."
+        )
+    if not has_sluurp:
+        raise ImportError("sluurp not install. Please install it to distribute stitching on slurm (pip install sluurm)")
+
+    slurm_configuration = configuration.slurm_config
+    n_jobs = slurm_configuration.n_jobs
+    stitching_type = configuration.stitching_type
+    # cleqn slurm configurqtion
+    slurm_configuration = slurm_configuration.to_dict()
+    clean_script = slurm_configuration.pop("clean_scripts", False)
+    slurm_configuration.pop("n_jobs", None)
+    # for now other_options is not handled
+    slurm_configuration.pop("other_options", None)
+
+    if "memory" in slurm_configuration and isinstance(slurm_configuration["memory"], str):
+        memory = slurm_configuration["memory"].lower().replace(" ", "")
+        memory = memory.rstrip("b").rstrip("g")
+        slurm_configuration["memory"] = memory
+    # handle slices if None
+    configuration.settle_inputs()
+
+    slice_sub_parts = split_slices(slices=configuration.slices, n_parts=n_jobs)
+    if isinstance(configuration, PreProcessedZStitchingConfiguration):
+        stitch_prefix = os.path.basename(os.path.splitext(configuration.output_file_path)[0])
+        configuration.output_file_path = os.path.abspath(configuration.output_file_path)
+
+    elif isinstance(configuration, PostProcessedZStitchingConfiguration):
+        stitch_prefix = os.path.basename(os.path.splitext(configuration.output_volume.file_path)[0])
+        configuration.output_volume.file_path = os.path.abspath(configuration.output_volume.file_path)
+    else:
+        raise TypeError(f"{type(configuration)} not handled")
+
+    for i_sub_part, slice_sub_part in enumerate(slice_sub_parts):
+        sub_configuration = copy.deepcopy(configuration)
+        # update slice
+        sub_configuration.slices = slice_sub_part
+        # remove slurm configuration because once on the partition we run it manually
+        sub_configuration.slurm_config = None
+
+        if isinstance(sub_configuration, PreProcessedZStitchingConfiguration):
+            original_output_file_path, file_extension = os.path.splitext(sub_configuration.output_file_path)
+            sub_configuration.output_file_path = os.path.join(
+                original_output_file_path,
+                os.path.basename(original_output_file_path) + f"_part_{i_sub_part}" + file_extension,
+            )
+            output_obj = NXtomoScan(
+                scan=sub_configuration.output_file_path,
+                entry=sub_configuration.output_data_path,
+            )
+        elif isinstance(sub_configuration, PostProcessedZStitchingConfiguration):
+            if isinstance(sub_configuration.output_volume, (HDF5Volume, MultiTIFFVolume)):
+                original_output_file_path, file_extension = os.path.splitext(sub_configuration.output_volume.file_path)
+                sub_configuration.output_volume.file_path = os.path.join(
+                    original_output_file_path,
+                    os.path.basename(original_output_file_path) + f"_part_{i_sub_part}" + file_extension,
+                )
+            elif isinstance(sub_configuration.output_volume, VolumeSingleFrameBase):
+                url = sub_configuration.output_volume.url
+                original_output_folder = url.file_path()
+                sub_part_url = DataUrl(
+                    file_path=os.path.join(
+                        original_output_folder,
+                        os.path.basename(original_output_folder) + f"_part_{i_sub_part}",
+                    ),
+                    data_path=url.data_path(),
+                    scheme=url.scheme(),
+                    data_slice=url.data_slice(),
+                )
+                sub_configuration.output_volume.url = sub_part_url
+            output_obj = sub_configuration.output_volume
+        else:
+            raise TypeError(f"{type(sub_configuration)} not handled")
+
+        working_directory = get_working_directory(output_obj)
+        if working_directory is not None:
+            script_dir = working_directory
+        else:
+            script_dir = "./"
+
+        # save sub part nabu configuration file
+        slurm_script_name = f"{stitch_prefix}_part_{i_sub_part}.sh"
+        nabu_script_name = f"{stitch_prefix}_part_{i_sub_part}.conf"
+
+        os.makedirs(script_dir, exist_ok=True)
+        default_config = get_default_stitching_config(stitching_type)
+        default_config.pop(SLURM_SECTION, None)
+        generate_nabu_configfile(
+            fname=os.path.join(script_dir, nabu_script_name),
+            default_config=default_config,
+            comments=True,
+            sections_comments=_SECTIONS_COMMENTS,
+            prefilled_values=sub_configuration.to_dict(),
+            options_level="advanced",
+        )
+
+        command = f"python3 -m nabu.app.stitching {os.path.join(script_dir, nabu_script_name)}"
+        script = (command,)
+
+        job = SBatchScriptJob(
+            slurm_config=slurm_configuration,
+            script=script,
+            script_path=os.path.join(script_dir, slurm_script_name),
+            clean_script=clean_script,
+            working_directory=working_directory,
+        )
+        job.overwrite = True
+        if yield_configuration:
+            yield job, sub_configuration
+        else:
+            yield job
+
+
+def split_slices(slices: Union[slice, tuple], n_parts: int):
+    if not isinstance(n_parts, int):
+        raise TypeError(f"n_parts should be an int. {type(n_parts)} provided")
+    if isinstance(slices, slice):
+        assert isinstance(slices.start, int), "slices.start must be an integer"
+        assert isinstance(slices.stop, int), "slices.stop must be an integer"
+        start = stop = slices.start
+        steps_size = int(numpy.ceil((slices.stop - slices.start) / n_parts))
+        while stop < slices.stop:
+            stop = min(start + steps_size, slices.stop)
+            yield slice(start, stop, slices.step)
+            start = stop
+    elif isinstance(slices, (tuple, list)):
+        start = stop = 0
+        steps_size = int(numpy.ceil(len(slices) / n_parts))
+        while stop < len(slices):
+            stop = min(start + steps_size, len(slices))
+            yield slices[start:stop]
+            start = stop
+    else:
+        raise TypeError(f"slices type ({type(slices)}) is not handled. Must be a slice or an Iterable")
+
+
+def get_working_directory(obj: TomoObject) -> Optional[str]:
+    """
+    return working directory for a specific TomoObject
+    """
+    if not isinstance(obj, TomoObject):
+        raise TypeError(f"obj should be an instance of {TomoObject}. {type(obj)} provided")
+    if isinstance(obj, (HDF5Volume, MultiTIFFVolume)):
+        if obj.file_path is None:
+            return None
+        else:
+            return os.path.abspath(os.path.dirname(obj.file_path))
+    elif isinstance(obj, VolumeSingleFrameBase):
+        if obj.data_url is not None:
+            return os.path.abspath(obj.data_url.file_path())
+        else:
+            return None
+    elif isinstance(obj, EDFTomoScan):
+        return obj.path
+    elif isinstance(obj, NXtomoScan):
+        if obj.master_file is None:
+            return None
+        else:
+            return os.path.abspath(os.path.dirname(obj.master_file))
+    else:
+        raise RuntimeError(f"obj type not handled ({type(obj)})")

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/nabu/stitching/tests/test_alignment.py

@@ -0,0 +1,99 @@
+import numpy
+import pytest
+from nabu.stitching.alignment import align_horizontally, PaddedRawData
+from nabu.testutils import get_data
+
+
+def test_alignment_axis_2():
+    """
+    test 'align_horizontally' function
+    """
+    dataset = get_data("chelsea.npz")["data"]  # shape is (300, 451)
+
+    # test if new_width < current_width: should raise an error
+    with pytest.raises(ValueError):
+        align_horizontally(dataset, alignment="center", new_width=10)
+
+    # test some use cases
+    res = align_horizontally(
+        dataset,
+        alignment="center",
+        new_width=600,
+        pad_mode="mean",
+    )
+    assert res.shape == (300, 600)
+    numpy.testing.assert_array_almost_equal(res[:, 74:-75], dataset)
+
+    res = align_horizontally(
+        dataset,
+        alignment="left",
+        new_width=600,
+        pad_mode="median",
+    )
+    assert res.shape == (300, 600)
+    numpy.testing.assert_array_almost_equal(res[:, :451], dataset)
+
+    res = align_horizontally(
+        dataset,
+        alignment="right",
+        new_width=600,
+        pad_mode="reflect",
+    )
+    assert res.shape == (300, 600)
+    numpy.testing.assert_array_almost_equal(res[:, -451:], dataset)
+
+
+def test_PaddedRawData():
+    """
+    test PaddedVolume class
+    """
+    data = numpy.linspace(
+        start=0,
+        stop=20 * 6 * 3,
+        dtype=numpy.int64,
+        num=20 * 6 * 3,
+    )
+    data = data.reshape((3, 6, 20))
+
+    padded_volume = PaddedRawData(data=data, axis_1_pad_width=(4, 1))
+
+    assert padded_volume.shape == (3, 6 + 4 + 1, 20)
+
+    numpy.testing.assert_array_equal(
+        padded_volume[:, 0, :],
+        numpy.zeros(shape=(3, 1, 20), dtype=numpy.int64),
+    )
+    numpy.testing.assert_array_equal(
+        padded_volume[:, 3, :],
+        numpy.zeros(shape=(3, 1, 20), dtype=numpy.int64),
+    )
+    numpy.testing.assert_array_equal(
+        padded_volume[:, 10, :],
+        numpy.zeros(shape=(3, 1, 20), dtype=numpy.int64),
+    )
+    assert padded_volume[:, 3, :].shape == (3, 1, 20)
+    numpy.testing.assert_array_equal(
+        padded_volume[:, 4, :],
+        data[:, 0:1, :],  # TODO: have a look, return a 3D array when a 2D expected...
+    )
+
+    with pytest.raises(ValueError):
+        padded_volume[:, 40, :]
+    with pytest.raises(ValueError):
+        padded_volume[:, 5:1, :]
+
+    arrays = (
+        numpy.zeros(shape=(3, 4, 20), dtype=numpy.int64),
+        data,
+        numpy.zeros(shape=(3, 1, 20), dtype=numpy.int64),
+    )
+    expected_volume = numpy.hstack(
+        arrays,
+    )
+    assert padded_volume[:, :, :].shape == padded_volume.shape
+    assert expected_volume.shape == padded_volume.shape
+
+    numpy.testing.assert_array_equal(
+        padded_volume[:, :, :],
+        expected_volume,
+    )

build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/build/lib/nabu/stitching/tests/test_config.py

@@ -0,0 +1,226 @@
+# coding: utf-8
+# /*##########################################################################
+#
+# Copyright (c) 2016-2017 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.
+#
+# ###########################################################################*/
+
+__authors__ = ["H. Payno"]
+__license__ = "MIT"
+__date__ = "10/05/2022"
+
+
+import os
+from tempfile import TemporaryDirectory
+import pytest
+from nabu.pipeline.config import (
+    generate_nabu_configfile,
+    _options_levels,
+    parse_nabu_config_file,
+)
+from nabu.stitching.overlap import OverlapStitchingStrategy
+from nabu.stitching import config as stiching_config
+
+
+_stitching_types = list(stiching_config.StitchingType.values())
+_stitching_types.append(None)
+
+
+def nabu_config_to_dict(nabu_config):
+    res = {}
+    for section, section_content in nabu_config.items():
+        res[section] = {}
+        for key, values in section_content.items():
+            res[section][key] = values["default"]
+    return res
+
+
+@pytest.mark.parametrize("stitching_type", _stitching_types)
+@pytest.mark.parametrize("option_level", _options_levels.keys())
+def test_stitching_config(stitching_type, option_level):
+    """
+    insure get_default_stitching_config is returning a dict
+    and is coherent with the configuration classes
+    """
+    with TemporaryDirectory() as output_dir:
+        nabu_dict = stiching_config.get_default_stitching_config(stitching_type)
+        config = nabu_config_to_dict(nabu_dict)
+        assert isinstance(config, dict)
+
+        assert "stitching" in config
+        assert "type" in config["stitching"]
+        stitching_type = stiching_config.StitchingType.from_value(config["stitching"]["type"])
+        if stitching_type is stiching_config.StitchingType.Z_POSTPROC:
+            assert isinstance(
+                stiching_config.dict_to_config_obj(config),
+                stiching_config.PostProcessedZStitchingConfiguration,
+            )
+        elif stitching_type is stiching_config.StitchingType.Z_PREPROC:
+            assert isinstance(
+                stiching_config.dict_to_config_obj(config),
+                stiching_config.PreProcessedZStitchingConfiguration,
+            )
+        else:
+            raise ValueError("not handled")
+
+        # dump configuration to file
+        output_file = os.path.join(output_dir, "config.conf")
+        generate_nabu_configfile(
+            fname=output_file,
+            default_config=nabu_dict,
+            comments=True,
+            sections_comments=stiching_config.SECTIONS_COMMENTS,
+            options_level=option_level,
+            prefilled_values={},
+        )
+
+        # load configuration from file
+        loaded_config = parse_nabu_config_file(output_file)
+        config_class_instance = stiching_config.dict_to_config_obj(loaded_config)
+        if stitching_type is stiching_config.StitchingType.Z_POSTPROC:
+            assert isinstance(
+                config_class_instance,
+                stiching_config.PostProcessedZStitchingConfiguration,
+            )
+        elif stitching_type is stiching_config.StitchingType.Z_PREPROC:
+            assert isinstance(
+                config_class_instance,
+                stiching_config.PreProcessedZStitchingConfiguration,
+            )
+
+        assert isinstance(config_class_instance.to_dict(), dict)
+
+
+@pytest.mark.parametrize("stitching_strategy", OverlapStitchingStrategy.values())
+@pytest.mark.parametrize("overwrite_results", (True, "False", 0, "1"))
+@pytest.mark.parametrize(
+    "axis_shifts",
+    (
+        "",
+        None,
+        "None",
+        "",
+        "skimage",
+        "nabu-fft",
+    ),
+)
+@pytest.mark.parametrize("axis_shifts_params", ("", {}, "window_size=200"))
+@pytest.mark.parametrize(
+    "slice_for_correlation",
+    (
+        "middle",
+        "3",
+    ),
+)
+@pytest.mark.parametrize("slices", ("middle", "0:26:2"))
+@pytest.mark.parametrize(
+    "input_scans",
+    (
+        "",
+        "hdf5:scan:/data/scan.hdf5?path=entry; hdf5:scan:/data/scan.hdf5?path=entry1",
+    ),
+)
+@pytest.mark.parametrize(
+    "slurm_config",
+    (
+        {
+            stiching_config.SLURM_MODULES_TO_LOADS: "tomotools",
+            stiching_config.SLURM_PREPROCESSING_COMMAND: "",
+            stiching_config.SLURM_CLEAN_SCRIPTS: True,
+            stiching_config.SLURM_MEM: 56,
+            stiching_config.SLURM_N_JOBS: 5,
+            stiching_config.SLURM_PARTITION: "my_partition",
+        },
+    ),
+)
+def test_PreProcessedZStitchingConfiguration(
+    stitching_strategy,
+    overwrite_results,
+    axis_shifts,
+    axis_shifts_params,
+    input_scans,
+    slice_for_correlation,
+    slices,
+    slurm_config,
+):
+    """
+    make sure configuration works well for PreProcessedZStitchingConfiguration
+    """
+    pre_process_config = stiching_config.PreProcessedZStitchingConfiguration.from_dict(
+        {
+            stiching_config.STITCHING_SECTION: {
+                stiching_config.CROSS_CORRELATION_SLICE_FIELD: slice_for_correlation,
+                stiching_config.AXIS_0_POS_PX: axis_shifts,
+                stiching_config.AXIS_1_POS_PX: axis_shifts,
+                stiching_config.AXIS_2_POS_PX: axis_shifts,
+                stiching_config.AXIS_0_PARAMS: axis_shifts_params,
+                stiching_config.AXIS_1_PARAMS: axis_shifts_params,
+                stiching_config.AXIS_2_PARAMS: axis_shifts_params,
+                stiching_config.STITCHING_STRATEGY_FIELD: stitching_strategy,
+            },
+            stiching_config.INPUTS_SECTION: {
+                stiching_config.INPUT_DATASETS_FIELD: input_scans,
+                stiching_config.STITCHING_SLICES: slices,
+            },
+            stiching_config.OUTPUT_SECTION: {
+                stiching_config.OVERWRITE_RESULTS_FIELD: overwrite_results,
+            },
+            stiching_config.Z_PRE_PROC_SECTION: {
+                stiching_config.DATA_FILE_FIELD: "my_file.nx",
+                stiching_config.DATA_PATH_FIELD: "entry",
+                stiching_config.NEXUS_VERSION_FIELD: None,
+            },
+            stiching_config.SLURM_SECTION: slurm_config,
+            stiching_config.NORMALIZATION_BY_SAMPLE_SECTION: {
+                stiching_config.NORMALIZATION_BY_SAMPLE_MARGIN: 1,
+                stiching_config.NORMALIZATION_BY_SAMPLE_SIDE: "right",
+                stiching_config.NORMALIZATION_BY_SAMPLE_ACTIVE_FIELD: True,
+                stiching_config.NORMALIZATION_BY_SAMPLE_METHOD: "mean",
+                stiching_config.NORMALIZATION_BY_SAMPLE_WIDTH: 31,
+            },
+        },
+    )
+
+    from_dict = stiching_config.PreProcessedZStitchingConfiguration.from_dict(pre_process_config.to_dict())
+    # workaround for scans because a new object is created each time
+    pre_process_config.settle_inputs
+    assert len(from_dict.input_scans) == len(pre_process_config.input_scans)
+    from_dict.input_scans = None
+    pre_process_config.input_scans = None
+    assert pre_process_config == from_dict
+
+
+def test_PostProcessedZStitchingConfiguration():
+    """
+    make sure configuration works well for PostProcessedZStitchingConfiguration
+    """
+    pass
+
+
+def test_description_dict():
+    """
+    make sure the description dict (used for generating the file) is working and generates a dict
+    """
+    assert isinstance(stiching_config.PreProcessedZStitchingConfiguration.get_description_dict(), dict)
+    assert isinstance(
+        stiching_config.PostProcessedZStitchingConfiguration.get_description_dict(),
+        dict,
+    )