Merge pull request #221 from DiamondLightSource/rescaletoint_C

Making rescale_to_int GPU only function

Author: dkazanc

httomolibgpu/misc/rescale.py

@@ -24,32 +24,32 @@
 from httomolibgpu import cupywrapper
 
 cp = cupywrapper.cp
-cupy_run = cupywrapper.cupy_run
 
 from typing import Literal, Optional, Tuple, Union
 
 from httomolibgpu.misc.supp_func import data_checker
 
+
 __all__ = [
     "rescale_to_int",
 ]
 
 
 def rescale_to_int(
-    data: Union[np.ndarray, cp.ndarray],
+    data: cp.ndarray,
     perc_range_min: float = 0.0,
     perc_range_max: float = 100.0,
     bits: Literal[8, 16, 32] = 8,
     glob_stats: Optional[Tuple[float, float, float, int]] = None,
-) -> Union[np.ndarray, cp.ndarray]:
+) -> cp.ndarray:
     """
     Rescales the data given as float32 type and converts it into the range of an unsigned integer type
     with the given number of bits. For more detailed information and examples, see :ref:`method_rescale_to_int`.
 
     Parameters
     ----------
-    data : Union[np.ndarray, cp.ndarray]
-        Input data as a numpy or cupy array (the function is cpu-gpu agnostic)
+    data : cp.ndarray
+        Input data as a cupy array
     perc_range_min: float, optional
         The lower cutoff point in the input data, in percent of the data range (defaults to 0).
         The lower bound is computed as min + perc_range_min/100*(max-min)
@@ -69,7 +69,7 @@ def rescale_to_int(
 
     Returns
     -------
-    Union[np.ndarray, cp.ndarray]
+    cp.ndarray
         The original data, clipped to the range specified with the perc_range_min and
         perc_range_max, and scaled to the full range of the output integer type
     """
@@ -82,18 +82,13 @@ def rescale_to_int(
 
     data = data_checker(data, verbosity=True, method_name="rescale_to_int")
 
-    if cupy_run:
-        xp = cp.get_array_module(data)
-    else:
-        import numpy as xp
-
     # get the min and max integer values of the output type
-    output_min = xp.iinfo(output_dtype).min
-    output_max = xp.iinfo(output_dtype).max
+    output_min = cp.iinfo(output_dtype).min
+    output_max = cp.iinfo(output_dtype).max
 
     if not isinstance(glob_stats, tuple):
-        min_value = float(xp.min(data))
-        max_value = float(xp.max(data))
+        min_value = float(cp.min(data))
+        max_value = float(cp.max(data))
     else:
         min_value = glob_stats[0]
         max_value = glob_stats[1]
@@ -102,32 +97,21 @@ def rescale_to_int(
     input_min = (perc_range_min * (range_intensity) / 100) + min_value
     input_max = (perc_range_max * (range_intensity) / 100) + min_value
 
+    factor = cp.float32(1.0)
     if (input_max - input_min) != 0.0:
-        factor = xp.float32((output_max - output_min) / (input_max - input_min))
-    else:
-        factor = 1.0
-
-    res = xp.empty(data.shape, dtype=output_dtype)
-    if xp.__name__ == "numpy":
-        if input_max == pow(2, 32):
-            input_max -= 1
-        res = np.copy(data.astype(float))
-        res[data.astype(float) < input_min] = int(input_min)
-        res[data.astype(float) > input_max] = int(input_max)
-        res -= input_min
-        res *= factor
-        res = output_dtype(res)
-    else:
-        rescale_kernel = cp.ElementwiseKernel(
-            "T x, raw T input_min, raw T input_max, raw T factor",
-            "O out",
-            """
-            T x_clean = isnan(x) || isinf(x) ? T(0) : x;
-            T x_clipped = x_clean < input_min ? input_min : (x_clean > input_max ? input_max : x_clean);
-            T x_rebased = x_clipped - input_min;
-            out = O(x_rebased * factor);
-            """,
-            "rescale_to_int",
-        )
-        rescale_kernel(data, input_min, input_max, factor, res)
+        factor = cp.float32((output_max - output_min) / (input_max - input_min))
+
+    res = cp.empty(data.shape, dtype=output_dtype)
+    rescale_kernel = cp.ElementwiseKernel(
+        "T x, raw T input_min, raw T input_max, raw T factor",
+        "O out",
+        """
+        T x_clean = isnan(x) || isinf(x) ? T(0) : x;
+        T x_clipped = x_clean < input_min ? input_min : (x_clean > input_max ? input_max : x_clean);
+        T x_rebased = x_clipped - input_min;
+        out = O(x_rebased * factor);
+        """,
+        "rescale_to_int",
+    )
+    rescale_kernel(data, input_min, input_max, factor, res)
     return res

httomolibgpu/misc/supp_func.py

@@ -162,7 +162,7 @@ def data_checker(
 ) -> bool:
     """
     Function that performs the variety of checks on input data, in some cases also correct the data and prints warnings.
-    Currently it checks for: the presence of infs and nans in data. 
+    Currently it checks for: the presence of infs and nans in data.
 
     Parameters
     ----------

tests/test_misc/test_rescale.py

@@ -14,14 +14,11 @@ def test_rescale_no_change():
     res_dev = rescale_to_int(
         data_dev, bits=8, glob_stats=(0.0, 255.0, 100.0, data.size)
     )
-    res_cpu = rescale_to_int(data, bits=8, glob_stats=(0.0, 255.0, 100.0, data.size))
 
     res = cp.asnumpy(res_dev).astype(np.float32)
 
     assert res_dev.dtype == np.uint8
-    assert res_cpu.dtype == np.uint8
     np.testing.assert_array_equal(res, data)
-    np.testing.assert_array_equal(res, res_cpu)
 
 
 @pytest.mark.parametrize("bits", [8, 16, 32])
@@ -31,28 +28,22 @@ def test_rescale_no_change_no_stats(bits: Literal[8, 16, 32]):
     data[13, 1] = (2**bits) - 1
     data_dev = cp.asarray(data)
     res_dev = rescale_to_int(data_dev, bits=bits)
-    res_cpu = rescale_to_int(data, bits=bits)
 
     res_dev_float32 = cp.asnumpy(res_dev).astype(np.float32)
 
     assert res_dev.dtype.itemsize == bits // 8
     np.testing.assert_array_equal(res_dev_float32, data)
-    assert res_cpu.dtype.itemsize == bits // 8
-    res_cpu_float32 = np.float32(res_cpu)
-    np.testing.assert_array_equal(res_dev_float32, res_cpu_float32)
 
 
 def test_rescale_double():
     data = np.ones((30, 50), dtype=np.float32)
 
     data_dev = cp.asarray(data)
     res_dev = rescale_to_int(data_dev, bits=8, glob_stats=(0, 2, 100, data.size))
-    res_cpu = rescale_to_int(data, bits=8, glob_stats=(0, 2, 100, data.size))
 
     res = cp.asnumpy(res_dev).astype(np.float32)
 
     np.testing.assert_array_almost_equal(res, 127.0)
-    np.testing.assert_array_almost_equal(res_cpu, 127.0)
 
 
 def test_rescale_handles_nan_inf():
@@ -63,25 +54,21 @@ def test_rescale_handles_nan_inf():
 
     data_dev = cp.asarray(data)
     res_dev = rescale_to_int(data_dev, bits=8, glob_stats=(0, 2, 100, data.size))
-    res_cpu = rescale_to_int(data, bits=8, glob_stats=(0, 2, 100, data.size))
 
     res = cp.asnumpy(res_dev).astype(np.float32)
 
     np.testing.assert_array_equal(res[0, 0:3], 0.0)
-    np.testing.assert_array_equal(res_cpu[0, 0:3], 0.0)
 
 
 def test_rescale_double_offset():
     data = np.ones((30, 50), dtype=np.float32) + 10
 
     data_dev = cp.asarray(data)
     res_dev = rescale_to_int(data_dev, bits=8, glob_stats=(10, 12, 100, data.size))
-    res_cpu = rescale_to_int(data, bits=8, glob_stats=(10, 12, 100, data.size))
 
     res = cp.asnumpy(res_dev).astype(np.float32)
 
     np.testing.assert_array_almost_equal(res, 127.0)
-    np.testing.assert_array_almost_equal(res_cpu, 127.0)
 
 
 @pytest.mark.parametrize("bits", [8, 16])
@@ -99,14 +86,6 @@ def test_rescale_double_offset_min_percentage(bits: Literal[8, 16, 32]):
         perc_range_max=90.0,
     )
 
-    res_cpu = rescale_to_int(
-        data,
-        bits=bits,
-        glob_stats=(10, 20, 100, data.size),
-        perc_range_min=10.0,
-        perc_range_max=90.0,
-    )
-
     res = cp.asnumpy(res_dev).astype(np.float32)
 
     max = (2**bits) - 1
@@ -116,22 +95,13 @@ def test_rescale_double_offset_min_percentage(bits: Literal[8, 16, 32]):
     assert res[0, 0] == 0.0
     assert res[0, 1] == max
 
-    res_cpu = res_cpu.astype(np.float32)
-    np.testing.assert_array_almost_equal(res_cpu[1:, :], num)
-    assert res_cpu[0, 0] == 0.0
-    assert res_cpu[0, 1] == max
-
 
 def test_tomo_data_scale(data):
-    data_cpu = data.get()
     res_dev = rescale_to_int(
         data.astype(cp.float32), perc_range_min=10, perc_range_max=90, bits=8
     )
-    res_cpu = rescale_to_int(data_cpu, perc_range_min=10, perc_range_max=90, bits=8)
     res = res_dev.get()
-    assert res_dev.dtype == np.uint8
-    assert res_dev.dtype == np.uint8
-    np.testing.assert_array_equal(res_cpu, res)
+    assert res.dtype == np.uint8
 
 
 @pytest.mark.perf