nabu 2025.1.0.dev14__py3-none-any.whl → 2025.1.0rc2__py3-none-any.whl

nabu/io/cast_volume.py

@@ -1,10 +1,12 @@
 import os
-
-from tomoscan.esrf.volume.singleframebase import VolumeSingleFrameBase
-from nabu.misc.utils import rescale_data
-from nabu.pipeline.params import files_formats
+import logging
+from glob import glob
+from shutil import rmtree
+import numpy
+from silx.io.utils import get_data
+from silx.io.url import DataUrl
 from tomoscan.volumebase import VolumeBase
-from tomoscan.scanbase import TomoScanBase
+from tomoscan.esrf.volume.singleframebase import VolumeSingleFrameBase
 from tomoscan.esrf.volume import (
     EDFVolume,
     HDF5Volume,
@@ -13,11 +15,10 @@ from tomoscan.esrf.volume import (
     TIFFVolume,
 )
 from tomoscan.io import HDF5File
-from silx.io.utils import get_data
-import numpy
-from silx.io.url import DataUrl
-from typing import Optional
-import logging
+from ..utils import first_generator_item
+from ..misc.utils import rescale_data
+from ..pipeline.params import files_formats
+from .reader import get_hdf5_file_all_virtual_sources, list_hdf5_entries
 
 _logger = logging.getLogger(__name__)
 
@@ -133,11 +134,12 @@ def cast_volume(
     output_data_type: numpy.dtype,
     data_min=None,
     data_max=None,
-    scan: Optional[TomoScanBase] = None,
+    scan=None,
     rescale_min_percentile=RESCALE_MIN_PERCENTILE,
     rescale_max_percentile=RESCALE_MAX_PERCENTILE,
     save=True,
     store=False,
+    remove_input_volume: bool = False,
 ) -> VolumeBase:
     """
     cast givent volume to output_volume of 'output_data_type' type
@@ -247,6 +249,11 @@ def cast_volume(
     except (OSError, KeyError):
         # if no metadata provided and or saved in disk or if some key are missing
         pass
+
+    if remove_input_volume:
+        _logger.info(f"Removing {input_volume.data_url.file_path()}")
+        remove_volume(input_volume, check=True)
+
     return output_volume
 
 
@@ -283,7 +290,7 @@ def clamp_and_rescale_data(
     return rescaled_data
 
 
-def find_histogram(volume: VolumeBase, scan: Optional[TomoScanBase] = None) -> Optional[DataUrl]:
+def find_histogram(volume: VolumeBase, scan=None):
     """
     Look for histogram of the provided url. If found one return the DataUrl of the nabu histogram
     """
@@ -330,7 +337,7 @@ def find_histogram(volume: VolumeBase, scan: Optional[TomoScanBase] = None) -> O
             data_path = getattr(scan, "entry/histogram/results/data", "entry/histogram/results/data")
         else:
 
-            def get_file_entries(file_path: str) -> Optional[tuple]:
+            def get_file_entries(file_path: str):
                 if os.path.exists(file_path):
                     with HDF5File(file_path, mode="r") as h5s:
                         return tuple(h5s.keys())
@@ -408,3 +415,83 @@ def _min_max_from_histo(url: DataUrl, rescale_min_percentile: int, rescale_max_p
         return _get_hst_saturations(
             hist, bins, numpy.float32(rescale_min_percentile), numpy.float32(rescale_max_percentile)
         )
+
+
+def _remove_volume_singleframe(volume, check=True):
+    volume_directory = volume.data_url.file_path()
+    if check:
+        volume_files = set(volume.browse_data_files())
+        files_names_pattern = os.path.join(volume_directory, "*." + volume.data_extension)
+        files_on_disk = set(glob(files_names_pattern))
+        # Don't check strict equality here, as some files on disk might be already removed.
+        # i.e, there should be no more files on disk than expected files in the volume
+        if not (files_on_disk.issubset(volume_files)):
+            raise RuntimeError(f"Unexpected files present in {volume_directory}: {files_on_disk - volume_files}")
+        # TODO also check for metadata file(s) ?
+    rmtree(volume_directory)
+
+
+def _remove_volume_multiframe(volume, check=True):
+    file_path = volume.data_url.file_path()
+    if check:
+        if not (os.path.isfile(file_path)):
+            raise RuntimeError(f"Expected a file: {file_path}")
+    os.remove(file_path)
+
+
+def _remove_volume_hdf5(volume, check=True):
+    file_path = volume.data_url.file_path()
+    entry = volume.data_url.data_path().lstrip("/").split("/")[0]
+
+    # Nabu HDF5 reconstructions have a folder alongside the HDF5 file, with the same prefix
+    # For example the HDF5 file "/path/to/rec.hdf5" has an associated directory "/path/to/rec"
+    associated_dir, _ = os.path.splitext(os.path.basename(file_path))
+    associated_dir_abs = os.path.join(os.path.dirname(file_path), associated_dir)
+
+    with HDF5File(file_path, "r") as f:
+        fdesc = f[entry]
+        virtual_sources = get_hdf5_file_all_virtual_sources(fdesc, return_only_filenames=True)
+
+    # TODO check if this is legitimate. Nabu reconstruction will only do one VS (for entry/reconstruction/results/data).
+    # Bliss/Lima do have multiple VS (flats/darks/projs), but we generally don't want to remove raw data ?
+    if len(virtual_sources) > 1:
+        raise ValueError("Found more than one virtual source - this looks weird. Interrupting.")
+    #
+    if len(virtual_sources) > 0:
+        h5path, virtual_source_files_paths = first_generator_item(virtual_sources[0].items())
+        if len(virtual_source_files_paths) == 1:
+            target_dir = os.path.dirname(virtual_source_files_paths[0])
+        else:
+            target_dir = os.path.commonpath(virtual_source_files_paths)
+        target_dir_abs = os.path.join(os.path.dirname(file_path), target_dir)
+        if check and (target_dir_abs != associated_dir_abs):
+            raise ValueError(
+                f"The virtual sources in {file_path}:{h5path} reference the directory {target_dir}, but expected was {associated_dir}"
+            )
+        if os.path.isdir(target_dir_abs):
+            rmtree(associated_dir_abs)
+    os.remove(file_path)
+
+
+def remove_volume(volume, check=True):
+    """
+    Remove files belonging to a volume, claim disk space.
+
+    Parameters
+    ----------
+    volume: tomoscan.esrf.volume
+        Volume object
+    check: bool, optional
+        Whether to check if the files that would be removed do not have extra other files ; interrupt the operation if so.
+
+    """
+    if isinstance(volume, (EDFVolume, JP2KVolume, TIFFVolume)):
+        _remove_volume_singleframe(volume, check=check)
+    elif isinstance(volume, MultiTIFFVolume):
+        _remove_volume_multiframe(volume, check=check)
+    elif isinstance(volume, HDF5Volume):
+        if len(list_hdf5_entries(volume.file_path)) > 1:
+            raise NotImplementedError("Removing a HDF5 volume with more than one entry is not supported")
+        _remove_volume_hdf5(volume, check=check)
+    else:
+        raise TypeError("Unknown type of volume")

nabu/io/reader.py

@@ -4,6 +4,7 @@ from math import ceil
 from multiprocessing.pool import ThreadPool
 from posixpath import sep as posix_sep, join as posix_join
 import numpy as np
+from h5py import Dataset
 from silx.io import get_data
 from silx.io.dictdump import h5todict
 from tomoscan.io import HDF5File
@@ -770,7 +771,7 @@ class NXDarksFlats:
         reduced_frames = [self._reduce_func[method](frames, axis=0) for frames in raw_frames]
         reader = getattr(self, "%s_reader" % what)
         if as_dict:
-            return {k: v for k, v in zip([s.start for s in reader._image_key_slices], reduced_frames)}  # noqa: C416
+            return {k: v for k, v in zip([s.start for s in reader._image_key_slices], reduced_frames)}
         return reduced_frames
 
     def get_raw_darks(self, force_reload=False, as_multiple_array=True):
@@ -987,6 +988,12 @@ def get_entry_from_h5_path(h5_path):
     return v[0] or v[1]
 
 
+def list_hdf5_entries(fname):
+    with HDF5File(fname, "r") as f:
+        entries = list(f.keys())
+    return entries
+
+
 def check_virtual_sources_exist(fname, data_path):
     with HDF5File(fname, "r") as f:
         if data_path not in f:
@@ -1006,6 +1013,30 @@ def check_virtual_sources_exist(fname, data_path):
     return True
 
 
+def get_hdf5_file_all_virtual_sources(file_path_or_obj, return_only_filenames=False):
+    result = []
+
+    def collect_vsources(name, obj):
+        if isinstance(obj, Dataset) and obj.is_virtual:
+            vs = obj.virtual_sources()
+            if return_only_filenames:
+                vs = [vs_.file_name for vs_ in vs]
+            result.append({name: vs})
+
+    _self_opened_file = False
+    if isinstance(file_path_or_obj, str):
+        fdesc = HDF5File(file_path_or_obj, "r")
+        _self_opened_file = True
+    else:
+        fdesc = file_path_or_obj
+
+    fdesc.visititems(collect_vsources)
+
+    if _self_opened_file:
+        fdesc.close()
+    return result
+
+
 def import_h5_to_dict(h5file, h5path, asarray=False):
     """
     Wrapper on top of silx.io.dictdump.dicttoh5 replacing "None" with None

nabu/io/tests/test_remove_volume.py

@@ -0,0 +1,152 @@
+from pathlib import Path
+from os import path, mkdir, rename
+import numpy as np
+import pytest
+from tomoscan.esrf.volume import (
+    EDFVolume,
+    HDF5Volume,
+    JP2KVolume,
+    MultiTIFFVolume,
+    TIFFVolume,
+)
+from tomoscan.esrf.volume.jp2kvolume import has_glymur
+from nabu.io.writer import merge_hdf5_files
+from nabu.io.cast_volume import remove_volume
+
+
+def test_remove_single_frame_volume(tmpdir):
+    """
+    Test volume removal for tiff, jp2 and EDF
+    """
+    # Have to use a not-too-small size because of jp2k
+    data = np.arange(10 * 40 * 50, dtype="f").reshape((10, 40, 50))
+
+    volume_classes = [EDFVolume, TIFFVolume, JP2KVolume]
+    if not (has_glymur):
+        volume_classes.pop()
+    for volume_cls in volume_classes:
+        ext = volume_cls.DEFAULT_DATA_EXTENSION
+        folder = path.join(tmpdir, f"{ext}_vol")
+        volume_basename = f"{ext}_basename"
+
+        vol_writer = volume_cls(folder=folder, volume_basename=volume_basename, overwrite=True, start_index=0)
+        vol_writer.data = data
+        vol_writer.save()
+
+        vol_reader = volume_cls(folder=folder, volume_basename=volume_basename)
+        assert path.isdir(folder), f"Expected to find a folder f{folder}"
+        remove_volume(vol_reader)
+        assert not (path.isdir(folder)), f"Expected to have removed the folder f{folder}"
+
+        vol_writer.save()
+        vol_reader = volume_cls(folder=folder, volume_basename=volume_basename)
+        Path(path.join(folder, f"unexpected.{ext}")).touch()
+        with pytest.raises(RuntimeError) as exc:
+            remove_volume(vol_reader, check=True)
+        assert "Unexpected files present" in str(exc.value), "Expected check to find extraneous files"
+
+
+def test_remove_multiframe_volume(tmpdir):
+    """
+    Test volume removal for "multiframe" formats (HDF5, tiff3D)
+    The HDF5 files considered in this test do not have virtual sources
+    """
+    data = np.arange(3 * 4 * 5, dtype="f").reshape((3, 4, 5))
+
+    for ext, volume_cls in {"h5": HDF5Volume, "tiff": MultiTIFFVolume}.items():
+        file_path = path.join(tmpdir, f"{ext}_vol.{ext}")
+
+        init_kwargs = {"file_path": file_path}
+        if ext == "h5":
+            init_kwargs["data_path"] = "entry"
+        vol_writer = volume_cls(**init_kwargs)
+        vol_writer.data = data
+        vol_writer.save()
+
+        vol_reader = volume_cls(**init_kwargs)
+        assert path.isfile(file_path), f"Expected to find a {ext} volume at {file_path}"
+        remove_volume(vol_reader)
+        assert not (path.isfile(file_path)), f"Expected to have removed f{file_path}"
+
+
+def test_remove_hdf5_multiple_entries(tmpdir):
+    data = np.arange(3 * 4 * 5, dtype="f").reshape((3, 4, 5))
+    file_path = path.join(tmpdir, "h5_vol.h5")
+    vol_writer_1 = HDF5Volume(file_path=file_path, data_path="entry0000")
+    vol_writer_1.data = data
+    vol_writer_1.save()
+    vol_writer_2 = HDF5Volume(file_path=file_path, data_path="entry0001")
+    vol_writer_2.data = data + 10
+    vol_writer_2.save()
+    vol_reader = HDF5Volume(file_path=file_path, data_path="entry0000")
+    with pytest.raises(NotImplementedError) as exc:
+        remove_volume(vol_reader, check=True)
+    assert "Removing a HDF5 volume with more than one entry is not supported" in str(
+        exc.value
+    ), "Expected an error message"
+
+
+def test_remove_nabu_hdf5_reconstruction(tmpdir):
+    """
+    Test removal of HDF5 reconstruction generated by nabu (i.e with virtual sources)
+    """
+
+    entry = "entry"
+    process_name = "reconstruction"
+
+    master_file_path = path.join(tmpdir, "sample_naburec.hdf5")
+    associated_dir = path.join(tmpdir, "sample_naburec")
+    if not (path.isdir(associated_dir)):
+        mkdir(associated_dir)
+
+    n_chunks = 5
+    local_files = []
+    for i in range(n_chunks):
+        fname = "sample_naburec_%06d.h5" % i
+        partial_rec_abspath = path.join(associated_dir, fname)
+        local_files.append(f"sample_naburec/{fname}")
+        # local_files.append(fname)
+        vol = HDF5Volume(partial_rec_abspath, data_path=f"{entry}/{process_name}")
+        vol.data = np.arange(3 * 4 * 5, dtype="f").reshape((3, 4, 5))
+        vol.save()
+
+    h5_path = f"{entry}/{process_name}/results/data"
+
+    merge_hdf5_files(
+        local_files,
+        h5_path,
+        master_file_path,
+        process_name,
+        output_entry=entry,
+        output_filemode="a",
+        processing_index=0,
+        config=None,
+        base_dir=path.dirname(associated_dir),
+        axis=0,
+        overwrite=True,
+    )
+
+    assert path.isfile(master_file_path), f"Expected to find the master file at {master_file_path}"
+    assert path.isdir(associated_dir)
+    for local_file in local_files:
+        partial_rec_file = path.join(tmpdir, local_file)
+        assert path.isfile(partial_rec_file), f"Expected to find partial file number {i} at {partial_rec_file}"
+
+    # Check that the virtual links are handled properly
+    #  sample_rec.hdf5   should reference sample_rec/sample_rec_{i}.h5
+    renamed_master_file_path = (
+        path.join(path.dirname(master_file_path), path.basename(master_file_path).split(".")[0]) + "_renamed" + ".h5"
+    )
+    rename(master_file_path, renamed_master_file_path)
+    h5_vol = HDF5Volume(file_path=renamed_master_file_path, data_path=f"{entry}/{process_name}")
+    with pytest.raises(ValueError) as exc:
+        remove_volume(h5_vol)
+    expected_error_message = f"The virtual sources in {renamed_master_file_path}:{process_name}/results/data reference the directory sample_naburec, but expected was sample_naburec_renamed"
+    assert str(exc.value) == expected_error_message
+
+    # Check removal in normal circumstances
+    rename(renamed_master_file_path, master_file_path)
+    h5_vol = HDF5Volume(file_path=master_file_path, data_path=f"{entry}/{process_name}")
+    remove_volume(h5_vol)
+    assert not (path.isfile(master_file_path)), f"Expected to find the master file at {master_file_path}"
+    assert not (path.isdir(associated_dir))

nabu/pipeline/config_validators.py

@@ -2,7 +2,7 @@
 import os
 
 path = os.path
-from ..utils import check_supported, is_writeable
+from ..utils import check_supported, deprecation_warning, is_writeable
 from .params import *  # noqa: F403
 
 """
@@ -97,12 +97,21 @@ def convert_to_bool_noerr(val):
     return res
 
 
-def name_range_checker(name, valid_names, descr, replacements=None):
+def name_range_checker(name, available_names, descr):
+    """
+    Check whether a parameter name is valid, against a list or dictionary of names.
+    """
     name = name.strip().lower()
-    if replacements is not None and name in replacements:
-        name = replacements[name]
-    valid = name in valid_names
-    assert valid, "Invalid %s '%s'. Available are %s" % (descr, name, str(valid_names))
+    valid = name in available_names
+    if isinstance(available_names, dict):
+        # handle replacements, eg. {"edge": "edges"}
+        name = available_names[name]
+        # we could use .keys() instead to be more permissive to the user
+        available_names_str = str(set(available_names.values()))
+    else:
+        # assuming list
+        available_names_str = str(available_names)
+    assert valid, "Invalid %s '%s'. Available are %s" % (descr, name, available_names_str)
     return name
 
 
@@ -338,9 +347,7 @@ def cor_validator(val):
         return val_float
     if len(val.strip()) == 0:
         return None
-    val = name_range_checker(
-        val.lower(), set(cor_methods.values()), "center of rotation estimation method", replacements=cor_methods
-    )
+    val = name_range_checker(val, cor_methods, "center of rotation estimation method")
     return val
 
 
@@ -351,9 +358,7 @@ def tilt_validator(val):
         return val_float
     if len(val.strip()) == 0:
         return None
-    val = name_range_checker(
-        val.lower(), set(tilt_methods.values()), "automatic detector tilt estimation method", replacements=tilt_methods
-    )
+    val = name_range_checker(val, tilt_methods, "automatic detector tilt estimation method")
     return val
 
 
@@ -395,53 +400,55 @@ def cor_slice_validator(val):
 
 
 @validator
-def flatfield_enabled_validator(val):
-    return name_range_checker(val, set(flatfield_modes.values()), "flatfield mode", replacements=flatfield_modes)
+def flatfield_validator(val):
+    ret = name_range_checker(val, flatfield_modes, "flatfield mode")
+    if ret in ["force-load", "force-compute"]:
+        deprecation_warning(
+            f"Using 'flatfield = {ret}' is deprecated since version 2025.1.0. Please use the parameter 'flatfield_loading_mode'",
+        )
+    return ret
+
+
+@validator
+def flatfield_loading_mode_validator(val):
+    return name_range_checker(val, flatfield_loading_mode, "flatfield mode")
 
 
 @validator
 def phase_method_validator(val):
-    return name_range_checker(
-        val, set(phase_retrieval_methods.values()), "phase retrieval method", replacements=phase_retrieval_methods
-    )
+    return name_range_checker(val, phase_retrieval_methods, "phase retrieval method")
 
 
 @validator
 def detector_distortion_correction_validator(val):
     return name_range_checker(
         val,
-        set(detector_distortion_correction_methods.values()),
+        detector_distortion_correction_methods,
         "detector_distortion_correction_methods",
-        replacements=detector_distortion_correction_methods,
     )
 
 
 @validator
 def unsharp_method_validator(val):
-    return name_range_checker(
-        val, set(unsharp_methods.values()), "unsharp mask method", replacements=phase_retrieval_methods
-    )
+    return name_range_checker(val, unsharp_methods, "unsharp mask method")
 
 
 @validator
 def padding_mode_validator(val):
-    return name_range_checker(val, set(padding_modes.values()), "padding mode", replacements=padding_modes)
+    return name_range_checker(val, padding_modes, "padding mode")
 
 
 @validator
 def reconstruction_method_validator(val):
-    return name_range_checker(
-        val, set(reconstruction_methods.values()), "reconstruction method", replacements=reconstruction_methods
-    )
+    return name_range_checker(val, reconstruction_methods, "reconstruction method")
 
 
 @validator
 def fbp_filter_name_validator(val):
     return name_range_checker(
         val,
-        set(fbp_filters.values()),
+        fbp_filters,
         "FBP filter",
-        replacements=fbp_filters,
     )
 
 
@@ -449,29 +456,24 @@ def fbp_filter_name_validator(val):
 def reconstruction_implementation_validator(val):
     return name_range_checker(
         val,
-        set(reco_implementations.values()),
+        reco_implementations,
         "Reconstruction method implementation",
-        replacements=reco_implementations,
     )
 
 
 @validator
 def optimization_algorithm_name_validator(val):
-    return name_range_checker(
-        val, set(optim_algorithms.values()), "optimization algorithm name", replacements=iterative_methods
-    )
+    return name_range_checker(val, optim_algorithms, "optimization algorithm name")
 
 
 @validator
 def output_file_format_validator(val):
-    return name_range_checker(val, set(files_formats.values()), "output file format", replacements=files_formats)
+    return name_range_checker(val, files_formats, "output file format")
 
 
 @validator
 def distribution_method_validator(val):
-    val = name_range_checker(
-        val, set(distribution_methods.values()), "workload distribution method", replacements=distribution_methods
-    )
+    val = name_range_checker(val, distribution_methods, "workload distribution method")
     # TEMP.
     if val != "local":
         raise NotImplementedError("Computation method '%s' is not implemented yet" % val)
@@ -481,9 +483,7 @@ def distribution_method_validator(val):
 
 @validator
 def sino_normalization_validator(val):
-    val = name_range_checker(
-        val, set(sino_normalizations.values()), "sinogram normalization method", replacements=sino_normalizations
-    )
+    val = name_range_checker(val, sino_normalizations, "sinogram normalization method")
     return val
 
 
@@ -491,9 +491,8 @@ def sino_normalization_validator(val):
 def sino_deringer_methods(val):
     val = name_range_checker(
         val,
-        set(rings_methods.values()),
+        rings_methods,
         "sinogram rings artefacts correction method",
-        replacements=rings_methods,
     )
     return val
 
@@ -556,7 +555,7 @@ def nonempty_string_validator(val):
 
 @validator
 def logging_validator(val):
-    return name_range_checker(val, set(log_levels.values()), "logging level", replacements=log_levels)
+    return name_range_checker(val, log_levels, "logging level")
 
 
 @validator