Merge pull request #446 from DiamondLightSource/integrate-padding-int…

…o-runner

Integrate padding into runner

httomo/loaders/standard_tomo_loader.py

@@ -390,7 +390,7 @@ def __init__(
         self.angles = angles
         self.preview = preview
 
-    def make_data_source(self) -> DataSetSource:
+    def make_data_source(self, padding: Tuple[int, int] = (0, 0)) -> DataSetSource:
         assert self.pattern in [Pattern.sinogram, Pattern.projection]
         loader = StandardTomoLoader(
             in_file=self.in_file,
@@ -402,6 +402,7 @@ def make_data_source(self) -> DataSetSource:
             preview_config=self.preview,
             slicing_dim=1 if self.pattern == Pattern.sinogram else 0,
             comm=self.comm,
+            padding=padding,
         )
         (self._angles_total, self._detector_y, self._detector_x) = loader.global_shape
         return loader

httomo/method_wrappers/generic.py

@@ -1,3 +1,5 @@
+from importlib import import_module
+
 import httomo.globals
 from httomo.block_interfaces import T, Block
 from httomo.runner.gpu_utils import get_gpu_id, gpumem_cleanup
@@ -83,7 +85,6 @@ def __init__(
         self._comm = comm
         self._module_path = module_path
         self._method_name = method_name
-        from importlib import import_module
 
         self._module = import_module(module_path)
         self._method: Callable = getattr(self._module, method_name)

httomo/method_wrappers/save_intermediate.py

@@ -70,13 +70,15 @@ def execute(self, block: T) -> T:
         # - we return just the input as it is
         self._gpu_time_info = GpuTimeInfo()
         with catchtime() as t:
-            data = block.data if block.is_cpu else xp.asnumpy(block.data)
+            data = (
+                block.data_unpadded if block.is_cpu else xp.asnumpy(block.data_unpadded)
+            )
         self._gpu_time_info.device2host += t.elapsed
 
         self._method(
             data,
             global_shape=block.global_shape,
-            global_index=block.global_index,
+            global_index=block.global_index_unpadded,
             slicing_dim=block.slicing_dim,
             file=self._file,
             frames_per_chunk=httomo.globals.FRAMES_PER_CHUNK,

httomo/methods_database/packages/external/httomolibgpu/httomolibgpu.yaml

@@ -5,7 +5,7 @@ misc:
       output_dims_change: False
       implementation: gpu_cupy
       save_result_default: False
-      padding: False
+      padding: True
       memory_gpu:
         multiplier: 2.1
         method: direct
@@ -14,7 +14,7 @@ misc:
       output_dims_change: False
       implementation: gpu_cupy
       save_result_default: False
-      padding: False
+      padding: True
       memory_gpu:
         multiplier: 2.1
         method: direct

httomo/methods_database/packages/external/httomolibgpu/supporting_funcs/misc/corr.py

@@ -0,0 +1,17 @@
+from typing import Tuple
+
+
+__all__ = [
+    "_calc_padding_remove_outlier",
+    "_calc_padding_median_filter",
+]
+
+
+def _calc_padding_remove_outlier(**kwargs) -> Tuple[int, int]:
+    kernel_size = kwargs["kernel_size"]
+    return (kernel_size // 2, kernel_size // 2)
+
+
+def _calc_padding_median_filter(**kwargs) -> Tuple[int, int]:
+    kernel_size = kwargs["kernel_size"]
+    return (kernel_size // 2, kernel_size // 2)

httomo/methods_database/query.py

@@ -1,3 +1,4 @@
+from importlib import import_module
 from types import ModuleType
 from typing import Callable, List, Literal, Optional, Tuple
 from pathlib import Path
@@ -140,7 +141,6 @@ def calculate_padding(self, **kwargs) -> Tuple[int, int]:
         return module_pad(**kwargs)
 
     def _import_supporting_funcs_module(self) -> ModuleType:
-        from importlib import import_module
 
         module_mem_path = "httomo.methods_database.packages.external."
         path = self.module_path.split(".")

httomo/runner/loader.py

@@ -1,4 +1,4 @@
-from typing import Protocol
+from typing import Protocol, Tuple
 
 from httomo.runner.dataset_store_interfaces import DataSetSource
 from httomo.utils import Pattern
@@ -13,8 +13,8 @@ class LoaderInterface(Protocol):
     method_name: str
     package_name: str = "httomo"
 
-    def make_data_source(self) -> DataSetSource:
-        """Create a dataset source that can produce blocks of data from the file.
+    def make_data_source(self, padding: Tuple[int, int]) -> DataSetSource:
+        """Create a dataset source that can produce padded blocks of data from the file.
 
         This will be called after the patterns and sections have been determined,
         just before the execution of the first section starts."""

httomo/runner/task_runner.py

@@ -22,7 +22,7 @@
 from httomo.runner.gpu_utils import get_available_gpu_memory, gpumem_cleanup
 from httomo.runner.monitoring_interface import MonitoringInterface
 from httomo.runner.pipeline import Pipeline
-from httomo.runner.section import Section, sectionize
+from httomo.runner.section import Section, determine_section_padding, sectionize
 from httomo.utils import (
     Pattern,
     _get_slicing_dim,
@@ -100,6 +100,10 @@ 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)
 
+        # 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,
@@ -218,7 +222,8 @@ def _load_datasets(self):
             self.pipeline.loader.pattern,
             self.pipeline.loader.method_name,
         )
-        self.source = self.pipeline.loader.make_data_source()
+        loader_padding = determine_section_padding(self._sections[0])
+        self.source = self.pipeline.loader.make_data_source(padding=loader_padding)
         self._log_task_end(
             "loader",
             start_time,

httomo/sweep_runner/param_sweep_runner.py

@@ -128,7 +128,7 @@ def prepare(self):
         """
         Load single block containing small number of sinogram slices in input data
         """
-        source = self._pipeline.loader.make_data_source()
+        source = self._pipeline.loader.make_data_source(padding=(0, 0))
 
         SINO_SLICING_DIM = 1
         no_of_middle_slices = source.global_shape[SINO_SLICING_DIM]

tests/method_wrappers/test_data_reducer.py

@@ -16,7 +16,9 @@ class FakeModule:
         def data_reducer(x):
             return 2 * x
 
-    mocker.patch("importlib.import_module", return_value=FakeModule)
+    mocker.patch(
+        "httomo.method_wrappers.generic.import_module", return_value=FakeModule
+    )
     wrp = make_method_wrapper(
         make_mock_repo(mocker),
         "mocked_module_path.morph",

tests/method_wrappers/test_dezinging.py

@@ -16,7 +16,9 @@ class FakeModule:
         def remove_outlier(x, axis="auto"):
             return 2 * x
 
-    mocker.patch("importlib.import_module", return_value=FakeModule)
+    mocker.patch(
+        "httomo.method_wrappers.generic.import_module", return_value=FakeModule
+    )
     wrp = make_method_wrapper(
         make_mock_repo(mocker),
         "mocked_module_path.prep",