adding memory estimator and tests for osem

httomo_backends/methods_database/packages/backends/httomolibgpu/httomolibgpu.yaml

@@ -236,6 +236,15 @@ recon:
       memory_gpu:
         multiplier: None
         method: module
+    OSEM3d_tomobar:
+      pattern: sinogram
+      output_dims_change: True
+      implementation: gpu_cupy
+      save_result_default: True
+      padding: True
+      memory_gpu:
+        multiplier: None
+        method: module        
   rotation:
     find_center_vo:
       pattern: projection

httomo_backends/methods_database/packages/backends/httomolibgpu/supporting_funcs/recon/algorithm.py

@@ -24,7 +24,7 @@
 from typing import Tuple
 import numpy as np
 from httomo_backends.cufft import CufftType, cufft_estimate_1d, cufft_estimate_2d
-from httomolibgpu.recon.algorithm import ADMM3d_tomobar, LPRec3d_tomobar
+from httomolibgpu.recon.algorithm import LPRec3d_tomobar
 from tomobar.supp.memory_estimator_helpers import DeviceMemStack
 
 __all__ = [
@@ -35,15 +35,18 @@
     "_calc_memory_bytes_CGLS3d_tomobar",
     "_calc_memory_bytes_FISTA3d_tomobar",
     "_calc_memory_bytes_ADMM3d_tomobar",
+    "_calc_memory_bytes_OSEM3d_tomobar",
     "_calc_output_dim_FBP2d_astra",
     "_calc_output_dim_FBP3d_tomobar",
     "_calc_output_dim_LPRec3d_tomobar",
     "_calc_output_dim_SIRT3d_tomobar",
     "_calc_output_dim_CGLS3d_tomobar",
     "_calc_output_dim_FISTA3d_tomobar",
     "_calc_output_dim_ADMM3d_tomobar",
+    "_calc_output_dim_OSEM3d_tomobar",    
     "_calc_padding_FISTA3d_tomobar",
     "_calc_padding_ADMM3d_tomobar",
+    "_calc_padding_OSEM3d_tomobar",
 ]
 
 
@@ -54,6 +57,9 @@ def _calc_padding_FISTA3d_tomobar(**kwargs) -> Tuple[int, int]:
 def _calc_padding_ADMM3d_tomobar(**kwargs) -> Tuple[int, int]:
     return (5, 5)
 
+def _calc_padding_OSEM3d_tomobar(**kwargs) -> Tuple[int, int]:
+    return (5, 5)
+
 
 def __calc_output_dim_recon(non_slice_dims_shape, **kwargs):
     """Function to calculate output dimensions for all reconstructors.
@@ -97,6 +103,9 @@ def _calc_output_dim_ADMM3d_tomobar(non_slice_dims_shape, **kwargs):
     return __calc_output_dim_recon(non_slice_dims_shape, **kwargs)
 
 
+def _calc_output_dim_OSEM3d_tomobar(non_slice_dims_shape, **kwargs):
+    return __calc_output_dim_recon(non_slice_dims_shape, **kwargs)
+
 def _calc_memory_bytes_FBP3d_tomobar(
     non_slice_dims_shape: Tuple[int, int],
     dtype: np.dtype,
@@ -627,6 +636,50 @@ def _calc_memory_bytes_FISTA3d_tomobar(
     return (tot_memory_bytes, 0)
 
 
+def _calc_memory_bytes_OSEM3d_tomobar(
+    non_slice_dims_shape: Tuple[int, int],
+    dtype: np.dtype,
+    **kwargs,
+) -> Tuple[int, int]:
+    detector_pad = 0
+    if "detector_pad" in kwargs:
+        detector_pad = kwargs["detector_pad"]
+    if detector_pad is True:
+        detector_pad = __estimate_detectorHoriz_padding(non_slice_dims_shape[1])
+    elif detector_pad is False:
+        detector_pad = 0
+
+    anglesnum = non_slice_dims_shape[0]
+    DetectorsLengthH_padded = non_slice_dims_shape[1] + 2 * detector_pad
+
+    # calculate the output shape
+    output_dims = _calc_output_dim_OSEM3d_tomobar(non_slice_dims_shape, **kwargs)
+    recon_data_size_original = (
+        np.prod(output_dims) * dtype.itemsize
+    )  # recon user-defined size
+
+    in_data_siz_pad = (anglesnum * DetectorsLengthH_padded) * dtype.itemsize
+    output_dims_larger_grid = (DetectorsLengthH_padded, DetectorsLengthH_padded)
+
+    out_data_size = np.prod(output_dims_larger_grid) * dtype.itemsize
+    normalisation = out_data_size
+    Ax = in_data_siz_pad
+    ratio = in_data_siz_pad
+    backproj = out_data_size
+
+    mlem_part = (
+        recon_data_size_original
+        + normalisation
+        + Ax
+        + ratio
+        + backproj
+        + out_data_size
+    )
+    regul_part = 8 * np.prod(output_dims_larger_grid) * dtype.itemsize
+
+    tot_memory_bytes = int(mlem_part + regul_part)
+    return (tot_memory_bytes, 0)
+
 def _calc_memory_bytes_ADMM3d_tomobar(
     non_slice_dims_shape: Tuple[int, int],
     dtype: np.dtype,

tests/test_httomolibgpu.py

@@ -34,6 +34,7 @@
     CGLS3d_tomobar,
     FISTA3d_tomobar,
     ADMM3d_tomobar,
+    OSEM3d_tomobar,
 )
 from httomolibgpu.misc.rescale import rescale_to_int
 
@@ -1010,6 +1011,103 @@ def test_recon_FISTA3d_tomobar_nonOS_memoryhook(
     assert percents_relative_maxmem <= 100
 
 
+@pytest.mark.cupy
+@pytest.mark.parametrize("slices", [3, 5])
+@pytest.mark.parametrize("recon_size_it", [2560])
+@pytest.mark.parametrize("padding", [0, 100, 200])
+def test_recon_MLEM3d_tomobar_nonOS_memoryhook(
+    slices, recon_size_it, padding, ensure_clean_memory
+):
+    angles_total = 901
+    detX_size = 2560
+    data = cp.random.random_sample((angles_total, slices, detX_size), dtype=np.float32)
+    kwargs = {}
+    kwargs["recon_size"] = recon_size_it
+
+    hook = MaxMemoryHook()
+    with hook:
+        recon_data = OSEM3d_tomobar(
+            cp.copy(data),
+            np.linspace(0.0 * np.pi / 180.0, 180.0 * np.pi / 180.0, data.shape[0]),
+            center=1200,
+            recon_size=recon_size_it,
+            iterations=1,
+            subsets_number=1,
+            regularisation_iterations=2,
+            nonnegativity=True,
+            detector_pad=padding,
+        )
+
+    # make sure estimator function is within range (80% min, 100% max)
+    max_mem = (
+        hook.max_mem
+    )  # the amount of memory in bytes needed for the method according to memoryhook
+
+    # now we estimate how much of the total memory required for this data
+    estimated_memory_bytes, subtract_bytes = _calc_memory_bytes_OSEM3d_tomobar(
+        (angles_total, detX_size), dtype=np.float32(), **kwargs
+    )
+    estimated_memory_mb = round(slices * estimated_memory_bytes / (1024**2), 2)
+    max_mem -= subtract_bytes
+    max_mem_mb = round(max_mem / (1024**2), 2)
+
+    # now we compare both memory estimations
+    difference_mb = abs(estimated_memory_mb - max_mem_mb)
+    percents_relative_maxmem = round((difference_mb / max_mem_mb) * 100)
+    # the estimated_memory_mb should be LARGER or EQUAL to max_mem_mb
+    # 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 <= 100
+
+@pytest.mark.cupy
+@pytest.mark.parametrize("slices", [3, 5])
+@pytest.mark.parametrize("recon_size_it", [2560])
+@pytest.mark.parametrize("padding", [0, 100, 200])
+def test_recon_OSEM3d_tomobar_memoryhook(
+    slices, recon_size_it, padding, ensure_clean_memory
+):
+    angles_total = 901
+    detX_size = 2560
+    data = cp.random.random_sample((angles_total, slices, detX_size), dtype=np.float32)
+    kwargs = {}
+    kwargs["recon_size"] = recon_size_it
+
+    hook = MaxMemoryHook()
+    with hook:
+        recon_data = OSEM3d_tomobar(
+            cp.copy(data),
+            np.linspace(0.0 * np.pi / 180.0, 180.0 * np.pi / 180.0, data.shape[0]),
+            center=1200,
+            recon_size=recon_size_it,
+            iterations=1,
+            subsets_number=12,
+            regularisation_iterations=2,
+            nonnegativity=True,
+            detector_pad=padding,
+        )
+
+    # make sure estimator function is within range (80% min, 100% max)
+    max_mem = (
+        hook.max_mem
+    )  # the amount of memory in bytes needed for the method according to memoryhook
+
+    # now we estimate how much of the total memory required for this data
+    estimated_memory_bytes, subtract_bytes = _calc_memory_bytes_OSEM3d_tomobar(
+        (angles_total, detX_size), dtype=np.float32(), **kwargs
+    )
+    estimated_memory_mb = round(slices * estimated_memory_bytes / (1024**2), 2)
+    max_mem -= subtract_bytes
+    max_mem_mb = round(max_mem / (1024**2), 2)
+
+    # now we compare both memory estimations
+    difference_mb = abs(estimated_memory_mb - max_mem_mb)
+    percents_relative_maxmem = round((difference_mb / max_mem_mb) * 100)
+    # the estimated_memory_mb should be LARGER or EQUAL to max_mem_mb
+    # 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 <= 100
+
+
 @pytest.mark.cupy
 @pytest.mark.parametrize("slices", [3, 5])
 @pytest.mark.parametrize("recon_size_it", [2560])