Merge branch 'main' into radway_69_iterative_lprec_estimator

Author: ferenc-gh

httomo/darks_flats.py

@@ -144,6 +144,38 @@ def get_separate(config: DarksFlatsFileConfig):
                 preview_config.detector_x.start : preview_config.detector_x.stop,
             ]
 
+    if (
+        darks_config.ignore
+        and not flats_config.ignore
+        and darks_config.file != flats_config.file
+    ):
+        flats = get_separate(flats_config)
+        darks = np.zeros(
+            (
+                1,
+                preview_config.detector_y.stop - preview_config.detector_y.start,
+                preview_config.detector_x.stop - preview_config.detector_x.start,
+            ),
+            dtype=flats.dtype,
+        )
+        return darks, flats
+
+    if (
+        not darks_config.ignore
+        and flats_config.ignore
+        and darks_config.file != flats_config.file
+    ):
+        darks = get_separate(darks_config)
+        flats = np.ones(
+            (
+                1,
+                preview_config.detector_y.stop - preview_config.detector_y.start,
+                preview_config.detector_x.stop - preview_config.detector_x.start,
+            ),
+            dtype=darks.dtype,
+        )
+        return darks, flats
+
     if darks_config.file != flats_config.file:
         darks = get_separate(darks_config)
         flats = get_separate(flats_config)

httomo/data/dataset_store.py

@@ -304,6 +304,10 @@ def __init__(
 
         source.finalize()
 
+    @property
+    def padding(self) -> Tuple[int, int]:
+        return self._padding
+
     @property
     def aux_data(self) -> AuxiliaryData:
         return self._aux_data

httomo/loaders/standard_tomo_loader.py

@@ -91,6 +91,10 @@ def __init__(
         self._log_info()
         weakref.finalize(self, self.finalize)
 
+    @property
+    def padding(self) -> Tuple[int, int]:
+        return self._padding
+
     @property
     def dtype(self) -> np.dtype:
         return self._data.dtype

httomo/runner/block_split.py

@@ -1,4 +1,3 @@
-import math
 from httomo.runner.dataset import DataSetBlock
 from httomo.runner.dataset_store_interfaces import DataSetSource
 
@@ -24,9 +23,31 @@ class BlockSplitter:
 
     def __init__(self, source: DataSetSource, max_slices: int):
         self._source = source
-        self._chunk_size = source.chunk_shape[source.slicing_dim]
+        self._chunk_size = (
+            source.chunk_shape[source.slicing_dim]
+            + source.padding[0]
+            + source.padding[1]
+        )
         self._max_slices = int(min(max_slices, self._chunk_size))
-        self._num_blocks = math.ceil(self._chunk_size / self._max_slices)
+
+        core_slices = self._max_slices - source.padding[0] - source.padding[1]
+        step = core_slices
+        num_blocks = 0
+        start_of_next_block = 0
+        while (
+            start_of_next_block + source.padding[0] + core_slices + source.padding[1]
+            < self._chunk_size
+        ):
+            num_blocks += 1
+            start_of_next_block += step
+
+        if (
+            start_of_next_block + source.padding[0] + core_slices + source.padding[1]
+            >= self._chunk_size
+        ):
+            num_blocks += 1
+
+        self._num_blocks = num_blocks
         self._current = 0
         assert self._source.slicing_dim in [
             0,
@@ -35,24 +56,23 @@ def __init__(self, source: DataSetSource, max_slices: int):
 
     @property
     def slices_per_block(self) -> int:
-        return self._max_slices
+        return self._max_slices - self._source.padding[0] - self._source.padding[1]
 
     def __len__(self):
         return self._num_blocks
 
     def __getitem__(self, idx: int) -> DataSetBlock:
         start = idx * self.slices_per_block
-        if start >= self._chunk_size:
+        if (
+            start
+            >= self._chunk_size - self._source.padding[0] - self._source.padding[1]
+        ):
             raise IndexError("Index out of bounds")
-        len = min(self.slices_per_block, self._chunk_size - start)
+        len = min(
+            self.slices_per_block,
+            self._chunk_size
+            - self._source.padding[0]
+            - self._source.padding[1]
+            - start,
+        )
         return self._source.read_block(start, len)
-
-    def __iter__(self):
-        return self
-
-    def __next__(self) -> DataSetBlock:
-        if self._current >= len(self):
-            raise StopIteration
-        v = self[self._current]
-        self._current += 1
-        return v

httomo/runner/dataset_store_interfaces.py

@@ -96,6 +96,13 @@ def finalize(self):
         to give implementations a chance to close files, free memory, etc."""
         ...  # pragma: no cover
 
+    @property
+    def padding(self) -> Tuple[int, int]:
+        """
+        Padding present in `DataSetBlock`'s produced by the `DataSetSource`.
+        """
+        ...  # pragma: no cover
+
 
 class DataSetSink(Protocol):
     @property

httomo/runner/task_runner.py

@@ -135,10 +135,6 @@ def _execute_section(self, section: Section, section_index: int = 0):
         slicing_dim_section: Literal[0, 1] = _get_slicing_dim(section.pattern) - 1  # type: ignore
         self.determine_max_slices(section, slicing_dim_section, self.source.aux_data.get_angles())
 
-        # Account for potential padding in number of max slices
-        padding = determine_section_padding(section)
-        section.max_slices -= padding[0] + padding[1]
-
         self._log_pipeline(
             f"Maximum amount of slices is {section.max_slices} for section {section_index}",
             level=logging.DEBUG,
@@ -260,9 +256,9 @@ def _execute_section_block(
     ) -> DataSetBlock:
         if_previous_block_is_on_gpu = False
         convert_gpu_block_to_cpu = False
+        if_current_block_is_on_gpu = False
 
         for ind, method in enumerate(section):
-            if_current_block_is_on_gpu = False
             if method.implementation == "gpu_cupy":
                 if_current_block_is_on_gpu = True
             if method.method_name == "calculate_stats" and if_previous_block_is_on_gpu:
@@ -410,7 +406,9 @@ def determine_max_slices(self, section: Section, slicing_dim: int, angles: np.nd
         assert len(section) > 0, "Section should contain at least 1 method"
 
         data_shape = self.source.chunk_shape
-        max_slices = data_shape[slicing_dim]
+        max_slices = (
+            data_shape[slicing_dim] + self.source.padding[0] + self.source.padding[1]
+        )
         # loop over all methods in section
         has_gpu = False
         for idx, m in enumerate(section):
@@ -463,6 +461,20 @@ def determine_max_slices(self, section: Section, slicing_dim: int, angles: np.nd
             max_slices_methods[idx] = min(max_slices, slices_estimated)
             non_slice_dims_shape = output_dims
 
+        if (
+            min(max_slices_methods)
+            < 1 + self.source.padding[0] + self.source.padding[1]
+        ):
+            padded_method = next(
+                method.method_name for method in section.methods if method.padding
+            )
+            err_str = (
+                "Unable to process data due to GPU memory limitations.\n"
+                f"Please remove method '{padded_method}' from the pipeline, or run on a "
+                "machine with more GPU memory."
+            )
+            raise ValueError(err_str)
+
         section.max_slices = min(max_slices_methods)
 
     def _pass_min_block_length_to_intermediate_data_wrapper(self, section: Section):

httomo/ui_layer.py

@@ -9,7 +9,7 @@
 import h5py
 from mpi4py.MPI import Comm
 
-from httomo.preview import PreviewConfig, PreviewDimConfig
+from httomo.preview import PreviewConfig
 from httomo.runner.method_wrapper import MethodWrapper
 from httomo.runner.pipeline import Pipeline
 
@@ -85,7 +85,6 @@ def build_pipeline(self) -> Pipeline:
             self._append_methods_list(
                 i, task_conf, methods_list, parameters, method_id_map
             )
-        fix_preview_y_if_smaller_than_padding(loader, methods_list)
         return Pipeline(loader=loader, methods=methods_list)
 
     def _append_methods_list(
@@ -169,27 +168,6 @@ def _setup_loader(self) -> LoaderInterface:
         return loader
 
 
-def fix_preview_y_if_smaller_than_padding(
-    loader: LoaderInterface, methods_list: List[MethodWrapper]
-) -> None:
-    vertical_preview_length = (
-        loader.preview.detector_y.stop - loader.preview.detector_y.start
-    ) // loader.comm.size
-    max_pad_value = 0
-    for _, m in enumerate(methods_list):
-        if m.padding:
-            max_pad_value = max(sum(m.calculate_padding()), max_pad_value)
-    if max_pad_value >= vertical_preview_length:
-        loader.preview = PreviewConfig(
-            angles=loader.preview.angles,
-            detector_y=PreviewDimConfig(
-                start=loader.preview.detector_y.start - max_pad_value // 2,
-                stop=loader.preview.detector_y.stop + max_pad_value // 2,
-            ),
-            detector_x=loader.preview.detector_x,
-        )
-
-
 def get_valid_ref_str(parameters: Dict[str, Any]) -> Dict[str, str]:
     """Find valid reference strings inside dictionary
 

tests/conftest.py

@@ -535,7 +535,7 @@ def compare_tif(files_list_to_compare: list, file_path_to_references: list):
             Image.open(tif_files_references[index])
         )
         res_norm = np.linalg.norm(res_images.flatten())
-        assert res_norm < 1e-6
+        assert res_norm < 1e-3
 
 
 def change_value_parameters_method_pipeline(