Merge pull request #737 from DiamondLightSource/frames_average

Frames/projection averaging feature

Author: dkazanc

httomo/method_wrappers/__init__.py

@@ -7,9 +7,11 @@
 # import all other wrappers to make sure they are available to the factory function
 # (add imports here when createing new wrappers)
 import httomo.method_wrappers.datareducer
+import httomo.method_wrappers.sino360_to_180
 import httomo.method_wrappers.dezinging
 import httomo.method_wrappers.distortion_correction
 import httomo.method_wrappers.seam_blender
+import httomo.method_wrappers.average_frames
 import httomo.method_wrappers.images
 import httomo.method_wrappers.reconstruction
 import httomo.method_wrappers.rotation

httomo/method_wrappers/average_frames.py

@@ -0,0 +1,37 @@
+from httomo.method_wrappers.generic import GenericMethodWrapper
+from httomo.block_interfaces import T
+import numpy as np
+
+
+class AverageFramesWrapper(GenericMethodWrapper):
+    """
+    Wrapper for frames/projection averaging.
+    """
+
+    @classmethod
+    def should_select_this_class(cls, module_path: str, method_name: str) -> bool:
+        return "average_projection_frames" in method_name
+
+    def _preprocess_data(self, block: T) -> T:
+        # when the angular preview is getting changed by averaging the angles should be changed accordingly
+        config_params = self._config_params
+        k = config_params["projection_averaging_factor"]
+
+        n_proj = block.data.shape[0]  # original data angular size
+        n_full = n_proj // k
+        remainder = n_proj % k
+
+        n_out = n_full + (remainder > 0)
+
+        averaged_angles = np.empty(n_out, dtype=block.angles.dtype)
+
+        if n_full:
+            averaged_angles[:n_full] = (
+                block.angles_radians[: n_full * k].reshape(n_full, k).mean(axis=1)
+            )
+
+        if remainder:
+            averaged_angles[-1] = block.angles_radians[n_full * k :].mean()
+
+        block.angles_radians = averaged_angles
+        return block

httomo/method_wrappers/reconstruction.py

@@ -12,24 +12,17 @@
 
 
 class ReconstructionWrapper(GenericMethodWrapper):
-    """Wraps reconstruction functions, limiting the length of the angles array
-    before calling the method."""
+    """Wraps reconstruction functions."""
 
     @classmethod
     def should_select_this_class(cls, module_path: str, method_name: str) -> bool:
         return module_path.endswith(".algorithm")
 
     def _preprocess_data(self, block: T) -> T:
-        # this is essential for the angles cutting below to be valid
         assert (
             self.pattern == Pattern.sinogram
         ), "reconstruction methods must be sinogram"
-
-        # for 360 degrees data the angular dimension will be truncated while angles are not.
-        # Truncating angles if the angular dimension has got a different size
-        datashape0 = block.data.shape[0]
-        if datashape0 != len(block.angles_radians):
-            block.angles_radians = block.angles_radians[0:datashape0]
+        assert len(block.angles_radians) == block.data.shape[0]
         self._input_shape = block.data.shape
         return super()._preprocess_data(block)
 

httomo/method_wrappers/sino360_to_180.py

@@ -0,0 +1,18 @@
+from httomo.method_wrappers.generic import GenericMethodWrapper
+from httomo.block_interfaces import T
+
+
+class Sino360to180Wrapper(GenericMethodWrapper):
+    """
+    Wrapper to perform extended FoV (360degrees) data conversion to a standard 180 degrees data.
+    The wrapper is responsible for changing the angles after the data has changed.
+    """
+
+    @classmethod
+    def should_select_this_class(cls, module_path: str, method_name: str) -> bool:
+        return "sino_360_to_180" in method_name
+
+    def _postprocess_data(self, block: T) -> T:
+        # for 360 degrees data the angular dimension is truncated so the angles should be changed in a similar fashion.
+        block.angles_radians = block.angles_radians[0 : block.data.shape[0]]
+        return block

httomo/runner/method_wrapper.py

@@ -194,6 +194,7 @@ def calculate_max_slices(
         data_dtype: np.dtype,
         slicing_dim: int,
         non_slice_dims_shape: Tuple[int, int],
+        angles: np.ndarray,
         available_memory: int,
     ) -> Tuple[int, int]:
         """If it runs on GPU, determine the maximum number of slices that can fit in the

tests/conftest.py

@@ -229,6 +229,11 @@ def FBP3d_tomobar_denoising():
     return "docs/source/pipelines_full/FBP3d_tomobar_denoising.yaml"
 
 
+@pytest.fixture
+def angles_averaging():
+    return "docs/source/pipelines_full/angles_averaging.yaml"
+
+
 @pytest.fixture
 def FISTA3d_tomobar():
     return "docs/source/pipelines_full/FISTA3d_tomobar.yaml"
@@ -339,6 +344,12 @@ def FBP2d_astra_i12_119647_npz():
     return np.load("tests/test_data/raw_data/i12/FBP2d_astra_i12_119647.npz")
 
 
+@pytest.fixture
+def angle_average_LPrec_i12_119647_npz():
+    # 10 slices numpy array
+    return np.load("tests/test_data/raw_data/i12/angle_average_LPrec_i12_119647.npz")
+
+
 @pytest.fixture
 def FBP3d_tomobar_TVdenoising_i13_177906_npz():
     # 10 slices numpy array

tests/method_wrappers/test_360_to_180.py

@@ -0,0 +1,48 @@
+from httomo.method_wrappers import make_method_wrapper
+from httomo.method_wrappers.sino360_to_180 import Sino360to180Wrapper
+from httomo.runner.auxiliary_data import AuxiliaryData
+from httomo.runner.dataset import DataSetBlock
+from ..testing_utils import make_mock_preview_config, make_mock_repo
+from httomo_backends.methods_database.query import Pattern
+
+import numpy as np
+from mpi4py import MPI
+from pytest_mock import MockerFixture
+
+
+def test_sino_360_to_180(mocker: MockerFixture):
+    GLOBAL_SHAPE = (10, 20, 30)
+    GLOBAL_SHAPE_MOD = (5, 20, 30)
+
+    class FakeModule:
+        def sino_360_to_180_tester(data):
+            np.testing.assert_array_equal(data, 1)
+            return data
+
+    mocker.patch(
+        "httomo.method_wrappers.generic.import_module", return_value=FakeModule
+    )
+    wrp = make_method_wrapper(
+        make_mock_repo(mocker, pattern=Pattern.sinogram),
+        "mocked_module_path.morph",
+        "sino_360_to_180_tester",
+        MPI.COMM_WORLD,
+        make_mock_preview_config(mocker),
+    )
+    assert isinstance(wrp, Sino360to180Wrapper)
+
+    aux_data = AuxiliaryData(angles=2.0 * np.ones(GLOBAL_SHAPE[0], dtype=np.float32))
+    data = np.ones(
+        GLOBAL_SHAPE_MOD, dtype=np.float32
+    )  # assuming the data already averaged here by factor of 2
+    input = DataSetBlock(
+        data[:, 0:3, :],
+        slicing_dim=1,
+        aux_data=aux_data,
+        chunk_shape=GLOBAL_SHAPE_MOD,
+        global_shape=GLOBAL_SHAPE_MOD,
+    )
+
+    wrp.execute(input)
+
+    assert aux_data.get_angles().shape[0] == 5

tests/method_wrappers/test_angle_average.py

@@ -0,0 +1,50 @@
+from httomo.method_wrappers import make_method_wrapper
+from httomo.method_wrappers.average_frames import AverageFramesWrapper
+from httomo.runner.auxiliary_data import AuxiliaryData
+from httomo.runner.dataset import DataSetBlock
+from ..testing_utils import make_mock_preview_config, make_mock_repo
+from httomo_backends.methods_database.query import Pattern
+
+import numpy as np
+from mpi4py import MPI
+from pytest_mock import MockerFixture
+
+
+def test_angle_averaging(mocker: MockerFixture):
+    GLOBAL_SHAPE = (10, 20, 30)
+
+    class FakeModule:
+        def average_projection_frames_tester(data, projection_averaging_factor):
+            np.testing.assert_array_equal(data, 1)
+            return data
+
+    mocker.patch(
+        "httomo.method_wrappers.generic.import_module", return_value=FakeModule
+    )
+    wrp = make_method_wrapper(
+        make_mock_repo(mocker, pattern=Pattern.sinogram),
+        "mocked_module_path.morph",
+        "average_projection_frames_tester",
+        MPI.COMM_WORLD,
+        make_mock_preview_config(mocker),
+        projection_averaging_factor=2,
+    )
+    assert isinstance(wrp, AverageFramesWrapper)
+
+    aux_data = AuxiliaryData(
+        angles=2.0 * np.ones(GLOBAL_SHAPE[0] + 10, dtype=np.float32)
+    )
+    data = np.ones(
+        GLOBAL_SHAPE, dtype=np.float32
+    )  # assuming the data already averaged here by factor of 2
+    input = DataSetBlock(
+        data[:, 0:3, :],
+        slicing_dim=1,
+        aux_data=aux_data,
+        chunk_shape=GLOBAL_SHAPE,
+        global_shape=GLOBAL_SHAPE,
+    )
+
+    wrp.execute(input)
+
+    assert aux_data.get_angles().shape[0] == 5

tests/method_wrappers/test_reconstruction.py

@@ -1,57 +1,13 @@
-import math
 from httomo.method_wrappers import make_method_wrapper
-from httomo.method_wrappers.reconstruction import ReconstructionWrapper
 from httomo.runner.auxiliary_data import AuxiliaryData
 from httomo.runner.dataset import DataSetBlock
 from ..testing_utils import make_mock_preview_config, make_mock_repo
 
-from httomo_backends.methods_database.query import Pattern
-
 import numpy as np
 from mpi4py import MPI
 from pytest_mock import MockerFixture
 
 
-def test_recon_handles_reconstruction_angle_reshape(mocker: MockerFixture):
-    GLOBAL_SHAPE = (10, 20, 30)
-
-    class FakeModule:
-        # we give the angles a different name on purpose
-        def recon_tester(data, theta):
-            np.testing.assert_array_equal(data, 1)
-            np.testing.assert_array_equal(theta, 2)
-            assert data.shape[0] == len(theta)
-            return data
-
-    mocker.patch(
-        "httomo.method_wrappers.generic.import_module", return_value=FakeModule
-    )
-    wrp = make_method_wrapper(
-        make_mock_repo(mocker, pattern=Pattern.sinogram),
-        "mocked_module_path.algorithm",
-        "recon_tester",
-        MPI.COMM_WORLD,
-        make_mock_preview_config(mocker),
-    )
-    assert isinstance(wrp, ReconstructionWrapper)
-
-    aux_data = AuxiliaryData(
-        angles=2.0 * np.ones(GLOBAL_SHAPE[0] + 10, dtype=np.float32)
-    )
-    data = np.ones(GLOBAL_SHAPE, dtype=np.float32)
-    input = DataSetBlock(
-        data[:, 0:3, :],
-        slicing_dim=1,
-        aux_data=aux_data,
-        chunk_shape=GLOBAL_SHAPE,
-        global_shape=GLOBAL_SHAPE,
-    )
-
-    wrp.execute(input)
-
-    assert aux_data.get_angles().shape[0] == GLOBAL_SHAPE[0]
-
-
 def test_recon_handles_reconstruction_axisswap(mocker: MockerFixture):
     class FakeModule:
         def recon_tester(data, theta):

tests/test_parallel_pipeline_big.py

@@ -176,6 +176,77 @@ def test_pipe_parallel_FBP3d_tomobar_k11_38730_in_memory_preview(
     assert res_norm < 1e-6
 
 
+# ########################################################################
+@pytest.mark.full_data_parallel
+def test_angles_averaging_LPRec_i12_119647_preview(
+    get_files: Callable,
+    cmd_mpirun,
+    i12_119647,
+    angles_averaging,
+    angle_average_LPrec_i12_119647_npz,
+    output_folder,
+):
+
+    change_value_parameters_method_pipeline(
+        angles_averaging,
+        method=[
+            "standard_tomo",
+            "average_projection_frames",
+        ],
+        key=[
+            "preview",
+            "projection_averaging_factor",
+        ],
+        value=[
+            {"detector_y": {"start": 800, "stop": 1200}},
+            12,
+        ],
+    )
+
+    cmd_mpirun.insert(9, i12_119647)
+    cmd_mpirun.insert(10, angles_averaging)
+    cmd_mpirun.insert(11, output_folder)
+
+    process = Popen(
+        cmd_mpirun, env=os.environ, shell=False, stdin=PIPE, stdout=PIPE, stderr=PIPE
+    )
+    output, error = process.communicate()
+    print(output)
+
+    files = get_files(output_folder)
+
+    #: check the generated reconstruction (hdf5 file)
+    h5_files = list(filter(lambda x: ".h5" in x, files))
+    assert len(h5_files) == 1
+
+    # load the pre-saved numpy array for comparison bellow
+    data_gt = angle_average_LPrec_i12_119647_npz["data"]
+    axis_slice = angle_average_LPrec_i12_119647_npz["axis_slice"]
+    slices, sizeX, sizeY = np.shape(data_gt)
+
+    step = axis_slice // (slices + 2)
+    # store for the result
+    data_result = np.zeros((slices, sizeX, sizeY), dtype=np.float32)
+
+    path_to_data = "data/"
+    h5_file_name = "LPRec3d_tomobar"
+    for file_to_open in h5_files:
+        if h5_file_name in file_to_open:
+            h5f = h5py.File(file_to_open, "r")
+            index_prog = step
+            for i in range(slices):
+                data_result[i, :, :] = h5f[path_to_data][:, index_prog, :]
+                index_prog += step
+            h5f.close()
+        else:
+            message_str = f"File name with {h5_file_name} string cannot be found."
+            raise FileNotFoundError(message_str)
+
+    residual_im = data_gt - data_result
+    res_norm = np.linalg.norm(residual_im.flatten()).astype("float32")
+    assert res_norm < 1e-6
+
+
 # ########################################################################