Simplify memory_gpu field in method database entry

httomo/method_wrappers/generic.py

@@ -146,7 +146,7 @@ def parameters(self) -> List[str]:
         return self._parameters
 
     @property
-    def memory_gpu(self) -> List[GpuMemoryRequirement]:
+    def memory_gpu(self) -> Optional[GpuMemoryRequirement]:
         return self._memory_gpu
 
     @property
@@ -414,7 +414,7 @@ def calculate_max_slices(
             return int(100e9), available_memory
 
         # if we have no information, we assume in-place operation with no extra memory
-        if len(self.memory_gpu) == 0:
+        if self.memory_gpu is None:
             return (
                 int(
                     available_memory
@@ -426,25 +426,22 @@ def calculate_max_slices(
         # NOTE: This could go directly into the methodquery / method database,
         # and here we just call calculated_memory_bytes
         memory_bytes_method = 0
-        for field in self.memory_gpu:
-            subtract_bytes = 0
-            # loop over the dataset names given in the library file and extracting
-            # the corresponding dimensions from the available datasets
-            if field.method == "direct":
-                assert field.multiplier is not None
-                # this calculation assumes a direct (simple) correspondence through multiplier
-                memory_bytes_method += int(
-                    field.multiplier
-                    * np.prod(non_slice_dims_shape)
-                    * data_dtype.itemsize
-                )
-            else:
-                (
-                    memory_bytes_method,
-                    subtract_bytes,
-                ) = self._query.calculate_memory_bytes(
-                    non_slice_dims_shape, data_dtype, **self.config_params
-                )
+        subtract_bytes = 0
+        if self.memory_gpu.method == "direct":
+            assert self.memory_gpu.multiplier is not None
+            # this calculation assumes a direct (simple) correspondence through multiplier
+            memory_bytes_method += int(
+                self.memory_gpu.multiplier
+                * np.prod(non_slice_dims_shape)
+                * data_dtype.itemsize
+            )
+        else:
+            (
+                memory_bytes_method,
+                subtract_bytes,
+            ) = self._query.calculate_memory_bytes(
+                non_slice_dims_shape, data_dtype, **self.config_params
+            )
 
         if memory_bytes_method == 0:
             return available_memory - subtract_bytes, available_memory

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

@@ -7,19 +7,17 @@ misc:
       save_result_default: False
       padding: False
       memory_gpu:
-        - datasets: [tomo]
-        - multipliers: [2.1]
-        - methods: [direct]
+        multiplier: 2.1
+        method: direct
     remove_outlier:
       pattern: all
       output_dims_change: False
       implementation: gpu_cupy
       save_result_default: False
       padding: False
       memory_gpu:
-        - datasets: [tomo]
-        - multipliers: [2.1]
-        - methods: [direct]
+        multiplier: 2.1
+        method: direct
   morph:
     sino_360_to_180:
       pattern: sinogram
@@ -28,19 +26,17 @@ misc:
       save_result_default: False
       padding: False
       memory_gpu:
-        - datasets: [tomo]
-        - multipliers: [None]
-        - methods: [module]
+        multiplier: None
+        method: module
     data_resampler:
       pattern: all
       output_dims_change: True
       implementation: gpu_cupy
       save_result_default: False
       padding: False
       memory_gpu:
-        - datasets: [tomo]
-        - multipliers: [None]
-        - methods: [module]
+        multiplier: None
+        method: module
   rescale:
     rescale_to_int:
       pattern: all
@@ -49,9 +45,8 @@ misc:
       save_result_default: False
       padding: False
       memory_gpu:
-        - datasets: [tomo]
-        - multipliers: [None]
-        - methods: [module]
+        multiplier: None
+        method: module
 prep:
   normalize:
     normalize:
@@ -61,9 +56,8 @@ prep:
       save_result_default: False
       padding: False
       memory_gpu:
-        - datasets: [tomo]
-        - multipliers: [None]
-        - methods: [module]
+        multiplier: None
+        method: module
   phase:
     paganin_filter_tomopy:
       pattern: projection
@@ -72,19 +66,17 @@ prep:
       save_result_default: False
       padding: False
       memory_gpu:
-        - datasets: [tomo]
-        - multipliers: [None]
-        - methods: [module]
+        multiplier: None
+        method: module
     paganin_filter_savu:
       pattern: projection
       output_dims_change: False
       implementation: gpu_cupy
       save_result_default: False
       padding: False
       memory_gpu:
-        - datasets: [tomo]
-        - multipliers: [None]
-        - methods: [module]
+        multiplier: None
+        method: module
   alignment:
     distortion_correction_proj_discorpy:
       pattern: projection
@@ -93,9 +85,8 @@ prep:
       save_result_default: False
       padding: False
       memory_gpu:
-        - datasets: [tomo]
-        - multipliers: [1.2]
-        - methods: [direct]
+        multiplier: 1.2
+        method: direct
   stripe:
     remove_stripe_based_sorting:
       pattern: sinogram
@@ -104,29 +95,26 @@ prep:
       save_result_default: False
       padding: False
       memory_gpu:
-        - datasets: [tomo]
-        - multipliers: [1.17]
-        - methods: [direct]
+        multiplier: 1.17
+        method: direct
     remove_stripe_ti:
       pattern: sinogram
       output_dims_change: False
       implementation: gpu_cupy
       save_result_default: False
       padding: False
       memory_gpu:
-        - datasets: [tomo]
-        - multipliers: [None]
-        - methods: [module]
+        multiplier: None
+        method: module
     remove_all_stripe:
       pattern: sinogram
       output_dims_change: False
       implementation: gpu_cupy
       save_result_default: False
       padding: False
       memory_gpu:
-        - datasets: [tomo]
-        - multipliers: [None]
-        - methods: [module]
+        multiplier: None
+        method: module
 recon:
   algorithm:
     FBP:
@@ -136,29 +124,26 @@ recon:
       save_result_default: True
       padding: False
       memory_gpu:
-        - datasets: [tomo]
-        - multipliers: [None]
-        - methods: [module]
+        multiplier: None
+        method: module
     SIRT:
       pattern: sinogram
       output_dims_change: True
       implementation: gpu_cupy
       save_result_default: True
       padding: False
       memory_gpu:
-        - datasets: [tomo]
-        - multipliers: [None]
-        - methods: [module]
+        multiplier: None
+        method: module
     CGLS:
       pattern: sinogram
       output_dims_change: True
       implementation: gpu_cupy
       save_result_default: True
       padding: False
       memory_gpu:
-        - datasets: [tomo]
-        - multipliers: [None]
-        - methods: [module]
+        multiplier: None
+        method: module
   rotation:
     find_center_vo:
       pattern: projection
@@ -167,26 +152,23 @@ recon:
       save_result_default: False
       padding: False
       memory_gpu:
-        - datasets: [tomo]
-        - multipliers: [0]
-        - methods: [direct]
+        multiplier: 0
+        method: direct
     find_center_360:
       pattern: projection
       output_dims_change: False
       implementation: gpu_cupy
       save_result_default: False
       padding: False
       memory_gpu:
-        - datasets: [tomo]
-        - multipliers: [0]
-        - methods: [direct]
+        multiplier: 0
+        method: direct
     find_center_pc:
       pattern: projection
       output_dims_change: False
       implementation: gpu_cupy
       save_result_default: False
       padding: False
       memory_gpu:
-        - datasets: [tomo]
-        - multipliers: [0]
-        - methods: [direct]
+        multiplier: 0
+        method: direct

httomo/methods_database/query.py

@@ -1,5 +1,5 @@
 from types import ModuleType
-from typing import Callable, List, Literal, Tuple
+from typing import Callable, List, Literal, Optional, Tuple
 from pathlib import Path
 import numpy as np
 
@@ -102,27 +102,18 @@ def padding(self) -> bool:
 
     def get_memory_gpu_params(
         self,
-    ) -> List[GpuMemoryRequirement]:
+    ) -> Optional[GpuMemoryRequirement]:
         p = get_method_info(self.module_path, self.method_name, "memory_gpu")
         if p is None or p == "None":
-            return []
+            return None
         if type(p) == list:
             # convert to dict first
-            dd: dict = dict()
+            d: dict = dict()
             for item in p:
-                dd |= item
+                d |= item
         else:
-            dd = p
-        # now iterate and make it into one
-        assert (
-            len(dd["datasets"]) == len(dd["multipliers"]) == len(dd["methods"])
-        ), "Invalid data"
-        return [
-            GpuMemoryRequirement(
-                dataset=d, multiplier=dd["multipliers"][i], method=dd["methods"][i]
-            )
-            for i, d in enumerate(dd["datasets"])
-        ]
+            d = p
+        return GpuMemoryRequirement(multiplier=d["multiplier"], method=d["method"])
 
     def calculate_memory_bytes(
         self, non_slice_dims_shape: Tuple[int, int], dtype: np.dtype, **kwargs

httomo/runner/method_wrapper.py

@@ -62,7 +62,7 @@ def parameters(self) -> List[str]:
         ...  # pragma: nocover
 
     @property
-    def memory_gpu(self) -> List[GpuMemoryRequirement]:
+    def memory_gpu(self) -> Optional[GpuMemoryRequirement]:
         """Memory requirements for GPU execution"""
         ...  # pragma: nocover
 

httomo/runner/methods_repository_interface.py

@@ -7,7 +7,6 @@
 
 @dataclass(frozen=True)
 class GpuMemoryRequirement:
-    dataset: Literal["data", "tomo", "darks", "flats"] = "tomo"
     multiplier: Optional[float] = 1.0
     method: Literal["direct", "module"] = "direct"
 
@@ -33,7 +32,7 @@ def get_implementation(self) -> Literal["cpu", "gpu", "gpu_cupy"]:
         """Check for implementation of the method"""
         ...  # pragma: no cover
 
-    def get_memory_gpu_params(self) -> List[GpuMemoryRequirement]:
+    def get_memory_gpu_params(self) -> Optional[GpuMemoryRequirement]:
         """Get the parameters for the GPU memory estimation"""
         ...  # pragma: no cover
 

httomo/runner/task_runner.py

@@ -334,7 +334,7 @@ def determine_max_slices(self, section: Section, slicing_dim: int):
 
         # loop over all methods in section
         for idx, m in enumerate(section):
-            if len(m.memory_gpu) == 0:
+            if m.memory_gpu is None:
                 max_slices_methods[idx] = max_slices
                 continue
 

httomo/utils.py

@@ -84,7 +84,9 @@ def log_exception(output: str) -> None:
 
 
 def _parse_preview(
-    preview: List[Dict[str, int]], data_shape: Tuple[int], data_indices: List[int]
+    preview: List[Dict[str, int]],
+    data_shape: Tuple[int],
+    data_indices: List[int],
 ) -> str:
     """Parse the python list that represents the preview parameter in the loader
     into a string that the helper loader functions in