Utils and tests for remove_stripe_fw

Author: mfep

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

@@ -107,6 +107,15 @@ prep:
       memory_gpu:
         multiplier: 1.17
         method: direct
+    remove_stripe_fw:
+      pattern: sinogram
+      output_dims_change: False
+      implementation: gpu_cupy
+      save_result_default: False
+      padding: False
+      memory_gpu:
+        multiplier: None
+        method: iterative
     remove_stripe_ti:
       pattern: sinogram
       output_dims_change: False

httomo_backends/methods_database/packages/backends/httomolibgpu/supporting_funcs/prep/stripe.py

@@ -25,9 +25,11 @@
 import numpy as np
 
 from httomo_backends.cufft import CufftType, cufft_estimate_1d
+from httomolibgpu.prep.stripe import DeviceMemStack, remove_stripe_fw
 
 
 __all__ = [
+    "_calc_memory_bytes_for_slices_remove_stripe_fw",
     "_calc_memory_bytes_remove_stripe_ti",
     "_calc_memory_bytes_remove_all_stripe",
     "_calc_memory_bytes_raven_filter",
@@ -53,6 +55,16 @@ def _calc_memory_bytes_remove_stripe_ti(
     return (tot_memory_bytes, gamma_mem)
 
 
+def _calc_memory_bytes_for_slices_remove_stripe_fw(
+    dims_shape: Tuple[int, int, int],
+    dtype: np.dtype,
+    **kwargs,
+) -> Tuple[int, int]:
+    mem_stack = DeviceMemStack()
+    remove_stripe_fw(dims_shape, mem_stack=mem_stack, **kwargs)
+    return mem_stack.highwater, 0
+
+
 def _calc_memory_bytes_remove_all_stripe(
     non_slice_dims_shape: Tuple[int, int],
     dtype: np.dtype,

tests/test_httomolibgpu.py

@@ -19,6 +19,7 @@
 from httomolibgpu.prep.alignment import distortion_correction_proj_discorpy
 from httomolibgpu.prep.stripe import (
     remove_stripe_based_sorting,
+    remove_stripe_fw,
     remove_stripe_ti,
     remove_all_stripe,
     raven_filter,
@@ -376,6 +377,36 @@ def test_remove_stripe_ti_memoryhook(slices, ensure_clean_memory):
     assert percents_relative_maxmem <= 20
 
 
+@pytest.mark.cupy
+@pytest.mark.parametrize("slices", [64, 129])
+def test_remove_stripe_fw_memoryhook(slices, ensure_clean_memory):
+    dim_x = 156
+    dim_y = 216
+    data = cp.random.random_sample((slices, dim_x, dim_y), dtype=np.float32)
+    hook = MaxMemoryHook()
+    with hook:
+        remove_stripe_fw(cp.copy(data)).get()
+    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_remove_stripe_fw(
+        (slices, dim_x, dim_y), dtype=np.float32()
+    )
+    estimated_memory_mb = round(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 1%
+    assert estimated_memory_mb >= max_mem_mb
+    assert percents_relative_maxmem <= 1
+
+
 @pytest.mark.cupy
 @pytest.mark.parametrize("projections", [180, 360, 720, 1080, 1440])
 def test_raven_filter_memoryhook(projections, ensure_clean_memory):