removing CPU counterpart from rescale to int and updating tests

Author: dkazanc

httomolibgpu/misc/rescale.py

@@ -30,28 +30,27 @@
 
 from httomolibgpu.misc.supp_func import data_checker
 
-from sofia.core.regularisers import rescale_to_int_C
 
 __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)
@@ -71,7 +70,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
     """
@@ -84,18 +83,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]
@@ -104,33 +98,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
-
-    if xp.__name__ == "numpy":
-        if input_max == pow(2, 32):
-            input_max -= 1
-        res = rescale_to_int_C(data, input_min, input_max, factor)
-        # 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:
-        res = xp.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)
+        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

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,14 @@ 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)
 
 
 @pytest.mark.perf