Merge pull request #363 from DiamondLightSource/360-stitching-memory-…

…estimation

Improve accuracy of 360 stitching memory estimation

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

@@ -23,14 +23,14 @@ misc:
   morph:
     sino_360_to_180:
       pattern: sinogram
-      output_dims_change: False
+      output_dims_change: True
       implementation: gpu_cupy
       save_result_default: False
       padding: False
       memory_gpu:
         - datasets: [tomo]
-        - multipliers: [2.2]
-        - methods: [direct]
+        - multipliers: [None]
+        - methods: [module]
     data_resampler:
       pattern: all
       output_dims_change: True

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

@@ -24,9 +24,13 @@
 from typing import Tuple
 import numpy as np
 
+from httomo.runner.output_ref import OutputRef
+
 __all__ = [
     "_calc_memory_bytes_data_resampler",
     "_calc_output_dim_data_resampler",
+    "_calc_memory_bytes_sino_360_to_180",
+    "_calc_output_dim_sino_360_to_180",
 ]
 
 
@@ -54,3 +58,63 @@ def _calc_memory_bytes_data_resampler(
 
     tot_memory_bytes = input_size + output_size + interpolator
     return (tot_memory_bytes, xi)
+
+
+def _calc_output_dim_sino_360_to_180(
+    non_slice_dims_shape: Tuple[int, int],
+    **kwargs,
+) -> Tuple[int, int]:
+    assert "overlap" in kwargs, "Expected overlap in method parameters"
+    overlap_side_output = kwargs["overlap"]
+    assert isinstance(
+        overlap_side_output, OutputRef
+    ), "Expected overlap to be in an OutputRef"
+    overlap: float = overlap_side_output.value
+
+    original_sino_width = non_slice_dims_shape[1]
+    stitched_sino_width = original_sino_width * 2 - math.ceil(overlap)
+    return non_slice_dims_shape[0] // 2, stitched_sino_width
+
+
+def _calc_memory_bytes_sino_360_to_180(
+    non_slice_dims_shape: Tuple[int, int],
+    dtype: np.dtype,
+    **kwargs,
+) -> Tuple[int, int]:
+    assert "overlap" in kwargs, "Expected overlap in method parameters"
+    overlap_side_output = kwargs["overlap"]
+    assert isinstance(
+        overlap_side_output, OutputRef
+    ), "Expected overlap to be in an OutputRef"
+    overlap: float = overlap_side_output.value
+
+    original_sino_width = non_slice_dims_shape[1]
+    stitched_sino_width = original_sino_width * 2 - math.ceil(overlap)
+    n = non_slice_dims_shape[0] // 2
+    stitched_non_slice_dims = (n, stitched_sino_width)
+
+    input_slice_size = int(np.prod(non_slice_dims_shape)) * dtype.itemsize
+    output_slice_size = int(np.prod(stitched_non_slice_dims)) * dtype.itemsize
+
+    summand_shape: Tuple[int, int] = (n, int(overlap))
+    # Multiplication between a subset of the original data (`float32`) and the 1D weights array
+    # (`float32`) causes a new array to be created that has dtype `float32` (for example,
+    # multiplications like `weights * data[:n, :, -overlap]`
+    summand_size = int(np.prod(summand_shape)) * np.float32().itemsize
+
+    total_memory_bytes = (
+        input_slice_size
+        + output_slice_size
+        # In both the `if` branch and the `else` branch checking the `rotation` variable value,
+        # in total, there are 4 copies of subsets of the `data` array that are made (note that
+        # the expressions below are referencing the `if` branch and are slightly different in
+        # the `else` branch):
+        # 1. fancy indexing: `data[n : 2 * n, :, overlap:][:, :, ::-1]`
+        # 2. multiplication: `weights * data[:n, :, :overlap]`
+        # 3. multiplication: `weights * data[n : 2 * n, :, :overlap]`
+        # 4. fancy indexing (performed on the result of 3):
+        # `(weights * data[n : 2 * n, :, :overlap])[:, :, ::-1]`
+        + 4 * summand_size
+    )
+
+    return total_memory_bytes, 0

tests/test_backends/test_httomolibgpu.py

@@ -5,7 +5,9 @@
 from mpi4py import MPI
 import numpy as np
 from numpy import uint16, float32
+from unittest import mock
 
+from pytest_mock import MockerFixture
 from numpy.testing import assert_allclose, assert_equal
 import os
 
@@ -16,9 +18,10 @@
 import cupy as cp
 
 from httomo.methods_database.query import get_method_info
+from httomo.runner.output_ref import OutputRef
 
 
-from httomolibgpu.misc.morph import data_resampler
+from httomolibgpu.misc.morph import data_resampler, sino_360_to_180
 from httomolibgpu.prep.normalize import normalize
 from httomolibgpu.prep.phase import paganin_filter_tomopy, paganin_filter_savu
 from httomolibgpu.prep.alignment import distortion_correction_proj_discorpy
@@ -39,6 +42,9 @@
 from httomo.methods_database.packages.external.httomolibgpu.supporting_funcs.misc.rescale import *
 from httomo.methods_database.packages.external.httomolibgpu.supporting_funcs.prep.normalize import *
 
+from tests.testing_utils import make_test_method
+
+
 module_mem_path = "httomo.methods_database.packages.external."
 
 
@@ -588,3 +594,50 @@ def test_rescale_to_int_memoryhook(
     # the resulting percent value should not deviate from max_mem on more than 20%
     assert estimated_memory_mb >= max_mem_mb
     assert percents_relative_maxmem <= 35
+
+
+@pytest.mark.cupy
+@pytest.mark.parametrize("slices", [3, 8, 30, 80])
+@pytest.mark.parametrize("det_x", [600, 2160])
+def test_sino_360_to_180_memoryhook(
+    ensure_clean_memory,
+    mocker: MockerFixture,
+    det_x: int,
+    slices: int,
+):
+    # Use a different overlap value for stitching based on the width of the 360 sinogram
+    overlap = 350 if det_x == 600 else 1950
+    shape = (1801, slices, det_x)
+    data = cp.random.random_sample(shape, dtype=np.float32)
+
+    # Run method to see actual memory usage
+    hook = MaxMemoryHook()
+    with hook:
+        sino_360_to_180(cp.copy(data), overlap)
+
+    # Call memory estimator to estimate memory usage
+    output_ref = OutputRef(
+        method=make_test_method(mocker),
+        mapped_output_name="overlap",
+    )
+    with mock.patch(
+        "httomo.runner.output_ref.OutputRef.value",
+        new_callable=mock.PropertyMock,
+    ) as mock_value_property:
+        mock_value_property.return_value = overlap
+        (estimated_bytes, subtract_bytes) = _calc_memory_bytes_sino_360_to_180(
+            non_slice_dims_shape=(shape[0], shape[2]),
+            dtype=np.float32(),
+            overlap=output_ref,
+        )
+    estimated_bytes *= slices
+
+    max_mem = hook.max_mem - subtract_bytes
+
+    # For the difference between the actual memory usage and estimated memory usage, calculate
+    # that as a percentage of the actual memory used
+    difference = abs(estimated_bytes - max_mem)
+    percentage_difference = round((difference / max_mem) * 100)
+
+    assert estimated_bytes >= max_mem
+    assert percentage_difference <= 35