Enable rotation wrapper to work with padded blocks

httomo/method_wrappers/rotation.py

@@ -147,15 +147,27 @@ def _run_method(self, block: T, args: Dict[str, Any]) -> T:
             # append to internal sinogram, until we have the last block
             if self.sino is None:
                 self.sino = np.empty(
-                    (block.chunk_shape[0], block.chunk_shape[2]), dtype=np.float32
+                    (block.chunk_shape_unpadded[0], block.chunk_shape_unpadded[2]),
+                    dtype=np.float32,
                 )
-            data = block.data[:, slice_for_cor, :]
+
+            # Extract core of block, in case padding slices are present
+            core_angles_start = 0
+            core_angles_stop = block.shape_unpadded[0]
+            if block.is_padded:
+                core_angles_start = block.padding[0]
+                core_angles_stop = core_angles_start + block.shape_unpadded[0]
+
+            data = block.data[core_angles_start:core_angles_stop, slice_for_cor, :]
+
             if block.is_gpu:
                 with catchtime() as t:
                     data = xp.asnumpy(data)
                 self._gpu_time_info.device2host += t.elapsed
             self.sino[
-                block.chunk_index[0] : block.chunk_index[0] + block.shape[0], :
+                block.chunk_index_unpadded[0] : block.chunk_index_unpadded[0]
+                + block.shape_unpadded[0],
+                :,
             ] = data
 
             if not block.is_last_in_chunk:  # exit if we didn't process all blocks yet

tests/method_wrappers/test_rotation.py

@@ -1,11 +1,11 @@
-from typing import Union
+from typing import List, Tuple, Union
 from unittest.mock import MagicMock
 import numpy as np
 from httomo.method_wrappers import make_method_wrapper
 from httomo.method_wrappers.rotation import RotationWrapper
 from httomo.runner.auxiliary_data import AuxiliaryData
 from httomo.runner.dataset import DataSetBlock
-from httomo.utils import Pattern, gpu_enabled, xp
+from httomo.utils import Pattern, gpu_enabled, make_3d_shape_from_shape, xp
 from ..testing_utils import make_mock_repo
 
 
@@ -29,7 +29,12 @@ def rotation_tester(data, ind=None):
         )
 
 
-def test_rotation_accumulates_blocks(mocker: MockerFixture):
+@pytest.mark.parametrize(
+    "padding",
+    [(0, 0), (2, 3)],
+    ids=["zero-padding", "non-zero-padding"],
+)
+def test_rotation_accumulates_blocks(mocker: MockerFixture, padding: Tuple[int, int]):
     GLOBAL_SHAPE = (10, 10, 30)
     global_data = np.arange(np.prod(GLOBAL_SHAPE), dtype=np.float32).reshape(
         GLOBAL_SHAPE
@@ -39,19 +44,39 @@ def test_rotation_accumulates_blocks(mocker: MockerFixture):
         darks=2.0 * np.ones((2, GLOBAL_SHAPE[1], GLOBAL_SHAPE[2]), np.float32),
         flats=3.0 * np.ones((2, GLOBAL_SHAPE[1], GLOBAL_SHAPE[2]), np.float32),
     )
+
+    chunk_shape_list: List[int] = list(GLOBAL_SHAPE)
+    if padding != (0, 0):
+        chunk_shape_list[0] += padding[0] + padding[1]
+
+    b1_data = global_data[0 : GLOBAL_SHAPE[0] // 2, :, :]
+    b1_block_start = 0
+    if padding != (0, 0):
+        b1_data = np.pad(b1_data, pad_width=(padding, (0, 0), (0, 0)))
+        b1_block_start -= padding[0]
     b1 = DataSetBlock(
-        data=global_data[0 : GLOBAL_SHAPE[0] // 2, :, :],
+        data=b1_data,
         aux_data=aux_data,
-        block_start=0,
+        block_start=b1_block_start,
+        chunk_start=0 if padding == (0, 0) else -padding[0],
         global_shape=GLOBAL_SHAPE,
-        chunk_shape=GLOBAL_SHAPE,
+        chunk_shape=make_3d_shape_from_shape(chunk_shape_list),
+        padding=padding,
     )
+
+    b2_data = global_data[GLOBAL_SHAPE[0] // 2 :, :, :]
+    b2_block_start = GLOBAL_SHAPE[0] // 2
+    if padding != (0, 0):
+        b2_data = np.pad(b2_data, pad_width=(padding, (0, 0), (0, 0)))
+        b2_block_start -= padding[0]
     b2 = DataSetBlock(
-        data=global_data[GLOBAL_SHAPE[0] // 2 :, :, :],
+        data=b2_data,
         aux_data=aux_data,
-        block_start=GLOBAL_SHAPE[0] // 2,
+        block_start=b2_block_start,
+        chunk_start=0 if padding == (0, 0) else -padding[0],
         global_shape=GLOBAL_SHAPE,
-        chunk_shape=GLOBAL_SHAPE,
+        chunk_shape=make_3d_shape_from_shape(chunk_shape_list),
+        padding=padding,
     )
 
     class FakeModule: